Create chat.py

chat.py

@@ -0,0 +1,133 @@
+# routes/chat.py
+import uuid
+from fastapi import APIRouter, HTTPException, Path
+from fastapi.responses import StreamingResponse
+from langchain_groq import ChatGroq
+from langchain.prompts import ChatPromptTemplate
+from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
+
+import config
+from models import ChatIDOut, MessageIn
+
+router = APIRouter(prefix="/chat", tags=["chat"])
+
+# ─── LLM & Prompt Setup ─────────────────────────────────────────────────
+def get_llm() -> ChatGroq:
+    if not config.CHATGROQ_API_KEY:
+        raise RuntimeError("CHATGROQ_API_KEY not set in environment")
+    return ChatGroq(
+        model="llama-3.3-70b-versatile",
+        temperature=0,
+        max_tokens=1024,
+        api_key=config.CHATGROQ_API_KEY
+    )
+
+llm = get_llm()
+
+SYSTEM_PROMPT = """
+You are an assistant specialized in solving quizzes. Your goal is to provide accurate,
+concise, and contextually relevant answers.
+"""
+qa_template = ChatPromptTemplate.from_messages(
+    [
+        ("system", SYSTEM_PROMPT),
+        ("user", "{question}"),
+    ]
+)
+
+# ─── MongoDB History Setup ────────────────────────────────────────────────
+chat_sessions: dict[str, MongoDBChatMessageHistory] = {}
+
+def create_history(session_id: str) -> MongoDBChatMessageHistory:
+    history = MongoDBChatMessageHistory(
+        session_id=session_id,
+        connection_string=config.CONNECTION_STRING,
+        database_name="Education_chatbot",
+        collection_name="chat_histories",
+    )
+    chat_sessions[session_id] = history
+    return history
+
+def get_history(session_id: str) -> MongoDBChatMessageHistory:
+    history = chat_sessions.get(session_id)
+    if not history:
+        raise HTTPException(status_code=404, detail="Chat session not found")
+    return history
+
+# ─── Summarization (to control token use) ─────────────────────────────────
+def summarize_if_needed(history: MongoDBChatMessageHistory, threshold: int = 10):
+    msgs = history.messages
+    if len(msgs) <= threshold:
+        return
+
+    summarization_prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a summarization assistant."),
+            ("user",
+                "Here is the chat history:\n\n{chat_history}\n\n"
+                "Summarize the above chat messages into a single concise message with key details."
+            ),
+        ]
+    )
+    text_history = "\n".join(
+        f"{'User' if m.type=='human' else 'Assistant'}: {m.content}"
+        for m in msgs
+    )
+    summary_chain = summarization_prompt | llm
+    summary = summary_chain.invoke({"chat_history": text_history})
+
+    history.clear()
+    history.add_ai_message(f"[Summary] {summary.content}")
+
+# ─── Endpoints ────────────────────────────────────────────────────────────
+
+@router.post("", response_model=ChatIDOut)
+async def create_chat():
+    """
+    Create a new chat session and return its ID.
+    """
+    session_id = str(uuid.uuid4())
+    create_history(session_id)
+    return ChatIDOut(chat_id=session_id)
+
+@router.post("/{chat_id}/message")
+async def post_message(
+    chat_id: str = Path(..., description="The chat session ID"),
+    payload: MessageIn = None
+):
+    """
+    Send a question and stream back the assistant's answer.
+    """
+    history = get_history(chat_id)
+    question = (payload and payload.question.strip()) or ""
+    if not question:
+        raise HTTPException(status_code=400, detail="Question cannot be empty")
+
+    # Summarize old turns if too long
+    summarize_if_needed(history)
+
+    # Build conversation for the LLM
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for msg in history.messages:
+        role = "user" if msg.type == "human" else "assistant"
+        messages.append({"role": role, "content": msg.content})
+    messages.append({"role": "user", "content": question})
+
+    # Persist user turn
+    history.add_user_message(question)
+
+    async def stream_generator():
+        full_response = ""
+        for chunk in llm.stream(messages=messages):
+            # adjust based on actual ChatGroq chunk schema
+            content = (
+                chunk.get("content")
+                or chunk.get("choices", [{}])[0].get("delta", {}).get("content")
+            )
+            if content:
+                yield content
+                full_response += content
+        # save final AI message
+        history.add_ai_message(full_response)
+
+    return StreamingResponse(stream_generator(), media_type="text/plain")