arittrabag/PharmAI_Navigator
0
1from __future__ import annotations
2
3import os
4import json
5import re
6from typing import Any, Dict, List, Optional
7from dotenv import load_dotenv
8
9from langchain_anthropic import ChatAnthropic
10from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage
11from langchain_core.tools import tool
12
13from langgraph.graph import StateGraph, START, END, MessagesState
14from langgraph.prebuilt import ToolNode, tools_condition
15
16from tools import (
17 tavily_search,
18 stub_evidence,
19 classify_query,
20 extract_entities,
21 normalize_evidence,
22 generate_graph_dot,
23 clinicaltrials_search,
24 render_dot_to_png_base64
25)
26
27# Load environment variables
28load_dotenv()
29
30# -----------------------------
31# LangChain Tool Wrappers
32# -----------------------------
33@tool("web_search")
34def web_search_tool(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
35 """Web search using Tavily. Returns a list of evidence dicts."""
36 ev = tavily_search(query=query, max_results=max_results)
37 return [e.model_dump() for e in ev]
38
39
40@tool("stub_evidence")
41def stub_evidence_tool(query: str) -> List[Dict[str, Any]]:
42 """Deterministic fallback evidence tool (offline/demo)."""
43 ev = stub_evidence(query=query)
44 return [e.model_dump() for e in ev]
45
46@tool("classify_query")
47def classify_query_tool(query: str) -> Dict[str, Any]:
48 """Classify query to decide which tools are needed."""
49 return classify_query(query)
50
51
52@tool("extract_entities")
53def extract_entities_tool(query: str) -> Dict[str, Optional[str]]:
54 """Extract drug and indication from query."""
55 return extract_entities(query)
56
57
58@tool("normalize_evidence")
59def normalize_evidence_tool(evidence: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
60 """Dedupe and clean evidence."""
61 return normalize_evidence(evidence)
62
63
64@tool("generate_graph_dot")
65def generate_graph_dot_tool(
66 title: str,
67 nodes: List[Dict[str, str]],
68 edges: List[Dict[str, str]],
69 rankdir: str = "LR",
70) -> str:
71 """
72 Generate Graphviz DOT.
73 IMPORTANT: Use this tool instead of writing DOT directly.
74 """
75 return generate_graph_dot(
76 title=title,
77 nodes=nodes,
78 edges=edges,
79 rankdir=rankdir,
80 )
81
82@tool("clinicaltrials_search")
83def clinicaltrials_search_tool(drug: str, indication: str, max_results: int = 5) -> List[Dict[str, Any]]:
84 """Search ClinicalTrials.gov (Tavily-based MVP)."""
85 ev = clinicaltrials_search(drug=drug, indication=indication, max_results=max_results)
86 return [e.model_dump() for e in ev]
87
88@tool("render_dot_to_png_base64")
89def render_dot_to_png_base64_tool(dot: str) -> Dict[str, Any]:
90 """Render DOT to PNG (base64). Optional dependency on graphviz."""
91 return render_dot_to_png_base64(dot)
92
93TOOLS = [
94 web_search_tool,
95 stub_evidence_tool,
96 classify_query_tool,
97 extract_entities_tool,
98 normalize_evidence_tool,
99 generate_graph_dot_tool,
100 clinicaltrials_search_tool,
101 render_dot_to_png_base64_tool
102]
103
104# -----------------------------
105# LangGraph State
106# -----------------------------
107class PharmAIState(MessagesState):
108 session_id: Optional[str]
109 user_query: str
110 decision_brief: str
111 citations: List[str]
112 confidence_score: float
113 tool_loops: int # safety counter
114 diagram_png_base64: Optional[str] # <-- add
115 diagram_dot: Optional[str] # <-- optional
116 intent: str # "simple" | "diligence" | "diagram"
117
118# -----------------------------
119# Guardrails + Prompts
120# -----------------------------
121SYSTEM_PROMPT = """You are PharmAI Navigator, an evidence-grounded diligence assistant for drug/asset evaluation.
122
123Your job:
124Turn a query like "Assess {Drug} for {Indication}" into a decision-grade brief OR structured output.
125
126CRITICAL TOOL USAGE RULES:
127- If the user asks for a diagram, flow, architecture, graph, visualization, or Graphviz:
128 → You MUST call `generate_graph_dot`.
129 → You MUST NOT write Graphviz DOT directly in your response.
130 → If the user asks for an image/PNG, call `render_dot_to_png_base64` AFTER you get DOT.
131- If the user asks for trials / phases / NCT IDs / endpoints:
132 → Prefer calling `extract_entities` then `clinicaltrials_search`.
133- If the user asks for factual claims (approvals, safety, pricing, patents, market):
134 → Prefer calling `web_search`.
135
136Guardrails (STRICT):
137- Do NOT invent specific facts (approval dates, trial names, endpoints, statistics, patent expiry).
138- Any concrete number/date/claim MUST be supported by tool evidence.
139- If evidence is insufficient, clearly list Evidence Gaps.
140- Be concise, structured, and decision-oriented.
141- Avoid medical advice; present as diligence/analysis.
142
143Simple Query Rule (CRITICAL):
144- If the user asks a simple definitional question ("what is", "define", "explain") and you can answer without external verification, do NOT call tools and respond directly.
145- Only use tools when you need current/specific data (trials, approvals, patents, market data).
146
147Citations policy:
148- The final response's "Citations" section is handled by the system.
149- Do NOT create your own citation list.
150"""
151
152FINAL_PROMPT = """Write the FINAL decision brief with these sections:
153
1541) Executive Recommendation (1–2 lines)
1552) Scientific Rationale (bullets)
1563) Clinical Evidence Snapshot (bullets)
1574) IP / Exclusivity Quick View (bullets)
1585) Market / SoC Snapshot (bullets)
1596) Key Risks + Next Actions (bullets)
160
161Rules:
162- If evidence is insufficient, include "Evidence Gaps" with bullets.
163- Do NOT add a citations section yourself; the system will append it.
164Return plain text only.
165"""
166
167# Placeholder detection to avoid wasting tokens on "Drug X / Indication Y"
168PLACEHOLDER_PATTERNS = [
169 r"\bdrug\s*x\b",
170 r"\bindication\s*y\b",
171 r"\bdrug\s*name\b",
172 r"\bindication\s*name\b",
173]
174def _looks_like_placeholder(q: str) -> bool:
175 ql = (q or "").strip().lower()
176 return any(re.search(p, ql) for p in PLACEHOLDER_PATTERNS)
177
178
179def _build_model() -> ChatAnthropic:
180 model_name = os.getenv("ANTHROPIC_MODEL", "claude-3-7-sonnet-latest")
181 return ChatAnthropic(
182 model=model_name,
183 temperature=0.2,
184 max_tokens=10000,
185 timeout=120,
186 streaming=False,
187 stop=None
188 ).bind_tools(TOOLS)
189
190
191# Safety cap to avoid endless tool loops
192MAX_TOOL_LOOPS = int(os.getenv("MAX_TOOL_LOOPS", "4"))
193
194
195def llm_call(state: PharmAIState) -> Dict[str, Any]:
196 """
197 Calls Claude with tool schemas attached.
198 Returns new messages to append into state["messages"].
199 """
200 llm = _build_model()
201 messages: List[BaseMessage] = state["messages"]
202
203 if not messages or not isinstance(messages[0], SystemMessage):
204 messages = [SystemMessage(content=SYSTEM_PROMPT)] + messages
205
206 tool_loops = state.get("tool_loops", 0)
207 if tool_loops >= MAX_TOOL_LOOPS:
208 # Stop tool-calling loop and force synthesis
209 stop_msg = HumanMessage(
210 content=(
211 "Stop calling tools now. Proceed to final synthesis using what you already have. "
212 "If evidence is insufficient, clearly list Evidence Gaps."
213 )
214 )
215 messages = messages + [stop_msg]
216
217 resp = llm.invoke(messages)
218 return {"messages": [resp]}
219
220
221# -----------------------------
222# Citations extraction (tool-only)
223# -----------------------------
224def _clean_url(u: str) -> str:
225 return u.strip().strip("),.]}\"'")
226
227def _extract_citations_from_messages(messages: List[BaseMessage]) -> List[str]:
228 """
229 Tool-only citation extraction (single source of truth):
230 - ONLY reads ToolMessage contents (actual tool outputs).
231 - If tool output is JSON (list/dict), pull `source` fields.
232 - Fallback: regex URL extraction from tool text.
233 """
234 citations: List[str] = []
235 url_re = re.compile(r"https?://[^\s\]\)\}\",']+")
236
237 for m in messages:
238 if not isinstance(m, ToolMessage):
239 continue
240
241 content = getattr(m, "content", None)
242 if not content:
243 continue
244
245 if isinstance(content, str):
246 parsed = None
247 try:
248 parsed = json.loads(content)
249 except Exception:
250 parsed = None
251
252 if isinstance(parsed, list):
253 for item in parsed:
254 if isinstance(item, dict):
255 src = item.get("source")
256 if isinstance(src, str) and src.startswith(("http://", "https://")):
257 citations.append(_clean_url(src))
258 elif isinstance(parsed, dict):
259 src = parsed.get("source")
260 if isinstance(src, str) and src.startswith(("http://", "https://")):
261 citations.append(_clean_url(src))
262
263 for u in url_re.findall(content):
264 citations.append(_clean_url(u))
265
266 # De-duplicate
267 seen = set()
268 out = []
269 for c in citations:
270 # drop clearly broken/truncated URLs
271 if len(c) < 12:
272 continue
273 if c not in seen:
274 seen.add(c)
275 out.append(c)
276 return out
277
278
279def _append_citations_section(brief_text: str, citations: List[str]) -> str:
280 """
281 Enforces "single source of truth":
282 - Removes any existing 'Citations' section the model may have produced
283 - Appends citations derived from tool outputs only
284 """
285 text = (brief_text or "").strip()
286
287 # Remove any model-generated citations section (best-effort)
288 # (handles '## Citations' or 'Citations' headers)
289 text = re.split(r"\n#{1,3}\s*Citations\s*\n|\nCitations\s*\n", text, maxsplit=1)[0].rstrip()
290
291 if citations:
292 lines = ["", "## Citations"]
293 for i, c in enumerate(citations, 1):
294 lines.append(f"{i}. {c}")
295 text = text + "\n" + "\n".join(lines)
296 else:
297 text = text + "\n\n## Citations\n- (No external sources retrieved.)"
298
299 return text
300
301def capture_diagram(state: PharmAIState) -> Dict[str, Any]:
302 # Find the last ToolMessage (most recent tool output)
303 last_tool = None
304 for m in reversed(state["messages"]):
305 if isinstance(m, ToolMessage):
306 last_tool = m
307 break
308
309 if not last_tool:
310 return {}
311
312 tool_name = getattr(last_tool, "name", "") or ""
313 content = getattr(last_tool, "content", "")
314
315 # If your render tool returns base64 string directly
316 if tool_name == "render_dot_to_png_base64":
317 return {"diagram_png_base64": content}
318
319 # If your generate_graph_dot returns dot string
320 if tool_name == "generate_graph_dot":
321 return {"diagram_dot": content}
322
323 return {}
324
325def route_after_tools(state: PharmAIState) -> str:
326 # If we already have the final diagram artifact, stop.
327 if state.get("diagram_png_base64"):
328 return END
329 return "bump_tool_loop"
330
331def preprocess(state: PharmAIState) -> Dict[str, Any]:
332 q = (state.get("user_query") or "").strip().lower()
333
334 if any(k in q for k in ["diagram", "flowchart", "architecture", "graphviz", "dot", "draw"]):
335 return {"intent": "diagram"}
336
337 if re.match(r"^(what is|define|explain)\b", q) and len(q) < 120:
338 return {"intent": "simple"}
339
340 return {"intent": "diligence"}
341
342def route_after_llm(state: PharmAIState):
343 # If query is simple, never call tools/synthesize
344 if state.get("intent") == "simple":
345 return "end_simple"
346
347 # If the model asked for tools, go tools
348 last = state["messages"][-1]
349 if getattr(last, "tool_calls", None):
350 return "tools"
351
352 return "synthesize"
353
354def end_simple(state: PharmAIState) -> Dict[str, Any]:
355 # Return the last assistant content as the final answer
356 last = state["messages"][-1]
357 text = getattr(last, "content", "") if isinstance(getattr(last, "content", ""), str) else str(getattr(last, "content", ""))
358 return {"decision_brief": text, "citations": []}
359
360
361# -----------------------------
362# Final Synthesis Node
363# -----------------------------
364def synthesize(state: PharmAIState) -> Dict[str, Any]:
365 # Fast guardrail: placeholders -> short response without tool burn
366 uq = state.get("user_query", "")
367 if _looks_like_placeholder(uq):
368 brief = (
369 "# FINAL DECISION BRIEF\n\n"
370 "I need the **actual drug name** and **specific indication** to perform diligence.\n\n"
371 "## Evidence Gaps\n"
372 "- Drug name (e.g., semaglutide)\n"
373 "- Indication (e.g., obesity)\n"
374 "- Trial/program context (if any)\n"
375 )
376 return {
377 "decision_brief": _append_citations_section(brief, []),
378 "citations": [],
379 "messages": [HumanMessage(content="(placeholder query detected; returned guardrail response)")],
380 }
381
382 llm = _build_model()
383 messages: List[BaseMessage] = state["messages"]
384 messages = messages + [HumanMessage(content=FINAL_PROMPT)]
385
386 resp = llm.invoke(messages)
387
388 tool_citations = _extract_citations_from_messages(state["messages"])
389 brief_text = resp.content if isinstance(resp.content, str) else str(resp.content)
390 brief_text = _append_citations_section(brief_text, tool_citations)
391
392 return {
393 "decision_brief": brief_text,
394 "citations": tool_citations,
395 "messages": [resp],
396 }
397
398
399# -----------------------------
400# Build + Compile Graph
401# -----------------------------
402def build_graph():
403 """
404 Graph with preprocessing and smart routing.
405 """
406 g = StateGraph(PharmAIState)
407
408 g.add_node("preprocess", preprocess)
409 g.add_node("llm_call", llm_call)
410 g.add_node("tools", ToolNode(TOOLS))
411 g.add_node("capture_diagram", capture_diagram)
412 g.add_node("bump_tool_loop", lambda s: {"tool_loops": s.get("tool_loops", 0) + 1})
413 g.add_node("synthesize", synthesize)
414 g.add_node("end_simple", end_simple)
415
416 g.add_edge(START, "preprocess")
417 g.add_edge("preprocess", "llm_call")
418
419 # After LLM: route based on intent and tool calls
420 g.add_conditional_edges(
421 "llm_call",
422 route_after_llm,
423 {
424 "tools": "tools",
425 "synthesize": "synthesize",
426 "end_simple": "end_simple",
427 },
428 )
429
430 # After tools: capture diagram data
431 g.add_edge("tools", "capture_diagram")
432
433 # After capture: check if we should stop (diagram complete) or continue
434 g.add_conditional_edges(
435 "capture_diagram",
436 route_after_tools,
437 {
438 END: END, # Stop if diagram is complete
439 "bump_tool_loop": "bump_tool_loop", # Continue otherwise
440 },
441 )
442
443 g.add_edge("bump_tool_loop", "llm_call")
444 g.add_edge("end_simple", END)
445 g.add_edge("synthesize", END)
446
447 return g.compile()
448
449# -----------------------------
450# Test execution
451# -----------------------------
452if __name__ == "__main__":
453 print("Building PharmAI Navigator graph...")
454 graph = build_graph()
455 print("Graph compiled successfully!")
456
457 # Test query designed to trigger generate_graph_dot tool
458 #test_query = "Assess semaglutide for obesity"
459 #test_query = "Assess donanemab for early Alzheimer’s disease. Retrieve key clinical trials, summarize efficacy and safety outcomes, normalize the evidence, and generate a system architecture graph showing how PharmAI Navigator evaluates this asset."
460 #test_query = "Create a DOT graph showing the relationship between Drug, Indication, Clinical Trials, FDA Approval, and Market Launch and render it as png"
461 test_query = "What is pembrolizumab?"
462 print(f"\nRunning test query: {test_query}")
463
464 result = graph.invoke({
465 "messages": [HumanMessage(content=test_query)],
466 "user_query": test_query,
467 "tool_loops": 0,
468 })
469
470 print("\n" + "=" * 60)
471 print("OUTPUT:")
472 print("=" * 60)
473 print(result.get("decision_brief", "No output"))
474
475 print("\n" + "=" * 60)
476 print("CITATIONS (tool-only):")
477 print("=" * 60)
478 for i, citation in enumerate(result.get("citations", []), 1):
479 print(f"{i}. {citation}")
480