diff --git a/CHANGELOG.md b/CHANGELOG.md index fbc1aed2..7bdb3a64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,11 @@ # Changelog ## [Unreleased] +### Changed + +- **Skills share a single `HaikuRAG` client per invocation** via the new `haiku.skills>=0.15.0` `lifespan` hook. The skill's sub-agent opens one read-only client on entry, all tool calls reuse it, and it closes on exit — replacing the old pattern of open/close around every `search` / `list_documents` / `get_document` call. +- **`max_searches` tracked on `RAGRunDeps.search_count`** instead of a module-level `ctx.run_id`-keyed dict. Eliminates a memory leak in long-running processes where old run ids were never evicted. + ## [0.41.0] - 2026-04-20 ### Added diff --git a/haiku_rag_slim/haiku/rag/skills/_deps.py b/haiku_rag_slim/haiku/rag/skills/_deps.py new file mode 100644 index 00000000..eb20f6df --- /dev/null +++ b/haiku_rag_slim/haiku/rag/skills/_deps.py @@ -0,0 +1,36 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from haiku.rag.config.models import AppConfig +from haiku.skills.state import SkillRunDeps + +if TYPE_CHECKING: + from haiku.rag.agents.analysis.sandbox import Sandbox + from haiku.rag.client import HaikuRAG + + +@dataclass +class RAGRunDeps(SkillRunDeps): + rag: "HaikuRAG | None" = None + search_count: int = 0 + + +@dataclass +class AnalysisRunDeps(RAGRunDeps): + sandbox: "Sandbox | None" = None + + +def make_rag_lifespan(db_path: Path, config: AppConfig): + @asynccontextmanager + async def lifespan(deps: RAGRunDeps) -> AsyncIterator[None]: + from haiku.rag.client import HaikuRAG + + async with HaikuRAG(db_path, config=config, read_only=True) as rag: + deps.rag = rag + deps.search_count = 0 + yield + + return lifespan diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 34bd977d..909ec65d 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/haiku_rag_slim/haiku/rag/skills/_tools.py @@ -5,9 +5,10 @@ from pydantic import BaseModel from pydantic_ai import RunContext from haiku.rag.agents.research.models import Citation +from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig +from haiku.rag.skills._deps import RAGRunDeps from haiku.rag.store.models.chunk import SearchResult -from haiku.skills.state import SkillRunDeps class CodeExecutionEntry(BaseModel): @@ -18,22 +19,13 @@ class CodeExecutionEntry(BaseModel): async def skill_search( - db_path: Path, - config: AppConfig, + rag: HaikuRAG, query: str, limit: int | None = None, document_filter: str | None = None, ) -> tuple[str, list[SearchResult]]: - from haiku.rag.client import HaikuRAG - - async with HaikuRAG(db_path, config=config, read_only=True) as rag: - results = await rag.search( - query, - limit=limit, - filter=document_filter, - ) - results = await rag.expand_context(results) - + results = await rag.search(query, limit=limit, filter=document_filter) + results = await rag.expand_context(results) formatted = "\n\n---\n\n".join( r.format_for_agent(rank=i + 1, total=len(results)) for i, r in enumerate(results) @@ -42,55 +34,55 @@ async def skill_search( async def skill_list_documents( - db_path: Path, - config: AppConfig, + rag: HaikuRAG, filter: str | None = None, ) -> list[dict[str, Any]]: - from haiku.rag.client import HaikuRAG - - async with HaikuRAG(db_path, config=config, read_only=True) as rag: - documents = await rag.list_documents(filter=filter) - return [ - { - "id": doc.id, - "title": doc.title, - "uri": doc.uri, - "metadata": doc.metadata, - "created_at": str(doc.created_at), - "updated_at": str(doc.updated_at), - } - for doc in documents - ] + documents = await rag.list_documents(filter=filter) + return [ + { + "id": doc.id, + "title": doc.title, + "uri": doc.uri, + "metadata": doc.metadata, + "created_at": str(doc.created_at), + "updated_at": str(doc.updated_at), + } + for doc in documents + ] async def skill_get_document( - db_path: Path, - config: AppConfig, + rag: HaikuRAG, query: str, ) -> dict[str, Any] | None: - from haiku.rag.client import HaikuRAG - - async with HaikuRAG(db_path, config=config, read_only=True) as rag: - document = await rag.resolve_document(query) - if document is None: - return None - return { - "id": document.id, - "content": document.content, - "title": document.title, - "uri": document.uri, - "metadata": document.metadata, - "created_at": str(document.created_at), - "updated_at": str(document.updated_at), - } + document = await rag.resolve_document(query) + if document is None: + return None + return { + "id": document.id, + "content": document.content, + "title": document.title, + "uri": document.uri, + "metadata": document.metadata, + "created_at": str(document.created_at), + "updated_at": str(document.updated_at), + } -def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> Any: +def _get_state(ctx: RunContext[RAGRunDeps], state_type: type[BaseModel]) -> Any: if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, state_type): return ctx.deps.state return None +def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG: + if ctx.deps is None or ctx.deps.rag is None: + raise RuntimeError( + "RAGRunDeps.rag is not set — skill lifespan must run before tools." + ) + return ctx.deps.rag + + def _register_citations(state: Any, citations: "list[Citation]") -> None: """Add citations to the index and record the turn's chunk IDs.""" chunk_ids = [] @@ -174,10 +166,9 @@ def create_skill_tools( if "search" in tool_names: max_searches = config.qa.max_searches - search_counts: dict[str, int] = {} async def search( - ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None + ctx: RunContext[RAGRunDeps], query: str, limit: int | None = None ) -> str: """Search the knowledge base using hybrid search (vector + full-text). @@ -187,9 +178,8 @@ def create_skill_tools( query: The search query. limit: Maximum number of results. """ - rid = ctx.run_id or "" - search_counts[rid] = search_counts.get(rid, 0) + 1 - if search_counts[rid] > max_searches: + ctx.deps.search_count += 1 + if ctx.deps.search_count > max_searches: return ( "Search limit reached. Answer the question using " "the results you already have." @@ -197,8 +187,7 @@ def create_skill_tools( state = _get_state(ctx, state_type) formatted, results = await skill_search( - db_path, - config, + _require_rag(ctx), query, limit=limit, document_filter=state.document_filter if state else None, @@ -212,36 +201,34 @@ def create_skill_tools( if "list_documents" in tool_names: async def list_documents( - ctx: RunContext[SkillRunDeps], + ctx: RunContext[RAGRunDeps], ) -> list[dict[str, Any]]: """List all documents in the knowledge base.""" state = _get_state(ctx, state_type) - result = await skill_list_documents( - db_path, - config, + return await skill_list_documents( + _require_rag(ctx), filter=state.document_filter if state else None, ) - return result tools["list_documents"] = list_documents if "get_document" in tool_names: async def get_document( - ctx: RunContext[SkillRunDeps], query: str + ctx: RunContext[RAGRunDeps], query: str ) -> dict[str, Any] | None: """Retrieve a document by ID, title, or URI. Args: query: Document ID, title, or URI to look up. """ - return await skill_get_document(db_path, config, query) + return await skill_get_document(_require_rag(ctx), query) tools["get_document"] = get_document if "execute_code" in tool_names: - async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str: + async def execute_code(ctx: RunContext[RAGRunDeps], code: str) -> str: """Execute Python code in a sandboxed interpreter. The code has access to search(), list_documents(), llm() functions @@ -291,7 +278,7 @@ def create_skill_tools( if "cite" in tool_names: - async def cite(ctx: RunContext[SkillRunDeps], chunk_ids: list[str]) -> str: + async def cite(ctx: RunContext[RAGRunDeps], chunk_ids: list[str]) -> str: """Register chunk IDs as citations for your answer. Call this after searching, with the chunk_id values from search diff --git a/haiku_rag_slim/haiku/rag/skills/analysis.py b/haiku_rag_slim/haiku/rag/skills/analysis.py index 63d2ea2c..c74637b3 100644 --- a/haiku_rag_slim/haiku/rag/skills/analysis.py +++ b/haiku_rag_slim/haiku/rag/skills/analysis.py @@ -60,6 +60,7 @@ def create_skill( config: haiku.rag AppConfig instance. If None, uses get_config(). """ from haiku.rag.config import get_config + from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan from haiku.rag.skills._tools import create_skill_extras, create_skill_tools if config is None: @@ -93,4 +94,6 @@ def create_skill( extras=extras, state_type=STATE_TYPE, state_namespace=STATE_NAMESPACE, + deps_type=AnalysisRunDeps, + lifespan=make_rag_lifespan(db_path, config), ) diff --git a/haiku_rag_slim/haiku/rag/skills/rag.py b/haiku_rag_slim/haiku/rag/skills/rag.py index a660c723..16fe1e38 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag.py +++ b/haiku_rag_slim/haiku/rag/skills/rag.py @@ -75,6 +75,7 @@ def create_skill( config: haiku.rag AppConfig instance. If None, uses get_config(). """ from haiku.rag.config import get_config + from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan from haiku.rag.skills._tools import create_skill_extras, create_skill_tools if config is None: @@ -103,4 +104,6 @@ def create_skill( extras=extras, state_type=STATE_TYPE, state_namespace=STATE_NAMESPACE, + deps_type=RAGRunDeps, + lifespan=make_rag_lifespan(db_path, config), ) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 7ffaf63d..dde74a44 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ dependencies = [ "docling-core>=2.71.0,<2.72", - "haiku.skills>=0.14.0", + "haiku.skills>=0.15.0", "httpx>=0.28.1", "jinja2>=3.1.0", "jsonpatch>=1.33", diff --git a/tests/skills/conftest.py b/tests/skills/conftest.py index 71a7a4c9..3bed1278 100644 --- a/tests/skills/conftest.py +++ b/tests/skills/conftest.py @@ -7,15 +7,20 @@ from pydantic_ai import RunContext from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.embeddings import EmbedderWrapper -from haiku.skills.state import SkillRunDeps +from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps VECTOR_DIM = 2560 -def _make_ctx(state=None): - """Create a mock RunContext with SkillRunDeps.""" +def _make_ctx(state=None, rag=None): + """Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState).""" + from haiku.rag.skills.analysis import AnalysisState + ctx = MagicMock(spec=RunContext) - ctx.deps = SkillRunDeps(state=state) + if isinstance(state, AnalysisState): + ctx.deps = AnalysisRunDeps(state=state, rag=rag) + else: + ctx.deps = RAGRunDeps(state=state, rag=rag) return ctx @@ -68,3 +73,10 @@ async def rag_db(temp_db_path): uri="test://ml-basics", ) return temp_db_path + + +@pytest.fixture +async def rag_client(rag_db): + """Yield an open read-only HaikuRAG client on the sample db.""" + async with HaikuRAG(rag_db, read_only=True) as rag: + yield rag diff --git a/tests/skills/test_analysis.py b/tests/skills/test_analysis.py index e2b66799..dffc5787 100644 --- a/tests/skills/test_analysis.py +++ b/tests/skills/test_analysis.py @@ -192,3 +192,28 @@ class TestExecuteCodeTool: ) assert "Error" in result assert "read-only" in result + + +class TestAnalysisLifespan: + async def test_opens_one_client_per_invocation(self, rag_db): + from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan + + config = AppConfig() + lifespan = make_rag_lifespan(rag_db, config) + deps = AnalysisRunDeps() + async with lifespan(deps): + assert deps.rag is not None + assert deps.rag.is_read_only + assert deps.search_count == 0 + docs = await deps.rag.list_documents() + assert len(docs) == 2 + + async def test_skill_has_lifespan_and_deps_type( + self, test_app_config, temp_db_path + ): + from haiku.rag.skills._deps import AnalysisRunDeps + from haiku.rag.skills.analysis import create_skill + + skill = create_skill(config=test_app_config, db_path=temp_db_path) + assert skill.deps_type is AnalysisRunDeps + assert skill.lifespan is not None diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py index 6fee0262..e73cb34d 100644 --- a/tests/skills/test_rag.py +++ b/tests/skills/test_rag.py @@ -157,50 +157,50 @@ class TestSkillExtras: class TestSearchTool: - async def test_search_returns_formatted_string(self, rag_db): + async def test_search_returns_formatted_string(self, rag_db, rag_client): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") - ctx = _make_ctx() + ctx = _make_ctx(rag=rag_client) result = await search(ctx, query="artificial intelligence") assert isinstance(result, str) assert len(result) > 0 - async def test_search_updates_state(self, rag_db): + async def test_search_updates_state(self, rag_db, rag_client): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") state = RAGState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, rag=rag_client) await search(ctx, query="artificial intelligence") assert "artificial intelligence" in state.searches results = state.searches["artificial intelligence"] assert len(results) > 0 assert isinstance(results[0], SearchResult) - async def test_search_applies_document_filter_from_state(self, rag_db): + async def test_search_applies_document_filter_from_state(self, rag_db, rag_client): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") state = RAGState(document_filter="title = 'AI Overview'") - ctx = _make_ctx(state) + ctx = _make_ctx(state, rag=rag_client) result = await search(ctx, query="artificial intelligence") assert "AI Overview" in result assert "ML Basics" not in result - async def test_search_without_state(self, rag_db): + async def test_search_without_state(self, rag_db, rag_client): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") - ctx = _make_ctx(state=None) + ctx = _make_ctx(state=None, rag=rag_client) result = await search(ctx, query="artificial intelligence") assert isinstance(result, str) - async def test_search_rate_limited(self, rag_db): + async def test_search_rate_limited(self, rag_db, rag_client): from haiku.rag.skills.rag import RAGState, create_skill config = AppConfig() @@ -208,69 +208,71 @@ class TestSearchTool: skill = create_skill(db_path=rag_db, config=config) search = _get_tool(skill, "search") state = RAGState() - ctx = _make_ctx(state) - ctx.run_id = "test-run" + ctx = _make_ctx(state, rag=rag_client) await search(ctx, query="first") await search(ctx, query="second") result = await search(ctx, query="third") assert "Search limit reached" in result + assert ctx.deps.search_count == 3 assert len(state.searches) == 2 class TestListDocumentsTool: - async def test_list_documents_returns_results(self, rag_db): + async def test_list_documents_returns_results(self, rag_db, rag_client): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) list_docs = _get_tool(skill, "list_documents") - ctx = _make_ctx() + ctx = _make_ctx(rag=rag_client) results = await list_docs(ctx) assert isinstance(results, list) assert len(results) == 2 - async def test_list_documents_applies_document_filter_from_state(self, rag_db): + async def test_list_documents_applies_document_filter_from_state( + self, rag_db, rag_client + ): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) list_docs = _get_tool(skill, "list_documents") state = RAGState(document_filter="title = 'AI Overview'") - ctx = _make_ctx(state) + ctx = _make_ctx(state, rag=rag_client) results = await list_docs(ctx) assert len(results) == 1 assert results[0]["title"] == "AI Overview" class TestGetDocumentTool: - async def test_get_document_by_title(self, rag_db): + async def test_get_document_by_title(self, rag_db, rag_client): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) get_doc = _get_tool(skill, "get_document") - ctx = _make_ctx() + ctx = _make_ctx(rag=rag_client) result = await get_doc(ctx, query="AI Overview") assert result is not None assert result["title"] == "AI Overview" - async def test_get_document_not_found(self, rag_db): + async def test_get_document_not_found(self, rag_db, rag_client): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) get_doc = _get_tool(skill, "get_document") - ctx = _make_ctx() + ctx = _make_ctx(rag=rag_client) result = await get_doc(ctx, query="nonexistent document xyz") assert result is None class TestCiteTool: - async def test_cite_registers_citations(self, rag_db): + async def test_cite_registers_citations(self, rag_db, rag_client): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") cite = _get_tool(skill, "cite") state = RAGState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, rag=rag_client) await search(ctx, query="artificial intelligence") chunk_ids = [ @@ -286,14 +288,14 @@ class TestCiteTool: assert len(state.citations[0]) == 2 assert all(cid in state.citation_index for cid in chunk_ids) - async def test_cite_deduplicates_in_index(self, rag_db): + async def test_cite_deduplicates_in_index(self, rag_db, rag_client): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") cite = _get_tool(skill, "cite") state = RAGState() - ctx = _make_ctx(state) + ctx = _make_ctx(state, rag=rag_client) await search(ctx, query="artificial intelligence") chunk_ids = [ @@ -316,3 +318,43 @@ class TestCiteTool: ctx = _make_ctx(state=None) result = await cite(ctx, chunk_ids=["nonexistent"]) assert "No state" in result + + +class TestLifespan: + async def test_opens_one_client_per_invocation(self, rag_db): + """Lifespan opens one HaikuRAG client, available on ctx.deps.rag throughout.""" + from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan + + config = AppConfig() + lifespan = make_rag_lifespan(rag_db, config) + deps = RAGRunDeps() + async with lifespan(deps): + assert deps.rag is not None + assert deps.rag.is_read_only + assert deps.search_count == 0 + docs = await deps.rag.list_documents() + assert len(docs) == 2 + # after exit the client has been closed; field still references it + assert deps.rag is not None + + async def test_search_count_resets_per_invocation(self, rag_db): + from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan + + config = AppConfig() + lifespan = make_rag_lifespan(rag_db, config) + + deps = RAGRunDeps(search_count=42) + async with lifespan(deps): + assert deps.search_count == 0 + + deps2 = RAGRunDeps(search_count=5) + async with lifespan(deps2): + assert deps2.search_count == 0 + + def test_skill_has_lifespan_and_deps_type(self, test_app_config, temp_db_path): + from haiku.rag.skills._deps import RAGRunDeps + from haiku.rag.skills.rag import create_skill + + skill = create_skill(config=test_app_config, db_path=temp_db_path) + assert skill.deps_type is RAGRunDeps + assert skill.lifespan is not None diff --git a/uv.lock b/uv.lock index 50d11777..c7287fb4 100644 --- a/uv.lock +++ b/uv.lock @@ -1570,7 +1570,7 @@ requires-dist = [ { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.21.1" }, { name = "docling", marker = "extra == 'docling'", specifier = ">=2.84.0" }, { name = "docling-core", specifier = ">=2.71.0,<2.72" }, - { name = "haiku-skills", specifier = ">=0.14.0" }, + { name = "haiku-skills", specifier = ">=0.15.0" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "jinja2", specifier = ">=3.1.0" }, { name = "jsonpatch", specifier = ">=1.33" }, @@ -1604,7 +1604,7 @@ provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jin [[package]] name = "haiku-skills" -version = "0.14.0" +version = "0.15.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ag-ui-protocol" }, @@ -1614,9 +1614,9 @@ dependencies = [ { name = "pyyaml" }, { name = "skills-ref" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/89/c4/82a6b82f70726a2e759aad4c6c553309f2cc2ca7f3c157b8a15fd14da709/haiku_skills-0.14.0.tar.gz", hash = "sha256:27074a171060a0ecae6b89c6b7756b3b8ed0dfb957a3f5076ab1d158e68feb82", size = 250637, upload-time = "2026-04-16T08:47:33.152Z" } +sdist = { url = "https://files.pythonhosted.org/packages/86/a1/e2bd00a72d002f9db1c53c068167ed436a457dae0f8996399f116c087f6a/haiku_skills-0.15.0.tar.gz", hash = "sha256:ce93e6846e05397f5d96c144f956edd395b9e7308cb5cef6213c49c183a08bbc", size = 252030, upload-time = "2026-04-22T09:00:13.775Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e5/7a/bc53abcbae8bf1aa013379f49f6690fbaff55f80a7b997824103a5f44384/haiku_skills-0.14.0-py3-none-any.whl", hash = "sha256:698d0012bcf06f43499c30aaf6a3b7b772c66737696bfa3f93a6b113fb0f6b17", size = 31613, upload-time = "2026-04-16T08:47:31.68Z" }, + { url = "https://files.pythonhosted.org/packages/9b/68/3df2c9761fc0b0592b60f4723c87adeda5f335ea0835b4cf77ca07784379/haiku_skills-0.15.0-py3-none-any.whl", hash = "sha256:a1771b16e0ffe7da775f28d38c704e029f8e8757791616d7f27e73feb2c16fd0", size = 32041, upload-time = "2026-04-22T09:00:12.824Z" }, ] [[package]]