File size: 3,000 Bytes
a6998ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
"""
Pytest configuration for agent testing.
"""
import os
import pytest
from typing import Dict, List, Optional
from gagent.agents import registry, BaseAgent, OllamaAgent, GeminiAgent, OpenAIAgent
@pytest.fixture
def agent_factory():
"""
Factory fixture to create agent instances with flexible configuration.
Returns:
Function that creates and returns an agent instance
"""
def _create_agent(
agent_type: str,
model_name: Optional[str] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
**kwargs,
) -> BaseAgent:
"""
Create an agent with the specified configuration.
Args:
agent_type: The type of agent to create
model_name: The model name to use
api_key: The API key to use
base_url: The base URL to use
**kwargs: Additional parameters for the agent
Returns:
An initialized agent instance
"""
# Get environment variables or defaults for any non-provided values
env_model = os.environ.get(f"{agent_type.upper()}_MODEL", "qwen3" if agent_type == "ollama" else None)
env_api_key = os.environ.get(f"{agent_type.upper()}_API_KEY", None)
env_base_url = os.environ.get(
f"{agent_type.upper()}_BASE_URL", "http://localhost:11434" if agent_type == "ollama" else None
)
return registry.get_agent(
agent_type=agent_type,
model_name=model_name or env_model,
api_key=api_key or env_api_key,
base_url=base_url or env_base_url,
**kwargs,
)
return _create_agent
@pytest.fixture
def ollama_agent(agent_factory) -> OllamaAgent:
"""Fixture to provide an Ollama agent."""
return agent_factory("ollama")
@pytest.fixture
def gemini_agent(agent_factory) -> GeminiAgent:
"""Fixture to provide a Gemini agent if environment variables are set."""
api_key = os.environ.get("GOOGLE_API_KEY", None)
if not api_key:
pytest.skip("GOOGLE_API_KEY environment variable not set")
return agent_factory("gemini")
@pytest.fixture
def openai_agent(agent_factory) -> OpenAIAgent:
"""Fixture to provide an OpenAI agent if environment variables are set."""
api_key = os.environ.get("OPENAI_API_KEY", None)
if not api_key:
pytest.skip("OPENAI_API_KEY environment variable not set")
return agent_factory("openai")
@pytest.fixture
def gaia_questions() -> List[Dict]:
"""Load GAIA questions for testing."""
import json
with open("exp/questions.json", "r") as f:
return json.load(f)
@pytest.fixture
def gaia_validation_data() -> Dict:
"""Load GAIA validation data."""
import json
validation_data = {}
with open("metadata.jsonl", "r") as f:
for line in f:
data = json.loads(line)
validation_data[data["task_id"]] = data["Final answer"]
return validation_data
|