feat: adding pydantic-ai tests

test-pydantic-ai/README.md

@@ -0,0 +1,132 @@
+# Pydantic AI Sentry Integration Tests
+
+Clean, focused tests for the Sentry Pydantic AI integration.
+
+## Prerequisites
+
+You need to have Python and `curl` installed.
+
+## Configure
+
+Set the following environment variables:
+- `SENTRY_DSN`
+- `OPENAI_API_KEY`
+- `ANTHROPIC_API_KEY` (for Anthropic provider tests)
+
+## Run Tests
+
+### Synchronous Tests
+```bash
+./run.sh
+```
+
+### Asynchronous Tests
+```bash
+./run_async.sh
+```
+
+## Test Structure
+
+The tests are organized into three main scenarios, each tested in both sync/async and streaming/non-streaming modes:
+
+### 1. **Simple Agent**
+- Basic agent without tools
+- Tests fundamental agent functionality
+- Demonstrates simple Q&A interactions
+
+### 2. **Agent with Tools**
+- Mathematical agent with calculation tools
+- Tools: `add()`, `multiply()`, `calculate_percentage()`
+- Tests tool integration and structured output
+- Returns `CalculationResult` with explanation
+
+### 3. **Two-Agent Workflow**
+- **Data Collector Agent**: Extracts and organizes data
+- **Data Analyzer Agent**: Analyzes data and provides insights
+- Demonstrates agent handoff and workflow patterns
+- Returns `AnalysisResult` with findings and recommendations
+
+### 4. **Anthropic Provider Tests**
+- **Anthropic Simple Agent**: Basic Claude agent without tools
+- **Anthropic Math Agent**: Claude agent with calculation tools
+- Tests both OpenAI and Anthropic providers for comparison
+- Uses `anthropic:claude-3-5-haiku-20241022` model
+
+## Test Modes
+
+Each scenario is tested in four different modes:
+
+| Mode | Sync/Async | Streaming | Description |
+|------|------------|-----------|-------------|
+| 1 | Sync | No | `agent.run_sync()` |
+| 2 | Sync | Yes | `agent.run_stream_sync()` |
+| 3 | Async | No | `await agent.run()` |
+| 4 | Async | Yes | `async with agent.run_stream()` |
+
+## Additional Features
+
+### Parallel Processing (Async Only)
+- Demonstrates running multiple agents concurrently
+- Uses `asyncio.gather()` for parallel execution
+- Shows scalable multi-agent patterns
+
+### Model Settings
+All agents use optimized model settings:
+- **Simple Agent**: Balanced settings (temp: 0.3)
+- **Math Agent**: Low temperature for precision (temp: 0.1)
+- **Data Collector**: Focused extraction (temp: 0.2)
+- **Data Analyzer**: Creative analysis (temp: 0.4)
+- **Anthropic Simple Agent**: Balanced settings (temp: 0.3)
+- **Anthropic Math Agent**: Low temperature for precision (temp: 0.1)
+
+## Sentry Integration
+
+The integration automatically creates Sentry spans for:
+- `gen_ai.pipeline` - Agent workflow execution (root span)
+- `gen_ai.invoke_agent` - Individual agent invocations
+- `gen_ai.chat` - Model requests (AI client calls)
+- `gen_ai.execute_tool` - Tool executions
+
+## File Structure
+
+```
+test-pydantic-ai/
+├── agents.py           # Agent definitions and tools
+├── main.py            # Synchronous tests
+├── main_async.py      # Asynchronous tests
+├── run.sh             # Run sync tests
+├── run_async.sh       # Run async tests
+├── pyproject.toml     # Dependencies
+└── README.md          # This file
+```
+
+## Output Example
+
+```
+🚀 Running Pydantic AI Synchronous Tests
+==================================================
+
+=== SIMPLE AGENT (Non-Streaming) ===
+Question: What is the capital of France?
+Answer: The capital of France is Paris.
+
+=== SIMPLE AGENT (Streaming) ===
+Question: Tell me a short story about a robot.
+Answer (streaming): Once upon a time, a curious robot named Zara discovered...
+
+=== AGENT WITH TOOLS (Non-Streaming) ===
+Task: Multi-step calculation
+Result: CalculationResult(result=126, operation='multi-step', explanation='...')
+
+=== TWO-AGENT WORKFLOW (Non-Streaming) ===
+Step 1: Data Collection
+Data Collector Result: Extracted sales data: [150, 200, 175, 225]
+
+Step 2: Data Analysis
+Data Analyzer Result: AnalysisResult(summary='...', key_findings=[...], recommendation='...')
+
+==================================================
+✅ All synchronous tests completed!
+```
+
+This clean structure focuses on the core functionality while thoroughly testing all aspects of the Pydantic AI integration with Sentry.

test-pydantic-ai/agents.py

@@ -0,0 +1,185 @@
+"""
+Simple, focused agents for testing Pydantic AI integration with Sentry.
+"""
+
+from pydantic_ai import Agent
+from pydantic import BaseModel
+
+
+class CalculationResult(BaseModel):
+    """Result from mathematical calculations."""
+    result: int
+    operation: str
+    explanation: str
+
+
+class AnalysisResult(BaseModel):
+    """Result from data analysis."""
+    summary: str
+    key_findings: list[str]
+    recommendation: str
+
+
+# Simple agent without tools
+simple_agent = Agent(
+    "openai:gpt-4o-mini",
+    name="simple_agent",
+    instructions="You are a helpful assistant. Provide clear, concise answers.",
+    model_settings={
+        "temperature": 0.3,
+        "max_tokens": 200,
+    },
+)
+
+
+# Agent with mathematical tools
+math_agent = Agent(
+    "openai:gpt-4o-mini",
+    name="math_agent",
+    instructions="You are a mathematical assistant. Use the available tools to perform calculations and return structured results.",
+    output_type=CalculationResult,
+    model_settings={
+        "temperature": 0.1,
+        "max_tokens": 300,
+    },
+)
+
+
+@math_agent.tool_plain
+def add(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+
+
+@math_agent.tool_plain
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers together."""
+    return a * b
+
+
+@math_agent.tool_plain
+def calculate_percentage(part: float, total: float) -> float:
+    """Calculate what percentage 'part' is of 'total'."""
+    if total == 0:
+        return 0.0
+    return (part / total) * 100
+
+
+# First agent in two-agent setup - data collector
+data_collector_agent = Agent(
+    "openai:gpt-4o-mini",
+    name="data_collector",
+    instructions="You collect and prepare data for analysis. Extract key numbers and organize information clearly.",
+    model_settings={
+        "temperature": 0.2,
+        "max_tokens": 400,
+    },
+)
+
+
+@data_collector_agent.tool_plain
+def extract_numbers(text: str) -> list[int]:
+    """Extract all numbers from a text string."""
+    import re
+    numbers = re.findall(r'\d+', text)
+    return [int(n) for n in numbers]
+
+
+@data_collector_agent.tool_plain
+def organize_data(items: list[str]) -> dict:
+    """Organize a list of items into categories."""
+    return {
+        "total_items": len(items),
+        "items": items,
+        "categories": list(set(item.split()[0] if item.split() else "unknown" for item in items))
+    }
+
+
+# Second agent in two-agent setup - data analyzer
+data_analyzer_agent = Agent(
+    "openai:gpt-4o-mini",
+    name="data_analyzer",
+    instructions="You analyze data provided by the data collector and provide insights and recommendations.",
+    output_type=AnalysisResult,
+    model_settings={
+        "temperature": 0.4,
+        "max_tokens": 500,
+    },
+)
+
+
+@data_analyzer_agent.tool_plain
+def calculate_statistics(numbers: list[int]) -> dict:
+    """Calculate basic statistics for a list of numbers."""
+    if not numbers:
+        return {"error": "No numbers provided"}
+
+    return {
+        "count": len(numbers),
+        "sum": sum(numbers),
+        "average": sum(numbers) / len(numbers),
+        "min": min(numbers),
+        "max": max(numbers),
+    }
+
+
+@data_analyzer_agent.tool_plain
+def identify_trends(numbers: list[int]) -> str:
+    """Identify trends in a sequence of numbers."""
+    if len(numbers) < 2:
+        return "Not enough data to identify trends"
+
+    differences = [numbers[i+1] - numbers[i] for i in range(len(numbers)-1)]
+
+    if all(d > 0 for d in differences):
+        return "Increasing trend"
+    elif all(d < 0 for d in differences):
+        return "Decreasing trend"
+    elif all(d == 0 for d in differences):
+        return "Constant values"
+    else:
+        return "Mixed trend"
+
+
+# Anthropic agents for provider testing
+anthropic_simple_agent = Agent(
+    "anthropic:claude-3-5-haiku-20241022",
+    name="anthropic_simple_agent",
+    instructions="You are a helpful assistant using Claude. Provide clear, concise answers.",
+    model_settings={
+        "temperature": 0.3,
+        "max_tokens": 200,
+    },
+)
+
+
+anthropic_math_agent = Agent(
+    "anthropic:claude-3-5-haiku-20241022",
+    name="anthropic_math_agent",
+    instructions="You are a mathematical assistant using Claude. Use the available tools to perform calculations and return structured results.",
+    output_type=CalculationResult,
+    model_settings={
+        "temperature": 0.1,
+        "max_tokens": 300,
+    },
+)
+
+
+@anthropic_math_agent.tool_plain
+def anthropic_add(a: int, b: int) -> int:
+    """Add two numbers together (Anthropic version)."""
+    return a + b
+
+
+@anthropic_math_agent.tool_plain
+def anthropic_multiply(a: int, b: int) -> int:
+    """Multiply two numbers together (Anthropic version)."""
+    return a * b
+
+
+@anthropic_math_agent.tool_plain
+def anthropic_calculate_percentage(part: float, total: float) -> float:
+    """Calculate what percentage 'part' is of 'total' (Anthropic version)."""
+    if total == 0:
+        return 0.0
+    return (part / total) * 100