kumardatascience/Multi-Source-RAG-AI-System-with-Query-Routing

1"""Chainlit chatbot with intelligent routing + live file upload support."""2 3import chainlit as cl4from router import ChatRouter5from rag import index_file_for_session, clear_session6 7 8workflow = ChatRouter(timeout=60)9 10 11@cl.on_chat_start12async def start():13    cl.user_session.set("history", [])14    await cl.Message(15        content="👋 Hi! I'm an intelligent assistant. Ask me anything, or upload a "16                "PDF/text file and ask questions about it. Watch the reasoning steps as I work!"17    ).send()18 19 20@cl.on_message21async def main(message: cl.Message):22    history = cl.user_session.get("history")23    session_id = cl.user_session.get("id") or "default"24 25    # 📎 Handle any uploaded files26    if message.elements:27        for element in message.elements:28            file_path = getattr(element, "path", None)29            if not file_path:30                continue31 32            async with cl.Step(33                name=f"📥 Indexing {element.name}",34                type="tool",35            ) as step:36                step.input = element.name37                try:38                    chunk_count = index_file_for_session(file_path, session_id)39                    step.output = f"Stored {chunk_count} chunks from **{element.name}**."40                except Exception as e:41                    step.output = f"⚠️ Failed to index {element.name}: {e}"42 43    history.append({"role": "user", "content": message.content})44 45    # Run the workflow with the session ID so RAG can find uploaded chunks46    result = await workflow.run(47        question=message.content,48        history=history,49        session_id=session_id,50    )51 52    # Stream the LLM reply53    reply = cl.Message(content="")54    full_response = ""55    async for chunk in result.stream:56        full_response += chunk57        await reply.stream_token(chunk)58    await reply.send()59 60    history.append({"role": "assistant", "content": full_response})61    cl.user_session.set("history", history)62 63 64@cl.on_chat_end65async def end():66    """Clean up the session's uploaded documents when the chat ends."""67    session_id = cl.user_session.get("id") or "default"68    clear_session(session_id)