Simplify InsightRecord, GapRecord
This commit is contained in:
parent
2060898018
commit
657245f686
3 changed files with 73 additions and 143 deletions
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue