src/core/graph.py

import re
import time

from langgraph.graph import StateGraph, END
from src.core.state import AgentState
from src.agents.agent_instances import (
    role_classifier, patient_llm, caregiver_llm, validator, safety_check,
    intent_classifier, diagnosis_assist, treatment_assist,
    monitoring_assist, general_assist, output_merger, research_agent,
    dietary_assist
)
from langchain_core.messages import AIMessage, ToolMessage
from langgraph.prebuilt import ToolNode
from src.tools.web_tools import web_search_tool
from src.utils.logger import setup_logger
from src.tools.fhir_memory import save_chat_as_fhir

logger = setup_logger("MedicalPipeline")

MAX_RECOVERY_ATTEMPTS = 3
EMERGENCY_PATTERNS = [
    r"\b(unconscious|passed out|not breathing|seizure|seizing|stroke|heart attack|chest pain|coma)\b",
    r"\b(dka|ketoacidosis|hypoglycemic emergency|hyperglycemic emergency|insulin overdose)\b",
]

def log_step(name: str, output: str = None):
    """Utility to log both to terminal and return a state update for the logs list."""
    logger.info(f"Executing: {name}")
    log_msg = f"➔ Executing Node: {name}"
    if output:
        log_msg += f"\nOutput: {output}"
    return {"logs": [log_msg]}


def _extract_message_content(message) -> str:
    if getattr(message, "content", None):
        return message.content
    return str(getattr(message, "tool_calls", ""))


def _set_latest_clinician_output(res: dict):
    latest_output = ""
    if res.get("messages"):
        latest_output = _extract_message_content(res["messages"][-1])

    if latest_output:
        res["clinician_outputs"] = [latest_output]
    else:
        res["clinician_outputs"] = []
    return res
def detect_emergency(state: AgentState) -> bool:
    last_message = state.get("messages")[-1] if state.get("messages") else None
    if not last_message:
        return False

    content = getattr(last_message, "content", "")
    if not content:
        return False

    normalized = content.lower()
    return any(re.search(pattern, normalized) for pattern in EMERGENCY_PATTERNS)


async def role_classifier_node(state: AgentState):
    res = await role_classifier.run(state)
    log = log_step("Role Classifier", f"Detected Role: {res.get('user_role')}")
    res.update(log)

    # Intent is only needed for the clinician pathway.
    if res.get('user_role') != "clinician":
        res['intent_type'] = "general"

    return res

from langchain_core.runnables.config import RunnableConfig

async def patient_llm_node(state: AgentState, config: RunnableConfig):
    res = await patient_llm.run(state, config=config)
    last_msg = res["messages"][-1]
    content = last_msg.content if getattr(last_msg, "content", "") else str(getattr(last_msg, "tool_calls", ""))
    log = log_step("Patient LLM", content)
    return {"messages": res["messages"], "logs": log["logs"]}

async def caregiver_llm_node(state: AgentState, config: RunnableConfig):
    res = await caregiver_llm.run(state, config=config)
    last_msg = res["messages"][-1]
    content = last_msg.content if getattr(last_msg, "content", "") else str(getattr(last_msg, "tool_calls", ""))
    log = log_step("Caregiver LLM", content)
    return {"messages": res["messages"], "logs": log["logs"]}

async def validator_node(state: AgentState):
    res = await validator.run(state)
    log = log_step("Response Validator", f"Valid: {res.get('is_valid')}")
    res.update(log)
    return res

async def safety_check_node(state: AgentState):
    res = await safety_check.run(state)
    log = log_step("Safety Check", f"Safe: {res.get('is_safe')}")
    res.update(log)
    return res

async def recovery_loop_node(state: AgentState):
    attempts = state.get("attempts", 0) + 1
    log = log_step("Recovery Loop", f"Attempt {attempts}")
    return {
        "attempts": attempts,
        "messages": [AIMessage(content="[RECOVERY] Let me try rephrasing or improving my previous response.")],
        "logs": log["logs"]
    }


async def emergency_response_node(state: AgentState):
    response = (
        "This appears to be a possible medical emergency. Seek urgent medical assistance now "
        "and do not delay care. If the person is unconscious, not breathing, or having a seizure, "
        "call emergency services immediately."
    )
    log = log_step("Emergency Fast Path", response)
    return {
        "messages": [AIMessage(content=response)],
        "logs": log["logs"],
        "is_valid": True,
        "is_safe": True,
    }


async def intent_classifier_node(state: AgentState):
    res = await intent_classifier.run(state)
    log = log_step("Intent Classifier", f"Intent: {res.get('intent_type')}")
    res.update(log)
    return res

async def diagnosis_assist_node(state: AgentState, config: RunnableConfig):
    res = _set_latest_clinician_output(await diagnosis_assist.run(state, config=config))
    log = log_step("Diagnosis Assistant", res.get('clinician_outputs', [""])[-1] if res.get('clinician_outputs') else "")
    res.update(log)
    clinician_outputs = (state.get("clinician_outputs") or []) + (res.get("clinician_outputs") or [])
    res["clinician_outputs"] = clinician_outputs
    return res

async def treatment_assist_node(state: AgentState, config: RunnableConfig):
    res = _set_latest_clinician_output(await treatment_assist.run(state, config=config))
    log = log_step("Treatment Assistant", res.get('clinician_outputs', [""])[-1] if res.get('clinician_outputs') else "")
    res.update(log)
    clinician_outputs = (state.get("clinician_outputs") or []) + (res.get("clinician_outputs") or [])
    res["clinician_outputs"] = clinician_outputs
    return res

async def monitoring_assist_node(state: AgentState, config: RunnableConfig):
    res = _set_latest_clinician_output(await monitoring_assist.run(state, config=config))
    log = log_step("Monitoring Assistant", res.get('clinician_outputs', [""])[-1] if res.get('clinician_outputs') else "")
    res.update(log)
    clinician_outputs = (state.get("clinician_outputs") or []) + (res.get("clinician_outputs") or [])
    res["clinician_outputs"] = clinician_outputs
    return res

async def general_assist_node(state: AgentState, config: RunnableConfig):
    res = _set_latest_clinician_output(await general_assist.run(state, config=config))
    log = log_step("General Clinical Assistant", res.get('clinician_outputs', [""])[-1] if res.get('clinician_outputs') else "")
    res.update(log)
    clinician_outputs = (state.get("clinician_outputs") or []) + (res.get("clinician_outputs") or [])
    res["clinician_outputs"] = clinician_outputs
    return res

async def merge_outputs_node(state: AgentState, config: RunnableConfig):
    # Prepare state with context about the specialist outputs for the OutputMerger
    clinician_outputs = state.get("clinician_outputs", [])
    if clinician_outputs:
        # Add specialist outputs summary to messages for context
        outputs_context = "\n\n".join([f"Specialist Output {i+1}:\n{output}" for i, output in enumerate(clinician_outputs)])
        state_with_context = dict(state)
        state_with_context["messages"] = state["messages"] + [AIMessage(content=outputs_context)]
        res = await output_merger.run(state_with_context, config=config)
    else:
        res = await output_merger.run(state, config=config)
    log = log_step("Output Merger", "Merged outputs successfully.")
    res.update(log)
    return res

async def research_agent_node(state: AgentState, config: RunnableConfig):
    res = await research_agent.run(state, config=config)
    log = log_step("Research Assistant", res.get('research_output', ''))
    res.update(log)
    return res

async def dietary_assist_node(state: AgentState, config: RunnableConfig):
    res = await dietary_assist.run(state, config=config)
    last_msg = res["messages"][-1]
    content = last_msg.content if getattr(last_msg, "content", "") else str(getattr(last_msg, "tool_calls", ""))
    log = log_step("Dietary Specialist", content)
    res.update(log)
    return res

async def tool_node_with_logging(state: AgentState):
    start_time = time.time()
    res = await tool_node.ainvoke(state)
    end_time = time.time()
    
    output_summary = f"{len(res)} tool(s) executed." if isinstance(res, list) else "Tool executed."
    log = log_step("Executing Tools (RAG/Web Search)", output_summary)
    
    metrics = {
        "agent": "ToolsNode",
        "tokens": 0, # Tools don't use tokens directly in their logic here
        "time": round(end_time - start_time, 3)
    }
    
    if isinstance(res, list):
        return {"messages": res, "logs": log["logs"], "metrics": [metrics]}
    res.update(log)
    res.update({"metrics": [metrics]})
    return res

async def persistence_node(state: AgentState):
    """Save the current chat history to Supabase in FHIR format."""
    patient_id = state.get("patient_id", "anonymous")
    session_id = state.get("session_id")

    # Convert LangChain messages to a simple list of dicts for the tool
    formatted_messages = []
    for msg in state["messages"]:
        role = "user" if msg.type == "human" else "assistant"
        formatted_messages.append({"role": role, "content": msg.content})

    # Persist patient-facing and caregiver-facing conversations for the active patient.
    if state.get("user_role") in ("patient", "caregiver"):
        start_time = time.time()
        res = save_chat_as_fhir.invoke({
            "patient_id": patient_id,
            "messages": formatted_messages,
            "session_id": session_id,
        })
        end_time = time.time()

        log = log_step("FHIR Persistence", res)

        metrics = {
            "agent": "PersistenceNode",
            "tokens": 0,
            "time": round(end_time - start_time, 3)
        }
        return {"logs": log["logs"], "metrics": [metrics]}

    return {}

# Define routing functions
def route_after_role(state: AgentState):
    role = state["user_role"]
    if role in {"patient", "caregiver"} and detect_emergency(state):
        return "emergency_response"
    if role == "patient":
        return "patient_llm"
    elif role == "caregiver":
        return "caregiver_llm"
    elif role == "clinician":
        return "intent_classifier"
    elif role == "researcher":
        return "research_agent"
    elif role == "dietary":
        return "dietary_assist"
    return END

def route_research_agent(state: AgentState):
    # Determine if the last message has tool calls
    last_message = state["messages"][-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools_node"
    return END

def route_patient_llm(state: AgentState):
    # Determine if the last message has tool calls
    last_message = state["messages"][-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools_node"
    return "validator"

def route_after_tools(state: AgentState):
    role = state.get("user_role")
    if role == "researcher":
        return "research_agent"
    elif role == "dietary":
        return "dietary_assist"
    elif role == "caregiver":
        return "caregiver_llm"
    return "patient_llm"

def route_after_validator(state: AgentState):
    if state.get("is_valid", False):
        return "safety_check"
    return "recovery_loop"

def route_after_safety(state: AgentState):
    if state.get("is_safe", False):
        return "persistence_node"
    return "recovery_loop"

def route_after_persistence(state: AgentState):
    return END

def route_after_recovery(state: AgentState):
    attempts = state.get("attempts", 0)
    if attempts >= MAX_RECOVERY_ATTEMPTS:
        return "persistence_node"

    role = state.get("user_role")
    if role == "dietary":
        return "dietary_assist"
    elif role == "caregiver":
        return "caregiver_llm"
    elif role == "patient":
        return "patient_llm"
    return "persistence_node"

def route_after_intent(state: AgentState):
    intent = state.get("intent_type", "general")
    if intent == "diagnosis":
        return "diagnosis_assist"
    elif intent == "treatment":
        return "treatment_assist"
    elif intent == "monitoring":
        return "monitoring_assist"
    else:
        return "general_assist"

from src.tools.dietary_tools import page_indexed_retrieval, search_guidelines, get_nutritional_data

# Updated Tool Node to include page indexing RAG
tools = [web_search_tool, page_indexed_retrieval, search_guidelines, get_nutritional_data]
tool_node = ToolNode(tools)

# Build the graph
builder = StateGraph(AgentState)

# Add nodes
builder.add_node("role_classifier", role_classifier_node)
builder.add_node("patient_llm", patient_llm_node)
builder.add_node("caregiver_llm", caregiver_llm_node)
builder.add_node("validator", validator_node)
builder.add_node("safety_check", safety_check_node)
builder.add_node("recovery_loop", recovery_loop_node)
builder.add_node("emergency_response", emergency_response_node)
builder.add_node("intent_classifier", intent_classifier_node)
builder.add_node("diagnosis_assist", diagnosis_assist_node)
builder.add_node("treatment_assist", treatment_assist_node)
builder.add_node("monitoring_assist", monitoring_assist_node)
builder.add_node("general_assist", general_assist_node)
builder.add_node("merge_outputs", merge_outputs_node)
builder.add_node("research_agent", research_agent_node)
builder.add_node("tools_node", tool_node_with_logging)
builder.add_node("dietary_assist", dietary_assist_node)
builder.add_node("persistence_node", persistence_node)

# Set entry point
builder.set_entry_point("role_classifier")

# Define edges
builder.add_conditional_edges("role_classifier", route_after_role, {
    "patient_llm": "patient_llm",
    "caregiver_llm": "caregiver_llm",
    "intent_classifier": "intent_classifier",
    "research_agent": "research_agent",
    "dietary_assist": "dietary_assist",
    "emergency_response": "emergency_response",
    END: END
})

# Patient and Caregiver Pathways
builder.add_conditional_edges("patient_llm", route_patient_llm, {
    "tools_node": "tools_node",
    "validator": "validator"
})
builder.add_conditional_edges("caregiver_llm", route_patient_llm, {
    "tools_node": "tools_node",
    "validator": "validator"
})
builder.add_conditional_edges("validator", route_after_validator, {
    "safety_check": "safety_check",
    "recovery_loop": "recovery_loop"
})
builder.add_conditional_edges("safety_check", route_after_safety, {
    "persistence_node": "persistence_node",
    "recovery_loop": "recovery_loop"
})
builder.add_edge("persistence_node", END)
builder.add_conditional_edges("recovery_loop", route_after_recovery, {
    "dietary_assist": "dietary_assist",
    "caregiver_llm": "caregiver_llm",
    "patient_llm": "patient_llm",
    "persistence_node": "persistence_node"
})

builder.add_edge("emergency_response", "persistence_node")

# Clinician Pathway
builder.add_conditional_edges("intent_classifier", route_after_intent, {
    "diagnosis_assist": "diagnosis_assist",
    "treatment_assist": "treatment_assist",
    "monitoring_assist": "monitoring_assist",
    "general_assist": "general_assist"
})
builder.add_edge("diagnosis_assist", "merge_outputs")
builder.add_edge("treatment_assist", "merge_outputs")
builder.add_edge("monitoring_assist", "merge_outputs")
builder.add_edge("general_assist", "merge_outputs")
builder.add_edge("merge_outputs", END)

# Researcher Pathway
builder.add_conditional_edges("research_agent", route_research_agent, {
    "tools_node": "tools_node",
    END: END
})

# Shared Tool Pathway
builder.add_conditional_edges("tools_node", route_after_tools, {
    "research_agent": "research_agent",
    "patient_llm": "patient_llm",
    "caregiver_llm": "caregiver_llm",
    "dietary_assist": "dietary_assist"
})

# Dietary Pathway
def route_dietary_assist(state: AgentState):
    # If the last message has tool calls, go to tools
    last_message = state["messages"][-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools_node"
    return "validator"

builder.add_conditional_edges("dietary_assist", route_dietary_assist, {
    "tools_node": "tools_node",
    "validator": "validator"
})

# Compile the graph
medical_pipeline = builder.compile()