graph-flow

repository·main·Indexed 43 days ago

https://github.com/a-agmon/rs-graph-llm

A type-safe, stateful graph workflow framework for building complex, interactive, and resumable AI agent workflows in Rust. Inspired by LangGraph and integrated with the Rig LLM ecosystem, it provides tools like GraphBuilder for structure definition, FlowRunner for execution, and SessionStorage for persistence via InMemory or Postgres implementations. It supports human-in-the-loop interactions through NextAction and ExecutionStatus, and includes built-in support for managing LLM chat history when the rig feature is enabled.

Tokens
21.7K
Snippets
56
Records
74
Agent score
56%

What's inside graph-flow

  1. Understand the Insurance Claims Workflow architecture

    main

    The Insurance Claims Workflow is a modular task-based system designed for processing insurance claims. It uses conditional routing to handle different insurance types and claim amounts.

    Workflow Flow:

    1. Initial Claim Query: Gathers initial information.
    2. Insurance Type Classifier: Routes the workflow to either Car or Apartment insurance paths.
    3. Insurance Details: Collects specific details based on the type (Car or Apartment).
    4. Smart Claim Validator & Approval Handler:
      • Claims < $1000: Automatically approved.
      • Claims ≥ $1000: Requires manual approval/decision before proceeding.
    5. Final Summary: The terminal endpoint that provides the outcome (Approved or Rejected).
    Initial Claim Query
           ↓
    Insurance Type Classifier
           ↓
        [Car | Apartment] Insurance Details
           ↓
    Smart Claim Validator & Approval Handler
      (Auto-approve <$1000 OR Wait for manual approval ≥$1000)
           ↓
       Final Summary
      (Approved OR Rejected)
  2. Understand the Recommendation Workflow

    main

    The service uses a multi-step workflow orchestrated by graph-flow to ensure high-quality recommendations. The process follows these stages:

    1. Query Refinement (QueryRefinementTask): Uses an LLM to rewrite the user's raw query into a semantically optimized version for vector search.
    2. Vector Search (VectorSearchTask): Generates embeddings for the refined query and performs a similarity search against the movies_with_vectors table, retrieving the top N movies (default 25).
    3. Answer Generation (AnswerGenerationTask): Uses the retrieved movie context and conversation history to generate a natural language recommendation using OpenRouter/GPT-4o-mini.
    4. Validation (ValidationTask): An LLM evaluates the generated answer for quality. If the answer is insufficient, it provides feedback.
    5. Retry Logic: If validation fails, the service automatically retries the generation (up to MAX_RETRIES times) using the validation feedback and chat history.
    6. Delivery (DeliveryTask): Returns the final validated recommendation and logs completion statistics.
  3. How the Medical Document Analysis workflow works

    main

    The service uses a graph-flow pattern consisting of 5 sequential tasks to process medical documents:

    1. PDF Extract: Converts PDF pages to high-resolution images (2000px width) using pdfium-render and performs OCR via GPT-4V to generate an initial summary.
    2. Human Review: Pauses the workflow to allow for human-in-the-loop feedback.
    3. Summary Integration: Combines the human feedback with the existing AI analysis.
    4. Research Search: Searches PubMed for relevant medical literature to augment the report.
    5. Final Report: Generates the comprehensive final medical analysis report.
  4. Control workflow execution with NextAction

    main

    Tasks control the flow of the graph by returning a NextAction variant in their TaskResult.

    VariantBehavior
    ContinueAdvance one edge, return control to the caller
    ContinueAndExecuteAdvance and keep executing until something pauses
    WaitForInputPark the workflow until new input arrives
    GoTo(task_id)Jump to a specific task
    EndComplete the workflow

    Guard rails:

    • GraphBuilder::with_task_timeout(duration): Sets a maximum duration for each task (default is 5 minutes).
    • GraphBuilder::with_max_execution_steps(n): Prevents infinite loops by turning a ContinueAndExecute cycle into an error after n steps.
  5. Understand ExecutionStatus and ExecutionResult

    main

    Every execution returns an ExecutionResult containing the task's response and an ExecutionStatus. The status tells you the current state of the workflow:

    • Paused { next_task_id, reason }: The task finished, and the workflow is ready to proceed to next_task_id on the next run (returned for Continue and GoTo).
    • WaitingForInput: The workflow is parked on the current task until you provide input and run it again.
    • Completed: The workflow has reached the End state.

    Failures are returned as Err(GraphError) from the execution call.

  6. Use GraphBuilder for immutable Graph construction

    main

    In graph-flow, the Graph object is immutable after it is created. All configuration and construction methods exist only on the GraphBuilder.

    • Builder Methods: Use add_task, add_edge, add_conditional_edge, and set_start_task on the GraphBuilder instance.
    • Result: Calling .build() consumes the builder and produces an immutable Graph. This design avoids lock overhead (like Mutex or DashMap) and prevents lock poisoning during lookups.
  7. Manage state with the `Context` object

    main

    The Context is a thread-safe, typed, and serialized key-value store used to pass data between tasks and maintain session state.

    Key capabilities:

    • Typed Storage: Use .set(key, value)? and .get::<T>(key) for type-safe state management. Note that .set() is fallible due to potential serialization errors.
    • Chat History: Built-in support for LLM conversation history via .add_user_message(string) and .add_assistant_message(string). History is stored in a ring buffer (default cap: 1000 messages).
    • Synchronous API: All Context methods are synchronous, making them safe to use in both async tasks and synchronous edge-condition closures.
    // Typed key-value state; set() is fallible (serialization can fail)
    context.set("claim_amount", 1500.0)?;
    let amount: f64 = context.get::<f64>("claim_amount").unwrap();
    
    // Works directly in edge conditions
    let ok = context.get::<bool>("condition").unwrap_or(false);
    
    // Built-in chat history (ring buffer, capped at 1000 messages by default)
    context.add_user_message("What's my claim status?".to_string());
    context.add_assistant_message("Your claim is being processed".to_string());
    let recent = context.get_last_messages(5);
  8. Production considerations for graph execution

    main

    When deploying graph-flow in production, keep the following behaviors in mind:

    • Concurrency & Optimistic Locking: Sessions use an optimistic-locking version. If two processes attempt to save the same session simultaneously, the losing attempt will return GraphError::SessionConflict. You should reload the session and retry. Note that task side effects from the losing attempt are not rolled back; ensure tasks are idempotent.
    • Timeouts & Limits: Use .with_task_timeout to prevent slow LLM calls from hanging the workflow. Use .with_max_execution_steps to prevent infinite loops in graphs containing cycles (e.g., via ContinueAndExecute).
    • Durability: Within a ContinueAndExecute chain, the session is only saved once the entire chain finishes. If the process crashes mid-chain, the session will be re-run from the first task of that chain upon recovery.
  9. Control graph advancement with NextAction

    main

    A task controls how the graph moves forward by returning a NextAction. This allows you to mix automated execution with interactive pauses.

    NextActionBehavior
    ContinueAdvance one edge, then return control to the caller (step-by-step/interactive mode).
    ContinueAndExecuteAdvance and keep executing tasks in the same call until a pause or end (fire-and-forget).
    WaitForInputPause at the current task until new input arrives.
    GoTo(task_id)Jump to a specific task.
    EndComplete the workflow.
    // Step-by-step: caller decides when the next task runs
    Ok(TaskResult::new(Some("Done".into()), NextAction::Continue))
    
    // Continuous: single runner.run() call executes the whole chain
    Ok(TaskResult::new(None, NextAction::ContinueAndExecute))
  10. Install graph-flow

    main

    Add graph-flow to your Cargo.toml. By default, it includes postgres support. If you need LLM integration helpers via Rig, enable the rig feature. If you want to exclude PostgreSQL, set default-features = false.

    [dependencies]
    graph-flow = { version = "0.8", features = ["rig"] }  # drop "rig" if you don't need LLM helpers
  11. Set up the Movie Recommendation Service

    main

    The Movie Recommendation Service is an Axum-based web service that provides AI-powered movie recommendations using a multi-step workflow. It requires two PostgreSQL databases: one for session state and one for movie data with vector embeddings.

    Prerequisites

    • PostgreSQL with pgvector extension installed.
    • An OpenRouter API key.

    Environment Variables

    You must configure the following environment variables:

    VariableRequirementDescription
    DATABASE_URLRequiredConnection string for the Sessions database
    MOVIES_DATABASE_URLRequiredConnection string for the Movies database
    OPENROUTER_API_KEYRequiredYour OpenRouter API key
    RUST_LOGOptionalLogging level (e.g., info)
    MAX_RETRIESOptionalMax validation retry attempts (default: 3)
    CHAT_HISTORY_LIMITOptionalMax chat messages to retain (default: 50)
    VECTOR_SEARCH_LIMITOptionalNumber of movies to retrieve (default: 25)

    Database Schema Setup

    Run the following SQL to prepare the movies database:

    CREATE TABLE movies_with_vectors (
        id SERIAL PRIMARY KEY,
        title TEXT NOT NULL,
        overview TEXT,
        vector vector(384)  -- Adjust dimension based on your embedding model
    );
    
    -- Add vector similarity search index
    CREATE INDEX movies_vector_idx ON movies_with_vectors USING ivfflat (vector vector_cosine_ops);
    # Install dependencies
    cargo build
    
    # Set environment variables
    export DATABASE_URL="postgresql://user:password@localhost:5432/sessions_db"
    export MOVIES_DATABASE_URL="postgresql://user:password@localhost:5432/movies_db"
    export OPENROUTER_API_KEY="your_api_key"
    
    # Run the service
    cargo run
  12. Install graph-flow and Rig integration

    main

    To use graph-flow with LLM capabilities, add it to your Cargo.toml along with rig-core and rig-agent. Enabling the rig feature in graph-flow allows for seamless integration between the workflow Context and Rig's message types.

    Note: PostgreSQL support is enabled by default via the postgres feature. If you want to exclude it, set default-features = false.

    [dependencies]
    graph-flow = { version = "0.8", features = ["rig"] }  # drop "rig" if you don't need LLM helpers
    rig-core = "0.42"                                     # provider clients + Message
    rig-agent = "0.42"                                    # the agent runtime (Chat, Prompt, .agent())