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
|
The server opens its own client and validates it on startup, so nothing
|
||||||
is opened here first.
|
is opened here first.
|
||||||
"""
|
"""
|
||||||
# The ref's own path, not `_path`: a URI-backed database has none, and
|
# The resolved scope, not a derived path and configuration: a path
|
||||||
# the local stand-in would override the URI in `_store_config`.
|
# would override a configured URI, and deriving drops the name results
|
||||||
|
# and citations carry.
|
||||||
server = create_mcp_server(
|
server = create_mcp_server(
|
||||||
self._one.db_path,
|
config=self.config, read_only=self.read_only, scope=self.scope
|
||||||
config=self._store_config,
|
|
||||||
read_only=self.read_only,
|
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
if transport == "stdio":
|
if transport == "stdio":
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,7 @@ import asyncio
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncIterator
|
||||||
from contextlib import AsyncExitStack, asynccontextmanager
|
from contextlib import AsyncExitStack, asynccontextmanager
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from fastmcp import FastMCP
|
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.tools.document import DocumentInfo
|
||||||
from haiku.rag.utils import format_citations
|
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:
|
def _decode_images(images_base64: list[str] | None) -> list[bytes] | None:
|
||||||
if not images_base64:
|
if not images_base64:
|
||||||
|
|
@ -25,8 +28,9 @@ def create_mcp_server(
|
||||||
db_path: Path | None = None,
|
db_path: Path | None = None,
|
||||||
config: AppConfig | None = None,
|
config: AppConfig | None = None,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
|
scope: "DatabaseScope | None" = None,
|
||||||
) -> FastMCP:
|
) -> FastMCP:
|
||||||
"""Create an MCP server with the specified database path.
|
"""Create an MCP server over one database.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
db_path: Path to the database file, or None to let `config` place it. A
|
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.
|
must pass None rather than a local stand-in.
|
||||||
config: Configuration to use.
|
config: Configuration to use.
|
||||||
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
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()
|
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
|
client: HaikuRAG | None = None
|
||||||
stack = AsyncExitStack()
|
stack = AsyncExitStack()
|
||||||
client_lock = asyncio.Lock()
|
client_lock = asyncio.Lock()
|
||||||
|
|
@ -51,7 +62,7 @@ def create_mcp_server(
|
||||||
async with client_lock:
|
async with client_lock:
|
||||||
if client is None:
|
if client is None:
|
||||||
client = await stack.enter_async_context(
|
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
|
return client
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -604,12 +604,42 @@ class TestMCPClientLifetime:
|
||||||
assert opens == 1
|
assert opens == 1
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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
|
self, monkeypatch
|
||||||
):
|
):
|
||||||
"""A path overrides `lancedb.uri`, so the server gets the database's own
|
"""A path would override a configured URI, and a derived configuration
|
||||||
path — None where a URI placed it — not the local stand-in a URI-backed
|
would lose the name, so `run_mcp` passes neither."""
|
||||||
ref resolves to for display."""
|
|
||||||
from haiku.rag.app import HaikuRAGApp
|
from haiku.rag.app import HaikuRAGApp
|
||||||
from haiku.rag.client.scope import DatabaseScope
|
from haiku.rag.client.scope import DatabaseScope
|
||||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||||
|
|
@ -623,9 +653,8 @@ class TestMCPClientLifetime:
|
||||||
async def run_stdio_async(self):
|
async def run_stdio_async(self):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def fake_create(db_path=None, config=None, read_only=False):
|
def fake_create(db_path=None, config=None, read_only=False, scope=None):
|
||||||
seen["db_path"] = db_path
|
seen.update(db_path=db_path, scope=scope)
|
||||||
seen["config"] = config
|
|
||||||
return _Server()
|
return _Server()
|
||||||
|
|
||||||
monkeypatch.setattr("haiku.rag.app.create_mcp_server", fake_create)
|
monkeypatch.setattr("haiku.rag.app.create_mcp_server", fake_create)
|
||||||
|
|
@ -636,9 +665,8 @@ class TestMCPClientLifetime:
|
||||||
await app.run_mcp(transport="stdio")
|
await app.run_mcp(transport="stdio")
|
||||||
|
|
||||||
assert seen["db_path"] is None
|
assert seen["db_path"] is None
|
||||||
[ref] = DatabaseScope.resolve(
|
[ref] = seen["scope"].databases
|
||||||
seen["config"], database_path=seen["db_path"]
|
assert ref.name == "prod"
|
||||||
).databases
|
|
||||||
assert ref.uri == "s3://bucket/prod.lancedb"
|
assert ref.uri == "s3://bucket/prod.lancedb"
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue