adding tests

Author: AlexAlbala

.gitignore

@@ -24,6 +24,7 @@ wheels/
 venv/
 env/
 ENV/
+.venv/
 
 # IDE
 .idea/

README.md

@@ -4,16 +4,27 @@
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![Python Versions](https://img.shields.io/pypi/pyversions/asktheapi-team-builder.svg)](https://pypi.org/project/asktheapi-team-builder/)
 
-A high-level Python library for easily building and managing autonomous agents networks that solve complex tasks using APIs defined with OpenAPI standard format. This package provides a clean, type-safe interface for creating, configuring, and running teams of agents that can work together to solve complex tasks.
+A high-level Python library for building and managing networks of autonomous agents that collaborate to solve complex tasks. It’s designed to work seamlessly with APIs defined using the OpenAPI standard. The library provides a clean, type-safe interface for creating, configuring, and running teams of agents, making it easy to orchestrate multi-agent workflows with minimal boilerplate.
 
 ## Features
 
-- 🚀 Easy creation of agent networks with custom tools and capabilities based on openAPI specification
-- 🤝 Team building with automatic coordination through a planning agent
-- 📡 Support for streaming agent interactions
-- 🔧 Built-in HTTP client for tool implementation
-- ✨ Pydantic models for type safety and validation
-- 🎯 Clean, intuitive API design
+- 🚀 Effortless Agent Network Creation
+Quickly build agent networks with custom tools and capabilities based on OpenAPI specifications.
+
+- 🤝 Team-Based Collaboration
+Easily define agent teams with automatic coordination handled by a built-in planning agent.
+
+- 📡 Streaming Interactions
+Stream agent communication in real-time for more dynamic and responsive workflows.
+
+- 🔧 Built-in HTTP Client
+Simplify tool implementation with an integrated HTTP client ready to call external APIs.
+
+- ✨ Type Safety with Pydantic
+Leverage Pydantic models for robust data validation and clear type definitions.
+
+- 🎯 Clean and Intuitive API
+Designed for developers—minimal boilerplate, maximum clarity.
 
 ## Installation
 
@@ -53,8 +64,10 @@ async def create_agents_from_spec():
         agents.append(agent_result)
 
     return agents
+```
 
-# 3. Build and run a team
+# 2. Build and run a team
+```python
 async def run_agent_team(agents: List[Agent], query: str):
     # Initialize team builder
     team_builder = TeamBuilder(
@@ -77,8 +90,9 @@ async def run_agent_team(agents: List[Agent], query: str):
     async for event in team_builder.run_team(team, messages, stream=True):
         if isinstance(event, ChatMessage):
             print(f"{event.source}: {event.content}")
-        
+```        
 # Example usage
+```python
 async def main():
     # Create agents from spec
     api_agents = await create_agents_from_spec()

pyproject.toml

@@ -8,7 +8,7 @@ package-dir = {"" = "src"}
 
 [project]
 name = "asktheapi-team-builder"
-version = "0.2.1"
+version = "0.2.2"
 authors = [
     { name = "Alex Albala" },
 ]

pytest.ini

@@ -0,0 +1,7 @@
+[pytest]
+asyncio_mode = auto
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*
+addopts = -v --tb=short 

setup.cfg

@@ -0,0 +1,7 @@
+[metadata]
+name = asktheapi-team-builder
+version = 0.2.2
+
+[options]
+packages = find:
+

src/asktheapi_team_builder/prompts/__init__.py

@@ -1,4 +1,4 @@
-from services.open_ai_service import OpenAIService
+from asktheapi_team_builder.services.open_ai_service import OpenAIService
 
 
 class LLMService():
@@ -7,4 +7,5 @@ def __init__(self, openai_service: OpenAIService, llm_headers: dict = {}):
         self.llm_headers = llm_headers
 
     async def chat_completion(self, model, messages, stream):
-        return await self.openai_service.completion_with_headers(model, messages, stream, self.llm_headers)
+        return await self.openai_service.completion_with_headers(
+            model=model, messages=messages, stream=stream, headers=self.llm_headers)

src/asktheapi_team_builder/services/__init__.py

@@ -0,0 +1,76 @@
+import os
+from autogen_core import FunctionCall
+from autogen_core.models import FunctionExecutionResult
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from asktheapi_team_builder.services.open_ai_service import OpenAIService
+from asktheapi_team_builder.services.llm_service import LLMService
+from asktheapi_team_builder.services.agent_evaluator_service import AgentEvaluatorService, AgentDTO, AgentToolDTO, TaskResult, TextMessage, ToolCallExecutionEvent, ToolCallRequestEvent
+
+@pytest.fixture(autouse=True)
+def setenvvar(monkeypatch):
+    with patch.dict(os.environ, clear=True):
+        envvars = {
+            "OPENAI_API_KEY": "mock-api-key",
+            "OPENAI_BASE_URL": "mock-base-url"
+        }
+        for k, v in envvars.items():
+            monkeypatch.setenv(k, v)
+        yield 
+        
+@pytest.fixture
+def mock_openai_service():
+    service = AsyncMock(spec=OpenAIService)
+    service.completion_with_headers.return_value = MagicMock(
+        choices=[MagicMock(message=MagicMock(content=['{"evaluation": []}']))]
+    )
+    return service
+
+@pytest.fixture
+def mock_llm_service():
+    return AsyncMock(spec=LLMService)
+
+@pytest.fixture
+def agent_evaluator_service(mock_llm_service):
+    service = AgentEvaluatorService()
+    service.llm_service = mock_llm_service
+    return service
+
+@pytest.fixture
+def sample_agent_tool():
+    return AgentToolDTO(
+        id="tool1",
+        name="test_tool",
+        description="A test tool",
+        method="GET",
+        path="/test",
+        jsonschema={},
+        auto_updated=False
+    )
+
+@pytest.fixture
+def sample_agent(sample_agent_tool):
+    return AgentDTO(
+        id="agent1",
+        name="test_agent",
+        system_prompt="You are a test agent",
+        description="A test agent",
+        base_url="http://test.com",
+        tools=[sample_agent_tool],
+        apispec_id="spec1",
+        auto_updated=False
+    )
+
+@pytest.fixture
+def sample_task_result():
+    return TaskResult(
+        messages=[
+            TextMessage(content="Test message", source="test_agent"),
+            ToolCallRequestEvent(content=[FunctionCall(
+                name="test_function", arguments="{}", id="call1")
+            ], source="test_agent"),
+            ToolCallExecutionEvent(content=[
+                FunctionExecutionResult(call_id="call1", is_error=False, content="Success", name="test_function")
+            ], source="test_agent")
+        ]
+    ) 

src/asktheapi_team_builder/services/llm_service.py

@@ -0,0 +1,62 @@
+import pytest
+from asktheapi_team_builder.services.agent_evaluator_service import AgentEvaluatorService
+from asktheapi_team_builder.types import TextMessage, ToolCallExecutionEvent
+from autogen_core.models import FunctionExecutionResult
+from unittest.mock import AsyncMock, MagicMock
+
+@pytest.mark.asyncio
+async def test_needs_evaluation_task_result_no_errors(agent_evaluator_service, sample_task_result):
+    # Act
+    result = agent_evaluator_service.needs_evaluation_task_result(sample_task_result)
+    
+    # Assert
+    assert result is False
+
+@pytest.mark.asyncio
+async def test_needs_evaluation_task_result_with_text_error(agent_evaluator_service, sample_task_result):
+    # Arrange
+    task_result = sample_task_result
+    task_result.messages.append(TextMessage(content="An error occurred", source="test_agent"))
+    
+    # Act
+    result = agent_evaluator_service.needs_evaluation_task_result(task_result)
+    
+    # Assert
+    assert result is True
+
+@pytest.mark.asyncio
+async def test_needs_evaluation_task_result_with_tool_error(agent_evaluator_service, sample_task_result):
+    # Arrange
+    task_result = sample_task_result
+    task_result.messages.append(ToolCallExecutionEvent(content=[
+                FunctionExecutionResult(call_id="call1", is_error=True, content="Error ocurred", name="test_function")
+            ], source="test_agent"))
+    
+    # Act
+    result = agent_evaluator_service.needs_evaluation_task_result(task_result)
+    
+    # Assert
+    assert result is True
+
+@pytest.mark.asyncio
+async def test_evaluate_task_result(agent_evaluator_service, sample_agent, sample_task_result, mock_llm_service):
+    # Arrange
+    mock_llm_service.chat_completion.return_value = AsyncMock(
+        choices=[MagicMock(message=MagicMock(content='{"evaluation": []}'))]
+    )
+    
+    # Act
+    result = await agent_evaluator_service.evaluate_task_result([sample_agent], sample_task_result)
+    
+    # Assert
+    assert result is not None
+    mock_llm_service.chat_completion.assert_called_once()
+    assert "test_agent" in str(mock_llm_service.chat_completion.call_args)
+
+@pytest.mark.asyncio
+async def test_evaluate_task_result_error_handling(agent_evaluator_service, sample_agent, sample_task_result, mock_llm_service):
+    # Arrange
+    mock_llm_service.chat_completion.side_effect = Exception("Evaluation error")
+    
+    # Act & Assert
+    await agent_evaluator_service.evaluate_task_result([sample_agent], sample_task_result)

tests/conftest.py

@@ -0,0 +1,44 @@
+import pytest
+from asktheapi_team_builder.services.llm_service import LLMService
+from unittest.mock import AsyncMock
+
+@pytest.fixture
+def llm_service(mock_openai_service):
+    return LLMService(mock_openai_service)
+
+@pytest.mark.asyncio
+async def test_chat_completion(llm_service, mock_openai_service):
+    # Arrange
+    model = "gpt-4"
+    messages = [{"role": "user", "content": "test"}]
+    stream = False
+    
+    # Act
+    result = await llm_service.chat_completion(model, messages, stream)
+    
+    # Assert
+    mock_openai_service.completion_with_headers.assert_called_once_with(
+        model=model,
+        messages=messages,
+        stream=stream,
+        headers={}
+    )
+    assert result is not None
+
+@pytest.mark.asyncio
+async def test_chat_completion_with_custom_headers():
+    # Arrange
+    custom_headers = {"custom-header": "value"}
+    mock_openai = AsyncMock()
+    service = LLMService(mock_openai, custom_headers)
+    
+    # Act
+    await service.chat_completion("gpt-4", [], False)
+    
+    # Assert
+    mock_openai.completion_with_headers.assert_called_once_with(
+        model="gpt-4",
+        messages=[],
+        stream=False,
+        headers=custom_headers
+    ) 

tests/services/test_agent_evaluator_service.py

@@ -0,0 +1,42 @@
+from openai import AsyncOpenAI
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from asktheapi_team_builder.services.open_ai_service import OpenAIService
+
+@pytest.fixture()
+def openai_service():
+    service = OpenAIService()
+    service.client = AsyncOpenAI()
+    service.client.chat.completions = AsyncMock()
+    return service
+
+@pytest.mark.asyncio
+async def test_completion_with_headers(openai_service):
+    # Arrange
+    model = "gpt-4"
+    messages = [{"role": "user", "content": "test"}]
+    stream = False
+    headers = {"test-header": "value"}
+    
+    # Act
+    result = await openai_service.completion_with_headers(model, messages, stream, headers)
+    
+    # Assert
+    openai_service.client.chat.completions.create.assert_called_once_with(
+        model=model,
+        messages=messages,
+        stream=stream,
+        extra_headers=headers,
+        response_format={'type': 'json_object'}
+    )
+    assert result is not None
+
+@pytest.mark.asyncio
+async def test_completion_with_headers_error_handling(openai_service):
+    # Arrange
+    openai_service.client.chat.completions.create.side_effect = Exception("API Error")
+    
+    # Act & Assert
+    with pytest.raises(Exception) as exc_info:
+        await openai_service.completion_with_headers("gpt-4", [], False, {})
+    assert str(exc_info.value) == "API Error"