Update main.py

main.py

@@ -1,5 +1,8 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+import json
+import asyncio
 from pydantic import BaseModel
 from typing import Dict, Any, Optional, List
 from contextlib import asynccontextmanager
@@ -9,16 +12,24 @@ from utils.config import settings
 # Import the LangGraph app
 from core.graph_workflow import app as graph_app
 
+print("🚀 Nexus Lex Backend: Initializing...")
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Startup: Ensure containers are running
-    print("Starting Qdrant container...")
-    start_qdrant_docker()
-    print("Starting Neo4j container...")
-    start_neo4j_docker()
-    
-    # Pre-loading models happens inside the services instantiated in graph_workflow
+    if settings.QDRANT_ENDPOINT:
+        print("📍 [STARTUP] Cloud Qdrant target detected.")
+    else:
+        print("⚠️ [STARTUP] No Qdrant endpoint found. Attempting local Docker lifecycle...")
+        start_qdrant_docker()
+        
+    if settings.NEO4J_URI:
+        print("📍 [STARTUP] Cloud Neo4j target detected.")
+    else:
+        print("⚠️ [STARTUP] No Neo4j URI found. Attempting local Docker lifecycle...")
+        start_neo4j_docker()
     
+    print("✅ [STARTUP] Initialization sequence complete. App is live.")
     yield
     
     # Shutdown: Stop containers
@@ -86,30 +97,65 @@ def clear_api_cache():
 
 @app.post("/search")
 async def search_amendments(request: SearchRequest):
-    try:
-        # Invoke LangGraph
-        initial_state = {"query": request.query, "retry_count": 0}
-        result = await graph_app.ainvoke(initial_state)
-        
-        # Format output from final state
-        return {
-            "query": result["query"],
-            "answer": result["draft_answer"].get("answer"),
-            "constitutional_status": result["draft_answer"].get("constitutional_status"),
-            "confidence": result["critique"].get("final_confidence"),
-            "sources": result["draft_answer"].get("sources"),
-            "quality_grade": result["critique"].get("quality_grade"),
-            "execution_trace": result["trace"],
-            # PASS RAW DATA FOR UI HIGHLIGHTING
-            "graph_nodes": result.get("graph_results", []),
-            "vector_chunks": [
-                {"id": c["id"], "text": c["text"], "metadata": c["metadata"]} 
-                for c in result.get("retrieved_chunks", [])
-            ]
-        }
-    except Exception as e:
-        print(f"Graph execution error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+    async def event_generator():
+        try:
+            initial_state = {"query": request.query, "retry_count": 0}
+            final_state = {}
+            
+            # Use astream to get updates after every node
+            # stream_mode="updates" gives us the incremental changes from each node
+            async for event in graph_app.astream(initial_state, stream_mode="updates"):
+                for node_name, state_update in event.items():
+                    # Update local state tracker
+                    final_state.update(state_update)
+                    
+                    # Yield a trace event for the UI terminal
+                    trace_msg = f">> {node_name.replace('_', ' ').title()} sequence complete."
+                    
+                    # Add descriptive context if available
+                    if node_name == "classify":
+                        entities = state_update.get("classification", {}).get("entities", {})
+                        arts = entities.get("articles", [])
+                        if arts: trace_msg = f">> Classified: Targeting Articles {', '.join(map(str, arts))}"
+                    elif node_name == "graph_plan":
+                        results = state_update.get("graph_results", [])
+                        trace_msg = f">> Neo4j: Retrieved {len(results)} legal relationships."
+                    elif node_name == "fetch_vector":
+                        chunks = state_update.get("raw_chunks", [])
+                        trace_msg = f">> Qdrant: Retrieved {len(chunks)} semantic text blocks."
+                    elif node_name == "reason":
+                        trace_msg = ">> Legal Reasoner: Synthesizing final opinion..."
+                        
+                    yield json.dumps({"type": "trace", "message": trace_msg}) + "\n"
+                    # Small sleep to ensure the UI feels sequential and "live"
+                    await asyncio.sleep(0.1)
+
+            # We need to ensure we have the full final state for the result
+            # LangGraph's astream "updates" doesn't necessarily give the final dictionary in one go
+            # So we format the accumulated final_state
+            
+            result_payload = {
+                "query": request.query,
+                "answer": final_state.get("draft_answer", {}).get("answer"),
+                "constitutional_status": final_state.get("draft_answer", {}).get("constitutional_status"),
+                "confidence": final_state.get("critique", {}).get("final_confidence"),
+                "sources": final_state.get("draft_answer", {}).get("sources"),
+                "quality_grade": final_state.get("critique", {}).get("quality_grade"),
+                "execution_trace": final_state.get("trace", []),
+                "graph_nodes": final_state.get("graph_results", []),
+                "vector_chunks": [
+                    {"id": c["id"], "text": c["text"], "metadata": c["metadata"]} 
+                    for c in final_state.get("retrieved_chunks", [])
+                ]
+            }
+            
+            yield json.dumps({"type": "result", "payload": result_payload}) + "\n"
+            
+        except Exception as e:
+            print(f"Streaming error: {e}")
+            yield json.dumps({"type": "error", "message": str(e)}) + "\n"
+
+    return StreamingResponse(event_generator(), media_type="application/x-ndjson")
 
 @app.get("/graph/explore")
 def get_graph_exploration(limit: int = 1000):