entity_extractor.py

"""
带上下文的实体/关系提取器
核心创新：P = (S, existing_entities, history, current)
实现 Mem0 风格的 CRUD 演化
"""

from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass, field, asdict
from llm_client import get_llm_client
from config import ENTITY_TYPES, RELATION_TYPES, STM_CONFIG, LPM_CONFIG


@dataclass
class EntityUpdate:
    """实体更新操作"""
    entity_name: str        # 规范化名称
    display_name: str       # 显示名称
    entity_type: str        # Person, Business, Location, Event, Activity, Artifact
    op: str                 # CREATE, UPDATE, DELETE, NOOP
    state_change: str       # 状态变化描述 (e.g., "banker → unemployed")
    context: str            # 原子笔记摘要

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class RelationUpdate:
    """关系更新操作"""
    source: str             # 源实体名称
    relation_type: str      # 关系类型
    target: str             # 目标实体名称
    op: str                 # CREATE, UPDATE, DELETE
    context: str            # 关系描述

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class ExtractionResult:
    """提取结果"""
    note_id: str
    entity_updates: List[EntityUpdate]
    relation_updates: List[RelationUpdate]
    raw_response: str = ""


# ==================== 提取 Prompt ====================

EXTRACTION_SYSTEM_PROMPT = """You are an entity and relationship extractor for conversational memory.
Your task is to extract entities and relationships, and determine CRUD operations based on context.

## Entity Types
- Person: People (Jon, Gina, Jean)
- Business: Businesses/ventures (dance studio, clothing store)
- Location: Places (Paris, Rome, gym)
- Event: Events (competition, festival, interview)
- Activity: Activities/hobbies (dancing, fashion)
- Artifact: Objects/creations (book, video, tattoo)

## Relation Types
OWNS, STARTED, WORKS_AT, LEFT, VISITED, LIVES_IN, PARTICIPATED_IN, HOSTED, ENJOYS, PRACTICES, CREATED, HAS, READING, RELATED_TO

## CRUD Operations
- CREATE: New entity or first mention
- UPDATE: State change of existing entity (e.g., "starting" → "opened")
- DELETE: Entity no longer valid (rare)
- NOOP: No change, just reference

## Rules
1. Entity names: lowercase with underscores (e.g., "jon", "dance_studio")
2. Check existing entities first - if mentioned entity matches an existing one (or its alias), use UPDATE not CREATE
3. State changes should capture transitions (e.g., "banker → unemployed", "preparing → opened")
4. Context should be a brief atomic summary of what happened
5. For cross-entity relationships, create edges (e.g., jon-ENJOYS->dancing, gina-ENJOYS->dancing)
6. Aliases: include common variations (e.g., "Jon's studio" → alias of "jon_dance_studio")

## Output JSON Format
{
  "entity_updates": [
    {
      "entity_name": "jon",
      "display_name": "Jon",
      "entity_type": "Person",
      "op": "UPDATE",
      "state_change": "employed → unemployed",
      "context": "Jon lost his job as a banker"
    }
  ],
  "relation_updates": [
    {
      "source": "jon",
      "relation_type": "STARTED",
      "target": "jon_dance_studio",
      "op": "CREATE",
      "context": "Jon decided to start a dance studio"
    }
  ]
}"""


EXTRACTION_PROMPT_TEMPLATE = """## Global Summary
{global_summary}

## Existing Entities (check before creating new ones!)
{existing_entities}

## Recent History (last {history_size} conversations)
{history}

## Current Conversation to Extract
Note ID: {note_id}
Date: {date}
User: {user}
Assistant: {assistant}

---
Extract entities and relationships from the CURRENT conversation.
- **MUST extract all mentioned entities** (people, places, businesses, events, activities, etc.)
- If entity already exists in the list above, use op="UPDATE"
- If entity doesn't exist, use op="CREATE"
- Create edges for cross-entity relationships
- Return valid JSON only. DO NOT return empty arrays if entities exist in conversation."""


class ContextAwareExtractor:
    """带上下文的实体关系提取器"""

    def __init__(self):
        self.llm = get_llm_client()
        self.history_window = STM_CONFIG["history_window"]
        self.max_entities = LPM_CONFIG["max_entities_in_prompt"]

    def extract(
        self,
        note: dict,
        global_summary: str,
        existing_entities: List[dict],
        history: List[dict]
    ) -> ExtractionResult:
        """
        带上下文提取实体和关系

        Args:
            note: 当前要提取的 Note
            global_summary: 全局摘要 S
            existing_entities: 已有实体列表 (from LPM)
            history: 最近的历史 Notes (last m)
        """
        # 格式化已有实体
        entities_str = self._format_existing_entities(existing_entities)

        # 格式化历史
        history_str = self._format_history(history)

        # 构建 prompt
        prompt = EXTRACTION_PROMPT_TEMPLATE.format(
            global_summary=global_summary or "No summary yet.",
            existing_entities=entities_str or "No entities yet.",
            history_size=len(history),
            history=history_str or "No history.",
            note_id=note.get("id", ""),
            date=note.get("session_date", ""),
            user=note.get("user", ""),
            assistant=note.get("assistant", "")
        )

        # 调用 LLM
        result = self.llm.call_json(prompt, EXTRACTION_SYSTEM_PROMPT)

        if not result:
            return ExtractionResult(
                note_id=note.get("id", ""),
                entity_updates=[],
                relation_updates=[],
                raw_response=""
            )

        # 解析结果
        entity_updates = []
        for e in result.get("entity_updates", []):
            try:
                entity_updates.append(EntityUpdate(
                    entity_name=e["entity_name"].lower().replace(" ", "_"),
                    display_name=e.get("display_name", e["entity_name"]),
                    entity_type=e.get("entity_type", "Person"),
                    op=e.get("op", "CREATE"),
                    state_change=e.get("state_change", ""),
                    context=e.get("context", "")
                ))
            except (KeyError, TypeError):
                continue

        relation_updates = []
        for r in result.get("relation_updates", []):
            try:
                relation_updates.append(RelationUpdate(
                    source=r["source"].lower().replace(" ", "_"),
                    relation_type=r.get("relation_type", "RELATED_TO"),
                    target=r["target"].lower().replace(" ", "_"),
                    op=r.get("op", "CREATE"),
                    context=r.get("context", "")
                ))
            except (KeyError, TypeError):
                continue

        return ExtractionResult(
            note_id=note.get("id", ""),
            entity_updates=entity_updates,
            relation_updates=relation_updates,
            raw_response=str(result)
        )

    def _format_existing_entities(self, entities: List[dict]) -> str:
        """格式化已有实体列表"""
        if not entities:
            return ""

        lines = []
        for e in entities[:self.max_entities]:
            summary = e.get('summary', '')  # 不截断，保留完整summary
            lines.append(f"- {e['name']} ({e['type']}): {summary}")

        return "\n".join(lines)

    def _format_history(self, history: List[dict]) -> str:
        """格式化历史对话"""
        if not history:
            return ""

        lines = []
        for h in history[-self.history_window:]:
            note_id = h.get("id", "")
            date = h.get("session_date", "")
            user = h.get("user", "")  # 不截断，保留完整对话
            assistant = h.get("assistant", "")  # 不截断，保留完整对话
            lines.append(f"[{note_id}] {date}: User: {user} | Assistant: {assistant}")

        return "\n".join(lines)


# ==================== 查询实体提取 ====================

QUERY_EXTRACTION_PROMPT = """Extract entity names from this question for knowledge graph search.

Question: {query}

Existing entities in the system:
{existing_entities}

Return JSON with:
- entities: list of entity names (use existing names if they match)
- keywords: key action/attribute words from the query (for context, not search)

Example:
{{"entities": ["jon", "paris"], "keywords": ["visited", "city"]}}"""


class QueryExtractor:
    """查询实体提取器"""

    def __init__(self):
        self.llm = get_llm_client()
        self._debug_file = None

    def _log(self, msg: str):
        """写入 debug 日志"""
        from pathlib import Path
        if self._debug_file is None:
            log_path = Path(__file__).parent / "debug_log.txt"
            self._debug_file = open(log_path, "a", encoding="utf-8")
        self._debug_file.write(msg + "\n")
        self._debug_file.flush()

    def extract(self, query: str, existing_entities: List[dict]) -> Tuple[List[str], List[str]]:
        """从查询中提取实体和关键词"""
        entities_str = "\n".join([
            f"- {e['name']} ({e['type']})"
            for e in existing_entities[:50]
        ]) or "No entities."

        prompt = QUERY_EXTRACTION_PROMPT.format(
            query=query,
            existing_entities=entities_str
        )

        self._log(f"[DEBUG-QE] Prompt entities list:\n{entities_str[:500]}...")

        result = self.llm.call_json(prompt)
        self._log(f"[DEBUG-QE] LLM raw response: {result}")

        if not result:
            self._log("[DEBUG-QE] ❌ LLM returned empty result!")
            return [], []

        entities = [e.lower().replace(" ", "_") for e in result.get("entities", [])]
        keywords = result.get("keywords", [])

        return entities, keywords


# ==================== 摘要生成 ====================

SUMMARY_PROMPT = """Generate a brief global summary of the conversation history based on the entities and their states.

Current Entities:
{entities}

Generate a 2-3 sentence summary that captures the main topics and entity states.
Focus on: who are the main people, what are they doing, key events.

Return plain text summary (no JSON)."""


class SummaryGenerator:
    """全局摘要生成器"""

    def __init__(self):
        self.llm = get_llm_client()

    def generate(self, entities: List[dict]) -> str:
        """生成全局摘要"""
        if not entities:
            return "No conversation history yet."

        entities_str = "\n".join([
            f"- {e['name']} ({e['type']}): {e.get('summary', 'No summary')}"
            for e in entities[:30]
        ])

        prompt = SUMMARY_PROMPT.format(entities=entities_str)
        result = self.llm.call(prompt)

        return result or "Conversation history available."


# ==================== Entity Summary 合并器 ====================

ENTITY_SUMMARY_MERGE_PROMPT = """Merge the old and new information about an entity into a comprehensive updated summary.

Entity: {entity_name} ({entity_type})

Old Summary:
{old_summary}

New Information:
{new_info}

Generate a comprehensive updated summary that:
1. Keeps ALL important information (especially relationships and state changes)
2. Preserves key relationships like STARTED, CREATED, PRACTICES, HOSTED, VISITED, etc.
3. Removes only truly redundant or outdated details
4. Captures state transitions clearly with full context
5. DO NOT truncate or abbreviate - preserve complete information

Return only the merged summary text (no JSON, no extra text)."""


class EntitySummaryMerger:
    """实体摘要合并器（使用 LLM 智能合并）"""

    def __init__(self):
        self.llm = get_llm_client()

    def merge(
        self,
        entity_name: str,
        entity_type: str,
        old_summary: str,
        new_info: str
    ) -> str:
        """
        使用 LLM 合并实体摘要

        Args:
            entity_name: 实体名称
            entity_type: 实体类型
            old_summary: 旧摘要
            new_info: 新信息（state_change 或 context）

        Returns:
            合并后的摘要（保留完整信息，不截断）
        """
        # 如果没有旧摘要，直接返回新信息
        if not old_summary or old_summary == "":
            return new_info  # 不截断，保留完整信息

        # 如果新信息已经在旧摘要中，不需要合并
        if new_info in old_summary:
            return old_summary  # 不截断，保留完整摘要

        # 调用 LLM 合并
        prompt = ENTITY_SUMMARY_MERGE_PROMPT.format(
            entity_name=entity_name,
            entity_type=entity_type,
            old_summary=old_summary,
            new_info=new_info
        )

        result = self.llm.call(prompt)

        if not result:
            # LLM 失败时回退到简单拼接
            combined = f"{old_summary}; {new_info}"
            return combined  # 不截断，保留完整信息

        return result.strip()  # 不截断，保留完整结果


if __name__ == "__main__":
    # 测试
    extractor = ContextAwareExtractor()

    test_note = {
        "id": "N54",
        "session_date": "2023-03-16",
        "user": "How's the studio coming along?",
        "assistant": "Great news! Jon's business is now officially open. He had a small opening ceremony yesterday."
    }

    existing = [
        {"name": "jon", "type": "Person", "summary": "Former banker, starting a dance studio"},
        {"name": "jon_dance_studio", "type": "Business", "summary": "Jon's dance studio, in preparation"},
    ]

    history = [
        {"id": "N52", "session_date": "2023-03-14", "user": "Is Jon ready for opening?", "assistant": "Almost! The renovation is done."},
        {"id": "N53", "session_date": "2023-03-15", "user": "What about the flooring?", "assistant": "The Marley flooring is installed."},
    ]

    result = extractor.extract(
        note=test_note,
        global_summary="Jon lost his job and is opening a dance studio. Gina is starting a clothing store.",
        existing_entities=existing,
        history=history
    )

    print("Entity Updates:")
    for e in result.entity_updates:
        print(f"  {e.op}: {e.entity_name} - {e.state_change}")

    print("\nRelation Updates:")
    for r in result.relation_updates:
        print(f"  {r.op}: {r.source} --{r.relation_type}--> {r.target}")