Spaces:
Sleeping
Sleeping
File size: 4,829 Bytes
cfcf570 | 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 | """
Chatbot API routes.
"""
import logging
from uuid import uuid4
from fastapi import APIRouter, Depends, HTTPException
from psycopg import Connection
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_tavily import TavilySearch
from app.core.chatbot.config import CHAT_HISTORY_TABLE
from app.core.chatbot.chains import general_chain, news_chain
from app.core.chatbot.database import db_connection_dependency
from app.schemas.chat import ChatRequest
from app.services.chatbot.history import delete_guest_session, get_message_history
from app.services.chatbot.search import (
extract_search_query,
format_search_results,
get_tavily_search,
should_use_search,
)
logger = logging.getLogger(__name__)
chat = APIRouter()
@chat.post("/chat")
async def chat_endpoint(
request: ChatRequest,
tavily_search: TavilySearch = Depends(get_tavily_search),
):
"""
Chat endpoint with intelligent query processing.
Flow:
1. LLM classifies if query needs web search
2. If NO search needed -> respond directly with LLM knowledge
3. If search needed -> extract optimized search query -> search with Tavily -> respond with sources
Returns:
response: The LLM's response
session_id: Session ID for conversation history
used_search: Whether web search was used
search_reason: Why search was/wasn't used
"""
try:
# Ensure every session has an id for history tracking (DB or in-memory)
session_id = request.session_id or str(uuid4())
save_to_db = request.save_to_db
# Step 1: LLM decides if search is needed
needs_search, search_reason = should_use_search(request.query)
logger.info("Search needed: %s, Reason: %s", needs_search, search_reason)
def get_history(sid: str):
return get_message_history(sid, save_to_db)
# Step 2: Take appropriate path
if not needs_search:
# Respond directly without search
chain_with_history = RunnableWithMessageHistory(
general_chain,
get_history,
input_messages_key="question",
history_messages_key="chat_history",
)
response = chain_with_history.invoke(
{"question": request.query},
config={"configurable": {"session_id": session_id}},
)
else:
# Step 2a: Extract optimized search query
optimized_query = extract_search_query(request.query)
logger.info("Optimized search query: %s", optimized_query)
# Step 2b: Search with optimized query
search_results = tavily_search.invoke(optimized_query)
formatted_results = format_search_results(search_results)
logger.debug("Search results: %s", search_results)
# Step 2c: Respond with sources
chain_with_history = RunnableWithMessageHistory(
news_chain,
get_history,
input_messages_key="question",
history_messages_key="chat_history",
)
response = chain_with_history.invoke(
{
"question": request.query,
"search_results": formatted_results,
},
config={"configurable": {"session_id": session_id}},
)
logger.debug("Response: %s", response.content)
return {
"response": {
"content": response.content,
},
"session_id": session_id,
"used_search": needs_search,
"search_reason": search_reason,
}
except HTTPException:
raise
except Exception as exc: # pragma: no cover - defensive server guard
logger.exception("Unhandled error in chat_endpoint")
raise HTTPException(status_code=500, detail=str(exc))
@chat.delete("/chat/{session_id}")
async def delete_chat_history(session_id: str, conn: Connection = Depends(db_connection_dependency)):
"""Delete chat history for a specific session."""
try:
# Delete from in-memory guest sessions
delete_guest_session(session_id)
# Delete from database
with conn.cursor() as cursor:
cursor.execute(
f'DELETE FROM "{CHAT_HISTORY_TABLE}" WHERE session_id = %s',
(session_id,),
)
deleted_count = cursor.rowcount
conn.commit()
return {
"success": True,
"deleted_count": deleted_count,
"session_id": session_id,
}
except Exception as e:
logger.exception("Error in delete_chat_history")
raise HTTPException(status_code=500, detail=str(e))
|