From 657245f686a96f2591f8deb2808b8bc9120d1127 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 12 Nov 2025 10:21:25 +0200 Subject: [PATCH] Simplify InsightRecord, GapRecord --- .../haiku/rag/graph/common/nodes.py | 2 +- .../haiku/rag/graph/research/dependencies.py | 118 +++++------------- .../haiku/rag/graph/research/models.py | 96 +++++++------- 3 files changed, 73 insertions(+), 143 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 1097193b..ba452a61 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -239,7 +239,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( if deps.agui_emitter: deps.agui_emitter.update_state(state) # Format the success message - if "{confidence}" in success_message_format: + if "{confidence" in success_message_format: message = success_message_format.format( sub_q=sub_q, confidence=answer.confidence ) diff --git a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py index fe65cc03..79a92cae 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py @@ -1,6 +1,6 @@ from collections.abc import Iterable -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, PrivateAttr from haiku.rag.client import HaikuRAG from haiku.rag.graph.common.models import SearchAnswer @@ -28,26 +28,29 @@ class ResearchContext(BaseModel): default_factory=list, description="Identified information gaps" ) + # Private dict indexes for O(1) lookups + _insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict) + _gaps_by_id: dict[str, GapRecord] = PrivateAttr(default_factory=dict) + + def model_post_init(self, __context: object) -> None: + """Build indexes after initialization.""" + self._insights_by_id = {ins.id: ins for ins in self.insights} + self._gaps_by_id = {gap.id: gap for gap in self.gaps} + def add_qa_response(self, qa: SearchAnswer) -> None: """Add a structured QA response (minimal context already included).""" self.qa_responses.append(qa) def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]: """Merge one or more insights into the shared context with deduplication.""" - merged: list[InsightRecord] = [] + for record in records: candidate = InsightRecord.model_validate(record) - existing = next( - (ins for ins in self.insights if ins.id == candidate.id), None - ) - if not existing: - existing = next( - (ins for ins in self.insights if ins.summary == candidate.summary), - None, - ) + existing = self._insights_by_id.get(candidate.id) if existing: + # Update existing insight existing.summary = candidate.summary existing.status = candidate.status if candidate.notes: @@ -60,36 +63,24 @@ class ResearchContext(BaseModel): ) merged.append(existing) else: - candidate = candidate.model_copy(deep=True) - if candidate.id is None: # pragma: no cover - defensive - raise ValueError( - "InsightRecord.id must be populated after validation" - ) - candidate_id: str = candidate.id - candidate.id = self._allocate_insight_id(candidate_id) - self.insights.append(candidate) - merged.append(candidate) + # Add new insight + new_insight = candidate.model_copy(deep=True) + self.insights.append(new_insight) + self._insights_by_id[new_insight.id] = new_insight + merged.append(new_insight) return merged def upsert_gaps(self, records: Iterable[GapRecord]) -> list[GapRecord]: """Merge one or more gap records into the shared context with deduplication.""" - merged: list[GapRecord] = [] + for record in records: candidate = GapRecord.model_validate(record) - existing = next((gap for gap in self.gaps if gap.id == candidate.id), None) - if not existing: - existing = next( - ( - gap - for gap in self.gaps - if gap.description == candidate.description - ), - None, - ) + existing = self._gaps_by_id.get(candidate.id) if existing: + # Update existing gap existing.description = candidate.description existing.severity = candidate.severity existing.blocking = candidate.blocking @@ -104,22 +95,19 @@ class ResearchContext(BaseModel): ) merged.append(existing) else: - candidate = candidate.model_copy(deep=True) - if candidate.id is None: # pragma: no cover - defensive - raise ValueError("GapRecord.id must be populated after validation") - candidate_id: str = candidate.id - candidate.id = self._allocate_gap_id(candidate_id) - self.gaps.append(candidate) - merged.append(candidate) + # Add new gap + new_gap = candidate.model_copy(deep=True) + self.gaps.append(new_gap) + self._gaps_by_id[new_gap.id] = new_gap + merged.append(new_gap) return merged def mark_gap_resolved( self, identifier: str, resolved_by: Iterable[str] | None = None ) -> GapRecord | None: - """Mark a gap as resolved by identifier (id or description).""" - - gap = self._find_gap(identifier) + """Mark a gap as resolved by identifier.""" + gap = self._gaps_by_id.get(identifier) if gap is None: return None @@ -131,7 +119,6 @@ class ResearchContext(BaseModel): def integrate_analysis(self, analysis: InsightAnalysis) -> None: """Apply an analysis result to the shared context.""" - merged_insights: list[InsightRecord] = [] if analysis.highlights: merged_insights = self.upsert_insights(analysis.highlights) @@ -141,9 +128,7 @@ class ResearchContext(BaseModel): analysis.gap_assessments = merged_gaps if analysis.resolved_gaps: resolved_by_list = ( - [ins.id for ins in merged_insights if ins.id is not None] - if merged_insights - else None + [ins.id for ins in merged_insights] if merged_insights else None ) for resolved in analysis.resolved_gaps: self.mark_gap_resolved(resolved, resolved_by=resolved_by_list) @@ -151,29 +136,6 @@ class ResearchContext(BaseModel): if question not in self.sub_questions: self.sub_questions.append(question) - def _allocate_insight_id(self, candidate_id: str) -> str: - taken: set[str] = set() - for ins in self.insights: - if ins.id is not None: - taken.add(ins.id) - return _allocate_sequential_id(candidate_id, taken) - - def _allocate_gap_id(self, candidate_id: str) -> str: - taken: set[str] = set() - for gap in self.gaps: - if gap.id is not None: - taken.add(gap.id) - return _allocate_sequential_id(candidate_id, taken) - - def _find_gap(self, identifier: str) -> GapRecord | None: - normalized = identifier.lower().strip() - for gap in self.gaps: - if gap.id is not None and gap.id == normalized: - return gap - if gap.description.lower().strip() == normalized: - return gap - return None - class ResearchDependencies(BaseModel): """Dependencies for research agents with multi-agent context.""" @@ -186,24 +148,4 @@ class ResearchDependencies(BaseModel): def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]: """Merge two iterables preserving order while removing duplicates.""" - - merged = list(existing) - seen = {item for item in existing if item} - for item in incoming: - if item and item not in seen: - merged.append(item) - seen.add(item) - return merged - - -def _allocate_sequential_id(candidate: str, taken: set[str]) -> str: - slug = candidate - if slug not in taken: - return slug - base = slug - counter = 2 - while True: - slug = f"{base}-{counter}" - if slug not in taken: - return slug - counter += 1 + return [k for k in dict.fromkeys([*existing, *incoming]) if k] diff --git a/haiku_rag_slim/haiku/rag/graph/research/models.py b/haiku_rag_slim/haiku/rag/graph/research/models.py index aae2f32a..a132bf4b 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/models.py +++ b/haiku_rag_slim/haiku/rag/graph/research/models.py @@ -1,19 +1,12 @@ -import re +import uuid from enum import Enum -from pydantic import BaseModel, Field, model_validator - -_SLUG_RE = re.compile(r"[^a-z0-9]+") +from pydantic import BaseModel, Field, field_validator -def _make_slug(text: str, prefix: str) -> str: - """Generate a lowercase slug with the given prefix as fallback.""" - - base = _SLUG_RE.sub("-", text.lower()).strip("-") - if not base: - base = prefix - # Trim overly long slugs but keep enough entropy for readability - return base[:48] +def _deduplicate_list(items: list[str]) -> list[str]: + """Remove duplicates while preserving order.""" + return list(dict.fromkeys(items)) class InsightStatus(str, Enum): @@ -28,48 +21,54 @@ class GapSeverity(str, Enum): HIGH = "high" -class InsightRecord(BaseModel): +class TrackedRecord(BaseModel): + """Base model for tracked entities with sources and metadata.""" + + model_config = {"validate_assignment": True} + + id: str = Field( + default_factory=lambda: str(uuid.uuid4())[:8], + description="Unique identifier for the record", + ) + supporting_sources: list[str] = Field( + default_factory=list, + description="Source identifiers backing this record", + ) + notes: str | None = Field( + default=None, + description="Optional elaboration or caveats", + ) + + @field_validator("supporting_sources", mode="before") + @classmethod + def deduplicate_sources(cls, v: list[str]) -> list[str]: + """Ensure supporting_sources has no duplicates.""" + return _deduplicate_list(v) if v else [] + + +class InsightRecord(TrackedRecord): """Structured insight with provenance and lifecycle metadata.""" - id: str | None = Field( - default=None, - description="Stable slug identifier for the insight (auto-generated if omitted)", - ) summary: str = Field(description="Concise description of the insight") status: InsightStatus = Field( default=InsightStatus.OPEN, description="Lifecycle status for the insight", ) - supporting_sources: list[str] = Field( - default_factory=list, - description="Source identifiers backing the insight", - ) originating_questions: list[str] = Field( default_factory=list, description="Research sub-questions that produced this insight", ) - notes: str | None = Field( - default=None, - description="Optional elaboration or caveats for the insight", - ) - @model_validator(mode="after") - def _set_defaults(self) -> "InsightRecord": - if not self.id: - self.id = _make_slug(self.summary, "insight") - self.id = self.id.lower() - self.supporting_sources = list(dict.fromkeys(self.supporting_sources)) - self.originating_questions = list(dict.fromkeys(self.originating_questions)) - return self + @field_validator("originating_questions", mode="before") + @classmethod + def deduplicate_questions(cls, v: list[str]) -> list[str]: + """Ensure originating_questions has no duplicates.""" + return _deduplicate_list(v) if v else [] -class GapRecord(BaseModel): +class GapRecord(TrackedRecord): """Structured representation of an identified research gap.""" - id: str | None = Field( - default=None, - description="Stable slug identifier for the gap (auto-generated if omitted)", - ) description: str = Field(description="Concrete statement of what is missing") severity: GapSeverity = Field( default=GapSeverity.MEDIUM, @@ -87,23 +86,12 @@ class GapRecord(BaseModel): default_factory=list, description="Insight IDs or notes explaining how the gap was closed", ) - supporting_sources: list[str] = Field( - default_factory=list, - description="Sources confirming the gap status (e.g., evidence of absence)", - ) - notes: str | None = Field( - default=None, - description="Optional clarification about the gap or follow-up actions", - ) - @model_validator(mode="after") - def _set_defaults(self) -> "GapRecord": - if not self.id: - self.id = _make_slug(self.description, "gap") - self.id = self.id.lower() - self.resolved_by = list(dict.fromkeys(self.resolved_by)) - self.supporting_sources = list(dict.fromkeys(self.supporting_sources)) - return self + @field_validator("resolved_by", mode="before") + @classmethod + def deduplicate_resolved_by(cls, v: list[str]) -> list[str]: + """Ensure resolved_by has no duplicates.""" + return _deduplicate_list(v) if v else [] class InsightAnalysis(BaseModel):