Spaces:
Running
Running
Update noRag.py
Browse files
noRag.py
CHANGED
@@ -1,6 +1,7 @@
|
|
1 |
# noRag.py
|
2 |
|
3 |
-
|
|
|
4 |
from pydantic import BaseModel
|
5 |
from groq import Groq
|
6 |
from pymongo import MongoClient
|
@@ -8,35 +9,45 @@ from config import CONNECTION_STRING, CHATGROQ_API_KEY, CUSTOM_PROMPT
|
|
8 |
|
9 |
router = APIRouter(prefix="/norag", tags=["noRag"])
|
10 |
|
11 |
-
#
|
12 |
client = Groq(api_key=CHATGROQ_API_KEY)
|
13 |
-
mongo
|
14 |
-
db
|
15 |
-
chats
|
16 |
|
17 |
SYSTEM_PROMPT = "You are a helpful assistant which helps people in their tasks."
|
18 |
|
|
|
19 |
class ChatRequest(BaseModel):
|
20 |
session_id: str
|
21 |
-
question:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
22 |
|
23 |
-
@router.post("/chat", summary="
|
24 |
async def chat_endpoint(req: ChatRequest):
|
25 |
-
# Fetch or create session in MongoDB
|
26 |
doc = chats.find_one({"session_id": req.session_id})
|
27 |
if not doc:
|
|
|
28 |
doc = {"session_id": req.session_id, "history": [], "summary": ""}
|
29 |
chats.insert_one(doc)
|
30 |
|
31 |
history, summary = doc["history"], doc["summary"]
|
32 |
|
33 |
-
#
|
34 |
if len(history) >= 10:
|
35 |
msgs = [f"{m['role']}: {m['content']}" for m in history]
|
36 |
combined = summary + "\n" + "\n".join(msgs)
|
37 |
sum_prompt = (
|
38 |
"Summarize the following chat history in one or two short sentences:\n\n"
|
39 |
-
|
|
|
40 |
)
|
41 |
sum_resp = client.chat.completions.create(
|
42 |
model="meta-llama/llama-4-scout-17b-16e-instruct",
|
@@ -49,15 +60,15 @@ async def chat_endpoint(req: ChatRequest):
|
|
49 |
summary = sum_resp.choices[0].message.content.strip()
|
50 |
history = []
|
51 |
|
52 |
-
#
|
53 |
-
|
54 |
full_prompt = CUSTOM_PROMPT.format(
|
55 |
context=SYSTEM_PROMPT,
|
56 |
-
chat_history=
|
57 |
-
question=req.question
|
58 |
)
|
59 |
|
60 |
-
#
|
61 |
resp = client.chat.completions.create(
|
62 |
model="meta-llama/llama-4-scout-17b-16e-instruct",
|
63 |
messages=[{"role": "user", "content": full_prompt}],
|
@@ -68,13 +79,13 @@ async def chat_endpoint(req: ChatRequest):
|
|
68 |
)
|
69 |
answer = resp.choices[0].message.content.strip()
|
70 |
|
71 |
-
#
|
72 |
-
history.append({"role": "user",
|
73 |
history.append({"role": "assistant", "content": answer})
|
74 |
chats.replace_one(
|
75 |
{"session_id": req.session_id},
|
76 |
{"session_id": req.session_id, "history": history, "summary": summary},
|
77 |
-
upsert=True
|
78 |
)
|
79 |
|
80 |
return {"session_id": req.session_id, "answer": answer, "summary": summary}
|
|
|
1 |
# noRag.py
|
2 |
|
3 |
+
import uuid
|
4 |
+
from fastapi import APIRouter
|
5 |
from pydantic import BaseModel
|
6 |
from groq import Groq
|
7 |
from pymongo import MongoClient
|
|
|
9 |
|
10 |
router = APIRouter(prefix="/norag", tags=["noRag"])
|
11 |
|
12 |
+
# clients
|
13 |
client = Groq(api_key=CHATGROQ_API_KEY)
|
14 |
+
mongo = MongoClient(CONNECTION_STRING)
|
15 |
+
db = mongo["edulearnai"]
|
16 |
+
chats = db["chats"]
|
17 |
|
18 |
SYSTEM_PROMPT = "You are a helpful assistant which helps people in their tasks."
|
19 |
|
20 |
+
|
21 |
class ChatRequest(BaseModel):
|
22 |
session_id: str
|
23 |
+
question: str
|
24 |
+
|
25 |
+
|
26 |
+
@router.post("/session", summary="Create a new chat session")
|
27 |
+
async def create_session():
|
28 |
+
session_id = str(uuid.uuid4())
|
29 |
+
chats.insert_one({"session_id": session_id, "history": [], "summary": ""})
|
30 |
+
return {"session_id": session_id}
|
31 |
+
|
32 |
|
33 |
+
@router.post("/chat", summary="Send a question to the assistant")
|
34 |
async def chat_endpoint(req: ChatRequest):
|
|
|
35 |
doc = chats.find_one({"session_id": req.session_id})
|
36 |
if not doc:
|
37 |
+
# if session not exist, create it
|
38 |
doc = {"session_id": req.session_id, "history": [], "summary": ""}
|
39 |
chats.insert_one(doc)
|
40 |
|
41 |
history, summary = doc["history"], doc["summary"]
|
42 |
|
43 |
+
# auto-summarize if too many turns
|
44 |
if len(history) >= 10:
|
45 |
msgs = [f"{m['role']}: {m['content']}" for m in history]
|
46 |
combined = summary + "\n" + "\n".join(msgs)
|
47 |
sum_prompt = (
|
48 |
"Summarize the following chat history in one or two short sentences:\n\n"
|
49 |
+
+ combined
|
50 |
+
+ "\n\nSummary:"
|
51 |
)
|
52 |
sum_resp = client.chat.completions.create(
|
53 |
model="meta-llama/llama-4-scout-17b-16e-instruct",
|
|
|
60 |
summary = sum_resp.choices[0].message.content.strip()
|
61 |
history = []
|
62 |
|
63 |
+
# build full prompt
|
64 |
+
hist_text = "\n".join(f"{m['role']}: {m['content']}" for m in history)
|
65 |
full_prompt = CUSTOM_PROMPT.format(
|
66 |
context=SYSTEM_PROMPT,
|
67 |
+
chat_history=hist_text or "(no prior messages)",
|
68 |
+
question=req.question,
|
69 |
)
|
70 |
|
71 |
+
# get answer
|
72 |
resp = client.chat.completions.create(
|
73 |
model="meta-llama/llama-4-scout-17b-16e-instruct",
|
74 |
messages=[{"role": "user", "content": full_prompt}],
|
|
|
79 |
)
|
80 |
answer = resp.choices[0].message.content.strip()
|
81 |
|
82 |
+
# persist
|
83 |
+
history.append({"role": "user", "content": req.question})
|
84 |
history.append({"role": "assistant", "content": answer})
|
85 |
chats.replace_one(
|
86 |
{"session_id": req.session_id},
|
87 |
{"session_id": req.session_id, "history": history, "summary": summary},
|
88 |
+
upsert=True,
|
89 |
)
|
90 |
|
91 |
return {"session_id": req.session_id, "answer": answer, "summary": summary}
|