src/core/graph_cdm.py

from langgraph.graph import StateGraph, END
from src.core.state import AgentState
from src.agents.agents import ResponseValidator, SafetyCheck
from src.tools.fhir_memory import (
    get_patient_summary_fhir, save_observation, ingest_fhir_bundle, 
    get_medications_by_patient, save_medication
)
from src.tools.web_tools import web_search_tool
from langgraph.prebuilt import ToolNode
from langchain_core.messages import ToolMessage
from langchain_core.runnables.config import RunnableConfig
try:
    from langgraph.runtime import Runtime
except ImportError:
    Runtime = None
import time
from src.utils.logger import setup_logger
from src.core.policy import PlatformGateService, audit_event, gate_update
from src.runtime.execution_context import ExecutionContext
from src.runtime.cdm_runtime import TrendAnalysis
from src.runtime.worker_runtime import WorkerRuntime
from src.skills.adapters import create_worker_adapter
from src.skills.registry import default_skill_registry

logger = setup_logger("CDMPipeline")


def _platform_agent(agent_type):
    return agent_type()

def log_step(name: str, output: str = None):
    logger.info(f"Executing: {name}")
    log_msg = f"➔ CDM Node: {name}"
    if output:
        log_msg += f"\nOutput: {output}"
    return {"logs": [log_msg]}

async def data_fetch_node(state: AgentState):
    """
    Fetch FHIR data and medications for CDM context.
    Bug 7.3: Include medication context in CDM pipeline.
    """
    patient_id = state.get("patient_id", "unknown")
    if patient_id == "unknown":
        return log_step("Data Fetch", "No Patient ID provided.")
    
    start_time = time.time()
    
    # Fetch FHIR summary
    summary = get_patient_summary_fhir.invoke({"patient_id": patient_id})
    
    # Bug 7.3: Fetch medications for context
    medications = get_medications_by_patient.invoke({"patient_id": patient_id})
    
    # If summary has active_medications field, populate it
    if isinstance(summary, dict) and "active_medications" in summary:
        summary["active_medications"] = medications if medications else []
    
    end_time = time.time()
    
    log_output = f"Retrieved summary and medications for {patient_id}"
    if medications:
        log_output += f" ({len(medications)} active medications)"
    
    log = log_step("FHIR Data Fetch", log_output)
    
    metrics = {
        "agent": "DataFetchNode",
        "tokens": 0,
        "time": round(end_time - start_time, 3)
    }
    
    return {
        "fhir_data": [summary] if isinstance(summary, dict) else [], 
        "current_medications": medications if isinstance(medications, list) else [],
        "logs": log["logs"],
        "metrics": [metrics]
    }

async def trend_analyzer_node(state: AgentState):
    registry = default_skill_registry()
    adapter = create_worker_adapter("cdm-trend-analysis", registry=registry)
    context = ExecutionContext(
        request_id=state.get("request_id") or "",
        trace_id=state.get("trace_id") or "",
        session_id=state.get("session_id") or "",
        patient_id=state.get("patient_id"),
        user_query=state["messages"][-1].content if state.get("messages") else "",
        user_role=state.get("user_role"),
    )
    worker_state = dict(state)
    if worker_state.get("user_role") in (None, "unknown"):
        worker_state["user_role"] = "patient"
    if state.get("runtime_mode", "skill") == "agent":
        result = await adapter.execute(worker_state)
    else:
        worker = WorkerRuntime(registry)
        try:
            result = await worker.execute_adapter(adapter, worker_state)
        finally:
            worker.flush_and_destroy()
    output = result.output or {}
    trend = TrendAnalysis(
        summary=output.get("trend_analysis", "") if isinstance(output, dict) else str(output),
        metrics=output.get("trend_data", {}) if isinstance(output, dict) else {},
    )
    metric_trend = dict(result.metrics) if result.metrics else {}
    metric_trend["skill_id"] = result.skill_id
    metric_trend["onloaded_skill"] = result.skill_id
    metric_trend["offloaded_skill"] = result.skill_id
    metric_trend["agent"] = "TrendAnalyzer"
    return {
        "trend_analysis": trend.summary,
        "trend_data": trend.metrics,
        "cdm_trend_analysis": trend.as_dict(),
        "logs": log_step("Trend Analyzer", trend.summary)["logs"],
        "metrics": [metric_trend],
        "active_skill_id": result.skill_id,
        "active_skill_version": result.skill_version,
        "skill_trace": (state.get("skill_trace") or []) + [
            audit_event({**state, "active_skill_id": result.skill_id, "active_skill_version": result.skill_version}, "skill_onloaded", agent="TrendAnalyzer"),
            audit_event({**state, "active_skill_id": result.skill_id, "active_skill_version": result.skill_version}, "skill_executed"),
            audit_event({**state, "active_skill_id": result.skill_id, "active_skill_version": result.skill_version}, "skill_offloaded", agent="TrendAnalyzer"),
        ],
    }

async def health_coach_node(state: AgentState, config: RunnableConfig = None):
    registry = default_skill_registry()
    adapter = create_worker_adapter("cdm-health-coach", registry=registry)
    worker_state = dict(state)
    if worker_state.get("user_role") in (None, "unknown"):
        worker_state["user_role"] = "patient"
    if state.get("runtime_mode", "skill") == "agent":
        result = await adapter.execute(worker_state, config=config)
    else:
        worker = WorkerRuntime(registry)
        try:
            result = await worker.execute_adapter(adapter, worker_state, config=config)
        finally:
            worker.flush_and_destroy()
    content = result.messages[-1].content if result.messages and getattr(result.messages[-1], "content", "") else "Tool calls generated."
    metric_coach = dict(result.metrics) if result.metrics else {}
    metric_coach["skill_id"] = result.skill_id
    metric_coach["onloaded_skill"] = result.skill_id
    metric_coach["offloaded_skill"] = result.skill_id
    metric_coach["agent"] = "HealthCoach"
    return {
        "messages": result.messages,
        "logs": log_step("Health Coach", content)["logs"],
        "metrics": [metric_coach],
        "active_skill_id": result.skill_id,
        "active_skill_version": result.skill_version,
        "skill_plan": [{"skill_id": result.skill_id, "reason": "CDM pipeline"}],
        "skill_trace": (state.get("skill_trace") or []) + [
            audit_event({**state, "active_skill_id": result.skill_id, "active_skill_version": result.skill_version}, "skill_onloaded", agent="HealthCoach"),
            audit_event({**state, "active_skill_id": result.skill_id, "active_skill_version": result.skill_version}, "skill_executed"),
            audit_event({**state, "active_skill_id": result.skill_id, "active_skill_version": result.skill_version}, "skill_offloaded", agent="HealthCoach"),
        ],
    }

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

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

# Tool routing
def route_health_coach(state: AgentState):
    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):
    return "health_coach"

def route_after_validator(state: AgentState):
    if state.get("is_valid", False):
        return "safety_check"
    return "health_coach" # Simple retry

# Setup tools
cdm_tools = [web_search_tool, save_observation, ingest_fhir_bundle, get_patient_summary_fhir, get_medications_by_patient, save_medication]
tool_node = ToolNode(cdm_tools)

def _ensure_tool_config(config: RunnableConfig = None) -> dict:
    call_config = dict(config) if config else {}
    configurable = dict(call_config.get("configurable", {}))
    if "__pregel_runtime" not in configurable and Runtime is not None:
        try:
            configurable["__pregel_runtime"] = Runtime()
        except Exception:
            pass
    call_config["configurable"] = configurable
    return call_config

async def tools_node_with_metrics(state: AgentState, config: RunnableConfig = None):
    start_time = time.time()
    call_config = _ensure_tool_config(config)
    res = await tool_node.ainvoke(state, config=call_config)
    end_time = time.time()
    
    tool_messages = res if isinstance(res, list) else (res.get("messages", []) if isinstance(res, dict) else [])
    tool_details = []
    for tm in tool_messages:
        if isinstance(tm, ToolMessage):
            tool_name = getattr(tm, "name", None) or "CDM Tool"
            out = str(getattr(tm, "content", ""))
            tool_details.append({
                "name": tool_name,
                "output": out[:200] if out else "",
                "time": round(end_time - start_time, 3)
            })

    names = [t["name"] for t in tool_details]
    tools_str = ", ".join(dict.fromkeys(names)) if names else "CDMTools"
    duration = round(end_time - start_time, 3)
    metrics = {
        "agent": "CDMToolsNode",
        "tool_name": tools_str,
        "tools": tool_details if tool_details else [{"name": tools_str, "time": duration}],
        "tokens": 0,
        "time": duration
    }
    
    # ToolNode returns a list of messages
    return {"messages": res, "metrics": [metrics]}

# Build CDM Graph
builder = StateGraph(AgentState)

builder.add_node("data_fetch", data_fetch_node)
builder.add_node("trend_analyzer", trend_analyzer_node)
builder.add_node("health_coach", health_coach_node)
builder.add_node("tools_node", tools_node_with_metrics)
builder.add_node("validator", validator_node)
builder.add_node("safety_check", safety_check_node)

builder.set_entry_point("data_fetch")
builder.add_edge("data_fetch", "trend_analyzer")
builder.add_edge("trend_analyzer", "health_coach")

builder.add_conditional_edges("health_coach", route_health_coach, {
    "tools_node": "tools_node",
    "validator": "validator"
})

builder.add_edge("tools_node", "health_coach")

builder.add_conditional_edges("validator", route_after_validator, {
    "safety_check": "safety_check",
    "health_coach": "health_coach"
})

builder.add_conditional_edges("safety_check", lambda x: "end" if x.get("is_safe") else "health_coach", {
    "end": END,
    "health_coach": "health_coach"
})

cdm_pipeline = builder.compile()