78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.skills.state import SkillRunDeps
|
|
|
|
if TYPE_CHECKING:
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.sandbox import Sandbox
|
|
|
|
|
|
@dataclass
|
|
class RAGRunDeps(SkillRunDeps):
|
|
rag: "HaikuRAG | None" = None
|
|
search_count: int = 0
|
|
|
|
|
|
@dataclass
|
|
class AnalysisRunDeps(RAGRunDeps):
|
|
sandbox: "Sandbox | None" = None
|
|
|
|
|
|
def _reset_invocation_state(state: Any) -> None:
|
|
"""Clear state fields scoped to a single invocation.
|
|
|
|
Keeps ``citation_index`` (accumulates resolved citations across the session
|
|
for lookup) and ``document_filter`` (session-level). Clears ``citations``,
|
|
``searches``, and (for analysis) ``executions``.
|
|
"""
|
|
if state is None:
|
|
return
|
|
citations = getattr(state, "citations", None)
|
|
if citations is not None:
|
|
citations.clear()
|
|
searches = getattr(state, "searches", None)
|
|
if searches is not None:
|
|
searches.clear()
|
|
executions = getattr(state, "executions", None)
|
|
if executions is not None:
|
|
executions.clear()
|
|
|
|
|
|
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
|
|
_reset_invocation_state(deps.state)
|
|
yield
|
|
|
|
return lifespan
|
|
|
|
|
|
def make_analysis_lifespan(db_path: Path, config: AppConfig):
|
|
@asynccontextmanager
|
|
async def lifespan(deps: AnalysisRunDeps) -> AsyncIterator[None]:
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
|
|
|
doc_filter = getattr(deps.state, "document_filter", None)
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
deps.rag = rag
|
|
deps.search_count = 0
|
|
deps.sandbox = Sandbox(
|
|
db_path=db_path,
|
|
config=config,
|
|
context=AnalysisContext(filter=doc_filter),
|
|
)
|
|
_reset_invocation_state(deps.state)
|
|
yield
|
|
|
|
return lifespan
|