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

examples/pydantic_ai/question_graph.py

@@ -0,0 +1,152 @@
+"""Example of a graph-based state machine for asking and evaluating questions.
+
+This demonstrates using pydantic_graph for complex control flow.
+"""
+
+from dataclasses import dataclass, field
+
+from pydantic import BaseModel
+from pydantic_ai import Agent, ModelMessage, format_as_xml
+from pydantic_graph import (
+    BaseNode,
+    End,
+    Graph,
+    GraphRunContext,
+)
+import pixie
+
+# Agent for asking questions
+ask_agent = Agent("openai:gpt-4o-mini", output_type=str)
+
+
+@dataclass
+class QuestionState:
+    question: str | None = None
+    answer: str | None = None
+    ask_agent_messages: list[ModelMessage] = field(default_factory=list)
+    evaluate_agent_messages: list[ModelMessage] = field(default_factory=list)
+
+
+@dataclass
+class Ask(BaseNode[QuestionState]):
+    """Generate a question using the AI."""
+
+    async def run(self, ctx: GraphRunContext[QuestionState]) -> "Answer":
+        agent_result = await ask_agent.run(
+            "Ask a simple question with a single correct answer.",
+            message_history=ctx.state.ask_agent_messages,
+        )
+        ctx.state.ask_agent_messages += agent_result.all_messages()
+        ctx.state.question = agent_result.output
+        ctx.state.answer = None
+        return Answer(agent_result.output)
+
+
+@dataclass
+class Answer(BaseNode[QuestionState]):
+    """Wait for user to provide an answer."""
+
+    question: str
+
+    async def run(self, ctx: GraphRunContext[QuestionState]) -> "Evaluate":
+        if ctx.state.answer is None:
+            raise ValueError("Answer need to be set after Ask() node in the main loop.")
+        return Evaluate(ctx.state.answer)
+
+
+class EvaluationOutput(BaseModel, use_attribute_docstrings=True):
+    correct: bool
+    """Whether the answer is correct."""
+    comment: str
+    """Comment on the answer, reprimand the user if the answer is wrong."""
+
+
+# Agent for evaluating answers
+evaluate_agent = Agent(
+    "openai:gpt-4o-mini",
+    output_type=EvaluationOutput,
+    system_prompt="Given a question and answer, evaluate if the answer is correct.",
+)
+
+
+@dataclass
+class Evaluate(BaseNode[QuestionState, None, str]):
+    """Evaluate the user's answer."""
+
+    answer: str
+
+    async def run(
+        self,
+        ctx: GraphRunContext[QuestionState],
+    ) -> End[str] | "Reprimand":
+        assert ctx.state.question is not None
+        agent_result = await evaluate_agent.run(
+            format_as_xml({"question": ctx.state.question, "answer": self.answer}),
+            message_history=ctx.state.evaluate_agent_messages,
+        )
+        ctx.state.evaluate_agent_messages += agent_result.all_messages()
+        if agent_result.output.correct:
+            return End(agent_result.output.comment)
+        else:
+            return Reprimand(agent_result.output.comment)
+
+
+@dataclass
+class Reprimand(BaseNode[QuestionState]):
+    """Tell the user they got it wrong and ask another question."""
+
+    comment: str
+
+    async def run(self, ctx: GraphRunContext[QuestionState]) -> Ask:
+        ctx.state.question = None
+        return Ask()
+
+
+# Create the question graph
+_question_graph = Graph(
+    nodes=(Ask, Answer, Evaluate, Reprimand), state_type=QuestionState
+)
+
+
+@pixie.app
+async def pydantic_ai_question_graph() -> pixie.PixieGenerator[str, str]:
+    """Interactive Q&A game using graph-based state machine.
+
+    The AI asks questions, the user answers, and the AI evaluates.
+    If wrong, user gets reprimanded and a new question is asked.
+    """
+
+    yield "🧠 Welcome to the Q&A Challenge!"
+    yield "I'll ask you questions and evaluate your answers.\n"
+
+    state = QuestionState()
+
+    # Use the graph's run_sync method which handles node execution
+    async with _question_graph.iter(Ask(), state=state) as graph_ctx:
+        while True:
+            node = await graph_ctx.next()
+
+            # Handle different node types
+            if isinstance(node, Ask):
+                yield "🤔 Generating a question..."
+
+            elif isinstance(node, Answer):
+                # Ask user for their answer
+                yield f"\n❓ Question: {node.question}"
+                yield "What is your answer?"
+                state.answer = yield pixie.InputRequired(str)
+                # Continue with the user's answer by passing it to the next iteration
+                # The graph will automatically move to Evaluate with this answer
+                continue
+
+            elif isinstance(node, Evaluate):
+                yield "⏳ Evaluating your answer..."
+
+            elif isinstance(node, Reprimand):
+                yield f"\n❌ {node.comment}"
+                yield "\n🔄 Let's try another question...\n"
+
+            elif isinstance(node, End):
+                yield f"\n✅ {node.data}"
+                yield "\n🎉 Congratulations! You got it right!"
+                return

examples/pydantic_ai/sql_gen.py

@@ -0,0 +1,182 @@
+"""Example demonstrating SQL generation using Pydantic AI.
+
+This shows multi-step workflow with validation and structured output.
+"""
+
+from dataclasses import dataclass
+from datetime import date
+from typing import Annotated, TypeAlias
+
+from annotated_types import MinLen
+from pydantic import BaseModel, Field
+from pydantic_ai import Agent, ModelRetry, RunContext, format_as_xml
+import pixie
+
+
+# Simulated database schema
+DB_SCHEMA = """
+CREATE TABLE records (
+    id bigint PRIMARY KEY,
+    created_at timestamptz,
+    start_timestamp timestamptz,
+    end_timestamp timestamptz,
+    trace_id text,
+    span_id text,
+    parent_span_id text,
+    level log_level,
+    span_name text,
+    message text,
+    attributes_json_schema text,
+    attributes jsonb,
+    tags text[],
+    is_exception boolean,
+    otel_status_message text,
+    service_name text
+);
+"""
+
+SQL_EXAMPLES = [
+    {
+        "request": "show me records where foobar is false",
+        "response": "SELECT * FROM records WHERE attributes->>'foobar' = false",
+    },
+    {
+        "request": 'show me records where attributes include the key "foobar"',
+        "response": "SELECT * FROM records WHERE attributes ? 'foobar'",
+    },
+    {
+        "request": "show me records from yesterday",
+        "response": "SELECT * FROM records WHERE start_timestamp::date > CURRENT_TIMESTAMP - INTERVAL '1 day'",
+    },
+    {
+        "request": 'show me error records with the tag "foobar"',
+        "response": "SELECT * FROM records WHERE level = 'error' and 'foobar' = ANY(tags)",
+    },
+]
+
+
+@dataclass
+class Deps:
+    """Dependencies for SQL generation (simulated database connection)."""
+
+    validate_sql: bool = True
+
+
+class Success(BaseModel):
+    """Response when SQL could be successfully generated."""
+
+    sql_query: Annotated[str, MinLen(1)]
+    explanation: str = Field(
+        "", description="Explanation of the SQL query, as markdown"
+    )
+
+
+class InvalidRequest(BaseModel):
+    """Response when the user input didn't include enough information to generate SQL."""
+
+    error_message: str
+
+
+Response: TypeAlias = Success | InvalidRequest
+
+agent = Agent[Deps, Response](
+    "openai:gpt-4o-mini",
+    output_type=Response,  # type: ignore
+    deps_type=Deps,
+)
+
+
+@agent.system_prompt
+async def system_prompt() -> str:
+    return f"""\
+Given the following PostgreSQL table of records, your job is to
+write a SQL query that suits the user's request.
+
+Database schema:
+
+{DB_SCHEMA}
+
+today's date = {date.today()}
+
+{format_as_xml(SQL_EXAMPLES)}
+"""
+
+
+@agent.output_validator
+async def validate_output(ctx: RunContext[Deps], output: Response) -> Response:
+    if isinstance(output, InvalidRequest):
+        return output
+
+    # Clean up the SQL query
+    output.sql_query = output.sql_query.replace("\\", "")
+    if not output.sql_query.upper().startswith("SELECT"):
+        raise ModelRetry("Please create a SELECT query")
+
+    # In a real scenario, we'd validate against an actual database
+    # For now, we'll do basic validation
+    if ctx.deps.validate_sql:
+        # Check for dangerous operations
+        upper_query = output.sql_query.upper()
+        if any(
+            danger in upper_query
+            for danger in ["DROP", "DELETE", "UPDATE", "INSERT", "TRUNCATE"]
+        ):
+            raise ModelRetry("Only SELECT queries are allowed")
+
+        # Check for required table reference
+        if "FROM RECORDS" not in upper_query and 'FROM "RECORDS"' not in upper_query:
+            raise ModelRetry("Query must reference the records table")
+
+    return output
+
+
+@pixie.app
+async def pydantic_ai_sql_gen(query: str) -> str:
+    """SQL generation agent with multi-step workflow and validation.
+
+    This example demonstrates:
+    - Dynamic system prompts with examples
+    - Output validation with ModelRetry
+    - Structured output with union types
+    - Agent dependencies for configuration
+
+    Args:
+        query: Natural language description of desired SQL query
+
+    Returns:
+        Generated SQL query or error message
+    """
+
+    deps = Deps(validate_sql=True)
+    result = await agent.run(query, deps=deps)
+
+    if isinstance(result.output, Success):
+        response = f"✅ SQL Query Generated:\n\n{result.output.sql_query}"
+        if result.output.explanation:
+            response += f"\n\n📝 Explanation:\n{result.output.explanation}"
+        return response
+    else:
+        return f"❌ Error: {result.output.error_message}"
+
+
+# For local testing
+async def test():
+    """Test function for local development."""
+    test_queries = [
+        "show me logs from yesterday with level 'error'",
+        "find all records where foobar is true",
+        "show me records from the last week",
+    ]
+
+    for query in test_queries:
+        print(f"\n{'='*60}")
+        print(f"Query: {query}")
+        print("=" * 60)
+        result = await pydantic_ai_sql_gen(query)
+        print(result)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(test())

examples/pydantic_ai/structured_output.py

@@ -0,0 +1,64 @@
+"""Simple example of using Pydantic AI to construct a Pydantic model from a text input.
+
+This example demonstrates:
+- Using Pydantic AI with structured output
+- Integration with Pixie SDK for observability
+
+Run with:
+    poetry run pixie
+
+Then query via GraphQL:
+    subscription {
+      run(name: "pydantic_model_example", inputData: "The windy city in the US of A.") {
+        runId
+        status
+        data
+      }
+    }
+"""
+
+import os
+from pydantic import BaseModel
+from pydantic_ai import Agent
+import pixie
+
+
+class MyModel(BaseModel):
+    """Model representing a city and country."""
+
+    city: str
+    country: str
+
+
+# Create the agent
+model = os.getenv("PYDANTIC_AI_MODEL", "openai:gpt-4o-mini")
+agent = Agent(model, output_type=MyModel)
+
+
+@pixie.app
+async def pydantic_ai_structured_output(query: str) -> MyModel:
+    """Extract city and country information from a text query.
+
+    Args:
+        query: Natural language text describing a city
+
+    Returns:
+        MyModel with extracted city and country
+    """
+
+    # Run the agent
+    agent_result = await agent.run(query)
+
+    # Return the structured output
+    return agent_result.output
+
+
+if __name__ == "__main__":
+    # For testing locally
+    import asyncio
+
+    async def test():
+        output = await pydantic_ai_structured_output("The windy city in the US of A.")
+        print(output)
+
+    asyncio.run(test())

examples/quickstart/chatbot.py

@@ -0,0 +1,25 @@
+from pydantic_ai import Agent
+
+import pixie
+
+agent = Agent(
+    name="Simple chatbot",
+    instructions="You are a helpful assistant.",
+    model="gpt-4o-mini",
+)
+
+
+@pixie.app
+async def example_chatbot():
+    """A simple chatbot using Pydantic-AI agent with GPT-4o-mini.
+
+    An OpenAI API key environment variable *(`OPENAI_API_KEY`)* is required to run this example.
+    """
+
+    yield "How can I help you today?"
+    messages = []
+    while True:
+        user_msg = yield pixie.InputRequired(str)
+        response = await agent.run(user_msg, message_history=messages)
+        messages = response.all_messages()
+        yield response.output

examples/quickstart/sleepy_poet.py

@@ -0,0 +1,96 @@
+"""An interactive agent that writes haikus while taking naps."""
+
+import asyncio
+from dataclasses import dataclass
+import logging
+from pydantic import BaseModel, Field
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.openai import OpenAIChatModelSettings
+
+import pixie
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PoetDeps:
+    """Dependencies for the sleepy poet agent."""
+
+    queue: asyncio.Queue
+    time_to_write: bool = False
+
+
+poet = Agent(
+    "openai:gpt-4o-mini",
+    system_prompt=(
+        "You are a sleepy poet. When given a topic, you first take a nap by calling the "
+        "`sleep_for_a_bit` tool once and once only, then write a haiku about the topic."
+    ),
+    deps_type=PoetDeps,
+    model_settings=OpenAIChatModelSettings(
+        parallel_tool_calls=False  # Ensure single sleep tool call per turn
+    ),
+)
+
+
+@poet.tool
+async def sleep_for_a_bit(ctx: RunContext[PoetDeps]) -> str:
+    """Sleep for 5 seconds before writing a haiku."""
+    if ctx.deps.time_to_write:
+        return "I've already napped, It's time to write the haiku now."
+
+    await ctx.deps.queue.put("Let me take a nap...")
+    for i in range(3):
+        await asyncio.sleep(1)
+        await ctx.deps.queue.put("z" * (i + 1))
+    ctx.deps.time_to_write = True
+    return "Poet napped for 3 seconds."
+
+
+class HaikuRequest(BaseModel):
+    """Input config for sleepy haiku agent."""
+
+    topic: str = Field(..., description="The topic to write haikus about.")
+    count: int = Field(3, description="Number of haikus to write.")
+
+
+@pixie.app
+async def example_sleepy_poet() -> pixie.PixieGenerator[str, HaikuRequest]:
+    """An interactive agent that writes haikus while taking naps.
+
+    This agent accepts a topic and number of haikus to write, then for each haiku,
+    it first calls a tool to sleep for a few seconds before composing the haiku.
+
+    Yields:
+        HaikuRequest: Requirement for user's input for the haiku request.
+        str: The generated haikus interleaved with sleep updates.
+    """
+
+    q = asyncio.Queue[str | None]()
+    deps = PoetDeps(queue=q)
+
+    yield "I'm the best poet in town! I can write haikus on any topic."
+
+    while True:
+        yield "What's your request?"
+        config = yield pixie.InputRequired(HaikuRequest)
+
+        async def write_haikus():
+            for i in range(0, config.count):
+                result = await poet.run(config.topic, deps=deps)
+                await q.put(f"### Haiku #{i+1}\n{result.output}\n")
+                deps.time_to_write = False  # Reset for next haiku
+
+            await q.put(None)  # Signal completion
+
+        task = asyncio.create_task(write_haikus())
+        try:
+            while True:
+                update = await q.get()
+                if update is None:
+                    break
+                yield update
+
+        except asyncio.CancelledError:
+            task.cancel()
+            return

examples/quickstart/weather_agent.py

@@ -0,0 +1,110 @@
+"""Interactive weather agent."""
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any
+from httpx import AsyncClient
+from pydantic import BaseModel
+from pydantic_ai import Agent, RunContext
+import pixie
+
+
+@dataclass
+class Deps:
+    """Dependencies for the weather agent."""
+
+    client: AsyncClient
+
+
+class LatLng(BaseModel):
+    """Latitude and longitude coordinates."""
+
+    lat: float
+    lng: float
+
+
+# Create the weather agent
+agent = Agent(
+    "openai:gpt-4o-mini",
+    instructions="Be concise, reply with one sentence.",
+    deps_type=Deps,
+    retries=2,
+)
+
+
+@agent.tool
+async def get_lat_lng(ctx: RunContext[Deps], location_description: str) -> LatLng:
+    """Get the latitude and longitude of a location.
+
+    Args:
+        ctx: The context.
+        location_description: A description of a location.
+    """
+    # NOTE: Uses demo endpoints that return random data
+    r = await ctx.deps.client.get(
+        "https://demo-endpoints.pydantic.workers.dev/latlng",
+        params={"location": location_description},
+    )
+    r.raise_for_status()
+    return LatLng.model_validate_json(r.content)
+
+
+@agent.tool
+async def get_weather(ctx: RunContext[Deps], lat: float, lng: float) -> dict[str, Any]:
+    """Get the weather at a location.
+
+    Args:
+        ctx: The context.
+        lat: Latitude of the location.
+        lng: Longitude of the location.
+    """
+    # NOTE: Uses demo endpoints that return random data
+    temp_response, descr_response = await asyncio.gather(
+        ctx.deps.client.get(
+            "https://demo-endpoints.pydantic.workers.dev/number",
+            params={"min": 10, "max": 30},
+        ),
+        ctx.deps.client.get(
+            "https://demo-endpoints.pydantic.workers.dev/weather",
+            params={"lat": lat, "lng": lng},
+        ),
+    )
+    temp_response.raise_for_status()
+    descr_response.raise_for_status()
+    return {
+        "temperature": f"{temp_response.text} °C",
+        "description": descr_response.text,
+    }
+
+
+@pixie.app
+async def example_weather_agent() -> pixie.PixieGenerator[str, str]:
+    """Interactive weather agent.
+
+    This agent interacts with the user to get weather information for a specified location.
+    It uses tools to fetch latitude/longitude and weather data, and supports a multi-turn
+    conversation with the user.
+
+    The agent will:
+    - Prompt the user for a location.
+    - Fetch latitude and longitude for the location.
+    - Retrieve weather data for the coordinates.
+    - Display the weather information to the user.
+
+    The conversation continues until the user enters a stop word ('exit', 'quit', or 'stop').
+    """
+
+    yield "Hi! I can help you find the weather for any location."
+    yield "Enter 'exit', 'quit', or 'stop' to end the conversation."
+
+    # Create HTTP client and dependencies
+    async with AsyncClient() as client:
+        deps = Deps(client=client)
+
+        user_message = None
+        stop_words = ["exit", "quit", "stop"]
+        while user_message not in stop_words:
+            yield "What location would you like the weather for?"
+            user_message = yield pixie.InputRequired(str)
+            result = await agent.run(user_message, deps=deps)
+            yield result.output