Give the MCP server the database, not a description of it
`run_mcp` derived a path and a configuration for the server, and deriving drops the configured name: MCP results and citations carried `source=None` where the same database named through any other path carried "alpha". `create_mcp_server` takes the resolved scope and opens through it, so the name survives. Passing a path and configuration still works and resolves to the same place.
This commit is contained in:
parent
d93cb5c891
commit
db7f0e0af6
3 changed files with 56 additions and 18 deletions
|
|
@ -923,12 +923,11 @@ class HaikuRAGApp:
|
|||
The server opens its own client and validates it on startup, so nothing
|
||||
is opened here first.
|
||||
"""
|
||||
# The ref's own path, not `_path`: a URI-backed database has none, and
|
||||
# the local stand-in would override the URI in `_store_config`.
|
||||
# The resolved scope, not a derived path and configuration: a path
|
||||
# would override a configured URI, and deriving drops the name results
|
||||
# and citations carry.
|
||||
server = create_mcp_server(
|
||||
self._one.db_path,
|
||||
config=self._store_config,
|
||||
read_only=self.read_only,
|
||||
config=self.config, read_only=self.read_only, scope=self.scope
|
||||
)
|
||||
try:
|
||||
if transport == "stdio":
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import asyncio
|
|||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
|
|
@ -12,6 +12,9 @@ from haiku.rag.store.models import Document, SearchResult
|
|||
from haiku.rag.tools.document import DocumentInfo
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
|
||||
|
||||
def _decode_images(images_base64: list[str] | None) -> list[bytes] | None:
|
||||
if not images_base64:
|
||||
|
|
@ -25,8 +28,9 @@ def create_mcp_server(
|
|||
db_path: Path | None = None,
|
||||
config: AppConfig | None = None,
|
||||
read_only: bool = False,
|
||||
scope: "DatabaseScope | None" = None,
|
||||
) -> FastMCP:
|
||||
"""Create an MCP server with the specified database path.
|
||||
"""Create an MCP server over one database.
|
||||
|
||||
Args:
|
||||
db_path: Path to the database file, or None to let `config` place it. A
|
||||
|
|
@ -34,8 +38,15 @@ def create_mcp_server(
|
|||
must pass None rather than a local stand-in.
|
||||
config: Configuration to use.
|
||||
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
||||
scope: The database, already resolved. Pass this rather than a derived
|
||||
path and configuration to keep its configured name, which results
|
||||
and citations carry as `source`.
|
||||
"""
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
|
||||
config = config if config is not None else get_config()
|
||||
if scope is None:
|
||||
scope = DatabaseScope.resolve(config, database_path=db_path)
|
||||
client: HaikuRAG | None = None
|
||||
stack = AsyncExitStack()
|
||||
client_lock = asyncio.Lock()
|
||||
|
|
@ -51,7 +62,7 @@ def create_mcp_server(
|
|||
async with client_lock:
|
||||
if client is None:
|
||||
client = await stack.enter_async_context(
|
||||
HaikuRAG(db_path, config=config, read_only=read_only)
|
||||
HaikuRAG._covering(scope, config, read_only=read_only)
|
||||
)
|
||||
return client
|
||||
|
||||
|
|
|
|||
|
|
@ -604,12 +604,42 @@ class TestMCPClientLifetime:
|
|||
assert opens == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_uri_backed_database_is_not_replaced_by_a_local_path(
|
||||
async def test_the_scope_decides_the_database_and_names_its_results(
|
||||
self, mcp_db, tmp_path
|
||||
):
|
||||
"""The scope is the selection, so the server reads the one database it
|
||||
names and results carry that name. The configuration alone would place
|
||||
every database it configures."""
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
||||
other = tmp_path / "beta.lancedb"
|
||||
async with HaikuRAG(other, create=True) as rag:
|
||||
await rag.create_document(
|
||||
"Zebras graze on the savannah.", title="Zebras", uri="test://zebras"
|
||||
)
|
||||
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(databases={"alpha": str(mcp_db), "beta": str(other)})
|
||||
)
|
||||
scope = DatabaseScope.resolve(config, database_name="alpha")
|
||||
|
||||
mcp = create_mcp_server(config=config, read_only=True, scope=scope)
|
||||
async with mcp._lifespan_manager():
|
||||
search = await _get_tool(mcp, "search_documents")
|
||||
results = await search(query="artificial intelligence")
|
||||
zebras = await search(query="zebras savannah")
|
||||
|
||||
assert results
|
||||
assert {r.source for r in results} == {"alpha"}
|
||||
assert all(r.source == "alpha" for r in zebras)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_command_hands_the_server_its_resolved_database(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""A path overrides `lancedb.uri`, so the server gets the database's own
|
||||
path — None where a URI placed it — not the local stand-in a URI-backed
|
||||
ref resolves to for display."""
|
||||
"""A path would override a configured URI, and a derived configuration
|
||||
would lose the name, so `run_mcp` passes neither."""
|
||||
from haiku.rag.app import HaikuRAGApp
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
|
@ -623,9 +653,8 @@ class TestMCPClientLifetime:
|
|||
async def run_stdio_async(self):
|
||||
return None
|
||||
|
||||
def fake_create(db_path=None, config=None, read_only=False):
|
||||
seen["db_path"] = db_path
|
||||
seen["config"] = config
|
||||
def fake_create(db_path=None, config=None, read_only=False, scope=None):
|
||||
seen.update(db_path=db_path, scope=scope)
|
||||
return _Server()
|
||||
|
||||
monkeypatch.setattr("haiku.rag.app.create_mcp_server", fake_create)
|
||||
|
|
@ -636,9 +665,8 @@ class TestMCPClientLifetime:
|
|||
await app.run_mcp(transport="stdio")
|
||||
|
||||
assert seen["db_path"] is None
|
||||
[ref] = DatabaseScope.resolve(
|
||||
seen["config"], database_path=seen["db_path"]
|
||||
).databases
|
||||
[ref] = seen["scope"].databases
|
||||
assert ref.name == "prod"
|
||||
assert ref.uri == "s3://bucket/prod.lancedb"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Reference in a new issue