refactor v0.3

app/main.py

@@ -133,12 +133,12 @@ def extract_user_query_fallback(data: Any) -> str:
 # Streaming Handlers for LangServe
 # ─────────────────────────────────────────
 
-async def streaming_flexible_handler(data):
+def streaming_flexible_handler(data):
     """Streaming handler that yields individual tokens for LangServe"""
+    logger.info(f"=== STREAMING HANDLER CALLED ===")
+    logger.info(f"streaming_flexible_handler received: {data} (type: {type(data)})")
+    
     try:
-        logger.info(f"=== STREAMING HANDLER START ===")
-        logger.info(f"streaming_flexible_handler received: {data} (type: {type(data)})")
-        
         # Extract query using the same logic as flexible_handler
         if hasattr(data, 'extract_query'):
             query = data.extract_query()
@@ -161,7 +161,8 @@ async def streaming_flexible_handler(data):
         
         if not query.strip():
             # Yield error message as individual characters
-            error_msg = "No query found. Please provide text in one of these fields: query, text, prompt, content, or messages array."
+            error_msg = "No query found. Please provide text."
+            logger.info(f"No query found, yielding error: '{error_msg}'")
             for char in error_msg:
                 yield char
             return
@@ -171,20 +172,18 @@ async def streaming_flexible_handler(data):
         full_text = result.result
         
         logger.info(f"Generated response: '{full_text}' (length: {len(full_text)})")
+        logger.info("Starting to yield individual characters...")
         
         # Yield individual characters - this is what LangServe expects for streaming
         for i, char in enumerate(full_text):
             logger.debug(f"Yielding character {i}: '{char}'")
             yield char
-            
-            # Small delay to simulate realistic streaming
-            await asyncio.sleep(0.02)
         
         logger.info("=== STREAMING HANDLER COMPLETE ===")
         
     except Exception as e:
         logger.error(f"Error in streaming_flexible_handler: {e}", exc_info=True)
-        error_msg = f"Error processing request: {str(e)}"
+        error_msg = f"Error: {str(e)}"
         for char in error_msg:
             yield char
 
@@ -195,6 +194,7 @@ async def streaming_flexible_handler(data):
 def flexible_handler(data) -> ChatFedOutput:
     """Primary handler for flexible input that accepts multiple formats"""
     try:
+        logger.info(f"=== NON-STREAMING HANDLER CALLED ===")
         logger.info(f"flexible_handler received: {data} (type: {type(data)})")
         
         # Handle both Pydantic model and raw dict input
@@ -240,7 +240,8 @@ def flexible_handler(data) -> ChatFedOutput:
             "extracted_query": query,
             "input_method": "flexible_handler",
             "input_type": str(type(data)),
-            "extracted_from": extracted_from
+            "extracted_from": extracted_from,
+            "handler_type": "NON_STREAMING"  # Added to identify which handler was used
         })
         return result