""" MedRAX API Module This module provides a FastAPI-based REST API for the MedRAX medical imaging AI assistant. It offers endpoints for processing medical images with text queries using the same agent architecture as the Gradio interface. The API supports: - Text-only queries - Single or multiple image inputs - Optional custom system prompts - Automatic thread management for each request - Tool execution and result aggregation """ import uuid import base64 from pathlib import Path from typing import List, Optional, Dict, Any import re import time from fastapi import FastAPI, HTTPException, UploadFile, File, Form from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from langchain_core.messages import AIMessage, ToolMessage # Import MedRAX components from medrax.agent import Agent class QueryRequest(BaseModel): """ Request model for text-only queries. Attributes: question (str): The question or query to ask the agent system_prompt (Optional[str]): Custom system prompt to override default thread_id (Optional[str]): Optional thread ID for conversation continuity """ question: str = Field(..., description="The question or query to ask the agent") system_prompt: Optional[str] = Field(None, description="Custom system prompt to override default") thread_id: Optional[str] = Field(None, description="Optional thread ID for conversation continuity") class QueryResponse(BaseModel): """ Response model for API queries. Attributes: response (str): The agent's text response thread_id (str): The thread ID used for this conversation tools_used (List[str]): List of tools that were executed processing_time (float): Time taken to process the request in seconds """ response: str = Field(..., description="The agent's text response") thread_id: str = Field(..., description="The thread ID used for this conversation") tools_used: List[str] = Field(..., description="List of tools that were executed") processing_time: float = Field(..., description="Time taken to process the request in seconds") class MedRAXAPI: """ FastAPI application wrapper for the MedRAX agent. This class provides a clean interface for creating and managing the API endpoints while maintaining separation of concerns from the core agent functionality. """ def __init__(self, agent: Agent, tools_dict: Dict[str, Any], temp_dir: str = "temp_api"): """ Initialize the MedRAX API. Args: agent (Agent): The initialized MedRAX agent tools_dict (Dict[str, Any]): Dictionary of available tools temp_dir (str): Directory for temporary file storage """ self.agent = agent self.tools_dict = tools_dict self.temp_dir = Path(temp_dir) self.temp_dir.mkdir(exist_ok=True) # Create FastAPI app self.app = FastAPI( title="MedRAX API", description="Medical Reasoning Agent for Chest X-ray Analysis", version="2.0.0", docs_url="/docs", redoc_url="/redoc", ) # Add CORS middleware self.app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Register routes self._register_routes() def _register_routes(self): """Register all API routes.""" @self.app.get("/health") async def health_check(): """Health check endpoint.""" return {"status": "healthy", "service": "MedRAX API"} @self.app.get("/tools") async def list_tools(): """List available tools.""" return {"available_tools": list(self.tools_dict.keys()), "total_count": len(self.tools_dict)} @self.app.post("/query", response_model=QueryResponse) async def query_text_only(request: QueryRequest): """ Process a text-only query without images. Args: request (QueryRequest): The query request Returns: QueryResponse: The agent's response """ return await self._process_query( question=request.question, system_prompt=request.system_prompt, thread_id=request.thread_id, images=None ) @self.app.post("/query-with-images", response_model=QueryResponse) async def query_with_images( question: str = Form(..., description="The question or query to ask the agent"), system_prompt: Optional[str] = Form(None, description="Custom system prompt to override default"), thread_id: Optional[str] = Form(None, description="Optional thread ID for conversation continuity"), images: List[UploadFile] = File(..., description="One or more medical images to analyze"), ): """ Process a query with one or more images. Args: question (str): The question or query to ask the agent system_prompt (Optional[str]): Custom system prompt to override default thread_id (Optional[str]): Optional thread ID for conversation continuity images (List[UploadFile]): List of uploaded image files Returns: QueryResponse: The agent's response """ # Validate image files if not images or len(images) == 0: raise HTTPException(status_code=400, detail="At least one image is required") # Validate file types allowed_types = {"image/jpeg", "image/jpg", "image/png", "image/bmp", "image/tiff", "application/dicom"} for image in images: if image.content_type not in allowed_types: raise HTTPException( status_code=400, detail=f"Unsupported file type: {image.content_type}. Allowed types: {allowed_types}", ) return await self._process_query( question=question, system_prompt=system_prompt, thread_id=thread_id, images=images ) async def _process_query( self, question: str, system_prompt: Optional[str] = None, thread_id: Optional[str] = None, images: Optional[List[UploadFile]] = None, ) -> QueryResponse: """ Internal method to process queries through the agent. Args: question (str): The question to ask system_prompt (Optional[str]): Custom system prompt thread_id (Optional[str]): Thread ID for conversation images (Optional[List[UploadFile]]): List of images Returns: QueryResponse: The processed response """ start_time = time.time() # Generate thread ID if not provided if not thread_id: thread_id = str(uuid.uuid4()) try: # Prepare messages messages = [] image_paths = [] # Handle image uploads if images: for i, image in enumerate(images): # Save uploaded file temporarily temp_path = self.temp_dir / f"{thread_id}_{i}_{image.filename}" with open(temp_path, "wb") as buffer: content = await image.read() buffer.write(content) image_paths.append(str(temp_path)) # Add image path for tools messages.append({"role": "user", "content": f"image_path: {temp_path}"}) # Add base64 encoded image for multimodal processing image_base64 = base64.b64encode(content).decode("utf-8") # Determine MIME type mime_type = "image/jpeg" # Default if image.content_type: mime_type = image.content_type elif temp_path.suffix.lower() in [".png"]: mime_type = "image/png" messages.append( { "role": "user", "content": [ { "type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_base64}"}, } ], } ) # Add text question messages.append({"role": "user", "content": [{"type": "text", "text": question}]}) # Process through agent workflow response_text = "" tools_used = [] # Temporarily update system prompt if provided original_prompt = None if system_prompt: original_prompt = self.agent.system_prompt self.agent.system_prompt = system_prompt try: async for chunk in self._stream_agent_response(messages, thread_id): if chunk.get("type") == "text": response_text += chunk.get("content", "") elif chunk.get("type") == "tool": tools_used.append(chunk.get("tool_name", "")) finally: # Restore original system prompt if original_prompt is not None: self.agent.system_prompt = original_prompt # Clean up temporary files for image_path in image_paths: try: Path(image_path).unlink(missing_ok=True) except Exception: pass # Ignore cleanup errors processing_time = time.time() - start_time return QueryResponse( response=response_text.strip(), thread_id=thread_id, tools_used=list(set(tools_used)), # Remove duplicates processing_time=processing_time, ) except Exception as e: # Clean up on error for image_path in image_paths: try: Path(image_path).unlink(missing_ok=True) except Exception: pass raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}") async def _stream_agent_response(self, messages: List[Dict], thread_id: str): """ Stream responses from the agent workflow. Args: messages (List[Dict]): Messages to process thread_id (str): Thread ID for the conversation Yields: Dict: Response chunks with type and content """ try: for chunk in self.agent.workflow.stream( {"messages": messages}, {"configurable": {"thread_id": thread_id}}, stream_mode="updates", ): if not isinstance(chunk, dict): continue for node_name, node_output in chunk.items(): if "messages" not in node_output: continue for msg in node_output["messages"]: if isinstance(msg, AIMessage) and msg.content: # Clean up temp paths from response clean_content = re.sub(r"temp[^\s]*", "", msg.content).strip() if clean_content: yield {"type": "text", "content": clean_content} elif isinstance(msg, ToolMessage): # Extract tool name from the message tool_call_id = msg.tool_call_id # We'll track tool usage but not include detailed output in API response yield {"type": "tool", "tool_name": "tool_executed"} except Exception as e: yield {"type": "error", "content": str(e)} def create_api(agent: Agent, tools_dict: Dict[str, Any], temp_dir: str = "temp_api") -> FastAPI: """ Create and configure the MedRAX FastAPI application. Args: agent (Agent): The initialized MedRAX agent tools_dict (Dict[str, Any]): Dictionary of available tools temp_dir (str): Directory for temporary file storage Returns: FastAPI: Configured FastAPI application """ api = MedRAXAPI(agent, tools_dict, temp_dir) return api.app