pixie-examples 0.1.1.dev3__py3-none-any.whl

examples/pydantic_ai/README.md

@@ -0,0 +1,246 @@
+# PydanticAI Examples with Pixie SDK
+
+This directory contains examples from the [PydanticAI Catalog](https://github.com/pydantic/pydantic-ai) integrated with Pixie SDK for observability and execution.
+
+## Examples from Catalog
+
+All examples are sourced from the official PydanticAI catalog and adapted for Pixie SDK.
+
+### 1. pydantic_model.py - Pydantic Model (Quickstart)
+
+Simple example demonstrating structured output extraction from text.
+
+**Pixie Handler:** `pydantic_model_example`
+**Documentation:** [pydantic-model](https://ai.pydantic.dev/examples/pydantic-model/)
+**Source:** [pydantic_model.py](https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/pydantic_model.py)
+
+- **Pattern:** Async function
+- **Input:** String query (e.g., "The windy city in the US of A.")
+- **Output:** Structured `MyModel` with city and country
+- **Features:**
+  - Basic agent setup with structured output
+  - Pydantic model validation
+  - Simple query → structured response
+
+### 2. bank_support.py - Bank Support (Multi-Turn Chatbot)
+
+Interactive bank support agent with tools and conversation history.
+
+**Pixie Handler:** `bank_support_agent`
+**Documentation:** [bank-support](https://ai.pydantic.dev/examples/bank-support/)
+**Source:** [bank_support.py](https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/bank_support.py)
+
+- **Pattern:** Async generator (multi-turn)
+- **Input:** None (interactive)
+- **Output:** Generator[str, str] for conversation
+- **Features:**
+  - Multi-turn conversation flow
+  - Tool usage (customer_balance)
+  - Dynamic system instructions
+  - SQLite database integration
+  - Interactive user input with `UserInputRequirement`
+
+### 3. flight_booking.py - Flight Booking (Multi-Agent)
+
+Multi-agent system for flight search, extraction, and seat selection.
+
+**Pixie Handler:** `flight_booking_example`
+**Documentation:** [flight-booking](https://ai.pydantic.dev/examples/flight-booking/)
+**Source:** [flight_booking.py](https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/flight_booking.py)
+
+- **Pattern:** Async generator (interactive)
+- **Input:** None (interactive)
+- **Output:** Generator[str, str] for booking flow
+- **Features:**
+  - **Agent Delegation:** search_agent delegates to extraction_agent
+  - **Programmatic Hand-off:** flight search → seat selection
+  - Multi-agent collaboration
+  - Usage limits and tracking
+  - Output validation with constraints
+  - Interactive confirmation workflow
+
+### 4. question_graph.py - Question Graph (Graph/State-Machine)
+
+Graph-based Q&A system using pydantic_graph for state machine control flow.
+
+**Pixie Handler:** `question_graph_example`
+**Documentation:** [question-graph](https://ai.pydantic.dev/examples/question-graph/)
+**Source:** [question_graph.py](https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/question_graph.py)
+
+- **Pattern:** Async generator (interactive with graph)
+- **Input:** None (interactive)
+- **Output:** Generator[str, str] for Q&A flow
+- **Features:**
+  - **pydantic_graph** for state machine control flow
+  - Graph nodes: Ask → Answer → Evaluate → Reprimand/End
+  - State management across nodes
+  - Conditional branching (correct/incorrect)
+  - Message history per agent
+
+**Graph Flow:**
+
+```
+Ask (generate question)
+  ↓
+Answer (get user input)
+  ↓
+Evaluate (check answer)
+  ↓
+  ├─ Correct → End (success)
+  └─ Wrong → Reprimand → Ask (retry)
+```
+
+### 5. sql_gen.py - SQL Generation (Multi-Step Workflow)
+
+SQL query generation with validation and structured output.
+
+**Pixie Handler:** `sql_gen_example`
+**Documentation:** [sql-gen](https://ai.pydantic.dev/examples/sql-gen/)
+**Source:** [sql_gen.py](https://github.com/pydantic/pydantic-ai/blob/main/examples/pydantic_ai_examples/sql_gen.py)
+
+- **Pattern:** Async function
+- **Input:** String query (natural language)
+- **Output:** String (SQL query or error message)
+- **Features:**
+  - Dynamic system prompts with SQL examples
+  - Output validation with `ModelRetry`
+  - Structured output with union types (`Success | InvalidRequest`)
+  - Multi-step validation workflow
+  - SQL safety checks (prevents dangerous operations)
+  - Schema-aware query generation
+
+## Setup
+
+### 1. Install Dependencies
+
+From the project root:
+
+```bash
+poetry install
+```
+
+### 2. Configure Environment Variables
+
+Copy `.env.example` to `.env` and add your API keys:
+
+```bash
+cp examples/pydantic_ai/.env.example examples/pydantic_ai/.env
+```
+
+Edit `.env` and add:
+
+```bash
+# Required for most examples
+OPENAI_API_KEY=sk-...
+
+# Optional
+LOGFIRE_TOKEN=...  # For observability
+```
+
+### 3. Start Pixie Server
+
+```bash
+poetry run pixie
+```
+
+The server will start at http://127.0.0.1:8000
+
+## Running Examples
+
+### Via GraphQL Playground
+
+Open http://127.0.0.1:8000/graphql and run subscriptions:
+
+#### Example 1: Pydantic Model
+
+```graphql
+subscription {
+  run(name: "pydantic_model_example", inputData: "The capital of France") {
+    runId
+    status
+    data
+  }
+}
+```
+
+#### Example 2: Bank Support (Interactive)
+
+```graphql
+subscription {
+  run(name: "bank_support_agent") {
+    runId
+    status
+    data
+  }
+}
+```
+
+#### Example 3: Flight Booking
+
+```graphql
+subscription {
+  run(name: "flight_booking_example") {
+    runId
+    status
+    data
+  }
+}
+```
+
+#### Example 4: Question Graph
+
+```graphql
+subscription {
+  run(name: "question_graph_example") {
+    runId
+    status
+    data
+  }
+}
+```
+
+#### Example 5: SQL Generation
+
+```graphql
+subscription {
+  run(name: "sql_gen_example", inputData: "show me error logs from yesterday") {
+    runId
+    status
+    data
+  }
+}
+```
+
+## Architecture Patterns
+
+### Async Functions vs Generators
+
+- **Async Function** (`async def → Type`): Single request-response (pydantic_model, sql_gen)
+- **Async Generator** (`async def → PixieGenerator[YieldType, SendType]`): Multi-turn (bank_support, flight_booking, question_graph)
+
+### Agent Patterns
+
+1. **Single Agent** - One agent handles workflow
+2. **Multi-Turn Chat** - Agent + user interaction loop
+3. **Agent Delegation** - One agent calls another
+4. **Programmatic Hand-off** - Sequential agent execution
+5. **Graph-Based Control** - pydantic_graph state machines
+
+## Integration with Pixie SDK
+
+All examples use:
+
+```python
+from pixie import app, PixieGenerator, UserInputRequirement
+
+@app
+async def my_handler(input: InputType) -> OutputType:
+    Agent.instrument_all()  # Enable observability
+    # ... implementation
+```
+
+## Resources
+
+- [PydanticAI Documentation](https://ai.pydantic.dev/)
+- [Pixie SDK Guide](../../.github/copilot-instructions.md)
+- [PydanticAI GitHub](https://github.com/pydantic/pydantic-ai)

examples/pydantic_ai/bank_support.py

@@ -0,0 +1,154 @@
+"""Bank support agent example using Pydantic AI.
+
+This example demonstrates:
+- Dynamic system prompts with agent instructions
+- Structured output types
+- Tools for querying dependencies
+- Integration with Pixie SDK
+
+Run with:
+    poetry run pixie
+
+Then query via GraphQL:
+    subscription {
+      run(name: "bank_support_agent", inputData: "What is my balance?") {
+        runId
+        status
+        data
+      }
+    }
+"""
+
+import sqlite3
+from dataclasses import dataclass
+from pydantic import BaseModel
+from pydantic_ai import Agent, RunContext
+import pixie
+
+
+@dataclass
+class DatabaseConn:
+    """A wrapper over the SQLite connection."""
+
+    sqlite_conn: sqlite3.Connection
+
+    async def customer_name(self, *, customer_id: int) -> str | None:
+        res = self.sqlite_conn.execute(
+            "SELECT name FROM customers WHERE id=?", (customer_id,)
+        )
+        row = res.fetchone()
+        if row:
+            return row[0]
+        return None
+
+    async def customer_balance(self, *, customer_id: int) -> float:
+        res = self.sqlite_conn.execute(
+            "SELECT balance FROM customers WHERE id=?", (customer_id,)
+        )
+        row = res.fetchone()
+        if row:
+            return row[0]
+        else:
+            raise ValueError("Customer not found")
+
+
+@dataclass
+class SupportDependencies:
+    customer_id: int
+    db: DatabaseConn
+
+
+class SupportOutput(BaseModel):
+    """Output model for bank support agent."""
+
+    support_advice: str
+    """Advice returned to the customer"""
+    block_card: bool
+    """Whether to block their card or not"""
+    risk: int
+    """Risk level of query"""
+
+
+# Create the support agent
+support_agent = Agent(
+    "openai:gpt-4o-mini",
+    deps_type=SupportDependencies,
+    output_type=SupportOutput,
+    instructions=(
+        "You are a support agent in our bank, give the "
+        "customer support and judge the risk level of their query. "
+        "Reply using the customer's name."
+    ),
+)
+
+
+@support_agent.instructions
+async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
+    customer_name = await ctx.deps.db.customer_name(customer_id=ctx.deps.customer_id)
+    return f"The customer's name is {customer_name!r}"
+
+
+@support_agent.tool
+async def customer_balance_tool(ctx: RunContext[SupportDependencies]) -> str:
+    """Returns the customer's current account balance."""
+    balance = await ctx.deps.db.customer_balance(
+        customer_id=ctx.deps.customer_id,
+    )
+    return f"${balance:.2f}"
+
+
+@pixie.app
+async def pydantic_ai_bank_support_agent() -> pixie.PixieGenerator[str, str]:
+    """Interactive bank support agent.
+
+    This agent helps customers with banking queries, can check balances,
+    and assesses risk levels of queries.
+
+    Yields:
+        str: Agent responses and support advice
+
+    Receives:
+        str: User queries via InputRequired
+    """
+
+    # Initialize database connection (in-memory for demo)
+    with sqlite3.connect(":memory:") as con:
+        cur = con.cursor()
+        cur.execute("CREATE TABLE customers(id, name, balance)")
+        cur.execute(
+            """
+            INSERT INTO customers VALUES
+                (123, 'John', 123.45),
+                (456, 'Jane', 567.89)
+        """
+        )
+        con.commit()
+
+        # Default customer
+        customer_id = 123
+        deps = SupportDependencies(
+            customer_id=customer_id, db=DatabaseConn(sqlite_conn=con)
+        )
+
+        yield "Welcome to Bank Support! I'm here to help you with your account."
+
+        while True:
+            # Get user query
+            user_query = yield pixie.InputRequired(str)
+
+            # Check for exit commands
+            if user_query.lower() in {"exit", "quit", "bye"}:
+                yield "Thank you for using our support service. Have a great day!"
+                break
+
+            # Run the agent
+            result = await support_agent.run(user_query, deps=deps)
+
+            # Format the response
+            output = result.output
+            response = f"""
+Support Advice: {output.support_advice}
+Card Status: {"🔒 BLOCKED" if output.block_card else "✓ Active"}
+Risk Level: {output.risk}/10
+"""
+            yield response.strip()

examples/pydantic_ai/flight_booking.py

@@ -0,0 +1,250 @@
+"""Example of a multi-agent flow where one agent delegates work to another.
+
+This demonstrates agent delegation and programmatic agent hand-off patterns.
+"""
+
+import datetime
+from dataclasses import dataclass
+from typing import Literal
+
+from pydantic import BaseModel, Field
+from pydantic_ai import (
+    Agent,
+    ModelMessage,
+    ModelRetry,
+    RunContext,
+    RunUsage,
+    UsageLimits,
+)
+import pixie
+
+
+class FlightDetails(BaseModel):
+    """Details of the most suitable flight."""
+
+    flight_number: str
+    price: int
+    origin: str = Field(description="Three-letter airport code")
+    destination: str = Field(description="Three-letter airport code")
+    date: datetime.date
+
+
+class NoFlightFound(BaseModel):
+    """When no valid flight is found."""
+
+
+@dataclass
+class Deps:
+    web_page_text: str
+    req_origin: str
+    req_destination: str
+    req_date: datetime.date
+
+
+# This agent is responsible for controlling the flow of the conversation.
+search_agent = Agent[Deps, FlightDetails | NoFlightFound](
+    "openai:gpt-4o-mini",
+    output_type=FlightDetails | NoFlightFound,  # type: ignore
+    retries=4,
+    system_prompt=(
+        "Your job is to find the cheapest flight for the user on the given date. "
+    ),
+)
+
+
+# This agent is responsible for extracting flight details from web page text.
+extraction_agent = Agent(
+    "openai:gpt-4o-mini",
+    output_type=list[FlightDetails],
+    system_prompt="Extract all the flight details from the given text.",
+)
+
+
+@search_agent.tool
+async def extract_flights(ctx: RunContext[Deps]) -> list[FlightDetails]:
+    """Get details of all flights."""
+    # we pass the usage to the search agent so requests within this agent are counted
+    agent_result = await extraction_agent.run(ctx.deps.web_page_text, usage=ctx.usage)
+    return agent_result.output
+
+
+@search_agent.output_validator
+async def validate_output(
+    ctx: RunContext[Deps], output: FlightDetails | NoFlightFound
+) -> FlightDetails | NoFlightFound:
+    """Procedural validation that the flight meets the constraints."""
+    if isinstance(output, NoFlightFound):
+        return output
+
+    errors: list[str] = []
+    if output.origin != ctx.deps.req_origin:
+        errors.append(
+            f"Flight should have origin {ctx.deps.req_origin}, not {output.origin}"
+        )
+    if output.destination != ctx.deps.req_destination:
+        errors.append(
+            f"Flight should have destination {ctx.deps.req_destination}, not {output.destination}"
+        )
+    if output.date != ctx.deps.req_date:
+        errors.append(f"Flight should be on {ctx.deps.req_date}, not {output.date}")
+
+    if errors:
+        raise ModelRetry("\n".join(errors))
+    else:
+        return output
+
+
+class SeatPreference(BaseModel):
+    row: int = Field(ge=1, le=30)
+    seat: Literal["A", "B", "C", "D", "E", "F"]
+
+
+class Failed(BaseModel):
+    """Unable to extract a seat selection."""
+
+
+# This agent is responsible for extracting the user's seat selection
+seat_preference_agent = Agent[None, SeatPreference | Failed](
+    "openai:gpt-4o-mini",
+    output_type=SeatPreference | Failed,
+    system_prompt=(
+        "Extract the user's seat preference. "
+        "Seats A and F are window seats. "
+        "Row 1 is the front row and has extra leg room. "
+        "Rows 14, and 20 also have extra leg room. "
+    ),
+)
+
+
+# in reality this would be downloaded from a booking site,
+# potentially using another agent to navigate the site
+flights_web_page = """
+1. Flight SFO-AK123
+- Price: $350
+- Origin: San Francisco International Airport (SFO)
+- Destination: Ted Stevens Anchorage International Airport (ANC)
+- Date: January 10, 2025
+
+2. Flight SFO-AK456
+- Price: $370
+- Origin: San Francisco International Airport (SFO)
+- Destination: Fairbanks International Airport (FAI)
+- Date: January 10, 2025
+
+3. Flight SFO-AK789
+- Price: $400
+- Origin: San Francisco International Airport (SFO)
+- Destination: Juneau International Airport (JNU)
+- Date: January 20, 2025
+
+4. Flight NYC-LA101
+- Price: $250
+- Origin: San Francisco International Airport (SFO)
+- Destination: Ted Stevens Anchorage International Airport (ANC)
+- Date: January 10, 2025
+
+5. Flight CHI-MIA202
+- Price: $200
+- Origin: Chicago O'Hare International Airport (ORD)
+- Destination: Miami International Airport (MIA)
+- Date: January 12, 2025
+
+6. Flight BOS-SEA303
+- Price: $120
+- Origin: Boston Logan International Airport (BOS)
+- Destination: Ted Stevens Anchorage International Airport (ANC)
+- Date: January 12, 2025
+
+7. Flight DFW-DEN404
+- Price: $150
+- Origin: Dallas/Fort Worth International Airport (DFW)
+- Destination: Denver International Airport (DEN)
+- Date: January 10, 2025
+
+8. Flight ATL-HOU505
+- Price: $180
+- Origin: Hartsfield-Jackson Atlanta International Airport (ATL)
+- Destination: George Bush Intercontinental Airport (IAH)
+- Date: January 10, 2025
+"""
+
+# restrict how many requests this app can make to the LLM
+usage_limits = UsageLimits(request_limit=15)
+
+
+@pixie.app
+async def pydantic_ai_flight_booking() -> pixie.PixieGenerator[str, str]:
+    """Multi-agent flight booking with search, extraction, and seat selection.
+
+    This example demonstrates:
+    - Agent delegation (search_agent -> extraction_agent)
+    - Programmatic agent hand-off (search -> seat selection)
+    - Interactive user input for flight confirmation and seat preference
+    """
+
+    deps = Deps(
+        web_page_text=flights_web_page,
+        req_origin="SFO",
+        req_destination="ANC",
+        req_date=datetime.date(2025, 1, 10),
+    )
+    message_history: list[ModelMessage] | None = None
+    usage: RunUsage = RunUsage()
+
+    yield "🛫 Welcome to Flight Booking! Searching for flights..."
+
+    # Flight search loop
+    while True:
+        agent_result = await search_agent.run(
+            f"Find me a flight from {deps.req_origin} to {deps.req_destination} on {deps.req_date}",
+            deps=deps,
+            usage=usage,
+            message_history=message_history,
+            usage_limits=usage_limits,
+        )
+
+        if isinstance(agent_result.output, NoFlightFound):
+            yield "❌ No flight found. Ending search."
+            break
+        else:
+            flight = agent_result.output
+            yield f"✈️ Flight found: {flight.flight_number}"
+            yield f"   Price: ${flight.price}"
+            yield f"   Route: {flight.origin} → {flight.destination}"
+            yield f"   Date: {flight.date}"
+            yield "\nDo you want to buy this flight, or keep searching? (buy/search)"
+
+            answer = yield pixie.InputRequired(str)
+            answer = answer.lower().strip()
+
+            if answer == "buy":
+                # Move to seat selection
+                yield "\n🪑 Great! Now let's select your seat..."
+                seat_message_history: list[ModelMessage] | None = None
+
+                while True:
+                    yield 'What seat would you like? (e.g., "row 1 seat A" or "window seat with leg room")'
+                    seat_input = yield pixie.InputRequired(str)
+
+                    seat_result = await seat_preference_agent.run(
+                        seat_input,
+                        message_history=seat_message_history,
+                        usage=usage,
+                        usage_limits=usage_limits,
+                    )
+
+                    if isinstance(seat_result.output, SeatPreference):
+                        seat = seat_result.output
+                        yield "\n✅ Booking confirmed!"
+                        yield f"   Flight: {flight.flight_number}"
+                        yield f"   Seat: Row {seat.row}, Seat {seat.seat}"
+                        yield f"   Total Cost: ${flight.price}"
+                        return
+                    else:
+                        yield "❌ Could not understand seat preference. Please try again."
+                        seat_message_history = seat_result.all_messages()
+            else:
+                yield "\n🔄 Searching for another flight..."
+                message_history = agent_result.all_messages(
+                    output_tool_return_content="Please suggest another flight"
+                )