life-echo/api/app/features/memory/llm_schemas.py

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