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

Make the shared-connection lock optional for direct tool calls

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-06-05 18:04:56 +03:00
parent e59ee56002
commit 4a65a69a27
No known key found for this signature in database
1 changed files with 21 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,4 +1,5 @@
import asyncio from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@ -85,11 +86,21 @@ def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
return ctx.deps.rag return ctx.deps.rag
def _require_lock(ctx: RunContext[RAGRunDeps]) -> asyncio.Lock: @asynccontextmanager
assert ctx.deps is not None and ctx.deps.rag_lock is not None, ( async def _serialized(ctx: RunContext[RAGRunDeps]) -> AsyncIterator[None]:
"RAGRunDeps.rag_lock is not set — skill lifespan must run before tools." """Serialize access to the shared connection through the run's lock.
)
return ctx.deps.rag_lock pydantic-ai runs a turn's tool calls concurrently and LanceDB's
per-connection state cannot take two in-flight operations at once. When no
lock is present (tools invoked directly, without a skill lifespan) there is
no concurrency to guard, so this is a no-op.
"""
lock = ctx.deps.rag_lock if ctx.deps is not None else None
if lock is None:
yield
else:
async with lock:
yield
def _register_citations(state: Any, citations: "list[Citation]") -> None: def _register_citations(state: Any, citations: "list[Citation]") -> None:
@ -203,7 +214,7 @@ def create_skill_tools(
) )
state = _get_state(ctx, state_type) state = _get_state(ctx, state_type)
async with _require_lock(ctx): async with _serialized(ctx):
formatted, results = await skill_search( formatted, results = await skill_search(
_require_rag(ctx), _require_rag(ctx),
query, query,
@ -230,7 +241,7 @@ def create_skill_tools(
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""List all documents in the knowledge base.""" """List all documents in the knowledge base."""
state = _get_state(ctx, state_type) state = _get_state(ctx, state_type)
async with _require_lock(ctx): async with _serialized(ctx):
return await skill_list_documents( return await skill_list_documents(
_require_rag(ctx), _require_rag(ctx),
filter=state.document_filter if state else None, filter=state.document_filter if state else None,
@ -248,7 +259,7 @@ def create_skill_tools(
Args: Args:
query: Document ID, title, or URI to look up. query: Document ID, title, or URI to look up.
""" """
async with _require_lock(ctx): async with _serialized(ctx):
return await skill_get_document(_require_rag(ctx), query) return await skill_get_document(_require_rag(ctx), query)
tools["get_document"] = get_document tools["get_document"] = get_document
@ -336,7 +347,7 @@ def create_skill_tools(
] ]
if missing: if missing:
async with _require_lock(ctx): async with _serialized(ctx):
rag = _require_rag(ctx) rag = _require_rag(ctx)
synthetic: list[SearchResult] = [] synthetic: list[SearchResult] = []
doc_cache: dict[str, Any] = {} doc_cache: dict[str, Any] = {}