Merge pull request #12 from Imaging-Plaza/feature/deep-wiki-mcp

Feature/deep wiki mcp

Author: qchapp

CHANGELOG.md

@@ -5,11 +5,17 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
+- **DeepWiki MCP integration**: Repository info tool now uses DeepWiki MCP server (https://mcp.deepwiki.com/sse) as primary source for GitHub repository documentation. DeepWiki provides fast, pre-indexed documentation access without API rate limits.
+- Automatic fallback to `repocards` library (replacing previous direct GitHub API implementation) when DeepWiki is unavailable or times out, ensuring robust repository information retrieval for both indexed and newly-created repositories.
+
+### Changed
+- Updated `pydantic-ai` dependency to include MCP support via `pydantic-ai[mcp]` extra.
+- Enhanced `RepoSummaryOutput` schema to include `source` field indicating whether data came from "deepwiki" or "repocards".
+- Repository info tool logs now track data source (DeepWiki vs repocards) for observability.
+- Replaced previous direct GitHub API implementation with `repocards` library as the fallback mechanism for repository information retrieval.
 - **YAML Model Configuration**: New `config.yaml` file for flexible model configuration supporting OpenAI, EPFL inference server, and any OpenAI-compatible API endpoints.
 - **Multi-Model Support**: Can now configure different models for agent (main reasoning & tool selection).
 - **Configuration Module**: New `utils/config.py` with Pydantic models for type-safe configuration loading and validation.
-
-### Changed
 - **Model Initialization**: Agent now uses configuration from `config.yaml`.
 - **API Client Creation**: OpenAI clients now support custom `base_url` for alternative API endpoints (EPFL, custom deployments).
 - **Dependency**: Added `pyyaml` to `pyproject.toml` dependencies.

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
   "pydantic>=2",
   "sentence-transformers",
   "openai>=1.30.0",
-  "pydantic-ai",
+  "pydantic-ai[mcp]",
   "requests",
   "python-dotenv",
   "gradio==5.42.0",
@@ -23,6 +23,7 @@ dependencies = [
   "imageio",
   "rdflib",
   "sparqlwrapper",
+  "repocards",
   "pyyaml",
 ]
 

src/ai_agent/agent/agent.py

@@ -13,11 +13,11 @@
 from utils.utils import _best_runnable_link
 from utils.config import get_config
 from .models import AgentToolSelection, ToolRunLog
-from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput, coerce_github_url_or_none
+from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput
 from .tools.rerank_tool import tool_rerank, RerankInput
 from .tools.search_tool import tool_search_tools, SearchToolsInput
 from .tools.gradio_space_tool import tool_run_example, RunExampleInput
-from .utils import AgentState, limit_tool_calls, cap_prepare
+from .utils import AgentState, limit_tool_calls, cap_prepare, coerce_github_url_or_none
 
 log = logging.getLogger("agent.core")
 
@@ -103,7 +103,7 @@ async def rerank(ctx: RunContext[AgentState], query: str, candidate_names: List[
 #     return out.model_dump(mode="python")
 
 @agent.tool(retries=0, prepare=cap_prepare)
-@limit_tool_calls("repo_info", cap=3)
+@limit_tool_calls("repo_info", cap=6)
 async def repo_info(ctx: RunContext[AgentState], url: str):
     norm_url = coerce_github_url_or_none(url)
     if not norm_url:
@@ -117,8 +117,13 @@ async def repo_info(ctx: RunContext[AgentState], url: str):
         return payload
 
     try:
-        out = tool_repo_summary(RepoSummaryInput(url=norm_url))
-        ctx.deps.tool_calls.append({"tool": "repo_info", "url": norm_url, "truncated": out.truncated})
+        out = await tool_repo_summary(RepoSummaryInput(url=norm_url))
+        ctx.deps.tool_calls.append({
+            "tool": "repo_info",
+            "url": norm_url,
+            "truncated": out.truncated,
+            "source": out.source
+        })
         return out.model_dump(mode="python")
     except Exception as e:
         ctx.deps.tool_calls.append({"tool": "repo_info", "url": norm_url, "error": str(e)})

src/ai_agent/agent/tools/deepwiki_tool.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Optional
+
+from pydantic import BaseModel
+from pydantic_ai.mcp import MCPServerSSE
+
+from .utils import _clip
+from ..utils import _coerce_owner_repo_ref
+
+log = logging.getLogger("agent.deepwiki")
+
+# DeepWiki MCP server endpoint (SSE transport)
+DEEPWIKI_SSE_URL = "https://mcp.deepwiki.com/sse"
+
+# Timeout for DeepWiki operations (seconds)
+DEEPWIKI_TIMEOUT = 60
+
+
+
+class DeepWikiInput(BaseModel):
+    """Input for DeepWiki operations."""
+    url: str  # GitHub repository URL or owner/repo format
+
+
+class DeepWikiContentsOutput(BaseModel):
+    """Output from read_wiki_contents."""
+    success: bool
+    contents: Optional[str] = None
+    error: Optional[str] = None
+    truncated: bool = False
+
+
+async def get_wiki_contents(input: DeepWikiInput) -> DeepWikiContentsOutput:
+    """
+    Fetch repo docs from DeepWiki MCP (SSE) and return a clipped string
+    to keep LLM token usage under control.
+    """
+    owner, repo, _ = _coerce_owner_repo_ref(input.url)
+    repo = f"{owner}/{repo}"
+
+    try:
+        server = MCPServerSSE(DEEPWIKI_SSE_URL)
+
+        async with server:
+            result = await asyncio.wait_for(
+                server.direct_call_tool("read_wiki_contents", {"repoName": repo}),
+                timeout=DEEPWIKI_TIMEOUT,
+            )
+
+            text = None
+            if isinstance(result, list):
+                text = "\n".join([p for p in result if isinstance(p, str)]) or None
+
+            if text and text.strip():
+                clipped_text, truncated = _clip(text.strip())
+                return DeepWikiContentsOutput(
+                    success=True,
+                    contents=clipped_text,
+                    truncated=truncated,
+                )
+
+            return DeepWikiContentsOutput(success=False, error="No content returned from DeepWiki")
+
+    except asyncio.TimeoutError:
+        log.warning(f"DeepWiki timed out after {DEEPWIKI_TIMEOUT}s for {repo}")
+        return DeepWikiContentsOutput(
+            success=False,
+            error=f"DeepWiki request timed out after {DEEPWIKI_TIMEOUT}s",
+        )
+    except Exception as e:
+        log.error(f"Failed to get wiki contents for {repo}: {e}")
+        return DeepWikiContentsOutput(
+            success=False,
+            error=f"Failed to connect to DeepWiki: {str(e)}",
+        )
+
+
+__all__ = [
+    "get_wiki_contents",
+    "DeepWikiInput",
+    "DeepWikiContentsOutput",
+]