Update api_server.py

api_server.py

@@ -4,6 +4,8 @@ from fastapi import FastAPI, HTTPException, Depends
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from backend.agent import run_agent
+from typing import Union
+
 
 app = FastAPI(title="Nursing Copilot API")
 
@@ -30,8 +32,7 @@ class PredictionRequest(BaseModel):
     model: str | None = "gpt2-nursing"
 
 class PredictionResponse(BaseModel):
-    completions: list[str]
-
+    completions: list[Union[str, dict]]
 
 # -----------------------------
 # API Endpoint