Merge pull request #350 from ggozad/fix/citations

flatten skill state `citations` from list[list[str]] to list[str]
This commit is contained in:
Yiorgis Gozadinos 2026-04-22 16:59:07 +03:00 committed by GitHub
commit 1202fe17e7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 20 additions and 20 deletions

View file

@ -1,6 +1,10 @@
# Changelog
## [Unreleased]
### Changed
- **BREAKING: Skill state `citations` is now `list[str]` instead of `list[list[str]]`.** With per-invocation state scoping (0.42), the outer list no longer tracked turn boundaries — it only grouped chunk ids per `cite` call within a single invocation, which has no downstream meaning. The field is now a flat, deduplicated list of chunk ids cited during the current invocation. Clients resolve each id through `citation_index` as before. Applies to both `RAGState` and `AnalysisState`.
## [0.42.0] - 2026-04-22
### Fixed

View file

@ -12,7 +12,7 @@ export interface Citation {
// Matches RAGState from the backend skill
export interface RAGState {
citation_index: Record<string, Citation>;
citations: string[][];
citations: string[];
document_filter: string | null;
searches: Record<string, unknown[]>;
}
@ -46,10 +46,7 @@ export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
}
export function getLatestCitations(state: RAGState): Citation[] {
const turns = state.citations;
if (turns.length === 0) return [];
const latestIds = turns[turns.length - 1];
return latestIds
return state.citations
.map((id) => state.citation_index[id])
.filter((c): c is Citation => c !== undefined);
}

View file

@ -33,14 +33,14 @@ class AnalysisState(BaseModel):
document_filter: str | None = None
executions: list[CodeExecutionEntry] = []
citation_index: dict[str, Citation] = {}
citations: list[list[str]] = []
citations: list[str] = []
searches: dict[str, list[SearchResult]] = {}
```
- **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.
- **citations** — Chunk IDs cited during the current invocation. Deduplicated; cleared at the start of each invocation.
- **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

@ -31,12 +31,12 @@ The skill manages a `RAGState` under the `"rag"` namespace:
```python
class RAGState(BaseModel):
citation_index: dict[str, Citation] = {}
citations: list[list[str]] = []
citations: list[str] = []
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = {}
```
- **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.
- **citations** — Chunk IDs registered via the `cite` tool during the current invocation. Deduplicated; cleared at the start of each invocation.
- **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

@ -83,8 +83,7 @@ def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
def _register_citations(state: Any, citations: "list[Citation]") -> None:
"""Add citations to the index and record the turn's chunk IDs."""
chunk_ids = []
"""Add citations to the index and record cited chunk IDs for this invocation."""
next_index = len(state.citation_index) + 1
for citation in citations:
cid = citation.chunk_id
@ -92,8 +91,8 @@ def _register_citations(state: Any, citations: "list[Citation]") -> None:
citation.index = next_index
next_index += 1
state.citation_index[cid] = citation
chunk_ids.append(cid)
state.citations.append(chunk_ids)
if cid not in state.citations:
state.citations.append(cid)
def create_skill_extras(

View file

@ -16,7 +16,7 @@ class AnalysisState(BaseModel):
document_filter: str | None = None
executions: list[CodeExecutionEntry] = Field(default_factory=list)
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[list[str]] = Field(default_factory=list)
citations: list[str] = Field(default_factory=list)
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)

View file

@ -30,7 +30,7 @@ def get_agent_preamble(config: AppConfig) -> str:
class RAGState(BaseModel):
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[list[str]] = Field(default_factory=list)
citations: list[str] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)

View file

@ -303,7 +303,7 @@ class TestAnalysisLifespan:
headings=[],
)
},
citations=[["c1"]],
citations=["c1"],
searches={"prior": []},
)
deps = AnalysisRunDeps(state=state)

View file

@ -284,8 +284,8 @@ class TestCiteTool:
result = await cite(ctx, chunk_ids=chunk_ids)
assert "Registered" in result
assert len(state.citations) == 1
assert len(state.citations[0]) == 2
assert len(state.citations) == 2
assert all(cid in state.citations for cid in chunk_ids)
assert all(cid in state.citation_index for cid in chunk_ids)
async def test_cite_deduplicates_in_index(self, rag_db, rag_client):
@ -308,7 +308,7 @@ class TestCiteTool:
await cite(ctx, chunk_ids=chunk_ids)
await cite(ctx, chunk_ids=chunk_ids)
assert len(state.citation_index) == 1
assert len(state.citations) == 2
assert len(state.citations) == 1
async def test_cite_without_state(self, rag_db):
from haiku.rag.skills.rag import create_skill
@ -394,7 +394,7 @@ class TestLifespan:
headings=[],
)
},
citations=[["c1"]],
citations=["c1"],
searches={"prior": []},
)
deps = RAGRunDeps(state=state)