Added main app files

.dockerignore

@@ -0,0 +1,33 @@
+# Python
+__pycache__
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+venv/
+ENV/
+
+# Testing
+.pytest_cache
+*.log
+
+# IDE
+.vscode
+.idea
+*.swp
+*.swo
+
+# Test files
+test_*.py
+test_*.ps1
+*_test.py
+fix_indexes.py
+
+# Documentation
+*.md
+MEMORY_COMPARISON.md
+
+# Output files
+*.png
+output.txt

Dockerfile

@@ -0,0 +1,30 @@
+# Use Python 3.11 slim image
+FROM python:3.11-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies for graphviz (optional but recommended)
+RUN apt-get update && apt-get install -y \
+    graphviz \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Expose port 7860 (HuggingFace Spaces default)
+EXPOSE 7860
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1
+ENV PORT=7860
+
+# Run the FastAPI application
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+import uuid
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from schemas import AgentRunRequest, AgentRunResponse, Message
+from memory_mongo import memory_store  # MongoDB-backed memory
+from graph import build_graph
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from fastapi.responses import StreamingResponse
+import json
+import time
+from fastapi.encoders import jsonable_encoder
+
+app = FastAPI(title="PharmAI Navigator (Agentic)", version="0.1.0")
+
+# CORS (HF Spaces + your Node proxy)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Compile graph once at startup
+GRAPH = build_graph()
+
+
+@app.get("/health")
+def health():
+    """Health check with MongoDB status."""
+    mongo_status = "connected"
+    session_count = 0
+    
+    try:
+        session_count = memory_store.get_session_count()
+    except Exception as e:
+        mongo_status = f"error: {str(e)}"
+    
+    return {
+        "status": "ok",
+        "mongodb": mongo_status,
+        "active_sessions": session_count
+    }
+
+
+@app.get("/session/{session_id}/history")
+def get_session_history(session_id: str):
+    """Get chat history for a session (for testing)."""
+    messages = memory_store.get(session_id)
+    return {
+        "session_id": session_id,
+        "message_count": len(messages),
+        "messages": [{"role": m.role, "content": m.content[:100] + "..." if len(m.content) > 100 else m.content} for m in messages]
+    }
+
+
+@app.delete("/session/{session_id}")
+def clear_session(session_id: str):
+    """Clear a session's history (for testing)."""
+    memory_store.clear(session_id)
+    return {"session_id": session_id, "status": "cleared"}
+
+
+@app.post("/admin/cleanup-sessions")
+def cleanup_old_sessions(days: int = 7):
+    """
+    Admin endpoint to manually cleanup old sessions.
+    (TTL index handles this automatically if configured)
+    """
+    try:
+        deleted = memory_store.cleanup_old_sessions(days=days)
+        return {
+            "status": "ok",
+            "deleted_sessions": deleted,
+            "days": days
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/test/echo")
+def test_echo(req: AgentRunRequest):
+    """
+    Lightweight test endpoint - no LLM calls, just tests memory.
+    Echoes back the query and shows session history.
+    """
+    session_id = req.session_id or str(uuid.uuid4())
+    
+    # Get prior history
+    prior = memory_store.get(session_id)
+    
+    # Append user message
+    memory_store.append(session_id, role="user", content=req.query)
+    
+    # Create fake response
+    fake_response = f"Echo: {req.query} (Session has {len(prior)} prior messages)"
+    
+    # Append assistant message
+    memory_store.append(session_id, role="assistant", content=fake_response)
+    
+    return {
+        "session_id": session_id,
+        "decision_brief": fake_response,
+        "prior_message_count": len(prior),
+        "current_message_count": len(memory_store.get(session_id)),
+        "citations": [],
+        "metadata": {"test_mode": True}
+    }
+
+
+@app.post("/run", response_model=AgentRunResponse)
+def run_agent(req: AgentRunRequest):
+    # 1) session handling
+    session_id = req.session_id or str(uuid.uuid4())
+
+    # 2) load prior history (for chat continuity)
+    prior = memory_store.get(session_id)
+
+    # Convert to LangChain message dict format for LangGraph MessagesState
+    # LangGraph expects state["messages"] as list of LC messages; we pass dict-like messages.
+    messages = []
+    for m in prior:
+        if m.role == "user":
+            messages.append(HumanMessage(content=m.content))
+        elif m.role == "assistant":
+            messages.append(AIMessage(content=m.content))
+        elif m.role == "system":
+            messages.append(SystemMessage(content=m.content))
+
+    # 3) append this user query to memory (pre-run)
+    memory_store.append(session_id, role="user", content=req.query)
+
+    # Append new user query as LangChain message
+    messages = messages + [HumanMessage(content=req.query)]
+
+    # 4) run graph (Mode A synchronous)
+    try:
+        final_state = GRAPH.invoke(
+            {
+                "session_id": session_id,
+                "user_query": req.query,
+                "messages": messages,
+            }
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Agent run failed: {str(e)}")
+
+    decision_brief = final_state.get("decision_brief") or final_state.get("final_decision") or ""
+    citations = final_state.get("citations") or []
+
+    # 5) save assistant response to memory (post-run)
+    if decision_brief.strip():
+        memory_store.append(session_id, role="assistant", content=decision_brief)
+
+    return AgentRunResponse(
+        session_id=session_id,
+        decision_brief=decision_brief,
+        confidence_score=final_state.get("confidence_score"),
+        citations=citations,
+        metadata={
+            "has_prior_messages": len(prior) > 0,
+        },
+    )

graph.py

@@ -0,0 +1,479 @@
+from __future__ import annotations
+
+import os
+import json
+import re
+from typing import Any, Dict, List, Optional
+from dotenv import load_dotenv
+
+from langchain_anthropic import ChatAnthropic
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.tools import tool
+
+from langgraph.graph import StateGraph, START, END, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
+
+from tools import (
+    tavily_search,
+    stub_evidence,
+    classify_query,
+    extract_entities,
+    normalize_evidence,
+    generate_graph_dot,
+    clinicaltrials_search,
+    render_dot_to_png_base64
+)
+
+# Load environment variables
+load_dotenv()
+
+# -----------------------------
+# LangChain Tool Wrappers
+# -----------------------------
+@tool("web_search")
+def web_search_tool(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+    """Web search using Tavily. Returns a list of evidence dicts."""
+    ev = tavily_search(query=query, max_results=max_results)
+    return [e.model_dump() for e in ev]
+
+
+@tool("stub_evidence")
+def stub_evidence_tool(query: str) -> List[Dict[str, Any]]:
+    """Deterministic fallback evidence tool (offline/demo)."""
+    ev = stub_evidence(query=query)
+    return [e.model_dump() for e in ev]
+
+@tool("classify_query")
+def classify_query_tool(query: str) -> Dict[str, Any]:
+    """Classify query to decide which tools are needed."""
+    return classify_query(query)
+
+
+@tool("extract_entities")
+def extract_entities_tool(query: str) -> Dict[str, Optional[str]]:
+    """Extract drug and indication from query."""
+    return extract_entities(query)
+
+
+@tool("normalize_evidence")
+def normalize_evidence_tool(evidence: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Dedupe and clean evidence."""
+    return normalize_evidence(evidence)
+
+
+@tool("generate_graph_dot")
+def generate_graph_dot_tool(
+    title: str,
+    nodes: List[Dict[str, str]],
+    edges: List[Dict[str, str]],
+    rankdir: str = "LR",
+) -> str:
+    """
+    Generate Graphviz DOT.
+    IMPORTANT: Use this tool instead of writing DOT directly.
+    """
+    return generate_graph_dot(
+        title=title,
+        nodes=nodes,
+        edges=edges,
+        rankdir=rankdir,
+    )
+
+@tool("clinicaltrials_search")
+def clinicaltrials_search_tool(drug: str, indication: str, max_results: int = 5) -> List[Dict[str, Any]]:
+    """Search ClinicalTrials.gov (Tavily-based MVP)."""
+    ev = clinicaltrials_search(drug=drug, indication=indication, max_results=max_results)
+    return [e.model_dump() for e in ev]
+
+@tool("render_dot_to_png_base64")
+def render_dot_to_png_base64_tool(dot: str) -> Dict[str, Any]:
+    """Render DOT to PNG (base64). Optional dependency on graphviz."""
+    return render_dot_to_png_base64(dot)
+
+TOOLS = [
+    web_search_tool,
+    stub_evidence_tool,
+    classify_query_tool,
+    extract_entities_tool,
+    normalize_evidence_tool,
+    generate_graph_dot_tool,
+    clinicaltrials_search_tool,
+    render_dot_to_png_base64_tool
+]
+
+# -----------------------------
+# LangGraph State
+# -----------------------------
+class PharmAIState(MessagesState):
+    session_id: Optional[str]
+    user_query: str
+    decision_brief: str
+    citations: List[str]
+    confidence_score: float
+    tool_loops: int                    # safety counter
+    diagram_png_base64: Optional[str]  # <-- add
+    diagram_dot: Optional[str]         # <-- optional
+    intent: str  # "simple" | "diligence" | "diagram"
+
+# -----------------------------
+# Guardrails + Prompts
+# -----------------------------
+SYSTEM_PROMPT = """You are PharmAI Navigator, an evidence-grounded diligence assistant for drug/asset evaluation.
+
+Your job:
+Turn a query like "Assess {Drug} for {Indication}" into a decision-grade brief OR structured output.
+
+CRITICAL TOOL USAGE RULES:
+- If the user asks for a diagram, flow, architecture, graph, visualization, or Graphviz:
+  → You MUST call `generate_graph_dot`.
+  → You MUST NOT write Graphviz DOT directly in your response.
+  → If the user asks for an image/PNG, call `render_dot_to_png_base64` AFTER you get DOT.
+- If the user asks for trials / phases / NCT IDs / endpoints:
+  → Prefer calling `extract_entities` then `clinicaltrials_search`.
+- If the user asks for factual claims (approvals, safety, pricing, patents, market):
+  → Prefer calling `web_search`.
+
+Guardrails (STRICT):
+- Do NOT invent specific facts (approval dates, trial names, endpoints, statistics, patent expiry).
+- Any concrete number/date/claim MUST be supported by tool evidence.
+- If evidence is insufficient, clearly list Evidence Gaps.
+- Be concise, structured, and decision-oriented.
+- Avoid medical advice; present as diligence/analysis.
+
+Simple Query Rule (CRITICAL):
+- 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.
+- Only use tools when you need current/specific data (trials, approvals, patents, market data).
+
+Citations policy:
+- The final response's "Citations" section is handled by the system.
+- Do NOT create your own citation list.
+"""
+
+FINAL_PROMPT = """Write the FINAL decision brief with these sections:
+
+1) Executive Recommendation (1–2 lines)
+2) Scientific Rationale (bullets)
+3) Clinical Evidence Snapshot (bullets)
+4) IP / Exclusivity Quick View (bullets)
+5) Market / SoC Snapshot (bullets)
+6) Key Risks + Next Actions (bullets)
+
+Rules:
+- If evidence is insufficient, include "Evidence Gaps" with bullets.
+- Do NOT add a citations section yourself; the system will append it.
+Return plain text only.
+"""
+
+# Placeholder detection to avoid wasting tokens on "Drug X / Indication Y"
+PLACEHOLDER_PATTERNS = [
+    r"\bdrug\s*x\b",
+    r"\bindication\s*y\b",
+    r"\bdrug\s*name\b",
+    r"\bindication\s*name\b",
+]
+def _looks_like_placeholder(q: str) -> bool:
+    ql = (q or "").strip().lower()
+    return any(re.search(p, ql) for p in PLACEHOLDER_PATTERNS)
+
+
+def _build_model() -> ChatAnthropic:
+    model_name = os.getenv("ANTHROPIC_MODEL", "claude-3-7-sonnet-latest")
+    return ChatAnthropic(
+        model=model_name,
+        temperature=0.2,
+        max_tokens=10000,
+        timeout=120,
+        streaming=False,
+        stop=None
+    ).bind_tools(TOOLS)
+
+
+# Safety cap to avoid endless tool loops
+MAX_TOOL_LOOPS = int(os.getenv("MAX_TOOL_LOOPS", "4"))
+
+
+def llm_call(state: PharmAIState) -> Dict[str, Any]:
+    """
+    Calls Claude with tool schemas attached.
+    Returns new messages to append into state["messages"].
+    """
+    llm = _build_model()
+    messages: List[BaseMessage] = state["messages"]
+
+    if not messages or not isinstance(messages[0], SystemMessage):
+        messages = [SystemMessage(content=SYSTEM_PROMPT)] + messages
+
+    tool_loops = state.get("tool_loops", 0)
+    if tool_loops >= MAX_TOOL_LOOPS:
+        # Stop tool-calling loop and force synthesis
+        stop_msg = HumanMessage(
+            content=(
+                "Stop calling tools now. Proceed to final synthesis using what you already have. "
+                "If evidence is insufficient, clearly list Evidence Gaps."
+            )
+        )
+        messages = messages + [stop_msg]
+
+    resp = llm.invoke(messages)
+    return {"messages": [resp]}
+
+
+# -----------------------------
+# Citations extraction (tool-only)
+# -----------------------------
+def _clean_url(u: str) -> str:
+    return u.strip().strip("),.]}\"'")
+
+def _extract_citations_from_messages(messages: List[BaseMessage]) -> List[str]:
+    """
+    Tool-only citation extraction (single source of truth):
+    - ONLY reads ToolMessage contents (actual tool outputs).
+    - If tool output is JSON (list/dict), pull `source` fields.
+    - Fallback: regex URL extraction from tool text.
+    """
+    citations: List[str] = []
+    url_re = re.compile(r"https?://[^\s\]\)\}\",']+")
+
+    for m in messages:
+        if not isinstance(m, ToolMessage):
+            continue
+
+        content = getattr(m, "content", None)
+        if not content:
+            continue
+
+        if isinstance(content, str):
+            parsed = None
+            try:
+                parsed = json.loads(content)
+            except Exception:
+                parsed = None
+
+            if isinstance(parsed, list):
+                for item in parsed:
+                    if isinstance(item, dict):
+                        src = item.get("source")
+                        if isinstance(src, str) and src.startswith(("http://", "https://")):
+                            citations.append(_clean_url(src))
+            elif isinstance(parsed, dict):
+                src = parsed.get("source")
+                if isinstance(src, str) and src.startswith(("http://", "https://")):
+                    citations.append(_clean_url(src))
+
+            for u in url_re.findall(content):
+                citations.append(_clean_url(u))
+
+    # De-duplicate
+    seen = set()
+    out = []
+    for c in citations:
+        # drop clearly broken/truncated URLs
+        if len(c) < 12:
+            continue
+        if c not in seen:
+            seen.add(c)
+            out.append(c)
+    return out
+
+
+def _append_citations_section(brief_text: str, citations: List[str]) -> str:
+    """
+    Enforces "single source of truth":
+    - Removes any existing 'Citations' section the model may have produced
+    - Appends citations derived from tool outputs only
+    """
+    text = (brief_text or "").strip()
+
+    # Remove any model-generated citations section (best-effort)
+    # (handles '## Citations' or 'Citations' headers)
+    text = re.split(r"\n#{1,3}\s*Citations\s*\n|\nCitations\s*\n", text, maxsplit=1)[0].rstrip()
+
+    if citations:
+        lines = ["", "## Citations"]
+        for i, c in enumerate(citations, 1):
+            lines.append(f"{i}. {c}")
+        text = text + "\n" + "\n".join(lines)
+    else:
+        text = text + "\n\n## Citations\n- (No external sources retrieved.)"
+
+    return text
+
+def capture_diagram(state: PharmAIState) -> Dict[str, Any]:
+    # Find the last ToolMessage (most recent tool output)
+    last_tool = None
+    for m in reversed(state["messages"]):
+        if isinstance(m, ToolMessage):
+            last_tool = m
+            break
+
+    if not last_tool:
+        return {}
+
+    tool_name = getattr(last_tool, "name", "") or ""
+    content = getattr(last_tool, "content", "")
+
+    # If your render tool returns base64 string directly
+    if tool_name == "render_dot_to_png_base64":
+        return {"diagram_png_base64": content}
+
+    # If your generate_graph_dot returns dot string
+    if tool_name == "generate_graph_dot":
+        return {"diagram_dot": content}
+
+    return {}
+
+def route_after_tools(state: PharmAIState) -> str:
+    # If we already have the final diagram artifact, stop.
+    if state.get("diagram_png_base64"):
+        return END
+    return "bump_tool_loop"
+
+def preprocess(state: PharmAIState) -> Dict[str, Any]:
+    q = (state.get("user_query") or "").strip().lower()
+
+    if any(k in q for k in ["diagram", "flowchart", "architecture", "graphviz", "dot", "draw"]):
+        return {"intent": "diagram"}
+
+    if re.match(r"^(what is|define|explain)\b", q) and len(q) < 120:
+        return {"intent": "simple"}
+
+    return {"intent": "diligence"}
+
+def route_after_llm(state: PharmAIState):
+    # If query is simple, never call tools/synthesize
+    if state.get("intent") == "simple":
+        return "end_simple"
+
+    # If the model asked for tools, go tools
+    last = state["messages"][-1]
+    if getattr(last, "tool_calls", None):
+        return "tools"
+
+    return "synthesize"
+
+def end_simple(state: PharmAIState) -> Dict[str, Any]:
+    # Return the last assistant content as the final answer
+    last = state["messages"][-1]
+    text = getattr(last, "content", "") if isinstance(getattr(last, "content", ""), str) else str(getattr(last, "content", ""))
+    return {"decision_brief": text, "citations": []}
+
+
+# -----------------------------
+# Final Synthesis Node
+# -----------------------------
+def synthesize(state: PharmAIState) -> Dict[str, Any]:
+    # Fast guardrail: placeholders -> short response without tool burn
+    uq = state.get("user_query", "")
+    if _looks_like_placeholder(uq):
+        brief = (
+            "# FINAL DECISION BRIEF\n\n"
+            "I need the **actual drug name** and **specific indication** to perform diligence.\n\n"
+            "## Evidence Gaps\n"
+            "- Drug name (e.g., semaglutide)\n"
+            "- Indication (e.g., obesity)\n"
+            "- Trial/program context (if any)\n"
+        )
+        return {
+            "decision_brief": _append_citations_section(brief, []),
+            "citations": [],
+            "messages": [HumanMessage(content="(placeholder query detected; returned guardrail response)")],
+        }
+
+    llm = _build_model()
+    messages: List[BaseMessage] = state["messages"]
+    messages = messages + [HumanMessage(content=FINAL_PROMPT)]
+
+    resp = llm.invoke(messages)
+
+    tool_citations = _extract_citations_from_messages(state["messages"])
+    brief_text = resp.content if isinstance(resp.content, str) else str(resp.content)
+    brief_text = _append_citations_section(brief_text, tool_citations)
+
+    return {
+        "decision_brief": brief_text,
+        "citations": tool_citations,
+        "messages": [resp],
+    }
+
+
+# -----------------------------
+# Build + Compile Graph
+# -----------------------------
+def build_graph():
+    """
+    Graph with preprocessing and smart routing.
+    """
+    g = StateGraph(PharmAIState)
+    
+    g.add_node("preprocess", preprocess)
+    g.add_node("llm_call", llm_call)
+    g.add_node("tools", ToolNode(TOOLS))
+    g.add_node("capture_diagram", capture_diagram)
+    g.add_node("bump_tool_loop", lambda s: {"tool_loops": s.get("tool_loops", 0) + 1})
+    g.add_node("synthesize", synthesize)
+    g.add_node("end_simple", end_simple)
+
+    g.add_edge(START, "preprocess")
+    g.add_edge("preprocess", "llm_call")
+    
+    # After LLM: route based on intent and tool calls
+    g.add_conditional_edges(
+        "llm_call",
+        route_after_llm,
+        {
+            "tools": "tools",
+            "synthesize": "synthesize",
+            "end_simple": "end_simple",
+        },
+    )
+
+    # After tools: capture diagram data
+    g.add_edge("tools", "capture_diagram")
+    
+    # After capture: check if we should stop (diagram complete) or continue
+    g.add_conditional_edges(
+        "capture_diagram",
+        route_after_tools,
+        {
+            END: END,  # Stop if diagram is complete
+            "bump_tool_loop": "bump_tool_loop",  # Continue otherwise
+        },
+    )
+    
+    g.add_edge("bump_tool_loop", "llm_call")
+    g.add_edge("end_simple", END)
+    g.add_edge("synthesize", END)
+    
+    return g.compile()
+
+# -----------------------------
+# Test execution
+# -----------------------------
+if __name__ == "__main__":
+    print("Building PharmAI Navigator graph...")
+    graph = build_graph()
+    print("Graph compiled successfully!")
+
+    # Test query designed to trigger generate_graph_dot tool
+    #test_query = "Assess semaglutide for obesity"
+    #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."
+    #test_query = "Create a DOT graph showing the relationship between Drug, Indication, Clinical Trials, FDA Approval, and Market Launch and render it as png"
+    test_query = "What is pembrolizumab?"
+    print(f"\nRunning test query: {test_query}")
+
+    result = graph.invoke({
+        "messages": [HumanMessage(content=test_query)],
+        "user_query": test_query,
+        "tool_loops": 0,
+    })
+
+    print("\n" + "=" * 60)
+    print("OUTPUT:")
+    print("=" * 60)
+    print(result.get("decision_brief", "No output"))
+
+    print("\n" + "=" * 60)
+    print("CITATIONS (tool-only):")
+    print("=" * 60)
+    for i, citation in enumerate(result.get("citations", []), 1):
+        print(f"{i}. {citation}")

memory.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+import time
+import threading
+from schemas import Message
+
+@dataclass
+class SessionMemory:
+    """In-memory chat history for a single session."""
+    messages: List[Message]
+    updated_at: float
+
+
+class MemoryStore:
+    """
+    Simple thread-safe in-memory store.
+
+    - session_id -> list[Message]
+    - trims to keep memory bounded
+    - includes basic TTL cleanup hooks (optional)
+    """
+
+    def __init__(self, max_messages: int = 30, ttl_seconds: Optional[int] = None):
+        self.max_messages = max_messages
+        self.ttl_seconds = ttl_seconds
+        self._lock = threading.Lock()
+        self._store: Dict[str, SessionMemory] = {}
+
+    def _now(self) -> float:
+        return time.time()
+
+    def get(self, session_id: str) -> List[Message]:
+        """Get messages for a session (returns empty list if new session)."""
+        if not session_id:
+            return []
+        with self._lock:
+            self._gc_locked()
+            if session_id not in self._store:
+                self._store[session_id] = SessionMemory(messages=[], updated_at=self._now())
+            return list(self._store[session_id].messages)
+
+    def append(self, session_id: str, role: str, content: str) -> None:
+        """Append a message and enforce trimming."""
+        if not session_id:
+            return
+        with self._lock:
+            self._gc_locked()
+            if session_id not in self._store:
+                self._store[session_id] = SessionMemory(messages=[], updated_at=self._now())
+
+            self._store[session_id].messages.append(Message(role=role, content=content))
+            self._store[session_id].updated_at = self._now()
+
+            # Trim oldest messages (keep most recent)
+            if len(self._store[session_id].messages) > self.max_messages:
+                overflow = len(self._store[session_id].messages) - self.max_messages
+                self._store[session_id].messages = self._store[session_id].messages[overflow:]
+
+    def set_messages(self, session_id: str, messages: List[Message]) -> None:
+        """Replace session history entirely (rarely needed, but handy)."""
+        if not session_id:
+            return
+        with self._lock:
+            self._store[session_id] = SessionMemory(
+                messages=messages[-self.max_messages :],
+                updated_at=self._now(),
+            )
+
+    def clear(self, session_id: str) -> None:
+        """Clear a single session."""
+        if not session_id:
+            return
+        with self._lock:
+            self._store.pop(session_id, None)
+
+    def _gc_locked(self) -> None:
+        """TTL cleanup (only runs if ttl_seconds is configured)."""
+        if not self.ttl_seconds:
+            return
+        cutoff = self._now() - self.ttl_seconds
+        expired = [sid for sid, mem in self._store.items() if mem.updated_at < cutoff]
+        for sid in expired:
+            self._store.pop(sid, None)
+
+
+# Global singleton (simple for HF Spaces demo)
+memory_store = MemoryStore(
+    max_messages=int(__import__("os").getenv("MAX_SESSION_MESSAGES", "30")),
+    ttl_seconds=int(__import__("os").getenv("SESSION_TTL_SECONDS", "0")) or None,
+)

memory_mongo.py

@@ -0,0 +1,260 @@
+"""
+MongoDB-backed session memory store.
+Replaces in-memory storage with persistent MongoDB storage.
+"""
+from __future__ import annotations
+from typing import List, Optional
+import os
+from datetime import datetime, timedelta
+from pymongo import MongoClient, ASCENDING
+from pymongo.errors import ConnectionFailure, OperationFailure
+from schemas import Message
+from dotenv import load_dotenv
+
+#load env vars
+load_dotenv()
+
+class MongoMemoryStore:
+    """
+    MongoDB-backed session memory store.
+    
+    Schema:
+    {
+        "_id": "session_id",
+        "messages": [
+            {"role": "user", "content": "..."},
+            {"role": "assistant", "content": "..."}
+        ],
+        "updated_at": datetime,
+        "created_at": datetime
+    }
+    """
+
+    def __init__(
+        self,
+        mongo_uri: Optional[str] = None,
+        database_name: str = "pharmai",
+        collection_name: str = "sessions",
+        max_messages: int = 30,
+        ttl_seconds: Optional[int] = None,
+    ):
+        self.max_messages = max_messages
+        self.ttl_seconds = ttl_seconds
+        
+        # Get MongoDB URI from env or parameter
+        self.mongo_uri = mongo_uri or os.getenv("MONGO_URI")
+        if not self.mongo_uri:
+            raise ValueError("MONGO_URI not found in environment variables")
+        
+        # Connect to MongoDB
+        try:
+            self.client = MongoClient(self.mongo_uri, serverSelectionTimeoutMS=5000)
+            # Test connection
+            self.client.admin.command('ping')
+            print(f"✅ MongoDB connected: {database_name}.{collection_name}")
+        except ConnectionFailure as e:
+            raise ConnectionError(f"Failed to connect to MongoDB: {e}")
+        
+        self.db = self.client[database_name]
+        self.collection = self.db[collection_name]
+        
+        # Create indexes
+        self._create_indexes()
+
+    def _create_indexes(self):
+        """Create indexes for performance and TTL."""
+        try:
+            # Get existing indexes
+            existing_indexes = self.collection.index_information()
+            
+            # TTL index - automatically delete old sessions
+            if self.ttl_seconds:
+                # Check if TTL index exists
+                ttl_exists = any(
+                    idx.get("expireAfterSeconds") is not None
+                    for idx in existing_indexes.values()
+                )
+                
+                if not ttl_exists:
+                    # Drop the basic updated_at index if it exists (without TTL)
+                    if "updated_at_1" in existing_indexes:
+                        self.collection.drop_index("updated_at_1")
+                    
+                    # Create TTL index
+                    self.collection.create_index(
+                        [("updated_at", ASCENDING)],
+                        expireAfterSeconds=self.ttl_seconds,
+                        name="session_ttl"
+                    )
+                    print(f"✅ Created TTL index (expires after {self.ttl_seconds}s)")
+            else:
+                # Just a regular index on updated_at (no TTL)
+                if "updated_at_1" not in existing_indexes and "session_ttl" not in existing_indexes:
+                    self.collection.create_index([("updated_at", ASCENDING)])
+                    print("✅ Created updated_at index")
+                    
+        except OperationFailure as e:
+            # Index creation failed, but continue anyway
+            print(f"⚠️  Index creation warning: {e}")
+            pass
+
+    def get(self, session_id: str) -> List[Message]:
+        """Get messages for a session."""
+        if not session_id:
+            return []
+        
+        try:
+            doc = self.collection.find_one({"_id": session_id})
+            if not doc:
+                return []
+            
+            # Convert dict messages to Message objects
+            messages = []
+            for msg in doc.get("messages", []):
+                messages.append(Message(
+                    role=msg.get("role", "user"),
+                    content=msg.get("content", "")
+                ))
+            
+            return messages
+        except OperationFailure as e:
+            print(f"Error getting session {session_id}: {e}")
+            return []
+
+    def append(self, session_id: str, role: str, content: str) -> None:
+        """Append a message to a session."""
+        if not session_id:
+            return
+        
+        now = datetime.utcnow()
+        message = {"role": role, "content": content}
+        
+        try:
+            # Try to update existing session
+            result = self.collection.update_one(
+                {"_id": session_id},
+                {
+                    "$push": {"messages": message},
+                    "$set": {"updated_at": now}
+                }
+            )
+            
+            # If session doesn't exist, create it
+            if result.matched_count == 0:
+                self.collection.insert_one({
+                    "_id": session_id,
+                    "messages": [message],
+                    "created_at": now,
+                    "updated_at": now
+                })
+            
+            # Trim old messages if needed
+            self._trim_messages(session_id)
+            
+        except OperationFailure as e:
+            print(f"Error appending to session {session_id}: {e}")
+
+    def _trim_messages(self, session_id: str) -> None:
+        """Keep only the most recent max_messages."""
+        try:
+            doc = self.collection.find_one({"_id": session_id})
+            if not doc:
+                return
+            
+            messages = doc.get("messages", [])
+            if len(messages) > self.max_messages:
+                # Keep only the most recent messages
+                trimmed = messages[-self.max_messages:]
+                self.collection.update_one(
+                    {"_id": session_id},
+                    {"$set": {"messages": trimmed}}
+                )
+        except OperationFailure as e:
+            print(f"Error trimming session {session_id}: {e}")
+
+    def set_messages(self, session_id: str, messages: List[Message]) -> None:
+        """Replace session history entirely."""
+        if not session_id:
+            return
+        
+        now = datetime.utcnow()
+        message_dicts = [{"role": m.role, "content": m.content} for m in messages]
+        
+        # Keep only most recent messages
+        if len(message_dicts) > self.max_messages:
+            message_dicts = message_dicts[-self.max_messages:]
+        
+        try:
+            self.collection.update_one(
+                {"_id": session_id},
+                {
+                    "$set": {
+                        "messages": message_dicts,
+                        "updated_at": now
+                    },
+                    "$setOnInsert": {"created_at": now}
+                },
+                upsert=True
+            )
+        except OperationFailure as e:
+            print(f"Error setting messages for session {session_id}: {e}")
+
+    def clear(self, session_id: str) -> None:
+        """Clear a single session."""
+        if not session_id:
+            return
+        
+        try:
+            self.collection.delete_one({"_id": session_id})
+        except OperationFailure as e:
+            print(f"Error clearing session {session_id}: {e}")
+
+    def cleanup_old_sessions(self, days: int = 7) -> int:
+        """
+        Manually cleanup sessions older than X days.
+        (TTL index handles this automatically if configured)
+        """
+        cutoff = datetime.utcnow() - timedelta(days=days)
+        try:
+            result = self.collection.delete_many({"updated_at": {"$lt": cutoff}})
+            return result.deleted_count
+        except OperationFailure as e:
+            print(f"Error cleaning up old sessions: {e}")
+            return 0
+
+    def get_session_count(self) -> int:
+        """Get total number of active sessions."""
+        try:
+            return self.collection.count_documents({})
+        except OperationFailure:
+            return 0
+
+    def close(self):
+        """Close MongoDB connection."""
+        if self.client:
+            self.client.close()
+
+
+# Create global singleton
+def create_memory_store() -> MongoMemoryStore:
+    """Factory function to create memory store based on configuration."""
+    try:
+        # Try MongoDB first
+        return MongoMemoryStore(
+            max_messages=int(os.getenv("MAX_SESSION_MESSAGES", "30")),
+            ttl_seconds=int(os.getenv("SESSION_TTL_SECONDS", "0")) or None,
+        )
+    except (ValueError, ConnectionError) as e:
+        print(f"⚠️  MongoDB not available: {e}")
+        print("⚠️  Falling back to in-memory storage")
+        
+        # Fallback to in-memory
+        from memory import MemoryStore
+        return MemoryStore(
+            max_messages=int(os.getenv("MAX_SESSION_MESSAGES", "30")),
+            ttl_seconds=int(os.getenv("SESSION_TTL_SECONDS", "0")) or None,
+        )
+
+
+# Global instance
+memory_store = create_memory_store()

requirements.txt

@@ -0,0 +1,32 @@
+# FastAPI and server
+fastapi
+uvicorn
+python-dotenv
+
+# LangChain and LangGraph
+langchain
+langchain-anthropic
+langchain-core
+langchain-community
+langgraph
+langgraph-checkpoint
+langsmith
+
+# Tools and utilities
+tavily-python
+pydantic
+pydantic-settings
+
+# MongoDB
+pymongo
+motor
+
+# HTTP client
+httpx
+aiohttp
+
+# Graph rendering (optional)
+graphviz
+
+# Other dependencies
+python-multipart

schemas.py

@@ -0,0 +1,98 @@
+# schemas.py
+from typing import List, Optional, Dict, Any
+from pydantic import BaseModel, Field
+from enum import Enum
+
+# Core Enums
+class AgentType(str, Enum):
+    PLANNER = "planner"
+    SCIENTIFIC = "scientific"
+    PATENT = "patent"
+    MARKET = "market"
+    SUPPLY = "supply"
+    SYNTHESIS = "synthesis"
+
+class EvidenceType(str, Enum):
+    LITERATURE = "literature"
+    CLINICAL_TRIAL = "clinical_trial"
+    PATENT = "patent"
+    MARKET = "market"
+    OTHER = "other"
+
+# API Schemas (FastAPI I/O)
+class AgentRunRequest(BaseModel):
+    """
+    Incoming request from Node.js backend or direct API call.
+    """
+    session_id: Optional[str] = Field(
+        default=None,
+        description="Optional session ID to maintain conversation state"
+    )
+    query: str = Field(
+        ...,
+        description="User query, e.g. 'Drug X for Indication Y'"
+    )
+
+class AgentRunResponse(BaseModel):
+    """
+    Final response returned by the agent system.
+    """
+    session_id: Optional[str]
+    decision_brief: str
+    confidence_score: Optional[float] = Field(
+        default=None,
+        description="Optional overall confidence score (0–1)"
+    )
+    citations: Optional[List[str]] = Field(
+        default=None,
+        description="List of citation identifiers or URLs"
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Extra debug or trace metadata"
+    )
+
+# Internal Agent State
+class Message(BaseModel):
+    """
+    Canonical message format passed between agents.
+    """
+    role: str  # system | user | assistant | tool
+    content: str
+
+class EvidenceItem(BaseModel):
+    """
+    A single piece of evidence produced by tools or agents.
+    """
+    type: EvidenceType
+    source: str
+    summary: str
+    confidence: Optional[float] = None
+    raw: Optional[Dict[str, Any]] = None
+
+class AgentOutput(BaseModel):
+    """
+    Output produced by a single agent.
+    """
+    agent: AgentType
+    text: str
+    evidence: Optional[List[EvidenceItem]] = None
+
+class AgentState(BaseModel):
+    """
+    LangGraph state object.
+    This is what flows between graph nodes.
+    """
+    session_id: Optional[str]
+    user_query: str
+
+    messages: List[Message] = Field(default_factory=list)
+
+    agent_outputs: Dict[AgentType, AgentOutput] = Field(
+        default_factory=dict,
+        description="Outputs from each agent"
+    )
+
+    final_decision: Optional[str] = None
+
+    confidence_score: Optional[float] = None

tools.py

@@ -0,0 +1,288 @@
+from typing import List, Dict, Any, Optional
+import os
+import uuid
+import re
+import base64
+from schemas import EvidenceItem, EvidenceType
+
+#Helper Functions
+def _etype(name: str, default: EvidenceType) -> EvidenceType:
+    """Return EvidenceType.<name> if it exists, else default (prevents breaking)."""
+    return getattr(EvidenceType, name, default)
+
+def _short(s: str, n: int = 700) -> str:
+    return (s or "")[:n]
+
+def _is_url(s: str) -> bool:
+    return isinstance(s, str) and s.startswith(("http://", "https://"))
+
+# Tool 1: Tavily Web Search (existing, unchanged)
+def tavily_search(query: str, max_results: int = 5) -> List[EvidenceItem]:
+    """
+    Uses Tavily API to perform web search.
+    Returns structured evidence items.
+    """
+    api_key = os.getenv("TAVILY_API_KEY")
+
+    if not api_key:
+        return [
+            EvidenceItem(
+                type=EvidenceType.OTHER,
+                source="tavily_disabled",
+                summary="Tavily API key not configured; search skipped.",
+                confidence=0.0,
+            )
+        ]
+
+    try:
+        from tavily import TavilyClient
+
+        client = TavilyClient(api_key=api_key)
+        results = client.search(
+            query=query,
+            max_results=max_results,
+            include_raw_content=False,
+        )
+
+        evidence: List[EvidenceItem] = []
+
+        for r in results.get("results", []):
+            evidence.append(
+                EvidenceItem(
+                    type=EvidenceType.LITERATURE,
+                    source=r.get("url", "unknown"),
+                    summary=r.get("content", "")[:500],
+                    confidence=0.6,
+                    raw=r,
+                )
+            )
+
+        return evidence
+
+    except Exception as e:
+        return [
+            EvidenceItem(
+                type=EvidenceType.OTHER,
+                source="tavily_error",
+                summary=f"Tavily search failed: {str(e)}",
+                confidence=0.0,
+            )
+        ]
+
+# Tool 2: Stub Evidence Generator (existing, unchanged)
+def stub_evidence(query: str) -> List[EvidenceItem]:
+    """
+    Deterministic fallback tool.
+    Useful for demos, offline mode, or testing agent logic.
+    """
+    return [
+        EvidenceItem(
+            type=EvidenceType.OTHER,
+            source="stub_tool",
+            summary=f"Stub evidence generated for query: '{query}'. "
+                    f"This indicates where real retrieval will plug in.",
+            confidence=0.2,
+            raw={
+                "id": str(uuid.uuid4()),
+                "note": "Replace with real retrieval later",
+            },
+        )
+    ]
+
+# Tool 3: Query Classifier (planner helper)
+def classify_query(query: str) -> Dict[str, Any]:
+    """
+    Lightweight classifier to help the agent decide
+    which tools (if any) are required.
+    """
+    q = (query or "").lower()
+    needs_graph = any(k in q for k in ["diagram", "graph", "graphviz", "dot", "flow", "architecture", "arch", "draw"])
+    needs_trials = any(k in q for k in ["trial", "clinical", "phase", "nct", "primary endpoint", "secondary endpoint"])
+    needs_facts = any(k in q for k in ["fda", "approval", "label", "patent", "exclusivity", "pricing", "aria", "safety", "market"])
+    needs_entities = any(k in q for k in ["evaluate", "assess", "analyze", "repurpose", "for "])
+    return {
+        "needs_graph": needs_graph,
+        "needs_clinical_trials": needs_trials,
+        "needs_web_search": needs_facts or needs_trials,
+        "needs_entity_extraction": needs_entities,
+    }
+
+# Tool 4: Entity Extraction (Drug / Indication)
+def extract_entities(query: str) -> Dict[str, Optional[str]]:
+    """
+    Minimal entity extractor for MVP.
+    """
+    text = (query or "").strip()
+    m = re.search(
+        r"(evaluate|assess|analyze)\s+(?P<drug>.+?)\s+for\s+(?P<indication>.+)",
+        text,
+        re.IGNORECASE,
+    )
+    if m:
+        return {
+            "drug": m.group("drug").strip(),
+            "indication": m.group("indication").strip(),
+        }
+    return {"drug": None, "indication": None}
+
+# Tool 5: Evidence Normalizer (dedupe + cleanup)
+def normalize_evidence(evidence: List[EvidenceItem]) -> List[EvidenceItem]:
+    """
+    Deduplicates evidence by source and trims noisy content.
+    """
+    seen = set()
+    cleaned: List[EvidenceItem] = []
+
+    for e in evidence:
+        if e.source in seen:
+            continue
+        seen.add(e.source)
+
+        cleaned.append(
+            EvidenceItem(
+                type=e.type,
+                source=e.source,
+                summary=(e.summary or "")[:800],
+                confidence=e.confidence,
+                raw=None,  # drop heavy payloads
+            )
+        )
+
+    return cleaned
+
+# Tool 6: Graph Generation (Graphviz DOT only)
+def generate_graph_dot(
+    title: str,
+    nodes: List[Dict[str, str]],
+    edges: List[Dict[str, str]],
+    rankdir: str = "LR",
+) -> str:
+    """
+    Generates Graphviz DOT code.
+    IMPORTANT: LLM must call this tool; never output DOT directly.
+    """
+    safe_title = (title or "PharmAI Graph").replace('"', "'")
+
+    lines = [
+        "digraph G {",
+        f"  rankdir={rankdir};",
+        '  labelloc="t";',
+        '  labeljust="c";',
+        f'  label=<<B><FONT POINT-SIZE="28">{safe_title}</FONT></B>>;',
+        "  node [shape=box, style=rounded];",
+        "",
+    ]
+
+    for n in nodes or []:
+        nid = n.get("id")
+        lbl = (n.get("label") or nid).replace('"', "'")
+        if nid:
+            lines.append(f'  {nid} [label="{lbl}"];')
+
+    lines.append("")
+
+    for e in edges or []:
+        src = e.get("from")
+        tgt = e.get("to")
+        lbl = e.get("label")
+        if src and tgt:
+            if lbl:
+                lines.append(f'  {src} -> {tgt} [label="{lbl}"];')
+            else:
+                lines.append(f"  {src} -> {tgt};")
+
+    lines.append("}")
+    return "\n".join(lines)
+
+#Tool 7: ClinicalTrials search (lightweight, Tavily-based)
+def clinicaltrials_search(drug: str, indication: str, max_results: int = 5) -> List[EvidenceItem]:
+    """
+    MVP approach:
+    - Uses Tavily to target ClinicalTrials.gov / NCT IDs
+    - Returns EvidenceItems for trial links + snippets
+    """
+    drug = (drug or "").strip()
+    indication = (indication or "").strip()
+
+    if not drug or not indication:
+        return [
+            EvidenceItem(
+                type=EvidenceType.OTHER,
+                source="clinicaltrials_search_invalid_input",
+                summary="Missing drug or indication for clinical trials search.",
+                confidence=0.0,
+            )
+        ]
+
+    query = f'site:clinicaltrials.gov ("{drug}") ("{indication}") NCT'
+    ev = tavily_search(query=query, max_results=max_results)
+
+    trial_type = _etype("CLINICAL_TRIAL", EvidenceType.LITERATURE)
+
+    out: List[EvidenceItem] = []
+    for e in ev:
+        # only keep plausible CT.gov results if possible
+        if _is_url(e.source) and "clinicaltrials.gov" in e.source:
+            out.append(
+                EvidenceItem(
+                    type=trial_type,
+                    source=e.source,
+                    summary=e.summary,
+                    confidence=max(0.55, float(e.confidence or 0.55)),
+                    raw=e.raw,
+                )
+            )
+
+    if out:
+        return out
+
+    # fallback: return whatever Tavily gave (still structured)
+    return [
+        EvidenceItem(
+            type=trial_type,
+            source=e.source,
+            summary=e.summary,
+            confidence=float(e.confidence or 0.4),
+            raw=e.raw,
+        )
+        for e in ev
+    ]
+
+#Tool 8 : DOT -> PNG
+def render_dot_to_png_base64(dot: str) -> Dict[str, Any]:
+    """
+    Converts DOT to PNG and returns base64 string.
+    - Uses python 'graphviz' package if available.
+    - If graphviz isn't installed in the environment, returns an error payload.
+    """
+    dot = (dot or "").strip()
+    if not dot:
+        return {"ok": False, "error": "Empty DOT string"}
+
+    try:
+        from graphviz import Source  # optional dependency
+
+        src = Source(dot)
+        png_bytes = src.pipe(format="png")
+        b64 = base64.b64encode(png_bytes).decode("utf-8")
+        return {"ok": True, "png_base64": b64}
+    except Exception as e:
+        return {
+            "ok": False,
+            "error": f"DOT->PNG render failed. Ensure `graphviz` Python package and system binaries are installed. Details: {str(e)}",
+        }
+
+# Tool Registry (extended, backward compatible)
+TOOL_REGISTRY: Dict[str, Any] = {
+    # existing
+    "web_search": tavily_search,
+    "stub_evidence": stub_evidence,
+
+    # new
+    "classify_query": classify_query,
+    "extract_entities": extract_entities,
+    "normalize_evidence": normalize_evidence,
+    "generate_graph_dot": generate_graph_dot,
+    "clinicaltrials_search": clinicaltrials_search,
+    "render_dot_to_png_base64": render_dot_to_png_base64
+}