mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 14:12:58 +00:00
arc53-machine
5d0992eef8
run_agent_headless records each unattended run under its endpoint; the scheduler passes its run id and the webhook worker its task id so Logs rows can find their trace, while the LLM's own request id stays untouched for quota counts. /api/search and MCP search_docs record their retrieval, and a graph build records every extraction call under one step.
62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
"""FastMCP server exposing DocsGPT retrieval over streamable HTTP.
|
|
|
|
Mounted at ``/mcp`` by ``docsgpt/asgi.py``. Bearer tokens are the
|
|
existing DocsGPT agent API keys — no new credential surface. The
|
|
``search_docs`` tool searches the caller's knowledge base.
|
|
|
|
The tool reads the ``Authorization`` header directly via
|
|
``get_http_headers(include={"authorization"})``. The ``include`` kwarg
|
|
is required: by default ``get_http_headers`` strips ``authorization``
|
|
(and a handful of other hop-by-hop headers) so they aren't forwarded
|
|
to downstream services — since we deliberately want the caller's
|
|
token, we opt it back in.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
from fastmcp import FastMCP
|
|
from fastmcp.server.dependencies import get_http_headers
|
|
|
|
from docsgpt.services.search_service import (
|
|
InvalidAPIKey,
|
|
SearchFailed,
|
|
search,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _extract_bearer_token() -> str | None:
|
|
auth = get_http_headers(include={"authorization"}).get("authorization", "")
|
|
parts = auth.split(None, 1)
|
|
if len(parts) != 2 or parts[0].lower() != "bearer" or not parts[1]:
|
|
return None
|
|
return parts[1]
|
|
|
|
|
|
mcp = FastMCP("docsgpt")
|
|
|
|
|
|
@mcp.tool
|
|
async def search_docs(query: str, chunks: int = 5) -> list[dict]:
|
|
"""Search the caller's DocsGPT knowledge base.
|
|
|
|
Authentication is via ``Authorization: Bearer <agent-api-key>`` on
|
|
the MCP request — the same opaque key that ``/api/search`` accepts
|
|
in its JSON body. Returns at most ``chunks`` hits, each a dict with
|
|
``text``, ``title``, ``source`` keys.
|
|
"""
|
|
api_key = _extract_bearer_token()
|
|
if not api_key:
|
|
raise PermissionError("Missing Bearer token")
|
|
try:
|
|
return await asyncio.to_thread(search, api_key, query, chunks, source="mcp")
|
|
except InvalidAPIKey as exc:
|
|
raise PermissionError("Invalid API key") from exc
|
|
except SearchFailed:
|
|
logger.exception("search_docs failed")
|
|
raise
|