Testing Multi-Agent Systems

from pydantic import BaseModel
from typing import List

# Define the contract between the Research Agent and the Writing Agent
class ResearchOutput(BaseModel):
    """Contract: what the Research Agent must produce."""
    topic: str
    findings: List[dict]  # Each finding has 'source', 'content', 'relevance'
    gaps: List[str]      # Topics that need more research
    confidence: float  # 0.0 to 1.0

class WritingInput(BaseModel):
    """Contract: what the Writing Agent expects to receive."""
    topic: str
    findings: List[dict]
    tone: str       # "formal", "casual", "technical"
    max_length: int  # words

def test_research_output_matches_writing_input():
    """Research Agent output must satisfy Writing Agent input contract."""
    # Run the Research Agent on a test task
    research_result = research_agent.run("Summarize recent advances in RAG")
    # Validate against the contract
    output = ResearchOutput(**research_result)
    assert len(output.findings) >= 1, "Must produce at least one finding"
    assert 0 <= output.confidence <= 1, "Confidence must be in [0, 1]"
    # Verify it can be transformed into the Writing Agent's expected input
    writing_input = WritingInput(
        topic=output.topic,
        findings=output.findings,
        tone="technical",
        max_length=2000,
    )
    assert writing_input  # Pydantic validation passed

import random
from unittest.mock import patch
class ChaosInjector:
    """Inject failures into agent tool calls for chaos testing."""
    def __init__(self, failure_rate: float = 0.3):
        self.failure_rate = failure_rate
        self.injected_failures = []
    def maybe_fail(self, tool_name: str, original_func):
        """Wrap a tool function with random failure injection."""
        async def chaos_wrapper(*args, **kwargs):
            if random.random() < self.failure_rate:
                failure_type = random.choice([
                    "timeout", "rate_limit", "server_error", "malformed_response"
                    ])
                self.injected_failures.append((tool_name, failure_type))
                if failure_type == "timeout":
                    raise TimeoutError(f"{tool_name} timed out")
                elif failure_type == "rate_limit":
                    raise RateLimitError(f"{tool_name} rate limited")
                elif failure_type == "server_error":
                    raise APIError(f"{tool_name} returned 500")
                elif failure_type == "malformed_response":
                    return "{{invalid json"
                    return await original_func(*args, **kwargs)
                    return chaos_wrapper
            def test_agent_resilience():
                """Chaos test: agent should handle random tool failures gracefully."""
                chaos = ChaosInjector(failure_rate=0.3)
                # Wrap all tools with chaos injection
                chaotic_tools = {
                    name: chaos.maybe_fail(name, tool.execute)
                    for name, tool in agent.tools.items()
                    }
                # Run the agent on a reference task
                result = agent.run(
                    "Analyze last month's sales data",
                    tools=chaotic_tools,
                    )
                # Verify graceful degradation
                assert result is not None, "Agent should produce some output even with failures"
                assert "error" not in result.lower() or "unavailable" in result.lower(), \
                "Error messages should be user-friendly"
                print(f"Injected {len(chaos.injected_failures)} failures: {chaos.injected_failures}")
                print(f"Agent output: {result[:200]}...")

# Lab starter: agent contract validation. Students fill in the TODOs.
from pydantic import BaseModel, Field
from typing import Literal

# 1) Define the contract the agent's tool calls must satisfy
class WeatherQuery(BaseModel):
    """TODO: extend with required and optional fields the agent must produce."""
    city: str = Field(..., description="City name; non-empty")
    units: Literal["c", "f"] = "c"

    def call_agent(prompt: str) -> dict:
        """TODO: call your agent and return its parsed JSON tool-call payload."""
        raise NotImplementedError

    def validate_tool_call(payload: dict) -> WeatherQuery:
        """TODO: parse `payload` into the WeatherQuery contract.
        Hint: use WeatherQuery.model_validate; let it raise on failure."""
        raise NotImplementedError

if __name__ == "__main__":
    prompt = "What's the weather in Tokyo in Fahrenheit?"
    payload = call_agent(prompt)
    contract = validate_tool_call(payload)
    print(f"Validated: {contract.model_dump()}")

# Full solution for the agent contract validation lab.
import json
from pydantic import BaseModel, Field, ValidationError
from typing import Literal
from openai import OpenAI

client = OpenAI()

class WeatherQuery(BaseModel):
    city: str = Field(..., min_length=1, description="City name; non-empty")
    country_code: str | None = Field(None, pattern=r"^[A-Z]{2}$",
                                      description="Optional ISO 3166-1 alpha-2 country code")
    units: Literal["c", "f"] = "c"

SYSTEM_PROMPT = (
    "You are a weather assistant. When the user asks about weather, respond with "
    "a JSON object {\"city\": ..., \"country_code\": ..., \"units\": \"c\" or \"f\"}. "
    "Nothing else."
)

def call_agent(prompt: str) -> dict:
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        response_format={"type": "json_object"},
        temperature=0.0,
    )
    return json.loads(resp.choices[0].message.content)

def validate_tool_call(payload: dict) -> WeatherQuery:
    try:
        return WeatherQuery.model_validate(payload)
    except ValidationError as e:
        # In production: log payload, return a structured error, ask agent to retry
        raise

def test_basic_call():
    payload = call_agent("What's the weather in Tokyo in Fahrenheit?")
    contract = validate_tool_call(payload)
    assert contract.city.lower() == "tokyo"
    assert contract.units == "f"

if __name__ == "__main__":
    test_basic_call()
    print("contract validated; tests passed")