import json from fastapi import HTTPException, Request from fastapi.responses import StreamingResponse from typing import List from langchain_core.documents import Document from ..schema.chat import ChatQueryRequest from ..schema.message import MessageCreate from ..schema.source import SourceCreate from ..core.chat_engine.query import ChatEngine from ..core.logger import SingletonLogger from ..controller import message as message_controller from ..controller import source as source_controller logger = SingletonLogger().get_logger() async def query_paper(user_id: int, payload: ChatQueryRequest, request: Request): """ Query a paper using the chat engine with conversation history. Automatically saves user query and assistant response to message table. Args: user_id: ID of the user making the query payload: ChatQueryRequest containing query and parameters request: FastAPI Request object to access app state Returns: StreamingResponse with SSE updates from the LangGraph execution Raises: HTTPException: If there's an error processing the query """ try: logger.info( f"User {user_id} querying paper {payload.paper_id} " f"in session {payload.session_id}: {payload.query}" ) # API keys are decrypted by APIKeyDecryptionMiddleware and stored in request.state graph = request.app.state.graph async def stream_and_save(): final_response = "" response_metadata = {} retrieved_docs: List[Document] = [] web_search_results: List[Document] = [] try: response_stream = ChatEngine.generate_response( graph=graph, query=payload.query, user_id=user_id, thread_id=payload.session_id, paper_id=payload.paper_id, model_name=payload.model_name, temperature=payload.temperature, max_tokens=payload.max_tokens, top_k=payload.top_k, use_web_search=payload.use_web_search, web_search_topic=payload.web_search_topic, request=request, ) async for chunk in response_stream: yield chunk try: if chunk.startswith("data: "): data_str = chunk[6:].strip() if data_str: data = json.loads(data_str) # Handle v2 streaming format if isinstance(data, dict): stream_type = data.get("type") stream_data = data.get("data") # LLM token streaming (new) if stream_type == "token": # Tokens are already streamed to client # Accumulate for final response if needed pass # State updates from nodes elif stream_type == "updates" and isinstance( stream_data, dict ): # Extract retrieved docs from rerank node if "rerank_docs_node" in stream_data: node_data = stream_data["rerank_docs_node"] if "retrieved_docs" in node_data: # Convert dict back to Document objects retrieved_docs = [ Document( page_content=doc.get( "page_content", "" ), metadata=doc.get( "metadata", {} ), ) for doc in node_data.get( "retrieved_docs", [] ) if isinstance(doc, dict) ] # Extract web search results from web crawl node if "web_crawl_node" in stream_data: node_data = stream_data["web_crawl_node"] if "web_search_results" in node_data: # Convert dict back to Document objects web_search_results = [ Document( page_content=doc.get( "page_content", "" ), metadata=doc.get( "metadata", {} ), ) for doc in node_data.get( "web_search_results", [] ) if isinstance(doc, dict) ] # Extract final response if "generate_response_node" in stream_data: node_data = stream_data[ "generate_response_node" ] if "response" in node_data: final_response = node_data["response"] if "response_metadata" in node_data: response_metadata = node_data[ "response_metadata" ] # Custom status messages elif stream_type == "custom": # Custom events are already streamed to client pass except (json.JSONDecodeError, KeyError) as e: logger.debug( f"Could not parse chunk for response extraction: {e}" ) continue if final_response: try: # Save the message first assistant_message_payload = MessageCreate( session_id=payload.session_id, user_id=user_id, content=[ {"role": "user", "content": payload.query}, {"role": "assistant", "content": final_response}, ], parent_message_id=None, model_used=payload.model_name, generation_metadata=response_metadata, ) assistant_message = await message_controller.create_message( user_id, assistant_message_payload ) logger.info( f"Saved assistant message with id: {assistant_message.id}" ) # Prepare and save sources sources_to_create: List[SourceCreate] = [] # Add retrieved document sources for doc in retrieved_docs: # Use all available metadata from the document doc_metadata = dict(doc.metadata) if doc.metadata else None sources_to_create.append( SourceCreate( message_id=assistant_message.id, source_text=doc.page_content[ :5000 ], # Limit text length source_type="document", source_url=doc.metadata.get("source", ""), metadata=doc_metadata, ) ) # Add web search result sources for doc in web_search_results: # Use all available metadata from web documents web_metadata = dict(doc.metadata) if doc.metadata else None sources_to_create.append( SourceCreate( message_id=assistant_message.id, source_text=doc.page_content[ :5000 ], # Limit text length source_type="web", source_url=doc.metadata.get("url", ""), metadata=web_metadata, ) ) # Save all sources in bulk if sources_to_create: saved_sources = await source_controller.create_sources( sources_to_create ) logger.info( f"Saved {len(saved_sources)} sources for message {assistant_message.id}" ) except Exception as e: logger.error( f"Failed to save assistant message or sources: {e}" ) except Exception as e: logger.error(f"Error in stream_and_save: {e}") yield f"data: {json.dumps({'error': str(e)})}\n\n" return StreamingResponse( stream_and_save(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) except Exception as e: logger.exception(f"Error querying paper: {e}") raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")