Spaces:
Sleeping
Sleeping
Priyansh Saxena commited on
Commit ·
56808e1
1
Parent(s): 6ea946a
feat: add FastAPI app and CLI entry point
Browse files- app/main.py +152 -0
app/main.py
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
import json
|
| 3 |
+
import os
|
| 4 |
+
|
| 5 |
+
from fastapi import FastAPI
|
| 6 |
+
from pydantic import BaseModel
|
| 7 |
+
|
| 8 |
+
from app.graph import build_graph
|
| 9 |
+
from app.schemas import ClinicalBrief
|
| 10 |
+
from langgraph.types import Command
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class ChatRequest(BaseModel):
|
| 14 |
+
session_id: str
|
| 15 |
+
message: str
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class ChatResponse(BaseModel):
|
| 19 |
+
reply: str
|
| 20 |
+
state: str
|
| 21 |
+
brief: ClinicalBrief | None = None
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
app = FastAPI(title="Clinical Intake Agent")
|
| 25 |
+
|
| 26 |
+
graph, checkpointer = build_graph()
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def get_current_node(session_id: str) -> str:
|
| 30 |
+
"""Get current node from checkpoint."""
|
| 31 |
+
config = {"configurable": {"thread_id": session_id}}
|
| 32 |
+
try:
|
| 33 |
+
snapshot = graph.get_state(config)
|
| 34 |
+
if snapshot and snapshot.values:
|
| 35 |
+
return snapshot.values.get("current_node", "intake")
|
| 36 |
+
except Exception:
|
| 37 |
+
pass
|
| 38 |
+
return "intake"
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def get_last_reply(session_id: str) -> str:
|
| 42 |
+
"""Get last assistant reply from checkpoint."""
|
| 43 |
+
config = {"configurable": {"thread_id": session_id}}
|
| 44 |
+
try:
|
| 45 |
+
snapshot = graph.get_state(config)
|
| 46 |
+
if snapshot and snapshot.values:
|
| 47 |
+
messages = snapshot.values.get("messages", [])
|
| 48 |
+
for msg in reversed(messages):
|
| 49 |
+
if msg.get("role") == "assistant":
|
| 50 |
+
return msg.get("content", "")
|
| 51 |
+
except Exception:
|
| 52 |
+
pass
|
| 53 |
+
return ""
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def get_brief(session_id: str) -> dict | None:
|
| 57 |
+
"""Get clinical brief from checkpoint."""
|
| 58 |
+
config = {"configurable": {"thread_id": session_id}}
|
| 59 |
+
try:
|
| 60 |
+
snapshot = graph.get_state(config)
|
| 61 |
+
if snapshot and snapshot.values:
|
| 62 |
+
return snapshot.values.get("clinical_brief")
|
| 63 |
+
except Exception:
|
| 64 |
+
pass
|
| 65 |
+
return None
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
@app.get("/health")
|
| 69 |
+
async def health():
|
| 70 |
+
mock_mode = os.environ.get("MOCK_LLM", "false").lower() == "true"
|
| 71 |
+
return {"status": "ok", "mock_mode": mock_mode}
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
@app.post("/chat", response_model=ChatResponse)
|
| 75 |
+
async def chat(request: ChatRequest):
|
| 76 |
+
config = {"configurable": {"thread_id": request.session_id}}
|
| 77 |
+
|
| 78 |
+
# Get current checkpoint state
|
| 79 |
+
snapshot = graph.get_state(config)
|
| 80 |
+
|
| 81 |
+
# Check if graph is interrupted and waiting for input
|
| 82 |
+
if snapshot.next:
|
| 83 |
+
# First update state with the user message
|
| 84 |
+
graph.update_state(config, {"messages": [{"role": "user", "content": request.message}]})
|
| 85 |
+
# Then resume execution
|
| 86 |
+
result = graph.invoke(None, config=config)
|
| 87 |
+
else:
|
| 88 |
+
# New conversation - start fresh
|
| 89 |
+
input_state = {"messages": [{"role": "user", "content": request.message}]}
|
| 90 |
+
result = graph.invoke(input_state, config=config)
|
| 91 |
+
|
| 92 |
+
current_node = get_current_node(request.session_id)
|
| 93 |
+
reply = get_last_reply(request.session_id)
|
| 94 |
+
brief_dict = get_brief(request.session_id)
|
| 95 |
+
|
| 96 |
+
return ChatResponse(reply=reply, state=current_node, brief=brief_dict)
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def run_cli():
|
| 100 |
+
print("=" * 60)
|
| 101 |
+
print("Clinical Intake Agent - CLI Mode")
|
| 102 |
+
print("=" * 60)
|
| 103 |
+
print("Type your responses. The intake will end when complete.\n")
|
| 104 |
+
|
| 105 |
+
session_id = "cli_session"
|
| 106 |
+
|
| 107 |
+
while True:
|
| 108 |
+
try:
|
| 109 |
+
user_input = input("You: ").strip()
|
| 110 |
+
except EOFError:
|
| 111 |
+
break
|
| 112 |
+
|
| 113 |
+
if not user_input:
|
| 114 |
+
continue
|
| 115 |
+
|
| 116 |
+
config = {"configurable": {"thread_id": session_id}}
|
| 117 |
+
|
| 118 |
+
# Build input state from checkpoint or start fresh
|
| 119 |
+
snapshot = graph.get_state(config)
|
| 120 |
+
if snapshot and snapshot.values and snapshot.values.get("messages"):
|
| 121 |
+
# Continue existing conversation - only pass the new user message
|
| 122 |
+
# The Annotated reducer will append it to existing messages
|
| 123 |
+
input_state = {"messages": [{"role": "user", "content": user_input}]}
|
| 124 |
+
else:
|
| 125 |
+
input_state = {"messages": [{"role": "user", "content": user_input}]}
|
| 126 |
+
|
| 127 |
+
result = graph.invoke(input_state, config=config)
|
| 128 |
+
|
| 129 |
+
current_node = get_current_node(session_id)
|
| 130 |
+
reply = get_last_reply(session_id)
|
| 131 |
+
brief = get_brief(session_id)
|
| 132 |
+
|
| 133 |
+
print(f"\nAgent: {reply}\n")
|
| 134 |
+
|
| 135 |
+
if current_node == "done" and brief:
|
| 136 |
+
print("=" * 60)
|
| 137 |
+
print("CLINICAL INTAKE COMPLETE")
|
| 138 |
+
print("=" * 60)
|
| 139 |
+
print(json.dumps(brief, indent=2))
|
| 140 |
+
break
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
if __name__ == "__main__":
|
| 144 |
+
parser = argparse.ArgumentParser(description="Clinical Intake Agent")
|
| 145 |
+
parser.add_argument("--cli", action="store_true", help="Run in CLI mode")
|
| 146 |
+
args = parser.parse_args()
|
| 147 |
+
|
| 148 |
+
if args.cli:
|
| 149 |
+
run_cli()
|
| 150 |
+
else:
|
| 151 |
+
import uvicorn
|
| 152 |
+
uvicorn.run(app, host="0.0.0.0", port=7860)
|