Priyansh Saxena commited on
Commit
56808e1
·
1 Parent(s): 6ea946a

feat: add FastAPI app and CLI entry point

Browse files
Files changed (1) hide show
  1. 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)