workflows-py/examples/visualization/resource_nodes_example.py

#!/usr/bin/env python3
"""
Example demonstrating resource nodes in workflow graph visualization.

This example shows how resources (dependencies injected via Annotated types)
are rendered in both Mermaid and Pyvis diagrams.

Run this script to generate:
- workflow_with_resources.mermaid - A Mermaid diagram file
- workflow_with_resources.html - An interactive Pyvis HTML visualization

You can view the Mermaid diagram at https://mermaid.live/ by pasting the contents.
The HTML file can be opened directly in any web browser.
"""

import argparse
from typing import Annotated

from llama_index.utils.workflow import (
    draw_all_possible_flows,
    draw_all_possible_flows_mermaid,
)
from workflows import Workflow, step
from workflows.events import Event, StartEvent, StopEvent
from workflows.resource import Resource

# --- Mock resource types ---


class DatabaseClient:
    """A database client for persistent storage."""

    def __init__(self, connection_string: str = "postgres://localhost/db"):
        self.connection_string = connection_string

    def query(self, sql: str) -> list:
        """Execute a SQL query."""
        return []


class CacheClient:
    """A cache client for fast data retrieval."""

    def __init__(self, host: str = "localhost", port: int = 6379):
        self.host = host
        self.port = port

    def get(self, key: str) -> str | None:
        """Get a value from cache."""
        return None

    def set(self, key: str, value: str) -> None:
        """Set a value in cache."""
        pass


class LLMClient:
    """A client for interacting with a large language model."""

    def __init__(self, api_key: str = "sk-..."):
        self.api_key = api_key

    async def complete(self, prompt: str) -> str:
        """Generate a completion for the given prompt."""
        return f"Response to: {prompt}"


# --- Resource factory functions ---


def get_database_client() -> DatabaseClient:
    """Factory function to create a database client.

    This function creates a PostgreSQL database client configured
    for the application's data storage needs.
    """
    return DatabaseClient(connection_string="postgres://localhost/myapp")


def get_cache_client() -> CacheClient:
    """Factory function to create a Redis cache client.

    Provides fast caching for frequently accessed data.
    """
    return CacheClient(host="localhost", port=6379)


def get_llm_client() -> LLMClient:
    """Factory function to create an LLM client.

    Creates a client for the language model API.
    """
    return LLMClient(api_key="sk-example-key")


# --- Event types ---


class QueryProcessedEvent(Event):
    """Event emitted after processing a user query."""

    query: str
    cached: bool = False


class ContextRetrievedEvent(Event):
    """Event emitted after retrieving context from the database."""

    context: str


class ResponseGeneratedEvent(Event):
    """Event emitted after generating an LLM response."""

    response: str


# --- Workflow with resources ---


class RAGWorkflow(Workflow):
    """A RAG (Retrieval-Augmented Generation) workflow demonstrating resource usage.

    This workflow shows how different steps can depend on shared resources
    like database clients, cache clients, and LLM clients.
    """

    @step
    async def process_query(
        self,
        ev: StartEvent,
        cache: Annotated[CacheClient, Resource(get_cache_client)],
    ) -> QueryProcessedEvent:
        """Process the incoming query, checking cache first."""
        query = getattr(ev, "query", "default query")

        # Check if query result is cached
        cached_result = cache.get(f"query:{query}")
        if cached_result:
            return QueryProcessedEvent(query=query, cached=True)

        return QueryProcessedEvent(query=query, cached=False)

    @step
    async def retrieve_context(
        self,
        ev: QueryProcessedEvent,
        db: Annotated[DatabaseClient, Resource(get_database_client)],
        cache: Annotated[CacheClient, Resource(get_cache_client)],
    ) -> ContextRetrievedEvent:
        """Retrieve relevant context from the database."""
        if ev.cached:
            context = "Cached context"
        else:
            # Query the database for relevant documents
            results = db.query(
                f"SELECT content FROM documents WHERE query = '{ev.query}'"
            )
            context = " ".join(str(r) for r in results) or "No context found"

            # Cache the result
            cache.set(f"context:{ev.query}", context)

        return ContextRetrievedEvent(context=context)

    @step
    async def generate_response(
        self,
        ev: ContextRetrievedEvent,
        llm: Annotated[LLMClient, Resource(get_llm_client)],
    ) -> ResponseGeneratedEvent:
        """Generate a response using the LLM with the retrieved context."""
        prompt = f"Context: {ev.context}\n\nGenerate a response."
        response = await llm.complete(prompt)
        return ResponseGeneratedEvent(response=response)

    @step
    async def finalize_response(
        self,
        ev: ResponseGeneratedEvent,
        cache: Annotated[CacheClient, Resource(get_cache_client)],
    ) -> StopEvent:
        """Finalize and cache the response."""
        # Cache the final response
        cache.set("last_response", ev.response)
        return StopEvent(result=ev.response)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate workflow visualizations with resource nodes"
    )
    parser.add_argument(
        "--output-dir",
        default=".",
        help="Directory to save output files (default: current directory)",
    )
    parser.add_argument(
        "--mermaid-only",
        action="store_true",
        help="Only generate Mermaid output (print to stdout)",
    )
    args = parser.parse_args()

    # Create the workflow
    workflow = RAGWorkflow()

    print("=" * 60)
    print("Workflow Graph Visualization with Resource Nodes")
    print("=" * 60)

    # Generate Mermaid diagram
    mermaid_file = f"{args.output_dir}/workflow_with_resources.mermaid"
    mermaid_output = draw_all_possible_flows_mermaid(
        workflow,
        filename="" if args.mermaid_only else mermaid_file,
    )

    print("\n--- Mermaid Diagram ---")
    print(mermaid_output)
    print()

    if not args.mermaid_only:
        print(f"Mermaid diagram saved to: {mermaid_file}")

        # Generate Pyvis HTML
        html_file = f"{args.output_dir}/workflow_with_resources.html"
        draw_all_possible_flows(workflow, filename=html_file)
        print(f"Interactive Pyvis diagram saved to: {html_file}")

    print("\n--- Resource Nodes Summary ---")
    print(
        """
The diagram shows:
- HEXAGON nodes (plum color): Resource dependencies
  - DatabaseClient: Database connection via get_database_client()
  - CacheClient: Cache connection via get_cache_client()
  - LLMClient: LLM API client via get_llm_client()

- Edge labels on resource connections show the variable name used in the step
  e.g., "db", "cache", "llm"

- Resources are deduplicated: CacheClient appears once even though
  it's used by multiple steps (process_query, retrieve_context, finalize_response)

To view the Mermaid diagram:
  1. Go to https://mermaid.live/
  2. Paste the diagram content above
  3. See the interactive visualization

To view the Pyvis diagram:
  1. Open workflow_with_resources.html in a web browser
  2. Hover over nodes to see metadata (type, getter, source location, docstring)
  3. Drag nodes to rearrange the layout
"""
    )


if __name__ == "__main__":
    main()