LangGraph & KirokuForms: HITL Examples

"""Pause a LangGraph run for a human, then carry on with their answer.

Uses LangGraph's own interrupt(): the graph stops, its state is checkpointed,
and the process is free while the reviewer takes however long they take. The run
resumes later, from anywhere, with what they actually submitted.
"""

from typing import Optional, TypedDict

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph

from kirokuforms import (
    KirokuFormsHITL,
    create_kiroku_interrupt_node,
    resume_with_answers,
)

client = KirokuFormsHITL(api_key="YOUR_KIROKU_API_KEY")


class TransactionState(TypedDict):
    transaction: dict
    review_data: Optional[dict]
    human_verification: Optional[dict]
    status: str


# The review step. `name` separates this step's cases from any other review node
# in the same run: it goes into the idempotency key, which is what stops the
# re-execution on resume from opening a second case and mailing the reviewer
# twice.
verify_transaction = create_kiroku_interrupt_node(
    client,
    name="transaction",
    title="Verify a high-value transaction",
    description="This one is over the threshold, so it needs a person.",
    fields=[
        {
            "type": "radio",
            "label": "Approve this transaction?",
            "name": "approved",
            "required": True,
            "options": [
                {"label": "Approve", "value": "yes"},
                {"label": "Reject", "value": "no"},
            ],
        },
        {
            "type": "textarea",
            "label": "Notes",
            "name": "notes",
            "required": False,
        },
    ],
)


def check_transaction(state: TransactionState) -> TransactionState:
    """Decide whether a person needs to look at this at all."""
    if state["transaction"]["amount"] > 1000:
        return {**state, "status": "needs_review"}
    return {**state, "status": "auto_approved"}


def needs_review(state: TransactionState) -> str:
    return "verify" if state["status"] == "needs_review" else "done"


def record_decision(state: TransactionState) -> TransactionState:
    """Read what the human said and turn it into a status."""
    answers = (state.get("human_verification") or {}).get("result") or {}
    approved = answers.get("approved") == "yes"
    return {**state, "status": "approved" if approved else "rejected"}


builder = StateGraph(TransactionState)
builder.add_node("check_transaction", check_transaction)
builder.add_node("verify", verify_transaction)
builder.add_node("record_decision", record_decision)

builder.add_edge(START, "check_transaction")
builder.add_conditional_edges(
    "check_transaction", needs_review, {"verify": "verify", "done": END}
)
builder.add_edge("verify", "record_decision")
builder.add_edge("record_decision", END)

# A checkpointer is required. Without one there is nothing to suspend into, and
# interrupt() has nowhere to put the state it is pausing.
graph = builder.compile(checkpointer=MemorySaver())

# A checkpointed run needs a thread_id: it is the key the state is saved under
# and the handle you pass back to resume this transaction later.
config = {"configurable": {"thread_id": "transaction-1001"}}

state = graph.invoke(
    {
        "transaction": {"id": "TX-1001", "customer": "Acme Ltd", "amount": 2500.00},
        "review_data": None,
        "human_verification": None,
        "status": "new",
    },
    config,
)

if "__interrupt__" in state:
    # Stopped. The payload carries the tokenized link to send the reviewer; in a
    # service this is where the process would end, and a webhook or a poll would
    # resume the thread later.
    payload = state["__interrupt__"][0].value
    print("Send this to the reviewer:", payload["form_url"])

    # Resumes with what they submitted, rather than with anything this script
    # assumes on their behalf.
    state = graph.invoke(
        resume_with_answers(client, payload["kiroku_task_id"], wait=True), config
    )

print("Final status:", state["status"])

from typing import TypedDict, Annotated, Sequence, Optional
import operator # Required for Annotated human_input setter
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver
from kirokuforms import KirokuFormsHITL, create_kiroku_interrupt_node, resume_with_answers
from langchain_core.tools import tool
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage # For explicit message construction
from langchain_openai import ChatOpenAI


# Define tools for the agent
@tool
def search_tool(query: str) -> str:
    """Searches for information based on the query."""
    # In a real scenario, this would call a search API
    return f"Search results for '{query}': Information found about {query}."

@tool
def financial_analysis_tool(data: dict) -> str:
    """Performs financial analysis on the provided data."""
    item = data.get("item", "the subject")
    # In a real scenario, this would perform calculations or call another service
    return f"Financial analysis complete for '{item}'. Recommendation: Based on the current data, proceed with caution."

tools = [search_tool, financial_analysis_tool]

# Agent state
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    human_input: Annotated[Optional[dict], operator.setitem] # KirokuForms interrupt handler adds 'human_input' after review

# Initialize LLM and KirokuForms interrupt handler
llm = ChatOpenAI(model="gpt-4-turbo-preview", api_key="YOUR_OPENAI_API_KEY") # Replace with your key
client = KirokuFormsHITL(api_key="YOUR_KIROKU_API_KEY") # Replace

# The oversight step. interrupt() suspends the graph, so nothing is held open
# while a person decides whether the agent may run the action it proposed.
#
# `data_key` names a dict on the state and generates one field per entry, so the
# reviewer sees this run's proposed action. Fields given with `fields=` are fixed
# when the node is built and cannot read state.
human_oversight_node = create_kiroku_interrupt_node(
    client,
    name="agent-oversight",
    title="Approve the action this agent wants to take",
    description="The agent proposed the action below. Approve it or reject it.",
    data_key="proposed_action",
)

# Define agent nodes
def call_model_node(state: AgentState) -> dict:
    response = llm.invoke(state["messages"], tools=tools)
    return {"messages": [response]}

def call_tools_node(state: AgentState) -> dict:
    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
        return {"messages": [HumanMessage(content="No tool calls found in the last AI message.")]} # Or handle as an error

    tool_results = []
    for tc in last_message.tool_calls:
        tool_name = tc["name"]
        try:
            # Find the tool in the global scope or a registered tool list
            selected_tool = globals().get(tool_name) # Or use a more robust tool registry
            if callable(selected_tool):
                tool_output = selected_tool(tc["args"])
                tool_results.append(ToolMessage(content=str(tool_output), tool_call_id=tc["id"]))
            else:
                tool_results.append(ToolMessage(content=f"Error: Tool '{tool_name}' not found or not callable.", tool_call_id=tc["id"]))
        except Exception as e:
            tool_results.append(ToolMessage(content=f"Error executing tool '{tool_name}': {e}", tool_call_id=tc["id"]))
            
    return {"messages": tool_results}


# Criteria for when to request human oversight
def needs_human_oversight(state: AgentState) -> bool:
    last_message = state["messages"][-1]
    if isinstance(last_message, AIMessage) and last_message.content and isinstance(last_message.content, str):
        critical_phrases = ["critical decision", "large investment", "significant financial impact", "major commitment"]
        if any(phrase in last_message.content.lower() for phrase in critical_phrases):
            return True
    # Example: if a specific tool (like financial_analysis_tool) was just called and produced output
    if isinstance(last_message, ToolMessage) and last_message.name == "financial_analysis_tool":
        return True
    return False

# Node to trigger human oversight
def prepare_oversight_payload(state: AgentState) -> dict:
    # Try to get the last AI message content or a summary
    agent_plan_summary = "No specific plan articulated by AI yet."
    if state["messages"]:
        # Iterate backwards to find the last AI message that is not a tool call response
        for msg in reversed(state["messages"]):
            if isinstance(msg, AIMessage) and msg.content:
                agent_plan_summary = msg.content
                break
            elif isinstance(msg, HumanMessage) and msg.content: # If last is human, show that too
                agent_plan_summary = f"Last user input: {msg.content}"
                break


    # Hand the proposed action to the oversight node, which suspends the graph.
    # This node only prepares what the reviewer sees; `human_oversight_node`
    # below is the one that stops the run.
    return {**state, "proposed_action": {"proposed_action": agent_plan_summary}}

# Node to process human feedback
def process_human_feedback_node(state: AgentState) -> dict:
    if state.get("human_input"):
        decision = state["human_input"].get("human_decision")
        feedback = state["human_input"].get("human_feedback", "")
        
        if decision == "approve":
            return {"messages": [HumanMessage(content=f"Human approved the plan. Original feedback: {feedback if feedback else 'None'}. Proceed.")]}
        elif decision == "modify":
            return {"messages": [HumanMessage(content=f"Human requested modification: '{feedback}'. Please revise the plan accordingly.")]}
        else: # Reject
            return {"messages": [HumanMessage(content=f"Human rejected the plan. Reason: '{feedback if feedback else 'None'}'. Halting operation.")]}
    # This case should ideally not be reached if human_input is guaranteed by the interrupt
    return {"messages": [HumanMessage(content="Error: Human feedback processing called without human_input in state.")]}


# Graph definition
agent_workflow_builder = StateGraph(AgentState) # Changed variable name for clarity
agent_workflow_builder.add_node("agent", call_model_node)
agent_workflow_builder.add_node("tools", call_tools_node)
agent_workflow_builder.add_node("prepare_oversight", prepare_oversight_payload)
agent_workflow_builder.add_node("human_oversight_trigger_node", human_oversight_node) # suspends the graph
agent_workflow_builder.add_node("process_human_feedback_node", process_human_feedback_node) # Renamed for clarity

# Conditional routing
def route_after_agent_logic(state: AgentState) -> str: # Renamed for clarity
    last_message = state["messages"][-1]
    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        return "tools"
    # Check for oversight AFTER LLM response (or tool use if that's a trigger)
    if needs_human_oversight(state): 
        return "prepare_oversight"
    return END # If no tools and no oversight needed

def route_after_human_feedback_logic(state: AgentState) -> str: # Renamed for clarity
    # Based on human feedback message, decide next step
    last_message_content = state["messages"][-1].content.lower() if state["messages"] and state["messages"][-1].content else ""
    if "rejected the plan" in last_message_content or "approved the plan" in last_message_content:
        return END
    # If "requested modification" or other cases that require more agent work
    return "agent" 

agent_workflow_builder.set_entry_point("agent")
agent_workflow_builder.add_conditional_edges("agent", route_after_agent_logic)
agent_workflow_builder.add_edge("tools", "agent") # Tools always return to agent for next step

# human_oversight_trigger_node calls interrupt; graph pauses. On resume, human_input is in state.
# The interrupt handler itself manages resumption to the *next* node in sequence if not conditional.
# Here, we want to explicitly route to process_human_feedback_node.
# The interrupt handler should add human_input, then this edge is followed.
agent_workflow_builder.add_edge("prepare_oversight", "human_oversight_trigger_node")
agent_workflow_builder.add_edge("human_oversight_trigger_node", "process_human_feedback_node")
agent_workflow_builder.add_conditional_edges("process_human_feedback_node", route_after_human_feedback_logic)


# Compile the agent
memory_agent = MemorySaver()
graph_agent = agent_workflow_builder.compile(checkpointer=memory_agent)

# Example Invocation:
# config_agent = {"configurable": {"thread_id": "agent-thread-1"}}
# initial_input_agent = {"messages": [HumanMessage(content="Analyze ACME Corp's financials for a large investment.")]}
# for event in graph_agent.stream(initial_input_agent, config_agent, stream_mode="values"):
#     print("--- Agent Stream Event ---")
#     for key, value in event.items():
#         if key == "messages":
#             print(f"Messages:")
#             for msg in value:
#                 print(f"  - Role: {msg.role if hasattr(msg, 'role') else type(msg)}, Content: {msg.content if hasattr(msg, 'content') else 'N/A'}")
#                 if hasattr(msg, 'tool_calls') and msg.tool_calls:
#                     print(f"    Tool Calls: {msg.tool_calls}")
#         else:
#             print(f"{key}: {value}")
#     print("--- End Agent Stream Event ---")
#     if "human_input" in event and event["human_input"]: # This check might be too early, human_input appears after resume
#         print(f"Human Input was processed: {event['human_input']}")

# To see the form URL (if HITL is triggered):
# response = graph_agent.invoke(initial_input_agent, config_agent)
# if response.get("hitl_form_url"): # Assuming interrupt handler adds this
#    print(f"KirokuForm URL for human review: {response.get('hitl_form_url')}")
# else:
#    print(f"Final Agent State: {response}")

from typing import TypedDict, Optional
from langgraph.graph import StateGraph, END
from kirokuforms import KirokuFormsHITL, create_kiroku_interrupt_node, resume_with_answers
import os
from fastapi import FastAPI, Request, BackgroundTasks
from pydantic import BaseModel
import uvicorn
import asyncio
import uuid # For generating unique thread_ids

# --- LangGraph Setup ---
# In production, use a persistent checkpointer (e.g., RedisSaver, SqliteSaver from langgraph.checkpoint.redis import RedisSaver)
# and a robust way to map KirokuForms task_id to LangGraph thread_id.
# For this example, we use a simple in-memory dictionary.
thread_task_map = {} 

# State for the workflow
class DocumentState(TypedDict):
    document_id: str
    document_content: str
    thread_id: str # Store thread_id in state for easier reference if needed
    hitl_task_id: Optional[str] # KirokuForms task ID
    hitl_form_url: Optional[str] # URL for the KirokuForm
    human_input: Optional[dict] # Data from KirokuForm submission
    status: str
    notes: Optional[str] # Notes from human review

# Initialize KirokuForms interrupt handler with webhook details
# Ensure your server is reachable at this webhook_url by KirokuForms service
WEBHOOK_BASE_URL = os.environ.get("WEBHOOK_BASE_URL", "http://localhost:8000") # Use ngrok for local dev
client = KirokuFormsHITL(
    api_key="YOUR_KIROKU_API_KEY", # Replace
    webhook_url=f"{WEBHOOK_BASE_URL}/webhook/kirokuforms_hitl_callback",
    webhook_secret="YOUR_SECURE_WEBHOOK_SECRET" # Replace
)

# The review step. interrupt() suspends the graph here, so this process is free
# while the reviewer takes however long they take; the webhook below is what
# wakes the thread up again.
#
# `data_key` names a dict on the state, and one field is generated per entry, so
# the reviewer sees this run's document rather than a fixed form. Fields passed
# with `fields=` are fixed when the node is built and cannot read state.
request_human_review_node = create_kiroku_interrupt_node(
    client,
    name="document-review",
    title="Async document review required",
    description="Please review this document and approve or reject it.",
    data_key="review_payload",
)

# Define workflow nodes
def start_document_processing_node(state: DocumentState) -> DocumentState:
    # Simulate initial processing
    print(f"Thread {state['thread_id']}: Starting processing for document {state['document_id']}")
    return {**state, "status": "processing_started"}

def prepare_review_payload(state: DocumentState) -> DocumentState:
    """Put what the reviewer needs to see where the review node can find it."""
    return {
        **state,
        "review_payload": {
            "document_id": state["document_id"],
            "document_content": state["document_content"][:2000],
        },
    }


def process_human_feedback_node(state: DocumentState) -> DocumentState:
    thread_id = state['thread_id']
    print(f"Thread {thread_id}: Processing human feedback for document {state['document_id']}")
    if state.get("human_input"):
        approved = state["human_input"].get("approved") == "yes"
        comments = state["human_input"].get("comments", "")
        final_status = "approved_by_human" if approved else "rejected_by_human"
        print(f"Thread {thread_id}: Document review result - Approved: {approved}, Comments: {comments}")
        return {**state, "status": final_status, "notes": comments}
    
    print(f"Thread {thread_id}: Error - No human_input found in state during feedback processing.")
    return {**state, "status": "error_no_human_input"}

# Build workflow
async_workflow_builder = StateGraph(DocumentState)
async_workflow_builder.add_node("start_processing", start_document_processing_node)
async_workflow_builder.add_node("prepare_review", prepare_review_payload)
async_workflow_builder.add_node("request_review", request_human_review_node) # suspends the graph
async_workflow_builder.add_node("process_feedback", process_human_feedback_node)

async_workflow_builder.set_entry_point("start_processing")
async_workflow_builder.add_edge("start_processing", "prepare_review")
async_workflow_builder.add_edge("prepare_review", "request_review")
# Graph pauses at request_review node (due to interrupt call with webhook).
# When webhook is received, it will resume graph execution which should proceed to process_feedback.
async_workflow_builder.add_edge("request_review", "process_feedback") 
async_workflow_builder.add_edge("process_feedback", END)

# Compile workflow
# For a real async system with webhooks, you MUST use a persistent checkpointer.
# Example: from langgraph.checkpoint.sqlite import SqliteSaver
# memory = SqliteSaver.from_conn_string(":memory:") # Or a file path
# graph_async = async_workflow_builder.compile(checkpointer=memory)
# For this simplified example, we will manage resumption manually without an explicit checkpointer in compile(),
# relying on the FastAPI handler to re-invoke with state. This is NOT robust for production.
# A proper checkpointer allows graph.get_state and graph.update_state.
# The KirokuForms interrupt handler is designed to work best with a checkpointer.
graph_async = async_workflow_builder.compile()


# --- FastAPI Application ---
app_fastapi = FastAPI()

class ProcessRequest(BaseModel):
    document_id: str
    content: str
    # thread_id: Optional[str] = None # Client can provide or we generate

class KirokuWebhookPayload(BaseModel):
    eventType: str
    taskId: str
    formId: str
    submissionId: str
    timestamp: str
    data: dict # Contains formData which has the human's input

async def run_graph_async_task(thread_id: str, initial_state: DocumentState):
    """Helper to run graph in background after initial invoke for webhook setup"""
    config = {"configurable": {"thread_id": thread_id}}
    # This first invoke runs until the interrupt (request_human_review_node)
    # The interrupt handler should populate hitl_task_id and hitl_form_url.
    # The graph execution pauses here.
    current_graph_state = graph_async.invoke(initial_state, config)
    
    kiroku_task_id = current_graph_state.get('hitl_task_id')
    form_url = current_graph_state.get('hitl_form_url')

    if kiroku_task_id:
        thread_task_map[kiroku_task_id] = thread_id
        print(f"POST /process_document (Thread: {thread_id}): Kiroku Task {kiroku_task_id} created. Form URL: {form_url}. Workflow paused awaiting webhook.")
    else:
        print(f"POST /process_document (Thread: {thread_id}): Error - Kiroku Task ID not found after interrupt call. State: {current_graph_state}")


@app_fastapi.post("/process_document")
async def start_doc_processing_endpoint(payload: ProcessRequest, background_tasks: BackgroundTasks):
    thread_id = str(uuid.uuid4()) # Generate a unique thread_id for this workflow instance
    
    initial_state = DocumentState(
        document_id=payload.document_id,
        document_content=payload.content,
        thread_id=thread_id,
        status="pending_start",
        hitl_task_id=None,
        hitl_form_url=None,
        human_input=None,
        notes=None
    )
    
    # Run the initial part of the graph in the background up to the HITL pause point
    # This ensures the API call returns quickly.
    background_tasks.add_task(run_graph_async_task, thread_id, initial_state)
    
    # The client will need to get the form_url through other means if not returned immediately,
    # or the system needs to notify the user. For this example, the form_url is logged by run_graph_async_task.
    return {
        "message": "Document processing initiated. Awaiting human review if required.",
        "thread_id": thread_id,
        "status_note": "Workflow will pause for human review; form URL will be logged by the server. Check server logs."
    }


@app_fastapi.post("/webhook/kirokuforms_hitl_callback")
async def kiroku_webhook_receiver_endpoint(payload: KirokuWebhookPayload, request: Request, background_tasks: BackgroundTasks):
    # IMPORTANT: Add robust webhook signature verification in production!
    # header_signature = request.headers.get("X-Kiroku-Signature")
    # raw_body = await request.body()
    # if not human_review_async_interrupt.verify_webhook_signature(raw_body, header_signature, "YOUR_SECURE_WEBHOOK_SECRET"):
    #     print("Webhook Error: Invalid signature")
    #     return {"status": "error", "detail": "Invalid signature"}, 401

    print(f"Webhook received - Kiroku Task ID: {payload.taskId}, Event: {payload.eventType}")

    if payload.eventType == "hitl.task.completed":
        kiroku_task_id = payload.taskId
        thread_id = thread_task_map.get(kiroku_task_id)

        if not thread_id:
            print(f"Webhook Error: No LangGraph thread_id found mapped for Kiroku Task ID: {kiroku_task_id}")
            return {"status": "error", "detail": "Thread ID not found for task."}

        human_form_data = payload.data.get("formData", {})
        print(f"Webhook for Thread {thread_id}: Human input received: {human_form_data}")
        
        # Prepare the input to resume the graph. This contains the human's submitted data.
        # The KirokuForms interrupt handler is designed to expect this in the 'human_input' field of the state.
        resume_input = {"human_input": human_form_data}
        
        config = {"configurable": {"thread_id": thread_id}}
        
        # Resume the graph by invoking it with the new input (human_form_data).
        # The graph's checkpointer (if configured) would load the state for thread_id,
        # merge this input, and continue execution from where it paused.
        # Since we're not using a persistent checkpointer directly in graph.compile for this simplified demo,
        # the 'invoke' here effectively continues the flow, and the state is managed in-memory by LangGraph for the thread.
        
        # To make this cleaner with a real checkpointer, you might do:
        # current_state = graph_async.get_state(config)
        # updated_values_for_state = {**current_state.values, "human_input": human_form_data}
        # graph_async.update_state(config, updated_values_for_state)
        # final_result = graph_async.invoke(None, config) # Invoke with None as input to continue with updated state

        # For simplicity with the current setup (no explicit checkpointer in compile):
        # The KirokuForms interrupt handler, when it resumes after a webhook,
        # typically handles updating the state internally before continuing the graph.
        # The following invoke passes the human_input which the already paused graph thread will pick up.
        
        final_result = graph_async.invoke(resume_input, config) # Pass human_input to be merged
        
        print(f"Webhook for Thread {thread_id}: Workflow resumed and completed. Final state: {final_result}")
        
        if kiroku_task_id in thread_task_map:
            del thread_task_map[kiroku_task_id] # Clean up map
        
        return {"status": "success_webhook_processed", "thread_id": thread_id, "final_status_from_graph": final_result.get("status")}
    
    return {"status": "webhook_event_ignored", "event_type": payload.eventType}

# To run this FastAPI app (save as, e.g., main.py):
# uvicorn main:app_fastapi --reload --port 8000
#
# And use ngrok for a public URL if testing KirokuForms webhooks from the cloud:
# ngrok http 8000
# Update WEBHOOK_BASE_URL with your ngrok https URL.