Simplify InsightRecord, GapRecord

This commit is contained in:
Yiorgis Gozadinos 2025-11-12 10:21:25 +02:00
parent 2060898018
commit 657245f686
No known key found for this signature in database
3 changed files with 73 additions and 143 deletions

View file

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

View file

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

View file

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