diff --git a/CHANGELOG.md b/CHANGELOG.md index 310f5ecd..bce01898 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ ### 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. +- `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. - `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]`. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index b9f0ff60..2f0bea64 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -129,6 +129,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): state: StateT | 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) + """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) resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) search_count: int = field(default=0, repr=False) @@ -347,6 +350,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): raise error async def _ensure_rag(self) -> HaikuRAG: + if self.borrowed_rag is not None: + return self.borrowed_rag if self.rag is None: async with self.resource_lock: if self.rag is None: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index e2de8868..b75fa06c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -1,13 +1,16 @@ from dataclasses import dataclass, field from functools import cache from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from pydantic_ai import RunContext, ToolFailed from pydantic_ai.messages import ToolReturn from pydantic_ai.toolsets import FunctionToolset +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + from haiku.rag.capabilities._base import ( CodeExecutionEntry, RAGCapabilityBase, @@ -162,6 +165,7 @@ def create_capability( config: AppConfig | None = None, *, defer_loading: bool = True, + rag: "HaikuRAG | None" = None, request_limit: int | None = 30, vision: bool | None = None, ) -> AnalysisCapability: @@ -180,6 +184,7 @@ def create_capability( return AnalysisCapability( db_path=resolve_db_path(db_path, config), config=config, + borrowed_rag=rag, state_type=AnalysisState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 79216a87..f8ea44e9 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -1,13 +1,16 @@ from dataclasses import dataclass from functools import cache from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from pydantic_ai import RunContext from pydantic_ai.messages import ToolReturn from pydantic_ai.toolsets import FunctionToolset +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, @@ -72,6 +75,7 @@ def create_capability( config: AppConfig | None = None, *, defer_loading: bool = True, + rag: "HaikuRAG | None" = None, request_limit: int | None = 20, vision: bool | None = None, ) -> RAGCapability: @@ -88,6 +92,7 @@ def create_capability( return RAGCapability( db_path=resolve_db_path(db_path, config), config=config, + borrowed_rag=rag, state_type=RAGState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index 9a739f74..c85fa1ab 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -63,6 +63,7 @@ async def ask( capability = create_capability( db_path=client.store.db_path, config=client._config, + rag=client, defer_loading=False, ) deps = _AgentDeps( @@ -115,6 +116,7 @@ async def analyze( capability = create_capability( db_path=client.store.db_path, config=client._config, + rag=client, defer_loading=False, ) deps = _AgentDeps( diff --git a/tests/capabilities/test_borrowed_client.py b/tests/capabilities/test_borrowed_client.py new file mode 100644 index 00000000..af1dcec6 --- /dev/null +++ b/tests/capabilities/test_borrowed_client.py @@ -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