Pipeline Composition

from pathlib import Path
import sys

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from shipit_agent import Agent, Pipeline, step, parallel
from examples.run_multi_tool_agent import build_llm_from_env
from IPython.display import display, Markdown

llm = build_llm_from_env("bedrock")

# Create specialized agents
researcher = Agent(
    llm=llm, prompt="You are a research expert. Provide key facts only, be concise."
)
writer = Agent(
    llm=llm, prompt="You are a technical writer. Write clear, engaging content."
)

# Chain them: research first, then write using the research
pipe = Pipeline.sequential(
    step("research", agent=researcher, prompt="Find 3 key facts about {topic}"),
    step(
        "write",
        agent=writer,
        prompt="Write a short paragraph using these facts:\n{research.output}",
    ),
)

result = pipe.run(topic="Python  asyncio library")

print("=== Research Output ===")
print(result.steps["research"].output[:300])
print("\n=== Final Written Output ===")
display(Markdown(result.output))

display(Markdown(result.output))

# Mix LLM agents with pure Python functions — no LLM call needed for deterministic transforms


def word_count(text: str) -> str:
    words = len(text.split())
    return f"Word count: {words}. Original text length: {len(text)} chars."


def uppercase(text: str) -> str:
    return text.upper()


pipe = Pipeline.sequential(
    step(
        "generate",
        agent=Agent(llm=llm, prompt="You are creative."),
        prompt="Write a haiku about {topic}",
    ),
    step("stats", fn=word_count),  # pure Python — no LLM call
    step("shout", fn=uppercase),  # another pure function
)

result = pipe.run(topic="coding")
print("Generated:", result.steps["generate"].output)
print("Stats:", result.steps["stats"].output)
print("Shouted:", result.output)

import time

# Run 3 agents in parallel, then merge results with a 4th agent
pros_agent = Agent(llm=llm, prompt="List only the pros/advantages. Be concise.")
cons_agent = Agent(llm=llm, prompt="List only the cons/disadvantages. Be concise.")
synthesizer = Agent(
    llm=llm, prompt="You are a balanced analyst. Synthesize both sides."
)

pipe = Pipeline(
    # These 2 steps run concurrently (parallel fan-out)
    parallel(
        step("pros", agent=pros_agent, prompt="Pros of {topic}"),
        step("cons", agent=cons_agent, prompt="Cons of {topic}"),
    ),
    # This step runs after both parallel steps complete
    step(
        "synthesis",
        agent=synthesizer,
        prompt="Combine these viewpoints:\n\nPros:\n{pros.output}\n\nCons:\n{cons.output}",
    ),
)

start = time.time()
result = pipe.run(topic="microservices architecture")
elapsed = time.time() - start

print(f"Completed in {elapsed:.1f}s\n")
print("=== Pros ===")
print(result.steps["pros"].output)
print("\n=== Cons ===")
print(result.steps["cons"].output)
print("\n=== Synthesis ===")
display(Markdown(result.output))

# Conditional routing: classify first, then route to the right handler

classifier = Agent(
    llm=llm,
    prompt="Classify the input as 'question' or 'statement'. Reply with just the word.",
)
question_agent = Agent(llm=llm, prompt="You answer questions helpfully and concisely.")
statement_agent = Agent(
    llm=llm, prompt="You acknowledge statements and add interesting context."
)

pipe = Pipeline.sequential(
    step("classify", agent=classifier, prompt="{input}"),
    step(
        "handle",
        router=lambda ctx: "question"
        if "question" in ctx["classify"].output.lower()
        else "statement",
        branches={
            "question": step("answer", agent=question_agent, prompt="{input}"),
            "statement": step("respond", agent=statement_agent, prompt="{input}"),
        },
    ),
)

# Test with a question
result = pipe.run(input="What is the speed of light?")
print("Classification:", result.steps["classify"].output.strip())
print("Route taken:", "question" if "answer" in result.steps else "statement")
display(Markdown(result.output))