Update noRag.py

noRag.py

@@ -1,120 +1,84 @@
-# noRag.py
+# norag_router.py
 
-import asyncio
-import os
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
 from groq import Groq
 from pymongo import MongoClient
-from config import (
-    CONNECTION_STRING,
-    CHATGROQ_API_KEY,
-    CUSTOM_PROMPT
-)
+from config import CONNECTION_STRING, CHATGROQ_API_KEY, CUSTOM_PROMPT
 
-# --- Setup Groq client and MongoDB ---
+# Create router under /norag
+router = APIRouter(prefix="/norag", tags=["noRag"])
+
+# Initialize Groq client and MongoDB
 client = Groq(api_key=CHATGROQ_API_KEY)
 mongo  = MongoClient(CONNECTION_STRING)
 db     = mongo["edulearnai"]
 chats  = db["chats"]
 
+# System prompt
 SYSTEM_PROMPT = "You are a helpful assistant which helps people in their tasks."
 
-# --- Session management in MongoDB ---
-def get_session(session_id: str) -> dict:
-    """Fetch or create a session doc with keys: session_id, history (list), summary (str)."""
-    doc = chats.find_one({"session_id": session_id})
+# Request model
+type ChatRequest(BaseModel):
+    session_id: str
+    question: str
+
+@router.post("/chat", summary="Ask a question to the noRag assistant")
+async def chat_endpoint(req: ChatRequest):
+    # Fetch or create session
+doc = chats.find_one({"session_id": req.session_id})
     if not doc:
-        doc = {"session_id": session_id, "history": [], "summary": ""}
+        doc = {"session_id": req.session_id, "history": [], "summary": ""}
         chats.insert_one(doc)
-    return doc
 
-def save_session(doc: dict):
-    """Overwrite the session document in Mongo."""
-    chats.replace_one({"session_id": doc["session_id"]}, doc)
+    history = doc["history"]
+    summary = doc["summary"]
 
-# --- History summarization ---
-async def summarize_history(prev_summary: str, history_msgs: list[str]) -> str:
-    """Ask the LLM to produce a short summary of the combined previous summary + new messages."""
-    combined = prev_summary + "\n" + "\n".join(history_msgs)
-    prompt = (
-        "Summarize the following chat history in one or two short sentences:\n\n"
-        f"{combined}\n\nSummary:"
-    )
-    resp = client.chat.completions.create(
-        model="meta-llama/llama-4-scout-17b-16e-instruct",
-        messages=[{"role": "user", "content": prompt}],
-        temperature=0.3,
-        max_completion_tokens=150,
-        top_p=1,
-        stream=False,
-    )
-    # the first (and only) completion
-    return resp.choices[0].message.content.strip()
-
-# --- Core chat logic ---
-async def chat(session_id: str, question: str):
-    session = get_session(session_id)
-    history = session["history"]
-    summary = session["summary"]
-
-    # If history is too long, summarize it and clear
+    # Summarize if history too long
     if len(history) >= 10:
-        msgs_to_summarize = [f"{m['role']}: {m['content']}" for m in history]
-        new_summary = await summarize_history(summary, msgs_to_summarize)
-        session["summary"] = new_summary
-        session["history"] = []
+        msgs = [f"{m['role']}: {m['content']}" for m in history]
+        combined = summary + "\n" + "\n".join(msgs)
+        sum_prompt = (
+            "Summarize the following chat history in one or two short sentences:\n\n"
+            + combined + "\n\nSummary:"
+        )
+        sum_resp = client.chat.completions.create(
+            model="meta-llama/llama-4-scout-17b-16e-instruct",
+            messages=[{"role": "user", "content": sum_prompt}],
+            temperature=0.3,
+            max_completion_tokens=150,
+            top_p=1,
+            stream=False,
+        )
+        summary = sum_resp.choices[0].message.content.strip()
         history = []
 
-    # Build the prompt
-    chat_history_text = "\n".join([f"{m['role']}: {m['content']}" for m in history])
+    # Build full prompt
+    chat_hist_text = "\n".join([f"{m['role']}: {m['content']}" for m in history])
     full_prompt = CUSTOM_PROMPT.format(
         context=SYSTEM_PROMPT,
-        chat_history=chat_history_text or "(no prior messages)",
-        question=question
+        chat_history=chat_hist_text or "(no prior messages)",
+        question=req.question
     )
 
-    # Call the model, streaming
-    completion = client.chat.completions.create(
+    # Call model
+    resp = client.chat.completions.create(
         model="meta-llama/llama-4-scout-17b-16e-instruct",
         messages=[{"role": "user", "content": full_prompt}],
         temperature=1,
         max_completion_tokens=1024,
         top_p=1,
-        stream=True,
+        stream=False,
     )
+    answer = resp.choices[0].message.content.strip()
 
-    # Print & accumulate the assistant’s reply
-    assistant_response = ""
-    print("Assistant:", end=" ", flush=True)
-    for chunk in completion:
-        delta = chunk.choices[0].delta.content or ""
-        print(delta, end="", flush=True)
-        assistant_response += delta
-    print()  # newline after done
-
-    # Persist the new exchange
-    session["history"].append({"role": "user",      "content": question})
-    session["history"].append({"role": "assistant", "content": assistant_response})
-    save_session(session)
-
-# --- CLI loop ---
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--session",
-        "-s",
-        default="default",
-        help="Session ID (used to key chat history in MongoDB)"
+    # Update session doc
+    history.append({"role": "user", "content": req.question})
+    history.append({"role": "assistant", "content": answer})
+    chats.replace_one(
+        {"session_id": req.session_id},
+        {"session_id": req.session_id, "history": history, "summary": summary},
+        upsert=True
     )
-    args = parser.parse_args()
 
-    print(f"Starting noRag chat (session={args.session}).  Type Ctrl+C to quit.\n")
-    try:
-        while True:
-            user_q = input("You: ")
-            if not user_q.strip():
-                continue
-            asyncio.run(chat(args.session, user_q))
-    except KeyboardInterrupt:
-        print("\nGoodbye!")
+    return {"session_id": req.session_id, "answer": answer, "summary": summary}