Merge pull request #349 from ggozad/fix/convert-urlparse
fix convert() misreading text content that starts with a URL
This commit is contained in:
commit
e736a8c73a
16 changed files with 562 additions and 174 deletions
11
CHANGELOG.md
11
CHANGELOG.md
|
|
@ -1,6 +1,17 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### Fixed
|
||||
|
||||
- **`create_document`, `update_document`, and rebuild (`RECHUNK` / full fallback) no longer misread URL-prefixed text as a URL to fetch.** These paths passed known-text content through `HaikuRAG.convert()`, which dispatches on `urlparse(source).scheme`; text whose first line was `https://...` (common for clipped web pages and notes) got handed to `httpx.get` and crashed with `httpx.InvalidURL` on embedded whitespace. Fixed by calling `converter.convert_text(...)` directly at those sites; `convert()` itself is unchanged for `create_document_from_source`.
|
||||
|
||||
### 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.
|
||||
- **Analysis sandbox persists variables across `execute_code` calls within one invocation.** Re-enables the incremental-exploration workflow (search in one call, process results in the next). Each new skill invocation constructs a fresh `Sandbox` via the analysis lifespan, so there is no cross-invocation leak.
|
||||
- **Skill state is scoped to the current invocation.** Lifespans now clear `citations`, `searches`, and (for analysis) `executions` at the start of each invocation, so state deltas sent to the AG-UI client reflect only the in-progress turn. `citation_index` is preserved across invocations so past-turn citation chunk ids remain resolvable, and `document_filter` is preserved as session-level config.
|
||||
|
||||
## [0.41.0] - 2026-04-20
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -37,10 +37,11 @@ class AnalysisState(BaseModel):
|
|||
searches: dict[str, list[SearchResult]] = {}
|
||||
```
|
||||
|
||||
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls.
|
||||
- **executions** — Each `execute_code` call appends a `CodeExecutionEntry` with code, stdout, stderr, and success status.
|
||||
- **citation_index** / **citations** — Same per-turn citation tracking as the RAG skill.
|
||||
- **searches** — Search results from both the `search` tool and sandbox-internal searches.
|
||||
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Persists across invocations as session-level configuration.
|
||||
- **executions** — Each `execute_code` call appends a `CodeExecutionEntry` with code, stdout, stderr, and success status. Cleared at the start of each invocation; mirrors the sandbox lifecycle (variables persist across calls within one invocation, a fresh sandbox is built per invocation).
|
||||
- **citation_index** — Citations indexed by chunk ID. Accumulates across invocations (same semantics as the RAG skill).
|
||||
- **citations** — Cleared at the start of each invocation; holds only the in-progress turn.
|
||||
- **searches** — Search results from both the `search` tool and sandbox-internal searches. Cleared at the start of each invocation.
|
||||
|
||||
## Usage with RAG Skill
|
||||
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ class RAGState(BaseModel):
|
|||
searches: dict[str, list[SearchResult]] = {}
|
||||
```
|
||||
|
||||
- **citation_index** — All citations indexed by chunk ID (deduplicated across turns).
|
||||
- **citations** — Per-turn lists of chunk IDs registered via the `cite` tool.
|
||||
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Set this to scope queries to specific documents.
|
||||
- **searches** — Search results keyed by query string.
|
||||
- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical turns' chunk IDs remain resolvable in the UI scrollback.
|
||||
- **citations** — Chunk IDs registered via the `cite` tool. Cleared at the start of each invocation; holds only the in-progress turn.
|
||||
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Persists across invocations as session-level configuration.
|
||||
- **searches** — Search results keyed by query string. Cleared at the start of each invocation.
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from pathlib import Path
|
|||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
import pydantic_monty
|
||||
from pydantic_monty import CallbackFile, MemoryFile, OSAccess
|
||||
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
|
@ -44,11 +44,12 @@ class Sandbox:
|
|||
and resolved asynchronously on the host.
|
||||
Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
|
||||
|
||||
Each ``execute()`` call runs in a fresh interpreter — variables do not
|
||||
persist between calls.
|
||||
The interpreter uses a REPL session — variables persist across
|
||||
``execute()`` calls within the same Sandbox instance.
|
||||
|
||||
sandbox = Sandbox(db_path, config, context)
|
||||
result = await sandbox.execute("print('hello')")
|
||||
result = await sandbox.execute("x = await search('query')")
|
||||
result = await sandbox.execute("print(x[0]['content'])") # x persists
|
||||
"""
|
||||
|
||||
_db_path: Path
|
||||
|
|
@ -56,6 +57,8 @@ class Sandbox:
|
|||
_context: AnalysisContext
|
||||
_search_results: "list[SearchResult]"
|
||||
_items_cache: dict[str, str] | None
|
||||
_repl: MontyRepl | None
|
||||
_vfs: OSAccess | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -68,6 +71,8 @@ class Sandbox:
|
|||
self._context = context
|
||||
self._search_results = []
|
||||
self._items_cache = None
|
||||
self._repl = None
|
||||
self._vfs = None
|
||||
|
||||
def _build_external_functions(self) -> dict[str, Any]:
|
||||
"""Build async external functions for the Monty interpreter."""
|
||||
|
|
@ -245,37 +250,46 @@ class Sandbox:
|
|||
|
||||
return OSAccess(files)
|
||||
|
||||
async def execute(self, code: str) -> SandboxResult:
|
||||
"""Execute Python code in the Monty interpreter."""
|
||||
external_fns = self._build_external_functions()
|
||||
vfs = await self._build_vfs()
|
||||
|
||||
input_names: list[str] = []
|
||||
inputs: dict[str, Any] | None = None
|
||||
if self._context.documents:
|
||||
input_names.append("documents")
|
||||
inputs = {
|
||||
"documents": [
|
||||
{
|
||||
"id": d.id,
|
||||
"title": d.title,
|
||||
"uri": d.uri,
|
||||
"content": d.content,
|
||||
}
|
||||
for d in self._context.documents
|
||||
]
|
||||
}
|
||||
|
||||
try:
|
||||
monty = pydantic_monty.Monty(
|
||||
code,
|
||||
inputs=input_names,
|
||||
async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]:
|
||||
"""Initialize the REPL session and VFS on first use."""
|
||||
if self._repl is None:
|
||||
self._vfs = await self._build_vfs()
|
||||
self._repl = MontyRepl(
|
||||
limits={
|
||||
"max_duration_secs": self._config.analysis.code_timeout,
|
||||
},
|
||||
)
|
||||
except (
|
||||
pydantic_monty.MontySyntaxError,
|
||||
pydantic_monty.MontyRuntimeError,
|
||||
) as e:
|
||||
return SandboxResult(stdout="", stderr=str(e), success=False)
|
||||
if self._context.documents:
|
||||
await pydantic_monty.run_repl_async(
|
||||
self._repl,
|
||||
"pass",
|
||||
inputs={
|
||||
"documents": [
|
||||
{
|
||||
"id": d.id,
|
||||
"title": d.title,
|
||||
"uri": d.uri,
|
||||
"content": d.content,
|
||||
}
|
||||
for d in self._context.documents
|
||||
]
|
||||
},
|
||||
external_functions=self._build_external_functions(),
|
||||
os=self._vfs,
|
||||
)
|
||||
repl = self._repl
|
||||
vfs = self._vfs
|
||||
if repl is None or vfs is None:
|
||||
raise RuntimeError("Sandbox initialization failed")
|
||||
return repl, vfs
|
||||
|
||||
async def execute(self, code: str) -> SandboxResult:
|
||||
"""Execute Python code in the Monty REPL.
|
||||
|
||||
Variables persist across calls within the same Sandbox instance.
|
||||
"""
|
||||
repl, vfs = await self._ensure_initialized()
|
||||
external_fns = self._build_external_functions()
|
||||
|
||||
stdout_lines: list[str] = []
|
||||
|
||||
|
|
@ -283,20 +297,19 @@ class Sandbox:
|
|||
stdout_lines.append(text)
|
||||
|
||||
max_chars = self._config.analysis.max_output_chars
|
||||
limits: pydantic_monty.ResourceLimits = {
|
||||
"max_duration_secs": self._config.analysis.code_timeout,
|
||||
}
|
||||
|
||||
try:
|
||||
output = await pydantic_monty.run_monty_async(
|
||||
monty,
|
||||
inputs=inputs,
|
||||
output = await pydantic_monty.run_repl_async(
|
||||
repl,
|
||||
code,
|
||||
external_functions=external_fns,
|
||||
limits=limits,
|
||||
print_callback=print_callback,
|
||||
os=vfs,
|
||||
)
|
||||
except pydantic_monty.MontyRuntimeError as e:
|
||||
except (
|
||||
pydantic_monty.MontySyntaxError,
|
||||
pydantic_monty.MontyRuntimeError,
|
||||
) as e:
|
||||
stdout = "".join(stdout_lines)
|
||||
if len(stdout) > max_chars:
|
||||
stdout = stdout[:max_chars] + "\n... (output truncated)"
|
||||
|
|
|
|||
|
|
@ -495,7 +495,8 @@ class HaikuRAG:
|
|||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
# Convert → Chunk → Embed using primitives
|
||||
docling_document = await self.convert(content, format=format)
|
||||
converter = get_converter(self._config)
|
||||
docling_document = await converter.convert_text(content, format=format)
|
||||
chunks = await self.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, self._config)
|
||||
|
||||
|
|
@ -1000,7 +1001,10 @@ class HaikuRAG:
|
|||
# Content provided without chunks - convert, chunk, and embed using primitives
|
||||
assert content is not None
|
||||
existing_doc.content = content
|
||||
converted_docling = await self.convert(existing_doc.content)
|
||||
converter = get_converter(self._config)
|
||||
converted_docling = await converter.convert_text(
|
||||
existing_doc.content, format="md"
|
||||
)
|
||||
existing_doc.set_docling(converted_docling)
|
||||
|
||||
new_chunks = await self.chunk(converted_docling)
|
||||
|
|
@ -1558,11 +1562,13 @@ class HaikuRAG:
|
|||
pending_docs: list[Document] = []
|
||||
pending_doc_ids: list[str] = []
|
||||
|
||||
converter = get_converter(self._config)
|
||||
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
|
||||
# Convert content to DoclingDocument
|
||||
docling_document = await self.convert(doc.content)
|
||||
# Convert stored markdown to DoclingDocument
|
||||
docling_document = await converter.convert_text(doc.content, format="md")
|
||||
|
||||
# Chunk and embed
|
||||
chunks = await self.chunk(docling_document)
|
||||
|
|
@ -1605,6 +1611,7 @@ class HaikuRAG:
|
|||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
pending_doc_ids: list[str] = []
|
||||
converter = get_converter(self._config)
|
||||
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
|
|
@ -1643,7 +1650,7 @@ class HaikuRAG:
|
|||
"Source missing for %s, re-embedding from content", doc.uri
|
||||
)
|
||||
|
||||
docling_document = await self.convert(doc.content)
|
||||
docling_document = await converter.convert_text(doc.content, format="md")
|
||||
chunks = await self.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, self._config)
|
||||
|
||||
|
|
|
|||
79
haiku_rag_slim/haiku/rag/skills/_deps.py
Normal file
79
haiku_rag_slim/haiku/rag/skills/_deps.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
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.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 _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.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
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
|
||||
|
|
@ -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,55 +201,52 @@ 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:
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps
|
||||
|
||||
async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
|
||||
async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
|
||||
"""Execute Python code in a sandboxed interpreter.
|
||||
|
||||
The code has access to search(), list_documents(), llm() functions
|
||||
and a virtual filesystem at /documents/ with document content and
|
||||
structure (metadata.json, content.txt, items.jsonl per document).
|
||||
|
||||
Use print() to output results. Each call runs in a fresh
|
||||
interpreter — variables do not persist between calls.
|
||||
Use print() to output results. Variables persist between calls
|
||||
within the same skill invocation.
|
||||
|
||||
Args:
|
||||
code: Python code to execute.
|
||||
"""
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
|
||||
state = _get_state(ctx, state_type)
|
||||
doc_filter = state.document_filter if state else None
|
||||
context = AnalysisContext(filter=doc_filter)
|
||||
sandbox = Sandbox(db_path=db_path, config=config, context=context)
|
||||
if ctx.deps is None or ctx.deps.sandbox is None:
|
||||
raise RuntimeError(
|
||||
"AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
|
||||
)
|
||||
sandbox = ctx.deps.sandbox
|
||||
result = await sandbox.execute(code)
|
||||
|
||||
state = _get_state(ctx, state_type)
|
||||
|
|
@ -291,7 +277,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
|
||||
|
|
|
|||
|
|
@ -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_analysis_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_analysis_lifespan(db_path, config),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ You solve complex analytical questions by writing and executing Python code agai
|
|||
## Tools
|
||||
|
||||
### execute_code
|
||||
Execute Python code in a sandboxed interpreter. Each call runs in a fresh interpreter — write self-contained code. Use `print()` to output results.
|
||||
Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results.
|
||||
|
||||
Inside the code, these functions are available (use `await`):
|
||||
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
|
||||
|
|
@ -93,7 +93,7 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
|
|||
|
||||
## Important
|
||||
|
||||
- Each `execute_code` call runs in a fresh interpreter — write self-contained code blocks
|
||||
- Variables persist between `execute_code` calls — you can search in one call and process results in the next
|
||||
- Use `print()` to output results — the output is your only feedback
|
||||
- Always execute code to answer questions — don't just describe what code would do
|
||||
- Use `await` for all async functions inside execute_code (search, list_documents, llm)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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, sandbox=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) or sandbox is not None:
|
||||
ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox)
|
||||
else:
|
||||
ctx.deps = RAGRunDeps(state=state, rag=rag)
|
||||
return ctx
|
||||
|
||||
|
||||
|
|
@ -68,3 +73,26 @@ 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
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sandbox_factory(rag_db, test_app_config):
|
||||
"""Build Sandbox instances bound to the sample db, optionally with a doc filter."""
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
|
||||
def _make(filter: str | None = None) -> Sandbox:
|
||||
return Sandbox(
|
||||
db_path=rag_db,
|
||||
config=test_app_config,
|
||||
context=AnalysisContext(filter=filter),
|
||||
)
|
||||
|
||||
return _make
|
||||
|
|
|
|||
|
|
@ -113,73 +113,75 @@ class TestDomainPreambleInAnalysisSkillInstructions:
|
|||
|
||||
|
||||
class TestExecuteCodeTool:
|
||||
async def test_execute_code_returns_output(self, rag_db):
|
||||
async def test_execute_code_returns_output(self, rag_db, sandbox_factory):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState()
|
||||
ctx = _make_ctx(state)
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory())
|
||||
result = await execute_code(ctx, code="print('hello')")
|
||||
assert "hello" in result
|
||||
|
||||
async def test_execute_code_updates_state(self, rag_db):
|
||||
async def test_execute_code_updates_state(self, rag_db, sandbox_factory):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState()
|
||||
ctx = _make_ctx(state)
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory())
|
||||
await execute_code(ctx, code="print('hello')")
|
||||
assert len(state.executions) == 1
|
||||
assert state.executions[0].code == "print('hello')"
|
||||
assert state.executions[0].success is True
|
||||
assert "hello" in state.executions[0].stdout
|
||||
|
||||
async def test_execute_code_reports_errors(self, rag_db):
|
||||
async def test_execute_code_reports_errors(self, rag_db, sandbox_factory):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState()
|
||||
ctx = _make_ctx(state)
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory())
|
||||
result = await execute_code(ctx, code="x = 1/0")
|
||||
assert "Error" in result
|
||||
assert "ZeroDivisionError" in result
|
||||
assert state.executions[0].success is False
|
||||
|
||||
async def test_execute_code_applies_document_filter(self, rag_db):
|
||||
async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState(document_filter="title = 'AI Overview'")
|
||||
ctx = _make_ctx(state)
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory(filter=state.document_filter))
|
||||
result = await execute_code(
|
||||
ctx, code="docs = await list_documents()\nprint(len(docs))"
|
||||
)
|
||||
assert "1" in result
|
||||
|
||||
async def test_execute_code_accumulates_search_results(self, rag_db):
|
||||
async def test_execute_code_accumulates_search_results(
|
||||
self, rag_db, sandbox_factory
|
||||
):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState()
|
||||
ctx = _make_ctx(state)
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory())
|
||||
await execute_code(
|
||||
ctx, code="results = await search('intelligence')\nprint(len(results))"
|
||||
)
|
||||
assert "_sandbox" in state.searches
|
||||
assert len(state.searches["_sandbox"]) > 0
|
||||
|
||||
async def test_execute_code_vfs_write_denied(self, rag_db):
|
||||
async def test_execute_code_vfs_write_denied(self, rag_db, sandbox_factory):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState()
|
||||
ctx = _make_ctx(state)
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory())
|
||||
result = await execute_code(
|
||||
ctx,
|
||||
code=(
|
||||
|
|
@ -192,3 +194,109 @@ class TestExecuteCodeTool:
|
|||
)
|
||||
assert "Error" in result
|
||||
assert "read-only" in result
|
||||
|
||||
async def test_execute_code_variables_persist_within_invocation(
|
||||
self, rag_db, sandbox_factory
|
||||
):
|
||||
"""Same sandbox across two calls → vars persist (one skill invocation)."""
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
state = AnalysisState()
|
||||
ctx = _make_ctx(state, sandbox=sandbox_factory())
|
||||
|
||||
await execute_code(ctx, code="x = 42")
|
||||
result = await execute_code(ctx, code="print(x * 2)")
|
||||
assert "84" in result
|
||||
|
||||
async def test_execute_code_isolated_across_invocations(
|
||||
self, rag_db, sandbox_factory
|
||||
):
|
||||
"""Different Sandbox instances → no cross-invocation leak."""
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
execute_code = _get_tool(skill, "execute_code")
|
||||
|
||||
ctx1 = _make_ctx(AnalysisState(), sandbox=sandbox_factory())
|
||||
await execute_code(ctx1, code="secret = 'do not leak'")
|
||||
|
||||
ctx2 = _make_ctx(AnalysisState(), sandbox=sandbox_factory())
|
||||
result = await execute_code(ctx2, code="print(secret)")
|
||||
assert not result.startswith("do not leak")
|
||||
assert "Error" in result or "NameError" in result
|
||||
|
||||
|
||||
class TestAnalysisLifespan:
|
||||
async def test_opens_client_and_sandbox_per_invocation(self, rag_db):
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
|
||||
|
||||
config = AppConfig()
|
||||
lifespan = make_analysis_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
|
||||
assert isinstance(deps.sandbox, Sandbox)
|
||||
docs = await deps.rag.list_documents()
|
||||
assert len(docs) == 2
|
||||
|
||||
async def test_lifespan_reads_document_filter_from_state(self, rag_db):
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
|
||||
from haiku.rag.skills.analysis import AnalysisState
|
||||
|
||||
config = AppConfig()
|
||||
lifespan = make_analysis_lifespan(rag_db, config)
|
||||
state = AnalysisState(document_filter="title = 'AI Overview'")
|
||||
deps = AnalysisRunDeps(state=state)
|
||||
async with lifespan(deps):
|
||||
assert deps.sandbox is not None
|
||||
assert deps.sandbox._context.filter == "title = 'AI Overview'"
|
||||
|
||||
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
|
||||
|
||||
async def test_lifespan_clears_executions_citations_searches(self, rag_db):
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
|
||||
from haiku.rag.skills._tools import CodeExecutionEntry
|
||||
from haiku.rag.skills.analysis import AnalysisState
|
||||
|
||||
config = AppConfig()
|
||||
lifespan = make_analysis_lifespan(rag_db, config)
|
||||
|
||||
state = AnalysisState(
|
||||
document_filter="title = 'AI Overview'",
|
||||
executions=[CodeExecutionEntry(code="prior", stdout="", success=True)],
|
||||
citation_index={
|
||||
"c1": Citation(
|
||||
index=1,
|
||||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
document_title="t",
|
||||
document_uri="u",
|
||||
content="x",
|
||||
page_numbers=[],
|
||||
headings=[],
|
||||
)
|
||||
},
|
||||
citations=[["c1"]],
|
||||
searches={"prior": []},
|
||||
)
|
||||
deps = AnalysisRunDeps(state=state)
|
||||
async with lifespan(deps):
|
||||
assert state.executions == []
|
||||
assert state.citations == []
|
||||
assert state.searches == {}
|
||||
assert "c1" in state.citation_index
|
||||
assert state.document_filter == "title = 'AI Overview'"
|
||||
|
|
|
|||
|
|
@ -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,75 @@ 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
|
||||
|
||||
async def test_lifespan_clears_citations_and_searches_but_keeps_index(self, rag_db):
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
|
||||
from haiku.rag.skills.rag import RAGState
|
||||
|
||||
config = AppConfig()
|
||||
lifespan = make_rag_lifespan(rag_db, config)
|
||||
|
||||
state = RAGState(
|
||||
document_filter="title = 'AI Overview'",
|
||||
citation_index={
|
||||
"c1": Citation(
|
||||
index=1,
|
||||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
document_title="t",
|
||||
document_uri="u",
|
||||
content="x",
|
||||
page_numbers=[],
|
||||
headings=[],
|
||||
)
|
||||
},
|
||||
citations=[["c1"]],
|
||||
searches={"prior": []},
|
||||
)
|
||||
deps = RAGRunDeps(state=state)
|
||||
async with lifespan(deps):
|
||||
assert state.citations == []
|
||||
assert state.searches == {}
|
||||
assert "c1" in state.citation_index # preserved for cross-turn lookup
|
||||
assert state.document_filter == "title = 'AI Overview'"
|
||||
|
|
|
|||
|
|
@ -1544,3 +1544,78 @@ async def test_sql_injection_is_blocked_with_escaping(temp_db_path):
|
|||
filter=f"title = '{injection_payload}'"
|
||||
)
|
||||
assert len(docs_unescaped) == 2 # SQL injection succeeds without escaping
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# URL-prefixed content regression tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _patch_embed_chunks(monkeypatch):
|
||||
async def fake_embed_chunks(chunks, config):
|
||||
for chunk in chunks:
|
||||
chunk.embedding = [0.0] * 2560
|
||||
return chunks
|
||||
|
||||
monkeypatch.setattr("haiku.rag.embeddings.embed_chunks", fake_embed_chunks)
|
||||
|
||||
|
||||
async def test_create_document_with_url_prefixed_content(temp_db_path, monkeypatch):
|
||||
"""Text whose first line is a URL must be stored as text, not fetched."""
|
||||
_patch_embed_chunks(monkeypatch)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
content = "https://example.com/foo\n\n# Heading\n\nBody text here."
|
||||
doc = await client.create_document(content=content, uri="test://url-prefixed")
|
||||
|
||||
assert doc.id is not None
|
||||
assert "example.com" in doc.content
|
||||
assert "Heading" in doc.content
|
||||
|
||||
|
||||
async def test_update_document_with_url_prefixed_content(temp_db_path, monkeypatch):
|
||||
"""update_document(content=...) with URL-prefixed text must not fetch it."""
|
||||
_patch_embed_chunks(monkeypatch)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(
|
||||
content="initial body", uri="test://update-url"
|
||||
)
|
||||
assert doc.id is not None
|
||||
|
||||
url_prefixed = "https://example.com/bar\n\n# New heading\n\nReplacement body."
|
||||
updated = await client.update_document(doc.id, content=url_prefixed)
|
||||
|
||||
assert "example.com" in updated.content
|
||||
assert "New heading" in updated.content
|
||||
|
||||
|
||||
async def test_rebuild_rechunk_with_url_prefixed_stored_content(
|
||||
temp_db_path, monkeypatch
|
||||
):
|
||||
"""RECHUNK rebuild must handle stored markdown whose first line is a URL."""
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
_patch_embed_chunks(monkeypatch)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(
|
||||
content="plain seed content", uri="file:///nonexistent/path.txt"
|
||||
)
|
||||
assert doc.id is not None
|
||||
|
||||
# Overwrite stored content to simulate markdown that starts with a URL,
|
||||
# bypassing the (also-affected) create_document path so this test
|
||||
# specifically exercises the rebuild path.
|
||||
doc.content = "https://example.com/baz\n\n# Stored\n\nStored body text."
|
||||
await client.document_repository.update(doc)
|
||||
|
||||
processed_ids = [
|
||||
doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK)
|
||||
]
|
||||
assert doc.id in processed_ids
|
||||
|
||||
doc_after = await client.document_repository.get_by_id(doc.id)
|
||||
assert doc_after is not None
|
||||
assert "example.com" in doc_after.content
|
||||
assert "Stored" in doc_after.content
|
||||
|
|
|
|||
8
uv.lock
8
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]]
|
||||
|
|
|
|||
Loading…
Reference in a new issue