cleanup

app/main.py

@@ -2,7 +2,7 @@ import gradio as gr
 from fastapi import FastAPI
 from langserve import add_routes
 from langgraph.graph import StateGraph, START, END
-from typing import Optional, Dict, Any, List, Literal
+from typing import Optional, Dict, Any
 from typing_extensions import TypedDict
 from pydantic import BaseModel
 from gradio_client import Client
@@ -11,26 +11,18 @@ import os
 from datetime import datetime
 import logging
 from contextlib import asynccontextmanager
-import io
-from PIL import Image
 import threading
 from langchain_core.runnables import RunnableLambda
 
-# Local imports
 from utils import getconfig
 
 config = getconfig("params.cfg")
-
 RETRIEVER = config.get("retriever", "RETRIEVER")
 GENERATOR = config.get("generator", "GENERATOR")
 
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Define langgraph state schema
 class GraphState(TypedDict):
     query: str
     context: str
@@ -41,7 +33,6 @@ class GraphState(TypedDict):
     year_filter: str
     metadata: Optional[Dict[str, Any]]
 
-# LangServe input/output schemas
 class ChatFedInput(TypedDict):
     query: str
     reports_filter: Optional[str]
@@ -55,24 +46,12 @@ class ChatFedOutput(TypedDict):
     result: str
     metadata: Dict[str, Any]
 
-# ChatUI specific schemas
-class ChatUIStreamInput(BaseModel):
-    text: str  # ChatUI sends input as "text" field
-    
-class ChatUIStreamOutput(BaseModel):
-    content: str
-
-class ChatMessage(BaseModel):
-    role: Literal["system", "user", "assistant"]
-    content: str
-
 class ChatUIInput(BaseModel):
-    messages: List[ChatMessage]
+    text: str
 
-# Retriever 
 def retrieve_node(state: GraphState) -> GraphState:
     start_time = datetime.now()
-    logger.info(f"Starting retrieval for query: {state['query'][:100]}...")
+    logger.info(f"Retrieval: {state['query'][:50]}...")
     
     try:
         client = Client(RETRIEVER)
@@ -88,30 +67,28 @@ def retrieve_node(state: GraphState) -> GraphState:
         duration = (datetime.now() - start_time).total_seconds()
         metadata = state.get("metadata", {})
         metadata.update({
-            "retrieval_duration_seconds": duration,
+            "retrieval_duration": duration,
             "context_length": len(context) if context else 0,
             "retrieval_success": True
         })
         
-        logger.info(f"Retrieval completed in {duration:.2f}s, context length: {len(context) if context else 0}")
         return {"context": context, "metadata": metadata}
         
     except Exception as e:
         duration = (datetime.now() - start_time).total_seconds()
-        logger.error(f"Retrieval failed after {duration:.2f}s: {str(e)}")
+        logger.error(f"Retrieval failed: {str(e)}")
         
         metadata = state.get("metadata", {})
         metadata.update({
-            "retrieval_duration_seconds": duration,
+            "retrieval_duration": duration,
             "retrieval_success": False,
             "retrieval_error": str(e)
         })
         return {"context": "", "metadata": metadata}
 
-# Generator
 def generate_node(state: GraphState) -> GraphState:
     start_time = datetime.now()
-    logger.info(f"Starting generation for query: {state['query'][:100]}...")
+    logger.info(f"Generation: {state['query'][:50]}...")
     
     try:
         client = Client(GENERATOR)
@@ -124,27 +101,25 @@ def generate_node(state: GraphState) -> GraphState:
         duration = (datetime.now() - start_time).total_seconds()
         metadata = state.get("metadata", {})
         metadata.update({
-            "generation_duration_seconds": duration,
+            "generation_duration": duration,
             "result_length": len(result) if result else 0,
             "generation_success": True
         })
         
-        logger.info(f"Generation completed in {duration:.2f}s, result length: {len(result) if result else 0}")
         return {"result": result, "metadata": metadata}
         
     except Exception as e:
         duration = (datetime.now() - start_time).total_seconds()
-        logger.error(f"Generation failed after {duration:.2f}s: {str(e)}")
+        logger.error(f"Generation failed: {str(e)}")
         
         metadata = state.get("metadata", {})
         metadata.update({
-            "generation_duration_seconds": duration,
+            "generation_duration": duration,
             "generation_success": False,
             "generation_error": str(e)
         })
-        return {"result": f"Error generating response: {str(e)}", "metadata": metadata}
+        return {"result": f"Error: {str(e)}", "metadata": metadata}
 
-# Build graph
 workflow = StateGraph(GraphState)
 workflow.add_node("retrieve", retrieve_node)
 workflow.add_node("generate", generate_node)
@@ -153,8 +128,7 @@ workflow.add_edge("retrieve", "generate")
 workflow.add_edge("generate", END)
 compiled_graph = workflow.compile()
 
-# Core processing function (shared by both Gradio and LangServe)
-def process_chatfed_query_core(
+def process_query_core(
     query: str,
     reports_filter: str = "",
     sources_filter: str = "",
@@ -164,13 +138,10 @@ def process_chatfed_query_core(
     user_id: Optional[str] = None,
     return_metadata: bool = False
 ):
-    """Core processing function used by both Gradio and LangServe interfaces."""
     start_time = datetime.now()
     if not session_id:
         session_id = f"session_{start_time.strftime('%Y%m%d_%H%M%S')}"
     
-    logger.info(f"Processing query in session {session_id}: {query[:100]}...")
-    
     try:
         initial_state = {
             "query": query,
@@ -183,8 +154,7 @@ def process_chatfed_query_core(
             "metadata": {
                 "session_id": session_id,
                 "user_id": user_id,
-                "start_time": start_time.isoformat(),
-                "orchestrator": "hybrid_gradio_langserve"
+                "start_time": start_time.isoformat()
             }
         }
         
@@ -193,13 +163,11 @@ def process_chatfed_query_core(
         
         final_metadata = final_state.get("metadata", {})
         final_metadata.update({
-            "total_duration_seconds": total_duration,
+            "total_duration": total_duration,
             "end_time": datetime.now().isoformat(),
             "pipeline_success": True
         })
         
-        logger.info(f"Query processing completed in {total_duration:.2f}s for session {session_id}")
-        
         if return_metadata:
             return {"result": final_state["result"], "metadata": final_metadata}
         else:
@@ -207,33 +175,22 @@ def process_chatfed_query_core(
         
     except Exception as e:
         total_duration = (datetime.now() - start_time).total_seconds()
-        logger.error(f"Pipeline failed after {total_duration:.2f}s for session {session_id}: {str(e)}")
+        logger.error(f"Pipeline failed: {str(e)}")
         
         if return_metadata:
             error_metadata = {
                 "session_id": session_id,
-                "total_duration_seconds": total_duration,
+                "total_duration": total_duration,
                 "pipeline_success": False,
                 "error": str(e)
             }
-            return {"result": f"Error processing query: {str(e)}", "metadata": error_metadata}
+            return {"result": f"Error: {str(e)}", "metadata": error_metadata}
         else:
-            return f"Error processing query: {str(e)}"
-
-# =============================================================================
-# GRADIO INTERFACE (MCP ENDPOINTS)
-# =============================================================================
+            return f"Error: {str(e)}"
 
-# Gradio wrapper functions for MCP compatibility
-def process_query_gradio(
-    query: str,
-    reports_filter: str = "",
-    sources_filter: str = "",
-    subtype_filter: str = "",
-    year_filter: str = ""
-) -> str:
-    """Gradio-compatible function that exposes MCP endpoints."""
-    return process_chatfed_query_core(
+def process_query_gradio(query: str, reports_filter: str = "", sources_filter: str = "", 
+                        subtype_filter: str = "", year_filter: str = "") -> str:
+    return process_query_core(
         query=query,
         reports_filter=reports_filter,
         sources_filter=sources_filter,
@@ -243,99 +200,7 @@ def process_query_gradio(
         return_metadata=False
     )
 
-def get_graph_visualization():
-    """Generate graph visualization for Gradio interface."""
-    try:
-        graph_png_bytes = compiled_graph.get_graph().draw_mermaid_png()
-        return Image.open(io.BytesIO(graph_png_bytes))
-    except Exception as e:
-        logger.error(f"Failed to generate graph visualization: {e}")
-        return None
-
-# Create Gradio interface
-def create_gradio_interface():
-    with gr.Blocks(title="ChatFed Orchestrator - MCP Endpoints") as demo:
-        gr.Markdown("# ChatFed Orchestrator")
-        gr.Markdown("**MCP Server Endpoints Available** - This interface provides MCP compatibility for ChatUI integration.")
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown("**Workflow Visualization**")
-                graph_display = gr.Image(
-                    value=get_graph_visualization(), 
-                    label="LangGraph Workflow",
-                    interactive=False,
-                    height=300
-                )
-                refresh_graph_btn = gr.Button("🔄 Refresh Graph", size="sm")
-                refresh_graph_btn.click(fn=get_graph_visualization, outputs=graph_display)
-                
-                gr.Markdown("**🔗 MCP Integration**")
-                gr.Markdown("MCP endpoints are active and ready for ChatUI integration.")
-                
-            with gr.Column(scale=2):
-                gr.Markdown("**MCP Endpoint Information**")
-                
-                with gr.Accordion("MCP Usage", open=True):
-                    gr.Markdown("""
-                    **MCP Server Endpoint:** Available at `/gradio_api/mcp/sse`
-                    
-                    **For ChatUI Integration:**
-                    ```python
-                    from gradio_client import Client
-                    
-                    # Connect to orchestrator MCP endpoint
-                    client = Client("https://your-space.hf.space")
-                    
-                    # Basic usage
-                    response = client.predict(
-                        query="your question",
-                        api_name="/process_query_gradio"
-                    )
-                    
-                    # With filters
-                    response = client.predict(
-                        query="your question",
-                        reports_filter="annual_reports",
-                        sources_filter="internal",
-                        year_filter="2024",
-                        api_name="/process_query_gradio"
-                    )
-                    ```
-                    """)
-        
-        with gr.Accordion("Test Interface", open=False):
-            # Test interface
-            with gr.Row():
-                with gr.Column():
-                    query_input = gr.Textbox(label="Query", lines=2, placeholder="Enter your question...")
-                    reports_filter_input = gr.Textbox(label="Reports Filter", placeholder="e.g., annual_reports")
-                    sources_filter_input = gr.Textbox(label="Sources Filter", placeholder="e.g., internal")
-                    subtype_filter_input = gr.Textbox(label="Subtype Filter", placeholder="e.g., financial")
-                    year_filter_input = gr.Textbox(label="Year Filter", placeholder="e.g., 2024")
-                    submit_btn = gr.Button("Submit", variant="primary")
-                
-                with gr.Column():
-                    output = gr.Textbox(label="Response", lines=10)
-            
-            submit_btn.click(
-                fn=process_query_gradio,
-                inputs=[query_input, reports_filter_input, sources_filter_input, subtype_filter_input, year_filter_input],
-                outputs=output
-            )
-    
-    return demo
-
-# =============================================================================
-# CHATUI STREAMING ADAPTER  
-# =============================================================================
-
-def chatui_streaming_adapter(data) -> str:
-    """
-    Adapter for ChatUI integration.
-    LangServe will automatically create /invoke and /stream endpoints.
-    ChatUI will use the /stream endpoint for streaming responses.
-    """
+def chatui_adapter(data) -> str:
     try:
         # Handle both dict and Pydantic model input
         if hasattr(data, 'text'):
@@ -343,58 +208,21 @@ def chatui_streaming_adapter(data) -> str:
         elif isinstance(data, dict) and 'text' in data:
             text = data['text']
         else:
-            # Log the actual input structure for debugging
             logger.error(f"Unexpected input structure: {data}")
             return "Error: Invalid input format. Expected 'text' field."
         
-        logger.info(f"ChatUI request: {text[:100]}...")
-        
-        # Process the query using your core function
-        result = process_chatfed_query_core(
+        result = process_query_core(
             query=text,
             session_id=f"chatui_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
             return_metadata=False
         )
-        
         return result
-            
     except Exception as e:
         logger.error(f"ChatUI error: {str(e)}")
-        return f"Error processing request: {str(e)}"
+        return f"Error: {str(e)}"
 
-def chatui_non_streaming_adapter(data) -> ChatUIStreamOutput:
-    """
-    Non-streaming adapter for ChatUI (fallback).
-    """
-    try:
-        # Handle both dict and Pydantic model input
-        if hasattr(data, 'text'):
-            text = data.text
-        elif isinstance(data, dict) and 'text' in data:
-            text = data['text']
-        else:
-            logger.error(f"Unexpected input structure: {data}")
-            return ChatUIStreamOutput(content="Error: Invalid input format. Expected 'text' field.")
-        
-        logger.info(f"ChatUI non-streaming request: {text[:100]}...")
-        
-        result = process_chatfed_query_core(
-            query=text,
-            session_id=f"chatui_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
-            return_metadata=False
-        )
-        return ChatUIStreamOutput(content=result)
-    except Exception as e:
-        logger.error(f"ChatUI adapter error: {str(e)}")
-        return ChatUIStreamOutput(content=f"Error processing request: {str(e)}")
-
-# =============================================================================
-# LANGSERVE API (TELEMETRY)
-# =============================================================================
-
-def process_chatfed_query_langserve(input_data: ChatFedInput) -> ChatFedOutput:
-    """LangServe function with full metadata return."""
-    result = process_chatfed_query_core(
+def process_query_langserve(input_data: ChatFedInput) -> ChatFedOutput:
+    result = process_query_core(
         query=input_data["query"],
         reports_filter=input_data.get("reports_filter", ""),
         sources_filter=input_data.get("sources_filter", ""),
@@ -406,147 +234,84 @@ def process_chatfed_query_langserve(input_data: ChatFedInput) -> ChatFedOutput:
     )
     return ChatFedOutput(result=result["result"], metadata=result["metadata"])
 
-def chatui_adapter(data: ChatUIInput):
-    """
-    Adapter to allow ChatUI to send full chat history.
-    We extract the latest user message for ChatFed.
-    """
-    last_user_msg = next(m.content for m in reversed(data.messages) if m.role == "user")
-    result = process_chatfed_query_core(query=last_user_msg)
-    return {"result": result, "metadata": {"source": "chatfed-langserve-adapter"}}
+def create_gradio_interface():
+    with gr.Blocks(title="ChatFed Orchestrator") as demo:
+        gr.Markdown("# ChatFed Orchestrator")
+        gr.Markdown("MCP endpoints available at `/gradio_api/mcp/sse`")
+        
+        with gr.Row():
+            with gr.Column():
+                query_input = gr.Textbox(label="Query", lines=2, placeholder="Enter your question...")
+                reports_filter_input = gr.Textbox(label="Reports Filter", placeholder="e.g., annual_reports")
+                sources_filter_input = gr.Textbox(label="Sources Filter", placeholder="e.g., internal")
+                subtype_filter_input = gr.Textbox(label="Subtype Filter", placeholder="e.g., financial")
+                year_filter_input = gr.Textbox(label="Year Filter", placeholder="e.g., 2024")
+                submit_btn = gr.Button("Submit", variant="primary")
+            
+            with gr.Column():
+                output = gr.Textbox(label="Response", lines=10)
+        
+        submit_btn.click(
+            fn=process_query_gradio,
+            inputs=[query_input, reports_filter_input, sources_filter_input, subtype_filter_input, year_filter_input],
+            outputs=output
+        )
+    
+    return demo
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    logger.info("🚀 Hybrid ChatFed Orchestrator starting up...")
-    logger.info("✅ LangGraph compiled successfully")
-    logger.info("🔗 MCP endpoints will be available via Gradio")
-    logger.info("📊 Enhanced API available via LangServe")
-    logger.info("🎯 ChatUI streaming integration enabled")
+    logger.info("ChatFed Orchestrator starting up...")
     yield
-    logger.info("🛑 Orchestrator shutting down...")
+    logger.info("Orchestrator shutting down...")
 
-# Create FastAPI app with docs disabled
 app = FastAPI(
-    title="ChatFed Orchestrator - Enhanced API",
+    title="ChatFed Orchestrator",
     version="1.0.0",
-    description="Enhanced API with observability. MCP endpoints available via Gradio interface.",
     lifespan=lifespan,
-    docs_url=None,  # Disable /docs endpoint
-    redoc_url=None  # Disable /redoc endpoint
+    docs_url=None,
+    redoc_url=None
 )
 
-# Health check
 @app.get("/health")
 async def health_check():
-    return {
-        "status": "healthy",
-        "mcp_endpoints": "available_via_gradio",
-        "enhanced_api": "available_via_langserve",
-        "chatui_integration": "enabled"
-    }
+    return {"status": "healthy"}
 
-# Add root endpoint
 @app.get("/")
 async def root():
     return {
         "message": "ChatFed Orchestrator API",
-        "version": "1.0.0",
         "endpoints": {
             "health": "/health",
             "chatfed": "/chatfed",
-            "chatfed-chatui": "/chatfed-chatui",
-            "chatfed-ui-stream": "/chatfed-ui-stream",  # New for ChatUI streaming
-            "chatfed-ui": "/chatfed-ui",  # New fallback
-            "process_query": "/process_query"
-        },
-        "gradio_interface": "http://localhost:7861/",
-        "mcp_endpoints": "http://localhost:7861/gradio_api/mcp/sse",
-        "note": "LangServe telemetry enabled - ChatUI integration available via /chatfed-ui-stream"
+            "chatfed-ui-stream": "/chatfed-ui-stream"
+        }
     }
 
-# =============================================================================
-# ADD LANGSERVE ROUTES
-# =============================================================================
-
-# Convert functions to Runnables
-process_chatfed_query_runnable = RunnableLambda(process_chatfed_query_langserve)
-chatui_adapter_runnable = RunnableLambda(chatui_adapter)
-chatui_streaming_runnable = RunnableLambda(chatui_streaming_adapter)
-chatui_non_streaming_runnable = RunnableLambda(chatui_non_streaming_adapter)
-
-# Add routes with explicit input/output schemas
+# LangServe routes
 add_routes(
     app,
-    process_chatfed_query_runnable,
+    RunnableLambda(process_query_langserve),
     path="/chatfed",
     input_type=ChatFedInput,
     output_type=ChatFedOutput
 )
 
-# Original ChatUI-compatible LangServe route
 add_routes(
     app,
-    chatui_adapter_runnable,
-    path="/chatfed-chatui",
-    input_type=ChatUIInput
-)
-
-# NEW: ChatUI streaming route (matches your ChatUI config)
-# LangServe will automatically create both /invoke and /stream endpoints
-add_routes(
-    app,
-    chatui_streaming_runnable,
+    RunnableLambda(chatui_adapter),
     path="/chatfed-ui-stream",
-    input_type=ChatUIStreamInput,
+    input_type=ChatUIInput,
     output_type=str,
     enable_feedback_endpoint=True,
     enable_public_trace_link_endpoint=True,
 )
 
-# NEW: ChatUI non-streaming fallback route
-add_routes(
-    app,
-    chatui_non_streaming_runnable,
-    path="/chatfed-ui",
-    input_type=ChatUIStreamInput,
-    output_type=ChatUIStreamOutput,
-    enable_feedback_endpoint=True,
-    enable_public_trace_link_endpoint=True,
-)
-
-# Backward compatibility endpoint
-@app.post("/process_query")
-async def process_query_endpoint(
-    query: str,
-    reports_filter: str = "",
-    sources_filter: str = "",
-    subtype_filter: str = "",
-    year_filter: str = "",
-    session_id: Optional[str] = None,
-    user_id: Optional[str] = None
-):
-    """Backward compatibility endpoint."""
-    return process_chatfed_query_core(
-        query=query,
-        reports_filter=reports_filter,
-        sources_filter=sources_filter,
-        subtype_filter=subtype_filter,
-        year_filter=year_filter,
-        session_id=session_id,
-        user_id=user_id,
-        return_metadata=False
-    )
-
-# =============================================================================
-# MAIN APPLICATION LAUNCHER
-# =============================================================================
-
 def run_gradio_server():
-    """Run Gradio server in a separate thread for MCP endpoints."""
     demo = create_gradio_interface()
     demo.launch(
         server_name="0.0.0.0",
-        server_port=7861,  # Different port from FastAPI
+        server_port=7861,
         mcp_server=True,   
         show_error=True,
         share=False,
@@ -554,24 +319,13 @@ def run_gradio_server():
     )
 
 if __name__ == "__main__":
-    # Start Gradio server in background thread for MCP endpoints
     gradio_thread = threading.Thread(target=run_gradio_server, daemon=True)
     gradio_thread.start()
-    logger.info("🔗 Gradio MCP server started on port 7861")
+    logger.info("Gradio MCP server started on port 7861")
     
-    # Start FastAPI server for enhanced API
     host = os.getenv("HOST", "0.0.0.0")
     port = int(os.getenv("PORT", "7860"))
     
-    logger.info(f"🚀 Starting FastAPI server on {host}:{port}")
-    logger.info("📊 Enhanced API with LangServe telemetry available")
-    logger.info("🔗 MCP endpoints available via Gradio on port 7861")
-    logger.info("🎯 ChatUI streaming integration ready at /chatfed-ui-stream")
+    logger.info(f"Starting FastAPI server on {host}:{port}")
     
-    uvicorn.run(
-        app,
-        host=host,
-        port=port,
-        log_level="info",
-        access_log=True
-    )
+    uvicorn.run(app, host=host, port=port, log_level="info", access_log=True)