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

Remove documents from state, fix tests

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-20 09:48:18 +03:00
parent d52f453c44
commit 68a9f191d2
No known key found for this signature in database
7 changed files with 24 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
app/frontend/lib/sessionStorage.ts
View file

@ -9,20 +9,12 @@ export interface Citation {
content: string;
}
export interface DocumentInfo {
id: string;
title: string;
uri: string;
created: string;
}
// Matches RAGState from the backend skill
export interface RAGState {
citation_index: Record<string, Citation>;
citations: string[][];
document_filter: string | null;
searches: Record<string, unknown[]>;
documents: DocumentInfo[];
}
export interface StoredMessage {
@ -50,7 +42,6 @@ export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
citations: state?.citations ?? [],
document_filter: state?.document_filter ?? null,
searches: state?.searches ?? {},
documents: state?.documents ?? [],
};
}

6
haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2
View file

@ -8,9 +8,6 @@ from haiku.skills.parser import parse_skill_md
{% if "cite" in tool_names %}
from haiku.rag.agents.research.models import Citation
{% endif %}
{% if "search" in tool_names or "list_documents" in tool_names or "get_document" in tool_names %}
from haiku.rag.tools.document import DocumentInfo
{% endif %}
{% if "search" in tool_names %}
from haiku.rag.store.models.chunk import SearchResult
{% endif %}
@ -38,9 +35,6 @@ class SkillState(BaseModel):
{% if "search" in tool_names %}
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
{% endif %}
{% if "search" in tool_names or "list_documents" in tool_names or "get_document" in tool_names %}
documents: list[DocumentInfo] = Field(default_factory=list)
{% endif %}
{% if "execute_code" in tool_names %}
executions: list[CodeExecutionEntry] = Field(default_factory=list)
{% endif %}

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

@ -7,7 +7,6 @@ from pydantic_ai import RunContext
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.skills.state import SkillRunDeps
@ -86,21 +85,6 @@ async def skill_get_document(
}
def update_documents_state(
documents_state: list[DocumentInfo],
doc_dicts: list[dict[str, Any]],
) -> None:
for doc_dict in doc_dicts:
doc_info = DocumentInfo(
id=str(doc_dict["id"]),
title=doc_dict["title"] or "Untitled",
uri=doc_dict.get("uri") or "",
created=doc_dict.get("created_at", ""),
)
if not any(d.id == doc_info.id for d in documents_state):
documents_state.append(doc_info)
def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> Any:
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, state_type):
return ctx.deps.state
@ -237,8 +221,6 @@ def create_skill_tools(
config,
filter=state.document_filter if state else None,
)
if state:
update_documents_state(state.documents, result)
return result
tools["list_documents"] = list_documents
@ -253,12 +235,7 @@ def create_skill_tools(
Args:
query: Document ID, title, or URI to look up.
"""
result = await skill_get_document(db_path, config, query)
if result is not None:
state = _get_state(ctx, state_type)
if state:
update_documents_state(state.documents, [result])
return result
return await skill_get_document(db_path, config, query)
tools["get_document"] = get_document

2
haiku_rag_slim/haiku/rag/skills/analysis.py
View file

@ -8,7 +8,6 @@ from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.skills._tools import CodeExecutionEntry
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
@ -19,7 +18,6 @@ class AnalysisState(BaseModel):
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[list[str]] = Field(default_factory=list)
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
documents: list[DocumentInfo] = Field(default_factory=list)
STATE_TYPE = AnalysisState

2
haiku_rag_slim/haiku/rag/skills/rag.py
View file

@ -7,7 +7,6 @@ from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
@ -34,7 +33,6 @@ class RAGState(BaseModel):
citations: list[list[str]] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
documents: list[DocumentInfo] = Field(default_factory=list)
STATE_TYPE = RAGState

12
tests/skills/test_rag.py
View file

@ -8,7 +8,6 @@ from haiku.rag.skills.rag import (
state_metadata,
)
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.skills.models import SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx
@ -213,17 +212,6 @@ class TestListDocumentsTool:
assert isinstance(results, list)
assert len(results) == 2
async def test_list_documents_updates_state(self, rag_db):
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()
ctx = _make_ctx(state)
await list_docs(ctx)
assert len(state.documents) == 2
assert isinstance(state.documents[0], DocumentInfo)
async def test_list_documents_applies_document_filter_from_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill

48
tests/test_skill_generator.py
View file

@ -23,9 +23,8 @@ class TestAvailableTools:
"list_documents",
"get_document",
"search",
"ask",
"research",
"analyze",
"execute_code",
"cite",
}
@ -43,7 +42,7 @@ class TestValidateTools:
validate_tools(["search"])
def test_valid_multiple_tools(self):
validate_tools(["list_documents", "get_document", "search", "ask"])
validate_tools(["list_documents", "get_document", "search", "cite"])
def test_valid_all_tools(self):
validate_tools(list(AVAILABLE_TOOLS))
@ -97,7 +96,7 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["list_documents", "get_document", "search", "ask"],
tool_names=["list_documents", "get_document", "search", "cite"],
)
assert result == tmp_path / "recipes-skill"
assert result.is_dir()
@ -112,7 +111,7 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="my-recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
assert result == tmp_path / "my-recipes-skill"
pkg = result / "my_recipes_skill"
@ -129,11 +128,11 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert '["search", "ask"]' in content
assert '["search", "cite"]' in content
def test_create_skill_tools_called(self, tmp_path):
render_templates(
@ -151,12 +150,12 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
assert '"search"' in content
assert '"ask"' in content
assert '"cite"' in content
def test_pyproject_toml(self, tmp_path):
render_templates(
@ -184,24 +183,23 @@ class TestRenderTemplates:
)
skill_md = tmp_path / "docs-skill" / "docs_skill" / "SKILL.md"
content = skill_md.read_text()
assert "**search**" in content
assert "**ask**" not in content
assert "**list_documents**" not in content
assert "**research**" not in content
assert "**analyze**" not in content
assert "### search" in content
assert "### cite" not in content
assert "### list_documents" not in content
assert "### execute_code" not in content
def test_skill_md_includes_all_selected(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "ask", "analyze"],
tool_names=["search", "execute_code", "cite"],
)
skill_md = tmp_path / "docs-skill" / "docs_skill" / "SKILL.md"
content = skill_md.read_text()
assert "**search**" in content
assert "**ask**" in content
assert "**analyze**" in content
assert "search" in content
assert "execute_code" in content
assert "cite" in content
def test_custom_preamble(self, tmp_path):
render_templates(
@ -226,23 +224,23 @@ class TestRenderTemplates:
content = init.read_text()
assert 'state_namespace="recipes"' in content
def test_analyze_state_fields(self, tmp_path):
def test_execute_code_state_fields(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "analyze"],
tool_names=["search", "execute_code"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "analyses" in content
assert "executions" in content
def test_imports_from_shared_tools(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask", "analyze"],
tool_names=["search", "execute_code", "cite"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
@ -329,7 +327,7 @@ class TestGenerateSkill:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
assert result == tmp_path / "recipes-skill"
assets = result / "recipes_skill" / "assets"
@ -479,7 +477,7 @@ class TestGenerateSkillRemote:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
config_path=config_file,
)
assets = result / "recipes_skill" / "assets"