[zephyrzilla]

Summary

This PR unifies the Workflow and Graph interfaces for workflow module to run/create GraSP pipelines.

How to Test the feature

Example 1:

import grasp

workflow = grasp.Workflow("tasks/examples/glaive_code_assistant")
workflow.run(num_records=1)

Example 2:

import grasp
from grasp.core.graph.functions.node_processor import NodePostProcessorWithState
from grasp.core.graph.grasp_state import GraspState
from grasp import GraspMessage


class TextExtractionPostProcessor(NodePostProcessorWithState):
    def apply(self, response: GraspMessage, state: GraspState) -> GraspState:
        content = response.message.content
        state["haiku_response"] = content
        return state


data = [
    {"text": "Hello, world!"},
    {"text": "Beautiful sunset today"},
    {"text": "Programming is fun"},
    {"text": "Nature is amazing"},
    {"text": "Coffee smells great"},
    {"text": "Books are knowledge"},
    {"text": "Music soothes soul"},
    {"text": "Travel broadens mind"},
    {"text": "Friends bring joy"},
    {"text": "Learning never stops"},
]

workflow = grasp.Workflow("my_workflow")

result = (
    workflow.source(data)
    .llm(
        model="gpt-4o",
        prompt="Create a Haiku: {text}",
        output="haiku_response",
        post_process=TextExtractionPostProcessor,
    )
    .sink("outputs/result.json")
    .run(num_records=1)
)

print(f"Processed {len(result) if result else 0} records")
if result:
    print(f"Sample result: {result[0]}")

Example 3 (Graph Creation):

from grasp import Workflow, DataSource, DataSink

graph = Workflow("simple_example")

graph.source(
    DataSource.memory(
        [
            {"id": 1, "text": "Machine learning transforms data into insights."},
            {"id": 2, "text": "Python provides powerful tools for AI development."},
        ]
    )
)

summarizer = (
    graph.add_llm_node("summarizer", "gpt-4o")
    .system_message("You create concise summaries.")
    .user_message("Summarize: {text}")
    .output_keys("summary")
)

enhancer = (
    graph.add_llm_node("enhancer", "gpt-4o")
    .system_message("You improve text clarity.")
    .user_message("Enhance: {summary}")
    .output_keys("enhanced")
)

graph.sequence("summarizer", "enhancer")

graph.sink(DataSink.disk("output/results_graph.jsonl"))

executable = graph.build()
results = executable.run()

print(results)

Example 4 (Conditional Edges and Output Mapper):

import random
from typing import Any

import grasp
from grasp.core.graph.functions.edge_condition import EdgeCondition
from grasp.core.graph.grasp_state import GraspState
from grasp.processors.output_record_generator import BaseOutputGenerator


class SimpleRouteCondition(EdgeCondition):
    """Simple condition that routes to different paths."""

    @staticmethod
    def apply(state: GraspState) -> str:
        if random.random() < 0.5:
            return "even_path"
        else:
            return "odd_path"


class ConditionalOutputGenerator(BaseOutputGenerator):
    """Custom output generator for conditional workflows."""

    def generate(self, state: GraspState) -> dict[str, Any]:
        """Generate clean output record."""
        messages = state.get("messages", [])
        final_response = (
            messages[-1].content
            if messages and hasattr(messages[-1], "content")
            else ""
        )

        return {
            "id": state.get("id"),
            "original_text": state.get("text"),
            "processing_path": "even" if state.get("id", 0) % 2 == 0 else "odd",
            "response": final_response,
            "workflow": "conditional_routing",
        }


graph = grasp.Workflow("simple_conditional")

graph.source(
    [
        {"id": 1, "text": "First message"},
        {"id": 2, "text": "Second message"},
        {"id": 3, "text": "Third message"},
        {"id": 4, "text": "Fourth message"},
    ]
)

classifier = (
    graph.add_llm_node("classifier", "gpt-4o")
    .system_message("You analyze messages.")
    .user_message("Analyzing: {text}")
    .temperature(0.3)
)

even_handler = (
    graph.add_llm_node("even_handler", "gpt-4o")
    .system_message("You handle even-numbered items.")
    .user_message("Processing even item #{id}: {text}")
    .temperature(0.3)
)

odd_handler = (
    graph.add_llm_node("odd_handler", "gpt-4o")
    .system_message("You handle odd-numbered items.")
    .user_message("Processing odd item #{id}: {text}")
    .temperature(0.3)
)

graph.add_edge("START", "classifier")

graph.add_conditional_edge(
    from_node="classifier",
    condition=SimpleRouteCondition,
    path_map={
        "even_path": "even_handler",
        "odd_path": "odd_handler",
    },
)

graph.add_edge("even_handler", "END")
graph.add_edge("odd_handler", "END")
(
    graph.preserve_fields("id", "question").output_metadata(
        {"workflow": "conditional", "version": "2.0"}
    )
)

graph.sink("output/conditional_results.jsonl")

executable = graph.build()
results = executable.run()

print(results)

Checklist

• Lint fixes and unit testing done
• End to end task testing
• Documentation updated

[amitsnow]

Same review as Sriram. Otherwise LGTM

[psriramsnc]

LGTM