52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""LLM JSON 输出校验（memory 富化）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
from pydantic import BaseModel, Field, field_validator
|
||
|
||
|
||
class ExtractedFactItem(BaseModel):
|
||
fact_type: str = "event"
|
||
subject: str | None = None
|
||
predicate: str | None = None
|
||
object_json: Any = None
|
||
confidence: float = Field(default=0.75, ge=0.0, le=1.0)
|
||
source_chunk_id: str | None = None
|
||
|
||
@field_validator("fact_type", mode="before")
|
||
@classmethod
|
||
def _coerce_fact_type(cls, v: object) -> str:
|
||
ft = str(v or "event").strip() or "event"
|
||
if ft not in ("person", "event", "relation", "place", "milestone"):
|
||
return "event"
|
||
return ft
|
||
|
||
|
||
class FactsExtractionPayload(BaseModel):
|
||
facts: list[ExtractedFactItem] = Field(default_factory=list)
|
||
|
||
|
||
class EnrichmentPayload(BaseModel):
|
||
"""单轮记忆富化：会话摘要 + 结构化事实（ingest 后一次 LLM 调用）。"""
|
||
|
||
summary: str = ""
|
||
facts: list[ExtractedFactItem] = Field(default_factory=list)
|
||
|
||
|
||
def facts_payload_to_dicts(payload: FactsExtractionPayload) -> list[dict]:
|
||
out: list[dict] = []
|
||
for item in payload.facts:
|
||
d = item.model_dump()
|
||
scid = d.get("source_chunk_id")
|
||
if scid is not None and not isinstance(scid, str):
|
||
d["source_chunk_id"] = str(scid)
|
||
out.append(d)
|
||
return out
|
||
|
||
|
||
def enrichment_payload_to_fact_dicts(payload: EnrichmentPayload) -> list[dict]:
|
||
"""将 EnrichmentPayload.facts 转为与 extract_facts 一致的字典列表。"""
|
||
return facts_payload_to_dicts(FactsExtractionPayload(facts=list(payload.facts)))
|