Merge pull request #349 from ggozad/fix/convert-urlparse

fix convert() misreading text content that starts with a URL
This commit is contained in:
Yiorgis Gozadinos 2026-04-22 14:56:53 +03:00 committed by GitHub
commit e736a8c73a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 562 additions and 174 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View 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

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

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

View file

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

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

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

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

View file

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

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

View file

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

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