wehub-resource-sync
280 lines
9.0 KiB
Python
280 lines
9.0 KiB
Python
import json
|
|
import time
|
|
from typing import List, Optional, Dict, Any
|
|
|
|
from fastapi import FastAPI, HTTPException, Query
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
|
|
from services.zep_service import zep_service
|
|
from services.llm_service import llm_service
|
|
from config.settings import settings
|
|
|
|
|
|
# -------------------------------
|
|
# Pydantic Models
|
|
# -------------------------------
|
|
|
|
class Message(BaseModel):
|
|
role: str
|
|
content: str
|
|
|
|
|
|
class ChatRequest(BaseModel):
|
|
messages: List[Message]
|
|
|
|
|
|
class SessionCreateRequest(BaseModel):
|
|
user_id: str
|
|
first_name: Optional[str] = "Demo"
|
|
last_name: Optional[str] = None
|
|
email: Optional[str] = None
|
|
|
|
|
|
# -------------------------------
|
|
# FastAPI App
|
|
# -------------------------------
|
|
|
|
app = FastAPI(
|
|
title="Zep + LLM Backend for Anam",
|
|
description="FastAPI backend that integrates Zep Cloud KG + your LLM for Anam avatar.",
|
|
version="1.0.0",
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"], # Allow all origins for local development
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
expose_headers=["*"],
|
|
)
|
|
|
|
|
|
# -------------------------------
|
|
# Health & Debug Endpoints
|
|
# -------------------------------
|
|
|
|
@app.get("/health")
|
|
async def health_check():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.get("/zep/test-graph")
|
|
async def test_graph(q: str = Query("what is Zep?")):
|
|
"""
|
|
Simple debug endpoint to verify Zep graph connectivity.
|
|
"""
|
|
try:
|
|
nodes = await zep_service.search_graph(
|
|
query=q,
|
|
user_id=settings.zep_docs_user_id,
|
|
limit=3,
|
|
scope="nodes",
|
|
)
|
|
return {"query": q, "results": nodes}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# -------------------------------
|
|
# Zep Session Helper Endpoint
|
|
# -------------------------------
|
|
|
|
@app.post("/zep/session")
|
|
async def create_zep_session(body: SessionCreateRequest):
|
|
"""
|
|
Helper endpoint to create a Zep user + thread (session) for a given user_id.
|
|
You can call this from Streamlit or any client instead of doing it inline.
|
|
"""
|
|
try:
|
|
user = await zep_service.create_user(
|
|
user_id=body.user_id,
|
|
first_name=body.first_name,
|
|
last_name=body.last_name,
|
|
email=body.email,
|
|
)
|
|
|
|
# 2) Create thread for this user
|
|
session_id = f"session-{body.user_id}"
|
|
thread = await zep_service.create_thread(
|
|
thread_id=session_id,
|
|
user_id=body.user_id,
|
|
)
|
|
|
|
return {
|
|
"user": user,
|
|
"thread": thread,
|
|
"session_id": session_id,
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# -------------------------------
|
|
# Core: LLM Stream Endpoint for Anam
|
|
# -------------------------------
|
|
|
|
@app.post("/llm/stream")
|
|
async def llm_stream(
|
|
payload: ChatRequest,
|
|
session_id: str = Query(..., description="Zep thread/session ID"),
|
|
):
|
|
"""
|
|
Streaming LLM endpoint with Zep KG integration.
|
|
This is what Anam's customLLMHandler() should call.
|
|
|
|
- Expects: JSON body { "messages": [ { "role": "...", "content": "..." }, ... ] }
|
|
- Requires: ?session_id=... query param (thread_id in Zep)
|
|
- Returns: SSE stream with chunks as { "text": "<token or chunk>" }
|
|
"""
|
|
|
|
messages: List[Message] = payload.messages
|
|
|
|
if not messages:
|
|
raise HTTPException(status_code=400, detail="No messages provided")
|
|
|
|
# Extract latest user message
|
|
user_message: Optional[str] = None
|
|
for m in reversed(messages):
|
|
if m.role == "user":
|
|
user_message = m.content
|
|
break
|
|
|
|
if not user_message:
|
|
raise HTTPException(status_code=400, detail="No user message found")
|
|
|
|
async def event_generator():
|
|
"""
|
|
Async generator that:
|
|
1) Saves the user message to Zep
|
|
2) Queries the Zep Knowledge Graph for docs context
|
|
3) Gets conversation context from Zep
|
|
4) Calls your LLM in streaming mode
|
|
5) Yields SSE 'data: { "text": ... }' lines to the client
|
|
6) Saves the final assistant message back into Zep
|
|
"""
|
|
full_response: str = ""
|
|
|
|
try:
|
|
print(f"\n{'='*60}")
|
|
print(f"NEW USER MESSAGE for session: {session_id}")
|
|
print(f"Message: {user_message}")
|
|
print(f"{'='*60}")
|
|
|
|
# 1) Save user message into Zep thread
|
|
await zep_service.add_message(
|
|
thread_id=session_id,
|
|
role="user",
|
|
content=user_message,
|
|
)
|
|
print("User message saved to Zep thread")
|
|
|
|
# 2) Build base system prompt
|
|
system_prompt = """
|
|
You are a helpful AI assistant. When provided with relevant information or context below,
|
|
use it to answer the user's question accurately and completely. Always use the information
|
|
provided to give thorough, detailed answers. Be conversational and friendly while being informative.
|
|
If the context contains the answer, state it clearly and add relevant details from the context.
|
|
""".strip()
|
|
|
|
# 3) Zep Knowledge Graph lookup
|
|
print(f"\nQUERYING KNOWLEDGE GRAPH...")
|
|
print(f" Query: '{user_message}'")
|
|
print(f" User ID: {settings.zep_docs_user_id}")
|
|
|
|
graph_start_time = time.time()
|
|
graph_results = await zep_service.search_graph(
|
|
query=user_message,
|
|
user_id=settings.zep_docs_user_id,
|
|
limit=3,
|
|
scope="nodes",
|
|
)
|
|
graph_end_time = time.time()
|
|
# Convert to milliseconds
|
|
graph_duration = (graph_end_time - graph_start_time) * 1000
|
|
|
|
print(
|
|
f"\nKNOWLEDGE GRAPH RESULTS: {len(graph_results)} nodes found")
|
|
print(f"Fetch time: {graph_duration:.2f}ms")
|
|
|
|
if graph_results:
|
|
print(f"\nRetrieved nodes:")
|
|
for i, node in enumerate(graph_results, 1):
|
|
print(f" {i}. {node.get('name', 'N/A')}")
|
|
summary_preview = node.get('summary', '')[:100]
|
|
print(f" Summary: {summary_preview}...")
|
|
|
|
docs_context = "\n\n".join(
|
|
[f"- {node['name']}: {node['summary']}" for node in graph_results]
|
|
)
|
|
system_prompt += f"\n\n=== RELEVANT INFORMATION ===\n{docs_context}\n\nUse the above information to answer the user's question."
|
|
print(
|
|
f"\nGRAPH CONTEXT ADDED TO PROMPT: {len(docs_context)} characters")
|
|
else:
|
|
print(
|
|
f"\nNO GRAPH RESULTS FOUND - LLM will respond without context")
|
|
|
|
# 4) Conversation context from this thread
|
|
print(f"\nFETCHING USER CONTEXT for thread: {session_id}")
|
|
user_context = await zep_service.get_user_context(thread_id=session_id)
|
|
if user_context:
|
|
system_prompt += f"\n\n=== CONVERSATION CONTEXT ===\n{user_context}"
|
|
print(f"USER CONTEXT ADDED: {len(user_context)} characters")
|
|
else:
|
|
print(f"No user context available yet (new conversation)")
|
|
|
|
# 5) Reformat messages for LLM
|
|
formatted_messages: List[Dict[str, Any]] = [
|
|
{"role": m.role, "content": m.content} for m in messages
|
|
]
|
|
|
|
# 6) Stream from LLM and send SSE chunks
|
|
print(f"\nSTREAMING LLM RESPONSE...")
|
|
print(f" Total prompt size: {len(system_prompt)} characters")
|
|
|
|
chunk_count = 0
|
|
async for chunk in llm_service.stream_chat_completion(
|
|
messages=formatted_messages,
|
|
system_prompt=system_prompt,
|
|
):
|
|
if not chunk:
|
|
continue
|
|
|
|
chunk_count += 1
|
|
full_response += chunk
|
|
payload = json.dumps({"content": chunk})
|
|
yield f"data: {payload}\n\n"
|
|
|
|
if chunk_count % 10 == 0:
|
|
print(f" Streamed {chunk_count} chunks so far...")
|
|
|
|
print(f" Total chunks streamed: {chunk_count}")
|
|
|
|
# 7) Save assistant message in Zep
|
|
if full_response.strip():
|
|
await zep_service.add_message(
|
|
thread_id=session_id,
|
|
role="assistant",
|
|
content=full_response,
|
|
)
|
|
print(
|
|
f"\nASSISTANT RESPONSE SAVED: {len(full_response)} characters")
|
|
print(f"{'='*60}\n")
|
|
|
|
except Exception as e:
|
|
# On error, stream an error message as SSE
|
|
err_payload = json.dumps({"content": "Error: " + str(e)})
|
|
yield f"data: {err_payload}\n\n"
|
|
|
|
# Return a streaming SSE response
|
|
return StreamingResponse(
|
|
event_generator(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|