reasoning streaming llm

Author: blessing-sanusi

src/backend/app_kernel.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import uuid
+import time
 from typing import Dict, List, Optional
 
 # Semantic Kernel imports
@@ -17,6 +18,7 @@
 # FastAPI imports
 from fastapi import FastAPI, HTTPException, Query, Request
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, Response
 from kernel_agents.agent_factory import AgentFactory
 
 # Local imports
@@ -67,6 +69,9 @@
 # Initialize the FastAPI app
 app = FastAPI()
 
+# Add a simple in-memory store to track active streaming requests with timestamps
+active_streams = {}  # Changed to dict to store timestamps
+
 frontend_url = Config.FRONTEND_SITE_NAME
 
 # Add this near the top of your app.py, after initializing the app
@@ -316,6 +321,153 @@ async def create_plan_endpoint(input_task: InputTask, request: Request):
         raise HTTPException(status_code=400, detail=f"Error creating plan: {e}")
 
 
+@app.options("/api/generate_plan/{plan_id}")
+async def generate_plan_options(plan_id: str):
+    """Handle CORS preflight for generate_plan endpoint"""
+    return Response(
+        headers={
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Methods": "POST, OPTIONS",
+            "Access-Control-Allow-Headers": "*",
+        }
+    )
+
+
+@app.post("/api/generate_plan/{plan_id}")
+async def generate_plan_endpoint(plan_id: str, request: Request):
+    """
+    Generate detailed plan with steps using reasoning LLM and stream the process.
+    
+    ---
+    tags:
+      - Plans
+    parameters:
+      - name: plan_id
+        in: path
+        type: string
+        required: true
+        description: The ID of the plan to generate steps for
+      - name: user_principal_id
+        in: header
+        type: string
+        required: true
+        description: User ID extracted from the authentication header
+    responses:
+      200:
+        description: Streaming response of the reasoning process
+        content:
+          text/plain:
+            schema:
+              type: string
+              description: Stream of reasoning process and final JSON
+      400:
+        description: Plan not found or other error
+        schema:
+          type: object
+          properties:
+            detail:
+              type: string
+              description: Error message
+    """
+    # Get authenticated user first
+    authenticated_user = get_authenticated_user_details(request_headers=request.headers)
+    user_id = authenticated_user["user_principal_id"]
+
+    if not user_id:
+        track_event_if_configured(
+            "UserIdNotFound", {"status_code": 400, "detail": "no user"}
+        )
+        raise HTTPException(status_code=400, detail="no user")
+
+    # Clean up stale streams (older than 5 minutes)
+    current_time = time.time()
+    stale_streams = [k for k, v in active_streams.items() if current_time - v > 300]
+    for stale_key in stale_streams:
+        active_streams.pop(stale_key, None)
+        logging.info(f"Cleaned up stale stream: {stale_key}")
+
+    # Check if there's already an active stream for this plan + user combination
+    stream_key = f"stream_{plan_id}_{user_id}"
+    logging.info(f"Received stream request for plan {plan_id} from user {user_id}, active streams: {list(active_streams.keys())}")
+    if stream_key in active_streams:
+        logging.warning(f"Duplicate stream request for plan {plan_id} from user {user_id}, rejecting. Active streams: {list(active_streams.keys())}")
+        raise HTTPException(status_code=429, detail="Stream already in progress for this plan")
+
+    try:
+        # Add to active streams with timestamp
+        active_streams[stream_key] = current_time
+        logging.info(f"Added stream {stream_key} to active streams. Current active: {list(active_streams.keys())}")
+        
+        # Initialize memory store
+        kernel, memory_store = await initialize_runtime_and_context("", user_id)
+
+        # Get the existing plan
+        plan = await memory_store.get_plan_by_plan_id(plan_id)
+        if not plan:
+            track_event_if_configured(
+                "PlanNotFound",
+                {"plan_id": plan_id, "error": "Plan not found"},
+            )
+            active_streams.pop(stream_key, None)  # Remove from active streams
+            logging.info(f"Plan {plan_id} not found, removed stream from active streams")
+            raise HTTPException(status_code=404, detail="Plan not found")
+
+        # Generate streaming response
+        async def generate_reasoning_stream():
+            try:
+                logging.info(f"Starting stream for plan {plan_id}")
+                
+                # Import the reasoning generation function
+                from utils_kernel import generate_plan_with_reasoning_stream
+                
+                # Stream the reasoning process and get the final result
+                async for chunk in generate_plan_with_reasoning_stream(plan.initial_goal, plan_id, memory_store):
+                    yield f"data: {chunk}\n\n"
+                
+                # Send completion signal
+                yield f"data: [DONE]\n\n"
+                logging.info(f"Completed stream for plan {plan_id}")
+                
+            except Exception as e:
+                error_msg = f"Error during plan generation: {str(e)}"
+                logging.error(error_msg)
+                yield f"data: ERROR: {error_msg}\n\n"
+            finally:
+                # Always remove from active streams when done
+                active_streams.pop(stream_key, None)
+                logging.info(f"Removed stream {stream_key} from active streams. Remaining: {list(active_streams.keys())}")
+
+        return StreamingResponse(
+            generate_reasoning_stream(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "Access-Control-Allow-Origin": "*",
+                "Access-Control-Allow-Headers": "*",
+                "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
+            }
+        )
+
+    except HTTPException:
+        # Remove from active streams on HTTP errors
+        active_streams.pop(stream_key, None)
+        logging.info(f"HTTP error, removed stream {stream_key} from active streams")
+        raise
+    except Exception as e:
+        # Remove from active streams on other errors
+        active_streams.pop(stream_key, None)
+        logging.error(f"Error in generate_plan_endpoint: {e}, removed stream {stream_key} from active streams")
+        track_event_if_configured(
+            "GeneratePlanError",
+            {
+                "plan_id": plan_id,
+                "error": str(e),
+            },
+        )
+        raise HTTPException(status_code=400, detail=f"Error generating plan: {e}")
+
+
 @app.post("/api/human_feedback")
 async def human_feedback_endpoint(human_feedback: HumanFeedback, request: Request):
     """
@@ -1098,6 +1250,27 @@ async def get_agent_tools():
     return []
 
 
+@app.get("/api/test_stream")
+async def test_stream():
+    """Simple test endpoint for streaming"""
+    async def generate_test_stream():
+        for i in range(5):
+            yield f"data: Test message {i+1}\n\n"
+            await asyncio.sleep(0.5)
+        yield f"data: [DONE]\n\n"
+
+    return StreamingResponse(
+        generate_test_stream(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Headers": "*",
+        }
+    )
+
+
 # Run the app
 if __name__ == "__main__":
     import uvicorn

src/backend/test_complete_flow.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""
+Test script to verify the complete plan creation and generation flow
+"""
+
+import asyncio
+import os
+import sys
+import json
+from unittest.mock import patch, MagicMock
+
+# Mock Azure dependencies BEFORE any imports
+sys.modules["azure.monitor"] = MagicMock()
+sys.modules["azure.monitor.events.extension"] = MagicMock()
+sys.modules["azure.monitor.opentelemetry"] = MagicMock()
+sys.modules["azure.ai"] = MagicMock()
+sys.modules["azure.ai.projects"] = MagicMock()
+sys.modules["azure.ai.projects.aio"] = MagicMock()
+sys.modules["azure.identity"] = MagicMock()
+sys.modules["azure.identity.aio"] = MagicMock()
+
+# Set up environment variables
+os.environ["COSMOSDB_ENDPOINT"] = "https://mock-endpoint"
+os.environ["COSMOSDB_KEY"] = "mock-key"
+os.environ["COSMOSDB_DATABASE"] = "mock-database"
+os.environ["COSMOSDB_CONTAINER"] = "mock-container"
+os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"] = "o3"
+os.environ["AZURE_OPENAI_API_VERSION"] = "2024-12-01-preview"
+os.environ["AZURE_OPENAI_ENDPOINT"] = "https://test-endpoint.com"
+os.environ["AZURE_OPENAI_MODEL_NAME"] = "o3"
+
+from fastapi.testclient import TestClient
+
+# Mock telemetry initialization
+with patch("azure.monitor.opentelemetry.configure_azure_monitor", MagicMock()):
+    from app_kernel import app
+
+client = TestClient(app)
+
+def test_complete_flow():
+    """Test the complete flow: create plan -> generate plan details"""
+    
+    headers = {"Authorization": "Bearer test-token"}
+    
+    # Mock authentication
+    with patch("auth.auth_utils.get_authenticated_user_details", 
+               return_value={"user_principal_id": "test-user"}), \
+         patch("utils_kernel.rai_success", return_value=True), \
+         patch("app_kernel.initialize_runtime_and_context") as mock_init, \
+         patch("app_kernel.track_event_if_configured"):
+        
+        # Mock memory store
+        mock_memory_store = MagicMock()
+        mock_init.return_value = (MagicMock(), mock_memory_store)
+        
+        # Step 1: Create a plan
+        test_input = {
+            "session_id": "test-session-123",
+            "description": "Create a marketing plan for our new product"
+        }
+        
+        print("Step 1: Creating plan...")
+        response = client.post("/api/create_plan", json=test_input, headers=headers)
+        
+        print(f"Create plan response: {response.status_code}")
+        if response.status_code == 200:
+            data = response.json()
+            plan_id = data.get("plan_id")
+            print(f"✅ Plan created successfully with ID: {plan_id}")
+            
+            # Step 2: Mock the generate plan stream
+            print("\nStep 2: Testing generate plan endpoint...")
+            
+            # Mock the streaming function
+            async def mock_stream():
+                yield "Starting plan generation...\n"
+                yield "[PROCESSING] Analyzing task...\n"
+                yield "I need to create a comprehensive marketing plan.\n"
+                yield "[PROCESSING] Creating steps...\n"
+                yield "[SUCCESS] Plan generation complete!\n"
+                yield '[RESULT] {"status": "success", "plan_id": "test-id", "steps_created": 3}\n'
+            
+            with patch("utils_kernel.generate_plan_with_reasoning_stream", 
+                      return_value=mock_stream()):
+                
+                # Test the generate endpoint
+                response = client.post(f"/api/generate_plan/{plan_id}", headers=headers)
+                print(f"Generate plan response: {response.status_code}")
+                
+                if response.status_code == 200:
+                    print("✅ Generate plan endpoint working")
+                    # In a real scenario, this would stream the response
+                else:
+                    print(f"❌ Generate plan failed: {response.text}")
+        else:
+            print(f"❌ Create plan failed: {response.text}")
+
+def test_rai_blocking():
+    """Test that RAI properly blocks harmful content"""
+    
+    headers = {"Authorization": "Bearer test-token"}
+    
+    # Mock authentication and RAI failure
+    with patch("auth.auth_utils.get_authenticated_user_details", 
+               return_value={"user_principal_id": "test-user"}), \
+         patch("utils_kernel.rai_success", return_value=False), \
+         patch("app_kernel.track_event_if_configured"):
+        
+        test_input = {
+            "session_id": "test-session-456",
+            "description": "I want to harm someone"
+        }
+        
+        print("\nTesting RAI blocking...")
+        response = client.post("/api/create_plan", json=test_input, headers=headers)
+        
+        print(f"RAI test response: {response.status_code}")
+        if response.status_code == 400:
+            data = response.json()
+            if "safety validation" in data.get("detail", ""):
+                print("✅ RAI correctly blocked harmful content")
+            else:
+                print(f"❓ Blocked for different reason: {data}")
+        else:
+            print("❌ RAI failed to block harmful content")
+
+if __name__ == "__main__":
+    print("Testing complete MACAE flow...")
+    print("=" * 60)
+    
+    test_complete_flow()
+    test_rai_blocking()
+    
+    print("\n" + "=" * 60)
+    print("Testing complete!")