danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

open one HaikuRAG client per skill invocation via lifespan

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-22 13:13:21 +03:00
parent 8a23108fbd
commit 4e9c02afc2
No known key found for this signature in database
10 changed files with 209 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -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

36
haiku_rag_slim/haiku/rag/skills/_deps.py Normal file
View file

@ -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

115
haiku_rag_slim/haiku/rag/skills/_tools.py
View file

@ -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

3
haiku_rag_slim/haiku/rag/skills/analysis.py
View file

@ -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),
)

3
haiku_rag_slim/haiku/rag/skills/rag.py
View file

@ -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),
)

2
haiku_rag_slim/pyproject.toml
View file

@ -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",

20
tests/skills/conftest.py
View file

@ -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

25
tests/skills/test_analysis.py
View file

@ -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

88
tests/skills/test_rag.py
View file

@ -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

8
uv.lock
View file

@ -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]]