3-ai-native-e2e-sample/backend/routers/medication.py

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import BingGroundingTool, FunctionTool, SubmitToolOutputsAction, RequiredFunctionToolCall, ToolOutput
from azure.identity import DefaultAzureCredential
from agents.medication_functions import medication_functions
import os
import logging
import json
import time
import asyncio
from pydantic import BaseModel
from typing import Optional

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/agents", tags=["medication"])

class MedicationInfo(BaseModel):
    name: str
    notes: Optional[str] = None

class MedicationAnalysis(BaseModel):
    analysis: str
    interactions: list[str]
    warnings: list[str]
    recommendations: list[str]

@router.post("/medication/analyze_stream")
async def analyze_medication_stream(info: MedicationInfo):
    async def event_generator():
        try:
            # Initialize client
            project_client = AIProjectClient.from_connection_string(
                credential=DefaultAzureCredential(),
                conn_str=os.environ["PROJECT_CONNECTION_STRING"]
            )
            logger.info(f"Starting streaming medication analysis for: {info.name}")
            
            # Get Bing connection and set up the tool
            bing_conn = project_client.connections.get(connection_name=os.environ["BING_CONNECTION_NAME"])
            if not bing_conn:
                yield f"data: {json.dumps({'type': 'error', 'content': 'No Bing connection found.'})}\n\n"
                return
            bing_tool = BingGroundingTool(connection_id=bing_conn.id)
            
            # Update system prompt with the new instructions:
            # The assistant must use Bing to search for medication data,
            # then call the analyze_medication_info function to format the response.
            system_prompt = (
                "You are a medication analysis assistant. Use Bing to retrieve accurate, up-to-date "
                "information about the medication specified by the user. Once you obtain Bing's results, "
                "call the analyze_medication_info function to structure your response as a JSON string. "
                "Return only a valid JSON string without any markdown or additional text."
            )
            
            # Configure function tools
            functions = FunctionTool(functions=medication_functions)
            
            # Create agent with both Bing and function tools
            agent = project_client.agents.create_agent(
                model=os.environ["MODEL_DEPLOYMENT_NAME"],
                name="medication-analysis-stream",
                instructions=system_prompt,
                tools=[*bing_tool.definitions, *functions.definitions],
                headers={"x-ms-enable-preview": "true"}
            )
            logger.info(f"Created agent with ID: {agent.id}")
            yield f"data: {json.dumps({'type': 'message', 'content': 'Agent created. Starting thread...'})}\n\n"
            
            # Create thread and initial message
            thread = project_client.agents.create_thread()
            message_content = f"Analyze the medication: {info.name}. {info.notes if info.notes else ''}"
            project_client.agents.create_message(
                thread_id=thread.id,
                role="user",
                content=message_content
            )
            yield f"data: {json.dumps({'type': 'message', 'content': 'Thread created and message sent.'})}\n\n"
            
            # Create and start the run
            run = project_client.agents.create_run(thread_id=thread.id, agent_id=agent.id)
            logger.info(f"Created run with ID: {run.id}")
            yield f"data: {json.dumps({'type': 'message', 'content': 'Run initiated. Processing...'})}\n\n"
            
            # Poll for run completion
            max_retries = 60
            retry_count = 0
            while run.status in ["queued", "in_progress", "requires_action"]:
                await asyncio.sleep(1)
                run = project_client.agents.get_run(thread_id=thread.id, run_id=run.id)
                logger.info(f"Current run status: {run.status}")
                yield f"data: {json.dumps({'type': 'message', 'content': f'Run status: {run.status}'})}\n\n"
                retry_count += 1
                if retry_count >= max_retries:
                    yield f"data: {json.dumps({'type': 'error', 'content': 'Run timed out.'})}\n\n"
                    return
                # Handle required function calls if any
                if run.status == "requires_action" and hasattr(run, "required_action") and isinstance(run.required_action, SubmitToolOutputsAction):
                    tool_calls = run.required_action.submit_tool_outputs.tool_calls
                    if not tool_calls:
                        logger.error("No tool calls provided")
                        break
                    tool_outputs = []
                    for tool_call in tool_calls:
                        if isinstance(tool_call, RequiredFunctionToolCall):
                            try:
                                output = functions.execute(tool_call)
                                tool_outputs.append(
                                    ToolOutput(
                                        tool_call_id=tool_call.id,
                                        output=output
                                    )
                                )
                                yield f"data: {json.dumps({'type': 'message', 'content': f'Executed function call: {tool_call.function.name}'})}\n\n"
                            except Exception as e:
                                logger.error(f"Error executing tool call: {e}")
                    if tool_outputs:
                        run = project_client.agents.submit_tool_outputs_to_run(
                            thread_id=thread.id,
                            run_id=run.id,
                            tool_outputs=tool_outputs
                        )
                        continue
            
            # Once the run is complete, look for the assistant message
            if run.status == "failed":
                logger.error(f"Run failed: {run.last_error}")
                yield f"data: {json.dumps({'type': 'error', 'content': f'Run failed: {run.last_error}'})}\n\n"
                return

            messages = project_client.agents.list_messages(thread_id=thread.id)
            final_result = None
            for msg in reversed(messages.data):
                if msg.role == "assistant":
                    # Loop through message content parts
                    for content in msg.content:
                        if hasattr(content, 'text') and content.text:
                            text = content.text.value
                            # If the response is formatted with markdown, extract JSON
                            if text.startswith('```json'):
                                text = text.split('```json')[1].split('```')[0].strip()
                            try:
                                final_result = json.loads(text)
                                break
                            except json.JSONDecodeError as e:
                                logger.error(f"JSON decode error: {e}")
                    if final_result:
                        break

            if final_result is None:
                yield f"data: {json.dumps({'type': 'error', 'content': 'No valid response received from AI'})}\n\n"
                return

            # Yield final result as a completed message
            yield f"data: {json.dumps({'done': True, 'type': 'final', 'content': final_result})}\n\n"
            
        except Exception as ex:
            logger.error(f"Exception in streaming analysis: {ex}")
            yield f"data: {json.dumps({'type': 'error', 'content': str(ex)})}\n\n"
    
    return StreamingResponse(event_generator(), media_type="text/event-stream")