Wave 1 Implementation Spec — Intake Restructuring¶

"""
Curaway — Layer State Management.

Tracks completion of 5 intake dimensions. Layers are independent —
completion is a high-water mark, not a linear pipeline.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any


class LayerName(str, Enum):
    INTENT = "intent_capture"
    MEDICAL = "medical_status"
    TRAVEL = "travel_readiness"
    LOGISTICS = "logistics"
    FINANCIAL = "financial_readiness"


class LayerStatus(str, Enum):
    NOT_STARTED = "not_started"
    IN_PROGRESS = "in_progress"
    COMPLETE = "complete"


@dataclass
class LayerState:
    status: LayerStatus = LayerStatus.NOT_STARTED
    completion: float = 0.0
    data: dict[str, Any] = field(default_factory=dict)
    summary_card_fired: bool = False


@dataclass
class PatientLayerState:
    """Full layer state for a patient's intake journey."""
    intent_capture: LayerState = field(default_factory=LayerState)
    medical_status: LayerState = field(default_factory=LayerState)
    travel_readiness: LayerState = field(default_factory=LayerState)
    logistics: LayerState = field(default_factory=LayerState)
    financial_readiness: LayerState = field(default_factory=LayerState)

    def get_layer(self, name: LayerName) -> LayerState:
        return getattr(self, name.value)

    def set_layer(self, name: LayerName, state: LayerState) -> None:
        setattr(self, name.value, state)

    def to_dict(self) -> dict:
        return {
            name.value: {
                "status": self.get_layer(name).status.value,
                "completion": self.get_layer(name).completion,
                "data": self.get_layer(name).data,
                "summary_card_fired": self.get_layer(name).summary_card_fired,
            }
            for name in LayerName
        }

    @classmethod
    def from_dict(cls, data: dict) -> "PatientLayerState":
        state = cls()
        for name in LayerName:
            if name.value in data:
                layer_data = data[name.value]
                setattr(state, name.value, LayerState(
                    status=LayerStatus(layer_data.get("status", "not_started")),
                    completion=layer_data.get("completion", 0.0),
                    data=layer_data.get("data", {}),
                    summary_card_fired=layer_data.get("summary_card_fired", False),
                ))
        return state

    def all_above_threshold(self, threshold: float, exclude: list[LayerName] | None = None) -> bool:
        """Check if all layers (except excluded) are above a completion threshold."""
        for name in LayerName:
            if exclude and name in exclude:
                continue
            if self.get_layer(name).completion < threshold:
                return False
        return True

# alembic/versions/xxxx_add_layer_state_to_cases.py

def upgrade():
    op.add_column('cases', sa.Column('layer_state', FlexibleJSON, nullable=True))

def downgrade():
    op.drop_column('cases', 'layer_state')

"""PFS — Patient Fitness Score (0-100). Pure function, deterministic."""

def compute_pfs(
    layer_state: dict,
    risk_items: list[dict],
    config: dict,
) -> dict:
    """Compute PFS from layer data + risk assessor output.

    Returns: {
        "score": 78,
        "band": "conditionally_ready",
        "transport_tier": "T1",
        "blocking_flags": [],
        "component_breakdown": {"medical": 85, "travel": 100, ...},
        "gaps": ["passport_timeline_tight"]
    }
    """

"""
Curaway — Intent Capture Extractor (Layer 1).

Pure function: (conversation_turn, patient_state) → structured_data_delta.
Uses structured output (JSON schema enforcement) via llm_gateway.invoke().
"""

from app.services.llm_gateway import invoke as llm_invoke
from app.services.prompt_loader import load_prompt


EXTRACTION_SCHEMA = {
    "type": "object",
    "properties": {
        "case_type": {"type": "string", "enum": ["elective", "selective", "urgent"]},
        "emotional_state": {
            "type": "object",
            "properties": {
                "readiness": {"type": "string"},
                "primary_fear": {"type": "string"},
                "primary_hope": {"type": "string"},
                "fear_intensity": {"type": "number"},
                "hope_intensity": {"type": "number"},
            }
        },
        # ... full schema from design Part 7
    }
}


async def extract_intent(
    message_text: str,
    current_layer_data: dict,
    patient_context: str,
) -> dict:
    """Extract Layer 1 signals from a conversation turn.

    Returns delta — only newly extracted signals, not the full layer.
    Caller merges delta into existing layer state.
    """
    prompt = load_prompt("intent_extraction", "en", version="v1")
    # ... build messages, call llm_invoke with expects_json=True
    # ... parse response, return delta

from typing import Annotated
from operator import methodcaller  # or a custom reducer


def merge_layer_deltas(current: dict, delta: dict) -> dict:
    """LangGraph reducer — merges per-extractor deltas into layer_state dict."""
    state = PatientLayerState.from_dict(current)
    # Apply delta — delta is {layer_name: {"completion": ..., "data_delta": ...}}
    for layer_name, layer_delta in delta.items():
        state.update_layer(
            LayerName(layer_name),
            completion=layer_delta.get("completion"),
            data_delta=layer_delta.get("data_delta"),
        )
    return state.to_dict()


class TriageState(TypedDict):
    layer_state: Annotated[dict, merge_layer_deltas]
    # ... other fields

"""
Curaway — Triage Agent (A3 Model).

Supervisor pattern: one conversational agent + parallel extractors.
Code nodes handle scoring and routing (deterministic).
LLM nodes handle conversation and extraction (non-deterministic).

Graph topology:
  triage_conversation → parallel_extractors → pfs_scorer → route_by_pfs
"""

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver


class TriageState(TypedDict):
    """State that flows through the triage graph."""
    messages: list[BaseMessage]
    layer_state: dict
    pfs: dict | None
    hss: dict | None
    fms: float | None
    active_layer: str
    response: str | None
    case_id: str
    patient_id: str
    tenant_id: str


def build_triage_graph(checkpointer=None):
    graph = StateGraph(TriageState)

    # Nodes
    graph.add_node("choose_layer", choose_next_layer)       # Code: adaptive routing
    graph.add_node("generate_response", generate_response)  # LLM: conversation
    graph.add_node("run_extractors", run_extractors)        # LLM: parallel extraction
    graph.add_node("compute_pfs", compute_pfs_node)         # Code: deterministic
    graph.add_node("check_gates", check_matching_gates)     # Code: threshold check

    # Edges
    graph.add_edge("choose_layer", "generate_response")
    graph.add_edge("generate_response", "run_extractors")
    graph.add_edge("run_extractors", "compute_pfs")
    graph.add_conditional_edges(
        "compute_pfs",
        route_by_pfs_band,
        {
            "continue": "choose_layer",   # PFS < threshold, keep collecting
            "matching_ready": END,         # PFS >= 60, layers >= 0.6
            "remediation": "choose_layer", # PFS 40-59, remediation prompts
            "not_ready": END,              # PFS < 40, honest redirect
        }
    )

    graph.set_entry_point("choose_layer")

    if checkpointer:
        return graph.compile(checkpointer=checkpointer)
    return graph.compile()

# In orchestrator, replace phase routing:
if case.status == "intake":
    # Old: route to identify_procedure / records_first / intake phases
    # New: route to triage agent
    from app.agents.triage_agent import run_triage_turn
    return await run_triage_turn(state)

def upgrade():
    # graph_node_audit_log
    op.create_table(
        'graph_node_audit_log',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('case_id', sa.String(36), nullable=False, index=True),
        sa.Column('turn_number', sa.Integer, nullable=False),
        sa.Column('node_name', sa.String(100), nullable=False),
        sa.Column('node_type', sa.String(20), nullable=False),  # llm/extractor/scorer/router
        sa.Column('input_hash', sa.String(64)),
        sa.Column('output_hash', sa.String(64)),
        sa.Column('state_before', FlexibleJSON),
        sa.Column('state_after', FlexibleJSON),
        sa.Column('decision_made', sa.String(100)),
        sa.Column('decision_reason', sa.Text),
        sa.Column('model_used', sa.String(50)),
        sa.Column('duration_ms', sa.Integer),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
    )

    # scoring_audit_log
    op.create_table(
        'scoring_audit_log',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('case_id', sa.String(36), nullable=False, index=True),
        sa.Column('score_type', sa.String(10), nullable=False),  # pfs/hss/fms
        sa.Column('score_value', sa.Float, nullable=False),
        sa.Column('config_version', sa.String(64)),
        sa.Column('inputs', FlexibleJSON),
        sa.Column('component_breakdown', FlexibleJSON),
        sa.Column('threshold_crossed', sa.Boolean, default=False),
        sa.Column('threshold_direction', sa.String(10)),
        sa.Column('previous_score', sa.Float),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
    )

    # extraction_audit_log
    op.create_table(
        'extraction_audit_log',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('case_id', sa.String(36), nullable=False, index=True),
        sa.Column('turn_number', sa.Integer, nullable=False),
        sa.Column('layer', sa.String(30), nullable=False),
        sa.Column('extracted_signals', FlexibleJSON),
        sa.Column('confidence_scores', FlexibleJSON),
        sa.Column('completion_before', sa.Float),
        sa.Column('completion_after', sa.Float),
        sa.Column('model_used', sa.String(50)),
        sa.Column('input_text', sa.Text),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
    )

"""Configure LangGraph checkpointing → Railway PostgreSQL."""

from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from app.config import settings


async def get_checkpointer():
    """Create checkpointer using DATABASE_URL_ADMIN (superuser, bypasses RLS).

    Thread ID = case_id. Checkpoints are case-scoped.
    """
    return AsyncPostgresSaver.from_conn_string(
        settings.database_url_admin or settings.database_url
    )