Update app.py

app.py

@@ -76,7 +76,6 @@ async def chat_endpoint(request: ChatRequest):
     Robust Streaming Endpoint that logs events to console.
     """
     config = {"configurable": {"thread_id": request.thread_id}}
-    
     inputs = {
         "query": request.query,
         "RAG": request.use_rag,
@@ -90,38 +89,69 @@ async def chat_endpoint(request: ChatRequest):
     async def event_generator():
         print(f"--- 🚀 Starting stream for {request.thread_id} ---")
         
-        # We iterate over ALL events
-        async for event in rag_app.astream_events(inputs, config=config, version="v1"):
-            
-            # [DEBUG] Print event types to your terminal so you can see what's happening
-            event_type = event.get("event")
+        try:
+            # Iterate over all events
+            async for event in rag_app.astream_events(inputs, config=config, version="v2"):
+                event_type = event.get("event")
+                name = event.get("name", "")
+                
+                # Debug: print all events to see what's happening
+                print(f"Event: {event_type} | Name: {name}")
+                
+                # Method 1: Check for on_chat_model_stream events (most reliable)
+                if event_type == "on_chat_model_stream":
+                    data = event.get("data", {})
+                    chunk = data.get("chunk")
+                    
+                    if chunk and hasattr(chunk, "content"):
+                        content = chunk.content
+                        if content:
+                            chunk_json = json.dumps({"content": content})
+                            yield f"data: {chunk_json}\n\n"
+                            print(f"✅ Yielded: {content[:50]}...")
+                
+                # Method 2: Alternative - check for any chunk with content
+                elif event_type in ["on_chain_stream", "on_llm_stream"]:
+                    data = event.get("data", {})
+                    chunk = data.get("chunk")
+                    
+                    # Handle different chunk types
+                    if chunk:
+                        content = None
+                        
+                        # AIMessageChunk
+                        if hasattr(chunk, "content"):
+                            content = chunk.content
+                        # String chunk
+                        elif isinstance(chunk, str):
+                            content = chunk
+                        # Dict with content
+                        elif isinstance(chunk, dict) and "content" in chunk:
+                            content = chunk["content"]
+                        
+                        if content:
+                            chunk_json = json.dumps({"content": content})
+                            yield f"data: {chunk_json}\n\n"
+                            print(f"✅ Yielded: {content[:50]}...")
             
-            # Logic: We don't care about the event name. 
-            # We only care: "Does this event have a chunk with text?"
-            data = event.get("data", {})
-            chunk = data.get("chunk")
+            print("--- ✅ Stream completed ---")
             
-            # Check if chunk exists and has .content attribute (standard LangChain message chunk)
-            if chunk and hasattr(chunk, "content"):
-                content = chunk.content
-                
-                # Filter out empty strings or weird artifacts
-                if content:
-                    # JSON encode the content
-                    chunk_json = json.dumps({"content": content})
-                    yield f"data: {chunk_json}\n\n"
-
+        except Exception as e:
+            print(f"❌ Error in stream: {e}")
+            error_json = json.dumps({"error": str(e)})
+            yield f"data: {error_json}\n\n"
+        
         # End of stream
         yield "data: [DONE]\n\n"
 
     return StreamingResponse(
-        event_generator(), 
+        event_generator(),
         media_type="text/event-stream",
         headers={
             "Cache-Control": "no-cache",
             "Connection": "keep-alive",
             "Content-Type": "text/event-stream",
-            "X-Accel-Buffering": "no", # Critical for HF
+            "X-Accel-Buffering": "no",
         },
     )