File size: 6,074 Bytes
41e79c8 8f6bb93 41e79c8 8a27bb1 8f6bb93 41e79c8 0b87551 8f6bb93 0b87551 8a27bb1 0b87551 8a27bb1 41e79c8 8a27bb1 8f6bb93 0b87551 8f6bb93 8a27bb1 41e79c8 8a27bb1 0b87551 8f6bb93 0b87551 8f6bb93 0b87551 8f6bb93 0b87551 8a27bb1 8f6bb93 8a27bb1 41e79c8 0b87551 8f6bb93 8a27bb1 8f6bb93 41e79c8 0b87551 8a27bb1 0b87551 8f6bb93 0b87551 8f6bb93 0b87551 8f6bb93 2b63102 98714ab 0b87551 8f6bb93 0b87551 8f6bb93 0b87551 8f6bb93 38b974d 8f6bb93 0b87551 8f6bb93 0b87551 8f6bb93 2b63102 8f6bb93 8a27bb1 41e79c8 8a27bb1 41e79c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | import os
import time
import shutil
import tempfile
import logging
from typing import List, Optional
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from src.agent import get_agent_executor, get_llm, file_processor
from src.evaluator import faithfulness_score, answer_relevance_score, accuracy_score
from src.memory import SessionStore
from src.utils import (
extract_content,
parse_tool_results,
extract_citations,
is_invalid_key,
retry_delay,
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(title="Agentic RAG Service", version="4.0")
sessions = SessionStore(max_turns=6, keep_recent=4)
class QueryRequest(BaseModel):
query: str
api_key: Optional[str] = None # BYOK: each visitor supplies their own key
session_id: Optional[str] = None # conversation memory key
reference: Optional[str] = None
class QueryResponse(BaseModel):
response: str
source: str # "rag" | "web" | "rag+web" | "unknown"
citations: List[str] = []
faithfulness: Optional[float] = None
answer_relevance: Optional[float] = None
accuracy: Optional[float] = None
def _summarizer(api_key: str):
"""Return a str->str summarizer backed by the user's own LLM."""
llm = get_llm(api_key)
def summarize(text: str) -> str:
prompt = (
"Condense the following conversation into a concise summary that preserves "
"key facts, entities, decisions, and the user's goals. Keep it under 150 words.\n\n"
f"{text}"
)
return extract_content(llm.invoke(prompt))
return summarize
@app.get("/")
async def root():
return {"status": "active", "service": "Agentic Knowledge Search",
"docs_url": "/docs", "uploaded_docs": file_processor.get_status()}
@app.post("/upload")
async def upload_files(files: List[UploadFile] = File(...)):
tmp_dir = tempfile.mkdtemp()
try:
saved_paths = []
for upload in files:
dest = os.path.join(tmp_dir, upload.filename)
with open(dest, "wb") as f:
shutil.copyfileobj(upload.file, f)
saved_paths.append(dest)
status = file_processor.process_files(saved_paths)
return JSONResponse({"status": status})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
shutil.rmtree(tmp_dir, ignore_errors=True)
@app.post("/reset")
async def reset_documents():
file_processor.reset()
return {"status": "Uploaded documents cleared."}
@app.post("/clear_memory")
async def clear_memory(session_id: str):
sessions.clear(session_id)
return {"status": "Conversation memory cleared."}
@app.post("/chat", response_model=QueryResponse)
async def chat(request: QueryRequest):
# BYOK: a key must be supplied with every request
if not request.api_key or not request.api_key.strip():
raise HTTPException(
status_code=400,
detail="Please enter your Google Gemini API key to ask a question.",
)
api_key = request.api_key.strip()
try:
agent = get_agent_executor(api_key)
except Exception as e:
raise HTTPException(status_code=400, detail=str(e))
memory = sessions.get(request.session_id or "default")
messages = memory.build_messages(request.query)
for attempt in range(3):
try:
logger.info(f"Query (attempt {attempt + 1}, session={request.session_id}): {request.query}")
result = agent.invoke({"messages": messages})
answer = extract_content(result["messages"][-1])
# Some models occasionally return an empty turn after a tool call; retry.
if not answer.strip():
logger.warning("Empty answer from model; retrying.")
if attempt < 2:
continue
return QueryResponse(
response="I couldn't generate a response for that. Please rephrase and try again.",
source="unknown",
)
source, tool_output = parse_tool_results(result["messages"])
citations = extract_citations(tool_output)
faith = faithfulness_score(answer, tool_output) if tool_output else None
relevance = answer_relevance_score(request.query, answer)
acc = accuracy_score(answer, request.reference) if request.reference else None
memory.add_turn(request.query, answer)
try:
memory.summarize_if_needed(_summarizer(api_key))
except Exception as e:
logger.warning(f"Summarization skipped: {e}")
return QueryResponse(
response=answer,
source=source,
citations=citations,
faithfulness=faith,
answer_relevance=relevance,
accuracy=acc,
)
except Exception as e:
error_str = str(e)
if is_invalid_key(error_str):
raise HTTPException(
status_code=401,
detail="Invalid or unauthorized API key. Please check your Google Gemini API key.",
)
if "RESOURCE_EXHAUSTED" not in error_str:
logger.error(f"Chat error: {error_str}")
raise HTTPException(status_code=500, detail=error_str)
delay = retry_delay(error_str)
if delay and delay <= 120 and attempt < 2:
logger.warning(f"Rate limited, retrying in {delay:.0f}s...")
time.sleep(delay + 1)
continue
raise HTTPException(
status_code=429,
detail="Daily API quota exhausted. Please wait until tomorrow or upgrade your Gemini API plan.",
)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|