refactored approach

app/main.py

@@ -1,6 +1,3 @@
-"""
-Complete ChatFed Orchestrator with flexible input handling for ChatUI compatibility
-"""
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
 import json, uuid
@@ -11,6 +8,7 @@ from typing import List, Literal, Optional, Dict, Any
 import gradio as gr
 from datetime import datetime
 import logging
+import asyncio
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -131,6 +129,43 @@ def extract_user_query_fallback(data: Any) -> str:
         logger.error(f"Error extracting query from {data}: {e}")
         return ""
 
+# ─────────────────────────────────────────
+# Streaming Generators
+# ─────────────────────────────────────────
+
+async def generate_streaming_response(query: str):
+    """Generate streaming response for ChatUI compatibility"""
+    try:
+        # Process the query
+        result = process_chatfed_query_core(query)
+        full_text = result.result
+        
+        # Emit metadata event
+        metadata = {
+            "run_id": str(uuid.uuid4()),
+            "timestamp": datetime.now().isoformat(),
+            **result.metadata
+        }
+        yield f"event: metadata\ndata: {json.dumps(metadata)}\n\n"
+        
+        # Stream tokens one by one
+        for i, char in enumerate(full_text):
+            # Ensure each token is a string
+            token_data = str(char)
+            yield f"event: data\ndata: {json.dumps(token_data)}\n\n"
+            
+            # Small delay to simulate realistic streaming
+            await asyncio.sleep(0.01)
+        
+        # End event
+        yield f"event: end\ndata: {json.dumps('[DONE]')}\n\n"
+        
+    except Exception as e:
+        logger.error(f"Error in streaming response: {e}")
+        error_data = {"error": str(e)}
+        yield f"event: error\ndata: {json.dumps(error_data)}\n\n"
+        yield f"event: end\ndata: {json.dumps('[ERROR]')}\n\n"
+
 # ─────────────────────────────────────────
 # Handlers - Multiple for Different Use Cases
 # ─────────────────────────────────────────
@@ -329,8 +364,6 @@ def create_gradio_interface():
     
     return demo
 
-
-
 # ─────────────────────────────────────────
 # App Startup
 # ─────────────────────────────────────────
@@ -341,30 +374,30 @@ app = FastAPI(
     version="1.0.0"
 )
 
-# # Add request logging middleware for debugging
-# @app.middleware("http")
-# async def log_requests(request: Request, call_next):
-#     """Log incoming requests for debugging"""
-#     if request.url.path.startswith("/chatfed") or request.url.path.startswith("/debug"):
-#         try:
-#             body = await request.body()
-#             logger.info(f"=== REQUEST DEBUG ===")
-#             logger.info(f"Path: {request.url.path}")
-#             logger.info(f"Method: {request.method}")
-#             logger.info(f"Headers: {dict(request.headers)}")
-#             logger.info(f"Body: {body.decode('utf-8') if body else 'Empty'}")
+# Add request logging middleware for debugging
+@app.middleware("http")
+async def log_requests(request: Request, call_next):
+    """Log incoming requests for debugging"""
+    if request.url.path.startswith("/chatfed") or request.url.path.startswith("/debug"):
+        try:
+            body = await request.body()
+            logger.info(f"=== REQUEST DEBUG ===")
+            logger.info(f"Path: {request.url.path}")
+            logger.info(f"Method: {request.method}")
+            logger.info(f"Headers: {dict(request.headers)}")
+            logger.info(f"Body: {body.decode('utf-8') if body else 'Empty'}")
             
-#             # Recreate request for next handler
-#             async def receive():
-#                 return {"type": "http.request", "body": body}
+            # Recreate request for next handler
+            async def receive():
+                return {"type": "http.request", "body": body}
             
-#             request._receive = receive
+            request._receive = receive
             
-#         except Exception as e:
-#             logger.error(f"Error logging request: {e}")
+        except Exception as e:
+            logger.error(f"Error logging request: {e}")
     
-#     response = await call_next(request)
-#     return response
+    response = await call_next(request)
+    return response
 
 # ─────────────────────────────────────────
 # LangServe Routes - Flexible input handling
@@ -375,7 +408,6 @@ add_routes(
     app,
     RunnableLambda(flexible_handler),
     path="/chatfed",
-    # Remove strict input type to allow both dicts and Pydantic models
     output_type=ChatFedOutput
 )
 
@@ -384,7 +416,6 @@ add_routes(
     app,
     RunnableLambda(chatui_handler),
     path="/chatfed-chatui",
-    # Remove strict input type to allow both dicts and Pydantic models
     output_type=ChatFedOutput
 )
 
@@ -393,16 +424,78 @@ add_routes(
     app,
     RunnableLambda(legacy_langserve_handler),
     path="/chatfed-strict",
-    # Remove strict input type to allow both dicts and Pydantic models
     output_type=ChatFedOutput
 )
 
-# ChatUI-compatible streaming route (yields tokens via SSE)
-add_routes(
-    app,
-    RunnableLambda(flexible_handler),
-    path="/chatfed-ui-stream"
-)
+# ─────────────────────────────────────────
+# Custom Streaming Endpoint for ChatUI
+# ─────────────────────────────────────────
+
+@app.post("/chatfed-ui-stream/stream")
+async def chatui_stream_endpoint(request: Request) -> StreamingResponse:
+    """
+    Proper streaming endpoint for ChatUI's langserve integration.
+    Returns Server-Sent Events with individual tokens.
+    """
+    try:
+        # Get the request payload
+        payload = await request.json()
+        logger.info(f"Stream endpoint received: {payload}")
+        
+        # Handle ChatUI's envelope format
+        if isinstance(payload.get("input"), dict):
+            input_data = payload["input"]
+        else:
+            input_data = payload
+        
+        # Extract query using flexible approach
+        try:
+            flexible_input = FlexibleChatInput(**input_data)
+            query = flexible_input.extract_query()
+        except Exception as e:
+            logger.warning(f"Failed to parse as FlexibleChatInput: {e}")
+            query = extract_user_query_fallback(input_data)
+        
+        if not query.strip():
+            # Return error stream
+            async def error_stream():
+                yield f"event: error\ndata: {json.dumps({'error': 'No valid query found'})}\n\n"
+                yield f"event: end\ndata: {json.dumps('[ERROR]')}\n\n"
+            
+            return StreamingResponse(
+                error_stream(),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                }
+            )
+        
+        # Return successful stream
+        return StreamingResponse(
+            generate_streaming_response(query),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            }
+        )
+        
+    except Exception as e:
+        logger.error(f"Error in stream endpoint: {e}")
+        
+        async def error_stream():
+            yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"
+            yield f"event: end\ndata: {json.dumps('[ERROR]')}\n\n"
+        
+        return StreamingResponse(
+            error_stream(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            }
+        )
 
 # ─────────────────────────────────────────
 # Additional Endpoints
@@ -423,6 +516,7 @@ async def root():
             "primary": "/chatfed (flexible input - use this for ChatUI)",
             "chatui": "/chatfed-chatui",
             "legacy": "/chatfed-strict (requires 'query' field)",
+            "streaming": "/chatfed-ui-stream/stream (proper SSE streaming for ChatUI)",
             "openai": "/v1/chat/completions",
             "simple": "/simple-chat",
             "gradio_ui": "/ui",
@@ -546,56 +640,6 @@ async def debug_input_endpoint(request: Request):
             "error": str(e),
             "raw_body": raw_body.decode('utf-8') if 'raw_body' in locals() else "failed_to_read"
         }
-    
-# @app.post("/chatfed-ui-stream/stream")
-# async def ui_stream_sse(request: Request) -> StreamingResponse:
-#     """
-#     Stream-friendly handler for ChatUI's langserve integration.
-#     Emits proper SSE: metadata → data (one char/token per event) → end.
-#     """
-#     try:
-#         payload = await request.json()
-#         # 1) Unwrap ChatUI envelope if needed
-#         if isinstance(payload.get("input"), dict):
-#             payload = payload["input"]
-
-#         # 2) Extract query (flexible or fallback)
-#         try:
-#             input_data = FlexibleChatInput(**payload)
-#             query = input_data.extract_query()
-#         except Exception:
-#             query = extract_user_query_fallback(payload)
-
-#         if not query.strip():
-#             # send a single “end” event with an error
-#             async def no_q():
-#                 yield 'event: end\ndata: ["error","No valid query found"]\n\n'
-#             return StreamingResponse(no_q(), media_type="text/event-stream")
-
-#         # 3) Generate full text
-#         full_text = process_chatfed_query_core(query).result
-
-#         # 4) Build SSE generator
-#         async def event_gen():
-#             # metadata
-#             meta = {"run_id": str(uuid.uuid4())}
-#             yield f"event: metadata\ndata: {json.dumps(meta)}\n\n"
-
-#             # one data event per character/token
-#             for ch in full_text:
-#                 yield f"event: data\ndata: {json.dumps(ch)}\n\n"
-
-#             # end
-#             # yield "event: end\ndata: [DONE]\n\n"
-#             yield f"event: end\ndata: {json.dumps('[DONE]')}\n\n"
-
-#         return StreamingResponse(event_gen(), media_type="text/event-stream")
-
-#     except Exception as e:
-#         logger.exception(f"ui_stream_sse error: {e}")
-#         async def crash():
-#             yield f'event: end\ndata: ["error","{str(e)}"]\n\n'
-#         return StreamingResponse(crash(), media_type="text/event-stream")
 
 # Mount Gradio at a specific path
 demo = create_gradio_interface()