Always create the shared-connection lock so serialization can't be skipped
This commit is contained in:
parent
458f365e29
commit
4450c1c908
3 changed files with 9 additions and 12 deletions
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
|
|
@ -18,8 +18,9 @@ class RAGRunDeps(SkillRunDeps):
|
|||
rag: "HaikuRAG | None" = None
|
||||
# pydantic-ai runs a turn's tool calls concurrently; LanceDB's per-connection
|
||||
# state cannot take two in-flight operations at once, so every use of ``rag``
|
||||
# (skill tools and the analysis sandbox) serializes through this lock.
|
||||
rag_lock: "asyncio.Lock | None" = None
|
||||
# (skill tools and the analysis sandbox) serializes through this lock. Always
|
||||
# present so serialization is never accidentally skipped.
|
||||
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||||
search_count: int = 0
|
||||
|
||||
|
||||
|
|
@ -55,7 +56,6 @@ def make_rag_lifespan(db_path: Path, config: AppConfig):
|
|||
|
||||
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
||||
deps.rag = rag
|
||||
deps.rag_lock = asyncio.Lock()
|
||||
deps.search_count = 0
|
||||
_reset_invocation_state(deps.state)
|
||||
yield
|
||||
|
|
@ -72,7 +72,6 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
|
|||
doc_filter = getattr(deps.state, "document_filter", None)
|
||||
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
||||
deps.rag = rag
|
||||
deps.rag_lock = asyncio.Lock()
|
||||
deps.search_count = 0
|
||||
sandbox = Sandbox(
|
||||
db_path=db_path,
|
||||
|
|
|
|||
|
|
@ -91,9 +91,9 @@ async def _serialized(ctx: RunContext[RAGRunDeps]) -> AsyncIterator[None]:
|
|||
"""Serialize access to the shared connection through the run's 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.
|
||||
per-connection state cannot take two in-flight operations at once. The lock
|
||||
is always present (``RAGRunDeps`` creates one by default); the no-op branch
|
||||
is a guard for a missing deps/lock.
|
||||
"""
|
||||
lock = ctx.deps.rag_lock if ctx.deps is not None else None
|
||||
if lock is None:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import asyncio
|
||||
import random
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
|
@ -18,11 +17,10 @@ def _make_ctx(state=None, rag=None, sandbox=None):
|
|||
from haiku.rag.skills.analysis import AnalysisState
|
||||
|
||||
ctx = MagicMock(spec=RunContext)
|
||||
lock = asyncio.Lock()
|
||||
if isinstance(state, AnalysisState) or sandbox is not None:
|
||||
ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox, rag_lock=lock)
|
||||
ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox)
|
||||
else:
|
||||
ctx.deps = RAGRunDeps(state=state, rag=rag, rag_lock=lock)
|
||||
ctx.deps = RAGRunDeps(state=state, rag=rag)
|
||||
return ctx
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue