Simplify InsightRecord, GapRecord

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
                     )

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(
+            existing = self._insights_by_id.get(candidate.id)
-                (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,
-                )
 
             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)
+                # Add new insight
-                if candidate.id is None:  # pragma: no cover - defensive
+                new_insight = candidate.model_copy(deep=True)
-                    raise ValueError(
+                self.insights.append(new_insight)
-                        "InsightRecord.id must be populated after validation"
+                self._insights_by_id[new_insight.id] = new_insight
-                    )
+                merged.append(new_insight)
-                candidate_id: str = candidate.id
-                candidate.id = self._allocate_insight_id(candidate_id)
-                self.insights.append(candidate)
-                merged.append(candidate)
 
         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)
+            existing = self._gaps_by_id.get(candidate.id)
-            if not existing:
-                existing = next(
-                    (
-                        gap
-                        for gap in self.gaps
-                        if gap.description == candidate.description
-                    ),
-                    None,
-                )
 
             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)
+                # Add new gap
-                if candidate.id is None:  # pragma: no cover - defensive
+                new_gap = candidate.model_copy(deep=True)
-                    raise ValueError("GapRecord.id must be populated after validation")
+                self.gaps.append(new_gap)
-                candidate_id: str = candidate.id
+                self._gaps_by_id[new_gap.id] = new_gap
-                candidate.id = self._allocate_gap_id(candidate_id)
+                merged.append(new_gap)
-                self.gaps.append(candidate)
-                merged.append(candidate)
 
         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)."""
+        """Mark a gap as resolved by identifier."""
-
+        gap = self._gaps_by_id.get(identifier)
-        gap = self._find_gap(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]
+                [ins.id for ins in merged_insights] if merged_insights else None
-                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."""
-
+    return [k for k in dict.fromkeys([*existing, *incoming]) if k]
-    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

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
+from pydantic import BaseModel, Field, field_validator
-
-_SLUG_RE = re.compile(r"[^a-z0-9]+")
 
 
-def _make_slug(text: str, prefix: str) -> str:
+def _deduplicate_list(items: list[str]) -> list[str]:
-    """Generate a lowercase slug with the given prefix as fallback."""
+    """Remove duplicates while preserving order."""
-
+    return list(dict.fromkeys(items))
-    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]
 
 
 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")
+    @field_validator("originating_questions", mode="before")
-    def _set_defaults(self) -> "InsightRecord":
+    @classmethod
-        if not self.id:
+    def deduplicate_questions(cls, v: list[str]) -> list[str]:
-            self.id = _make_slug(self.summary, "insight")
+        """Ensure originating_questions has no duplicates."""
-        self.id = self.id.lower()
+        return _deduplicate_list(v) if v else []
-        self.supporting_sources = list(dict.fromkeys(self.supporting_sources))
-        self.originating_questions = list(dict.fromkeys(self.originating_questions))
-        return self
 
 
-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")
+    @field_validator("resolved_by", mode="before")
-    def _set_defaults(self) -> "GapRecord":
+    @classmethod
-        if not self.id:
+    def deduplicate_resolved_by(cls, v: list[str]) -> list[str]:
-            self.id = _make_slug(self.description, "gap")
+        """Ensure resolved_by has no duplicates."""
-        self.id = self.id.lower()
+        return _deduplicate_list(v) if v else []
-        self.resolved_by = list(dict.fromkeys(self.resolved_by))
-        self.supporting_sources = list(dict.fromkeys(self.supporting_sources))
-        return self
 
 
 class InsightAnalysis(BaseModel):