Merge pull request #1 from AlexAlbala/feature/mcp_support

Feature/mcp support

Author: AlexAlbala

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "pydantic>=2.0.0",
     "openai>=1.63.2",
     "pyyaml>=6.0.0",
+    "fastmcp>=0.4.1",
 ]
 
 [project.urls]

src/asktheapi_team_builder/__init__.py

@@ -1,7 +1,7 @@
-from .team_builder import TeamBuilder
-from .agent_builder import AgentBuilder
-from .models import Agent, Tool, Message
-from .api_spec_handler import (
+from .core.team_builder import TeamBuilder
+from .core.agent_builder import AgentBuilder
+from .core.models import Agent, Tool, Message
+from .core.api_spec_handler import (
     APISpecHandler,
     APISpecClassification,
     APISpecClassificationResult,

src/asktheapi_team_builder/core/__init__.py

@@ -14,11 +14,19 @@ class APISpecClassification(BaseModel):
 class APISpecClassificationResult(BaseModel):
     specs: List[APISpecClassification]
 
+class APISpecAgentToolResult(BaseModel):
+    name: str
+    description: str
+    jsonschema: dict
+    path: str
+    method: str
+    
 class APISpecAgentResult(BaseModel):
     name: str
     description: str
     system_prompt: str
     user_prompt: str
+    tools: List[APISpecAgentToolResult]
 
 class APISpecHandler:
     def __init__(self, headers: dict = {}):

src/asktheapi_team_builder/agent_builder.py renamed to src/asktheapi_team_builder/core/agent_builder.py

@@ -0,0 +1,36 @@
+from fastmcp import FastMCP
+from pydantic import BaseModel
+
+from asktheapi_team_builder.core.api_spec_handler import APISpecHandler
+from asktheapi_team_builder.core.tool_builder import build_tool_function
+
+class MCPConfig(BaseModel):
+    transport: str = "sse"
+    port: int = 8000
+    name: str = "asktheapi_mcp"
+    
+class MCPService():
+    def __init__(self, mcp_config: MCPConfig):
+        self.mcp = FastMCP(name=mcp_config.name, port=mcp_config.port)
+        self.mcp_config = mcp_config
+        self.api_spec_handler = APISpecHandler()
+        
+    async def _create_from_spec(self, url_spec: str, headers: dict = {}):
+        spec_content = await self.api_spec_handler.download_url_spec(url_spec)
+        classification_result = await self.api_spec_handler.classify_spec(spec_content)
+        for group_spec in classification_result.specs:
+            agent_result = await self.api_spec_handler.generate_agent_for_group(
+                group_spec,
+                spec_content
+            )
+            for tool in agent_result.tools:
+                self.mcp.add_tool(build_tool_function(agent_result, tool, headers))
+        
+        return self.mcp
+        
+    async def _run_mcp(self):
+        self.mcp.run(self.mcp_config.transport)
+        
+    async def start_from_spec(self, url_spec: str, headers: dict = {}):
+        await self._create_from_spec(url_spec, headers)
+        await self._run_mcp()

src/asktheapi_team_builder/api_spec_handler.py renamed to src/asktheapi_team_builder/core/api_spec_handler.py

@@ -0,0 +1,98 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from asktheapi_team_builder.services.mcp_service import MCPService, MCPConfig
+from asktheapi_team_builder.core.api_spec_handler import APISpecHandler
+
+@pytest.fixture
+def mock_api_spec_handler():
+    handler = AsyncMock(spec=APISpecHandler)
+    handler.download_url_spec.return_value = "mock_spec_content"
+    handler.classify_spec.return_value = MagicMock(
+        specs=[
+            MagicMock(
+                group_name="test_group",
+                endpoints=["/test"]
+            )
+        ]
+    )
+    handler.generate_agent_for_group.return_value = MagicMock(
+        tools=[
+            MagicMock(
+                name="test_tool",
+                description="A test tool",
+                method="GET",
+                path="/test",
+                jsonschema={}
+            )
+        ]
+    )
+    return handler
+
+@pytest.fixture
+def mock_fast_mcp():
+    with patch("asktheapi_team_builder.services.mcp_service.FastMCP") as mock:
+        instance = MagicMock()
+        mock.return_value = instance
+        yield instance
+
+@pytest.fixture
+def mcp_service(mock_fast_mcp, mock_api_spec_handler):
+    with patch("asktheapi_team_builder.services.mcp_service.APISpecHandler", return_value=mock_api_spec_handler):
+        config = MCPConfig(transport="sse", port=8000, name="test_mcp")
+        service = MCPService(config)
+        return service
+
+@pytest.mark.asyncio
+async def test_mcp_service_initialization(mcp_service, mock_fast_mcp):
+    # Assert
+    assert mcp_service.mcp_config.name == "test_mcp"
+    assert mcp_service.mcp_config.port == 8000
+    assert mcp_service.mcp_config.transport == "sse"
+
+@pytest.mark.asyncio
+async def test_create_from_spec(mcp_service, mock_api_spec_handler, mock_fast_mcp):
+    # Arrange
+    url_spec = "http://example.com/spec"
+    headers = {"Authorization": "Bearer token"}
+    
+    # Act
+    result = await mcp_service._create_from_spec(url_spec, headers)
+    
+    # Assert
+    mock_api_spec_handler.download_url_spec.assert_called_once_with(url_spec)
+    mock_api_spec_handler.classify_spec.assert_called_once_with("mock_spec_content")
+    mock_api_spec_handler.generate_agent_for_group.assert_called_once()
+    assert result == mock_fast_mcp
+    # Verify that add_tool was called
+    assert mock_fast_mcp.add_tool.called
+
+@pytest.mark.asyncio
+async def test_run_mcp(mcp_service, mock_fast_mcp):
+    # Act
+    await mcp_service._run_mcp()
+    
+    # Assert
+    mock_fast_mcp.run.assert_called_once_with("sse")
+
+@pytest.mark.asyncio
+async def test_start_from_spec(mcp_service, mock_api_spec_handler):
+    # Arrange
+    url_spec = "http://example.com/spec"
+    headers = {"Authorization": "Bearer token"}
+    
+    # Act
+    await mcp_service.start_from_spec(url_spec, headers)
+    
+    # Assert
+    mock_api_spec_handler.download_url_spec.assert_called_once_with(url_spec)
+    mcp_service.mcp.run.assert_called_once_with("sse")
+
+@pytest.mark.asyncio
+async def test_mcp_config_defaults():
+    # Act
+    config = MCPConfig()
+    
+    # Assert
+    assert config.transport == "sse"
+    assert config.port == 8000
+    assert config.name == "asktheapi_mcp"