sources

app.py

@@ -1,7 +1,10 @@
 import gradio as gr
 import asyncio
 import logging
-from utils.generator import generate, generate_streaming
+import json
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+from utils.generator import generate_streaming, generate
 
 # Configure logging
 logging.basicConfig(
@@ -14,6 +17,100 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# ---------------------------------------------------------------------
+# FastAPI app for ChatUI endpoints
+# ---------------------------------------------------------------------
+app = FastAPI(title="ChatFed Generator", version="1.0.0")
+
+@app.post("/generate")
+async def generate_endpoint(request: Request):
+    """
+    Non-streaming generation endpoint for ChatUI format.
+    
+    Expected request body:
+    {
+        "query": "user question",
+        "context": [...] // list of retrieval results
+    }
+    
+    Returns ChatUI format:
+    {
+        "answer": "response with citations [1][2]",
+        "sources": [{"link": "doc://...", "title": "..."}]
+    }
+    """
+    try:
+        body = await request.json()
+        query = body.get("query", "")
+        context = body.get("context", [])
+        
+        result = await generate(query, context, chatui_format=True)
+        return result
+        
+    except Exception as e:
+        logger.exception("Generation endpoint failed")
+        return {"error": str(e)}
+
+@app.post("/generate/stream")
+async def generate_stream_endpoint(request: Request):
+    """
+    Streaming generation endpoint for ChatUI format.
+    
+    Expected request body:
+    {
+        "query": "user question", 
+        "context": [...] // list of retrieval results
+    }
+    
+    Returns Server-Sent Events in ChatUI format:
+    event: data
+    data: "response chunk"
+    
+    event: sources  
+    data: {"sources": [...]}
+    
+    event: end
+    """
+    try:
+        body = await request.json()
+        query = body.get("query", "")
+        context = body.get("context", [])
+        
+        async def event_stream():
+            async for event in generate_streaming(query, context, chatui_format=True):
+                event_type = event["event"]
+                event_data = event["data"]
+                
+                if event_type == "data":
+                    yield f"event: data\ndata: {json.dumps(event_data)}\n\n"
+                elif event_type == "sources":
+                    yield f"event: sources\ndata: {json.dumps(event_data)}\n\n"
+                elif event_type == "end":
+                    yield f"event: end\ndata: {{}}\n\n"
+                elif event_type == "error":
+                    yield f"event: error\ndata: {json.dumps(event_data)}\n\n"
+        
+        return StreamingResponse(
+            event_stream(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "Access-Control-Allow-Origin": "*",
+                "Access-Control-Allow-Headers": "*",
+            }
+        )
+        
+    except Exception as e:
+        logger.exception("Streaming endpoint failed")
+        async def error_stream():
+            yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"
+        
+        return StreamingResponse(
+            error_stream(),
+            media_type="text/event-stream"
+        )
+
 # ---------------------------------------------------------------------
 # Wrapper function to handle async streaming for Gradio
 # ---------------------------------------------------------------------
@@ -22,7 +119,7 @@ def generate_streaming_wrapper(query: str, context: str):
     logger.info(f"Starting generation request - Query length: {len(query)}, Context length: {len(context)}")
     
     async def _async_generator():
-        async for chunk in generate_streaming(query, context):
+        async for chunk in generate_streaming(query, context, chatui_format=False):
             yield chunk
     
     # Create a new event loop for this thread
@@ -81,13 +178,15 @@ ui = gr.Interface(
     api_name="generate"
 )
 
+# Mount Gradio app to FastAPI
+app = gr.mount_gradio_app(app, ui, path="/gradio")
+
 # Launch with MCP server enabled
 if __name__ == "__main__":
+    import uvicorn
     logger.info("Starting ChatFed Generation Module server")
-    logger.info("Server will be available at http://0.0.0.0:7860")
-    ui.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        # mcp_server=True,
-        show_error=True
-    )
+    logger.info("FastAPI server will be available at http://0.0.0.0:7860")
+    logger.info("Gradio UI will be available at http://0.0.0.0:7860/gradio")
+    logger.info("ChatUI endpoints: /generate (non-streaming), /generate/stream (streaming)")
+    
+    uvicorn.run(app, host="0.0.0.0", port=7860)

utils/generator.py

@@ -86,8 +86,8 @@ def extract_relevant_fields(retrieval_results: List[Dict[str, Any]]) -> List[Dic
     Returns:
         List of processed objects with only relevant fields
     """
-
-    retrieval_results = ast.literal_eval(retrieval_results)
+    if isinstance(retrieval_results, str):
+        retrieval_results = ast.literal_eval(retrieval_results)
 
     processed_results = []
     
@@ -191,7 +191,11 @@ def build_messages(question: str, context: str) -> list:
     """
     system_content = (
         "You are an expert assistant. Answer the USER question using only the "
-        "CONTEXT provided. If the context is insufficient say 'I don't know.'"
+        "CONTEXT provided. When referencing information from the context, use inline "
+        "citations in square brackets like [1], [2], etc. to reference the document "
+        "numbers shown in the context. Use multiple citations when information comes "
+        "from multiple documents, like [1][2]. If the context is insufficient, say "
+        "'I don't know.'"
     )
     
     user_content = f"### CONTEXT\n{context}\n\n### USER QUESTION\n{question}"
@@ -201,8 +205,7 @@ def build_messages(question: str, context: str) -> list:
         HumanMessage(content=user_content)
     ]
 
-    
-async def generate(query: str, context: Union[str, List[Dict[str, Any]]]) -> str:
+async def generate(query: str, context: Union[str, List[Dict[str, Any]]], chatui_format: bool = False) -> Union[str, Dict[str, Any]]:
     """
     Generate an answer to a query using provided context through RAG.
     
@@ -211,42 +214,79 @@ async def generate(query: str, context: Union[str, List[Dict[str, Any]]]) -> str
     
     Args:
         query (str): User query
-        context (list): List of retrieval result objects (dictionaries)
+        context (Union[str, List[Dict[str, Any]]]): Context as string or list of retrieval results
+        chatui_format (bool): If True, return ChatUI format with sources
+        
     Returns:
-        str: The generated answer based on the query and context
+        Union[str, Dict]: The generated answer or ChatUI format response
     """
     if not query.strip():
-        return "Error: Query cannot be empty"
+        return {"error": "Query cannot be empty"} if chatui_format else "Error: Query cannot be empty"
+    
+    processed_results = []
     
     # Handle both string context (for Gradio UI) and list context (from retriever)
     if isinstance(context, list):
         if not context:
-            return "Error: No retrieval results provided"
+            return {"error": "No retrieval results provided"} if chatui_format else "Error: No retrieval results provided"
         
         # Process the retrieval results
         processed_results = extract_relevant_fields(context)
         formatted_context = format_context_from_results(processed_results)
         
         if not formatted_context.strip():
-            return "Error: No valid content found in retrieval results"
+            return {"error": "No valid content found in retrieval results"} if chatui_format else "Error: No valid content found in retrieval results"
     
     elif isinstance(context, str):
         if not context.strip():
-            return "Error: Context cannot be empty"
+            return {"error": "Context cannot be empty"} if chatui_format else "Error: Context cannot be empty"
         formatted_context = context
     
     else:
-        return "Error: Context must be either a string or list of retrieval results"
+        return {"error": "Context must be either a string or list of retrieval results"} if chatui_format else "Error: Context must be either a string or list of retrieval results"
     
     try:
         messages = build_messages(query, formatted_context)
         answer = await _call_llm(messages)
-        return answer
+        
+        if chatui_format:
+            # Return ChatUI format
+            result = {"answer": answer}
+            if processed_results:
+                # Extract sources for ChatUI
+                sources = []
+                for result_item in processed_results:
+                    filename = result_item.get('filename', 'Unknown')
+                    page = result_item.get('page', 'Unknown')
+                    year = result_item.get('year', 'Unknown')
+                    
+                    # Create link using doc:// scheme
+                    link = f"doc://{filename}"
+                    
+                    # Create descriptive title
+                    title_parts = [filename]
+                    if page != 'Unknown':
+                        title_parts.append(f"Page {page}")
+                    if year != 'Unknown':
+                        title_parts.append(f"({year})")
+                    
+                    title = " - ".join(title_parts)
+                    
+                    sources.append({
+                        "link": link,
+                        "title": title
+                    })
+                
+                result["sources"] = sources
+            return result
+        else:
+            return answer
+            
     except Exception as e:
         logging.exception("Generation failed")
-        return f"Error: {str(e)}"
+        return {"error": str(e)} if chatui_format else f"Error: {str(e)}"
 
-async def generate_streaming(query: str, context: Union[str, List[Dict[str, Any]]]) -> AsyncGenerator[str, None]:
+async def generate_streaming(query: str, context: Union[str, List[Dict[str, Any]]], chatui_format: bool = False) -> AsyncGenerator[Union[str, Dict[str, Any]], None]:
     """
     Generate a streaming answer to a query using provided context through RAG.
     
@@ -256,18 +296,27 @@ async def generate_streaming(query: str, context: Union[str, List[Dict[str, Any]
     Args:
         query (str): User query
         context (Union[str, List[Dict[str, Any]]]): Context as string or list of retrieval results
+        chatui_format (bool): If True, yield ChatUI format events
         
     Yields:
-        str: Streaming chunks of the generated answer
+        Union[str, Dict]: Streaming chunks or ChatUI format events
     """
     if not query.strip():
-        yield "Error: Query cannot be empty"
+        if chatui_format:
+            yield {"event": "error", "data": {"error": "Query cannot be empty"}}
+        else:
+            yield "Error: Query cannot be empty"
         return
     
+    processed_results = []
+    
     # Handle both string context (for Gradio UI) and list context (from retriever)
     if isinstance(context, list):
         if not context:
-            yield "Error: No retrieval results provided"
+            if chatui_format:
+                yield {"event": "error", "data": {"error": "No retrieval results provided"}}
+            else:
+                yield "Error: No retrieval results provided"
             return
         
         # Process the retrieval results
@@ -275,23 +324,72 @@ async def generate_streaming(query: str, context: Union[str, List[Dict[str, Any]
         formatted_context = format_context_from_results(processed_results)
         
         if not formatted_context.strip():
-            yield "Error: No valid content found in retrieval results"
+            if chatui_format:
+                yield {"event": "error", "data": {"error": "No valid content found in retrieval results"}}
+            else:
+                yield "Error: No valid content found in retrieval results"
             return
     
     elif isinstance(context, str):
         if not context.strip():
-            yield "Error: Context cannot be empty"
+            if chatui_format:
+                yield {"event": "error", "data": {"error": "Context cannot be empty"}}
+            else:
+                yield "Error: Context cannot be empty"
             return
         formatted_context = context
     
     else:
-        yield "Error: Context must be either a string or list of retrieval results"
+        if chatui_format:
+            yield {"event": "error", "data": {"error": "Context must be either a string or list of retrieval results"}}
+        else:
+            yield "Error: Context must be either a string or list of retrieval results"
         return
     
     try:
         messages = build_messages(query, formatted_context)
+        
+        # Stream the text response
         async for chunk in _call_llm_streaming(messages):
-            yield chunk
+            if chatui_format:
+                yield {"event": "data", "data": chunk}
+            else:
+                yield chunk
+        
+        # Send sources at the end if available and in ChatUI format
+        if chatui_format and processed_results:
+            sources = []
+            for result in processed_results:
+                filename = result.get('filename', 'Unknown')
+                page = result.get('page', 'Unknown')
+                year = result.get('year', 'Unknown')
+                
+                # Create link using doc:// scheme
+                link = f"doc://{filename}"
+                
+                # Create descriptive title
+                title_parts = [filename]
+                if page != 'Unknown':
+                    title_parts.append(f"Page {page}")
+                if year != 'Unknown':
+                    title_parts.append(f"({year})")
+                
+                title = " - ".join(title_parts)
+                
+                sources.append({
+                    "link": link,
+                    "title": title
+                })
+            
+            yield {"event": "sources", "data": {"sources": sources}}
+        
+        # Send end event for ChatUI format
+        if chatui_format:
+            yield {"event": "end", "data": {}}
+        
     except Exception as e:
         logging.exception("Streaming generation failed")
-        yield f"Error: {str(e)}"
+        if chatui_format:
+            yield {"event": "error", "data": {"error": str(e)}}
+        else:
+            yield f"Error: {str(e)}"