Lend the caller's client to the capability

`client.ask` and `client.analyze` built their capability from a db_path, so the
capability opened a second connection to the database the client already had
open, once per call.

Ownership is now explicit rather than inferred. `rag` stays the connection the
capability opened and must close; `borrowed_rag` is a caller's, which
`_ensure_rag` prefers and `_close` never touches. Two fields rather than a flag,
so closing a borrowed connection is not expressible.

`for_run` still clears `rag` per run, since a run owns what it opens. It leaves
`borrowed_rag` alone: that connection belongs to the caller and outlives the run.
This commit is contained in:
Yiorgis Gozadinos 2026-08-18 14:41:11 +03:00
parent 8a4d488a72
commit 0882dc9fed
No known key found for this signature in database
6 changed files with 152 additions and 2 deletions

View file

@ -11,6 +11,7 @@
### Changed ### Changed
- Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged. - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged.
- `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs.
- The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift. - The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift.
- `import_documents` embeds chunks across the whole batch in one pass instead of per document. - `import_documents` embeds chunks across the whole batch in one pass instead of per document.
- `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`. - `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`.

View file

@ -129,6 +129,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
state: StateT | None = field(default=None, repr=False) state: StateT | None = field(default=None, repr=False)
outer_state: dict[str, Any] | None = field(default=None, repr=False) outer_state: dict[str, Any] | None = field(default=None, repr=False)
rag: HaikuRAG | None = field(default=None, repr=False) rag: HaikuRAG | None = field(default=None, repr=False)
"""A connection this capability opened, and must close."""
borrowed_rag: HaikuRAG | None = field(default=None, repr=False)
"""A caller's connection, reused and never closed here."""
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
search_count: int = field(default=0, repr=False) search_count: int = field(default=0, repr=False)
@ -347,6 +350,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
raise error raise error
async def _ensure_rag(self) -> HaikuRAG: async def _ensure_rag(self) -> HaikuRAG:
if self.borrowed_rag is not None:
return self.borrowed_rag
if self.rag is None: if self.rag is None:
async with self.resource_lock: async with self.resource_lock:
if self.rag is None: if self.rag is None:

View file

@ -1,13 +1,16 @@
from dataclasses import dataclass, field from dataclasses import dataclass, field
from functools import cache from functools import cache
from pathlib import Path from pathlib import Path
from typing import Any from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai import RunContext, ToolFailed from pydantic_ai import RunContext, ToolFailed
from pydantic_ai.messages import ToolReturn from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.toolsets import FunctionToolset
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.capabilities._base import ( from haiku.rag.capabilities._base import (
CodeExecutionEntry, CodeExecutionEntry,
RAGCapabilityBase, RAGCapabilityBase,
@ -162,6 +165,7 @@ def create_capability(
config: AppConfig | None = None, config: AppConfig | None = None,
*, *,
defer_loading: bool = True, defer_loading: bool = True,
rag: "HaikuRAG | None" = None,
request_limit: int | None = 30, request_limit: int | None = 30,
vision: bool | None = None, vision: bool | None = None,
) -> AnalysisCapability: ) -> AnalysisCapability:
@ -180,6 +184,7 @@ def create_capability(
return AnalysisCapability( return AnalysisCapability(
db_path=resolve_db_path(db_path, config), db_path=resolve_db_path(db_path, config),
config=config, config=config,
borrowed_rag=rag,
state_type=AnalysisState, state_type=AnalysisState,
state_namespace=STATE_NAMESPACE, state_namespace=STATE_NAMESPACE,
instruction_text=instructions(), instruction_text=instructions(),

View file

@ -1,13 +1,16 @@
from dataclasses import dataclass from dataclasses import dataclass
from functools import cache from functools import cache
from pathlib import Path from pathlib import Path
from typing import Any from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai import RunContext from pydantic_ai import RunContext
from pydantic_ai.messages import ToolReturn from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.toolsets import FunctionToolset
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.capabilities._base import ( from haiku.rag.capabilities._base import (
RAGCapabilityBase, RAGCapabilityBase,
resolve_db_path, resolve_db_path,
@ -72,6 +75,7 @@ def create_capability(
config: AppConfig | None = None, config: AppConfig | None = None,
*, *,
defer_loading: bool = True, defer_loading: bool = True,
rag: "HaikuRAG | None" = None,
request_limit: int | None = 20, request_limit: int | None = 20,
vision: bool | None = None, vision: bool | None = None,
) -> RAGCapability: ) -> RAGCapability:
@ -88,6 +92,7 @@ def create_capability(
return RAGCapability( return RAGCapability(
db_path=resolve_db_path(db_path, config), db_path=resolve_db_path(db_path, config),
config=config, config=config,
borrowed_rag=rag,
state_type=RAGState, state_type=RAGState,
state_namespace=STATE_NAMESPACE, state_namespace=STATE_NAMESPACE,
instruction_text=instructions(), instruction_text=instructions(),

View file

@ -63,6 +63,7 @@ async def ask(
capability = create_capability( capability = create_capability(
db_path=client.store.db_path, db_path=client.store.db_path,
config=client._config, config=client._config,
rag=client,
defer_loading=False, defer_loading=False,
) )
deps = _AgentDeps( deps = _AgentDeps(
@ -115,6 +116,7 @@ async def analyze(
capability = create_capability( capability = create_capability(
db_path=client.store.db_path, db_path=client.store.db_path,
config=client._config, config=client._config,
rag=client,
defer_loading=False, defer_loading=False,
) )
deps = _AgentDeps( deps = _AgentDeps(

View file

@ -0,0 +1,132 @@
import pytest
from haiku.rag.capabilities.rag import create_capability
from haiku.rag.client import HaikuRAG
@pytest.mark.asyncio
async def test_a_borrowed_client_is_reused_not_reopened(temp_db_path, monkeypatch):
"""A capability handed a client must not open a second connection to the
same database."""
from haiku.rag.store.engine import Store
async with HaikuRAG(temp_db_path, create=True) as client:
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
capability = create_capability(
db_path=client.store.db_path, config=client._config, rag=client
)
assert await capability._ensure_rag() is client
assert opens == 0
@pytest.mark.asyncio
async def test_closing_never_closes_a_borrowed_client(temp_db_path):
"""`_close` owns only what it opened. Closing the caller's client would be a
use-after-close for the caller."""
async with HaikuRAG(temp_db_path, create=True) as client:
capability = create_capability(
db_path=client.store.db_path, config=client._config, rag=client
)
await capability._ensure_rag()
await capability._close()
# Still usable by its owner.
assert await client.list_documents() == []
@pytest.mark.asyncio
async def test_a_borrowed_client_survives_for_run(temp_db_path):
"""for_run clears the owned connection per run; a borrowed one is the
caller's and carries into the run copy."""
from tests.capabilities.test_capabilities import Deps, make_context
async with HaikuRAG(temp_db_path, create=True) as client:
capability = create_capability(
db_path=client.store.db_path, config=client._config, rag=client
)
run_capability = await capability.for_run(make_context(Deps()))
assert run_capability is not capability
assert run_capability.rag is None
assert run_capability.borrowed_rag is client
assert await run_capability._ensure_rag() is client
@pytest.mark.asyncio
async def test_ask_hands_its_client_to_the_capability(temp_db_path, monkeypatch):
"""`ask` built the capability from a db_path alone, so the capability opened
its own connection to a database the client already had open."""
from haiku.rag.capabilities import rag as rag_capability
from haiku.rag.store.engine import Store
real = rag_capability.create_capability
built = {}
def spy(**kwargs):
built["capability"] = real(**kwargs)
raise RuntimeError("stop before running the agent")
async with HaikuRAG(temp_db_path, create=True) as client:
monkeypatch.setattr(rag_capability, "create_capability", spy)
with pytest.raises(RuntimeError, match="stop before running the agent"):
await client.ask("anything")
capability = built["capability"]
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
assert await capability._ensure_rag() is client
assert opens == 0
@pytest.mark.asyncio
async def test_analyze_hands_its_client_to_the_capability(temp_db_path, monkeypatch):
from haiku.rag.capabilities import analysis as analysis_capability
from haiku.rag.store.engine import Store
real = analysis_capability.create_capability
built = {}
def spy(**kwargs):
built["capability"] = real(**kwargs)
raise RuntimeError("stop before running the agent")
async with HaikuRAG(temp_db_path, create=True) as client:
monkeypatch.setattr(analysis_capability, "create_capability", spy)
with pytest.raises(RuntimeError, match="stop before running the agent"):
await client.analyze("anything")
capability = built["capability"]
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
assert await capability._ensure_rag() is client
assert opens == 0