Spaces:
Running
Running
Update chat.py
Browse files
chat.py
CHANGED
@@ -11,7 +11,7 @@ from models import ChatIDOut, MessageIn
|
|
11 |
|
12 |
router = APIRouter(prefix="/chat", tags=["chat"])
|
13 |
|
14 |
-
# βββ LLM & Prompt Setup
|
15 |
def get_llm() -> ChatGroq:
|
16 |
if not config.CHATGROQ_API_KEY:
|
17 |
raise RuntimeError("CHATGROQ_API_KEY not set in environment")
|
@@ -35,7 +35,7 @@ qa_template = ChatPromptTemplate.from_messages(
|
|
35 |
]
|
36 |
)
|
37 |
|
38 |
-
# βββ MongoDB History Setup
|
39 |
chat_sessions: dict[str, MongoDBChatMessageHistory] = {}
|
40 |
|
41 |
def create_history(session_id: str) -> MongoDBChatMessageHistory:
|
@@ -54,7 +54,7 @@ def get_history(session_id: str) -> MongoDBChatMessageHistory:
|
|
54 |
raise HTTPException(status_code=404, detail="Chat session not found")
|
55 |
return history
|
56 |
|
57 |
-
# βββ Summarization (to control token use)
|
58 |
def summarize_if_needed(history: MongoDBChatMessageHistory, threshold: int = 10):
|
59 |
msgs = history.messages
|
60 |
if len(msgs) <= threshold:
|
@@ -79,7 +79,7 @@ def summarize_if_needed(history: MongoDBChatMessageHistory, threshold: int = 10)
|
|
79 |
history.clear()
|
80 |
history.add_ai_message(f"[Summary] {summary.content}")
|
81 |
|
82 |
-
# βββ Endpoints
|
83 |
|
84 |
@router.post("", response_model=ChatIDOut)
|
85 |
async def create_chat():
|
@@ -118,8 +118,9 @@ async def post_message(
|
|
118 |
|
119 |
async def stream_generator():
|
120 |
full_response = ""
|
121 |
-
|
122 |
-
|
|
|
123 |
content = (
|
124 |
chunk.get("content")
|
125 |
or chunk.get("choices", [{}])[0].get("delta", {}).get("content")
|
|
|
11 |
|
12 |
router = APIRouter(prefix="/chat", tags=["chat"])
|
13 |
|
14 |
+
# βββ LLM & Prompt Setup ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
15 |
def get_llm() -> ChatGroq:
|
16 |
if not config.CHATGROQ_API_KEY:
|
17 |
raise RuntimeError("CHATGROQ_API_KEY not set in environment")
|
|
|
35 |
]
|
36 |
)
|
37 |
|
38 |
+
# βββ MongoDB History Setup βββββββββββββββββββββββββββββββββββββββββββββββββββ
|
39 |
chat_sessions: dict[str, MongoDBChatMessageHistory] = {}
|
40 |
|
41 |
def create_history(session_id: str) -> MongoDBChatMessageHistory:
|
|
|
54 |
raise HTTPException(status_code=404, detail="Chat session not found")
|
55 |
return history
|
56 |
|
57 |
+
# βββ Summarization (to control token use) ββββββββββββββββββββββββββββββββββββ
|
58 |
def summarize_if_needed(history: MongoDBChatMessageHistory, threshold: int = 10):
|
59 |
msgs = history.messages
|
60 |
if len(msgs) <= threshold:
|
|
|
79 |
history.clear()
|
80 |
history.add_ai_message(f"[Summary] {summary.content}")
|
81 |
|
82 |
+
# βββ Endpoints ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
83 |
|
84 |
@router.post("", response_model=ChatIDOut)
|
85 |
async def create_chat():
|
|
|
118 |
|
119 |
async def stream_generator():
|
120 |
full_response = ""
|
121 |
+
# Pass messages list as the first positional argument to .stream()
|
122 |
+
for chunk in llm.stream(messages):
|
123 |
+
# adapt to ChatGroq response structure
|
124 |
content = (
|
125 |
chunk.get("content")
|
126 |
or chunk.get("choices", [{}])[0].get("delta", {}).get("content")
|