meeting_memory/meeting_memory/extractor.py

import json
import logging
import re
import sys
from enum import Enum
from typing import Any, List, Optional

from openai import OpenAI
from pydantic import BaseModel, Field

from meeting_memory.config import config
from meeting_memory.prompts import extract_entities as prompt_extract_entities
from meeting_memory.prompts import extract_facts as prompt_extract_facts
from meeting_memory.prompts import resolve_entities as prompt_dedupe_nodes
from meeting_memory.prompts import resolve_facts as prompt_dedupe_edges
from meeting_memory.prompts import summarize_entity as prompt_summarize

logger = logging.getLogger(__name__)

client = OpenAI(
    api_key=config.llm.api_key or None,
    base_url=config.llm.base_url if config.llm.base_url else None,
)


class EntityType(str, Enum):
    DEPARTMENT = 'Department'
    PROJECT = 'Project'
    METRIC = 'Metric'
    PERSON = 'Person'
    SYSTEM = 'System'
    DOCUMENT = 'Document'
    PARTICIPANT = 'participant'
    UNKNOWN = 'Unknown'


# Normalization map: legacy LLM output → canonical type
_ENTITY_TYPE_ALIASES = {
    '组织': 'Department',
    'organization': 'Department',
    '部门': 'Department',
    '指标': 'Metric',
    'kpi': 'Metric',
    '项目': 'Project',
}


def _canonical_entity_type(raw: str) -> str:
    normalized = raw.strip()
    if normalized in _ENTITY_TYPE_ALIASES:
        return _ENTITY_TYPE_ALIASES[normalized]
    for member in EntityType:
        if member.value.lower() == normalized.lower():
            return member.value
    return EntityType.UNKNOWN.value


def _neo4j_labels(entity_type: str) -> list[str]:
    canonical = _canonical_entity_type(entity_type)
    labels = ['Entity']
    if canonical != EntityType.UNKNOWN.value:
        labels.append(canonical)
    return labels


class Entity(BaseModel):
    name: str
    entity_type: str = EntityType.UNKNOWN.value
    description: str = ''


class Relation(BaseModel):
    source_entity_name: str
    target_entity_name: str
    relation_type: str
    fact: str = ''
    valid_at: str = ''
    invalid_at: str = ''
    evidence: str = ''
    qualifiers: List[str] = Field(default_factory=list)
    confidence: float = 0.0


class ActionItem(BaseModel):
    task: str
    assignee: str = ''
    deadline: str = ''
    status: str = '待办'
    priority: str = '中'


class Decision(BaseModel):
    content: str
    proposer: str = ''
    status: str = '已决'


class MeetingMetric(BaseModel):
    metric_name: str
    value: str
    target: str = ''
    owner: str = ''
    trend: str = ''
    unit: str = ''


class DepartmentInfo(BaseModel):
    name: str
    description: str = ''
    projects: List[str] = Field(default_factory=list)


class MeetingExtraction(BaseModel):
    title: str
    date: str = ''
    participants: List[str] = Field(default_factory=list)
    agenda: List[str] = Field(default_factory=list)
    entities: List[Entity] = Field(default_factory=list)
    relations: List[Relation] = Field(default_factory=list)
    action_items: List[ActionItem] = Field(default_factory=list)
    decisions: List[Decision] = Field(default_factory=list)
    metrics: List[MeetingMetric] = Field(default_factory=list)
    departments: List[DepartmentInfo] = Field(default_factory=list)
    summary: str = ''


def _call_llm(
    messages: list[dict],
    response_model: type | None = None,
    stream: bool = False,
    max_tokens: int | None = None,
) -> Any:
    kwargs = {
        'model': config.llm.model,
        'messages': messages,
        'max_tokens': max_tokens or config.llm.max_tokens,
        'temperature': config.llm.temperature,
    }
    if response_model is not None:
        kwargs['response_format'] = {'type': 'json_object'}
    if stream:
        kwargs['stream'] = True

    if not stream:
        response = client.chat.completions.create(**kwargs)
        content = response.choices[0].message.content
        if content is None:
            raise ValueError('LLM returned empty response')
        return content

    kwargs['stream'] = True
    response = client.chat.completions.create(**kwargs)
    chunks: List[str] = []
    print('\n[LLM] 开始流式输出：')
    for event in response:
        if not event.choices:
            continue
        delta = event.choices[0].delta.content
        if not delta:
            continue
        chunks.append(delta)
        sys.stdout.write(delta)
        sys.stdout.flush()
    print('\n[LLM] 输出结束')
    return ''.join(chunks)


def _try_parse_json(content: str) -> dict | list:
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        logger.warning('JSON parsing failed; trying to repair extracted block')
        match = re.search(r'\{.*\}|\[.*\]', content, re.DOTALL)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError as exc:
                logger.error('Repaired JSON still failed to parse: %s', exc)
        raise


def _normalize_string(name: str) -> str:
    return re.sub(r'[\s]+', ' ', name.strip().lower())


def _format_episodes_for_context(episodes: list[dict] | None) -> str:
    if not episodes:
        return ''
    return '\n'.join(
        f'[Episode {i}] {ep.get("content", "")}'
        for i, ep in enumerate(episodes)
    )


# ===== Step 1: 实体节点抽取 =====

def extract_entities_from_text(
    text: str,
    previous_episodes: list[dict] | None = None,
    entity_types: list[dict] | None = None,
    stream: bool = False,
) -> list[dict]:
    context = {
        'episode_content': text,
        'previous_episodes': previous_episodes or [],
        'entity_types': entity_types or [],
    }
    messages = prompt_extract_entities(context)
    content = _call_llm(messages, stream=stream)
    try:
        data = _try_parse_json(content)
    except Exception as exc:
        logger.error('Failed to parse entity extraction result: %s', exc)
        return []
    if isinstance(data, dict):
        data = data.get('entities', data.get('extracted_entities', []))
    if not isinstance(data, list):
        return []
    result = []
    for item in data:
        if isinstance(item, dict) and item.get('name', '').strip():
            result.append({
                'name': item['name'].strip(),
                'entity_type': item.get('entity_type', 'Entity'),
                'description': item.get('description', ''),
                'evidence': item.get('evidence', ''),
            })
    return result


# ===== Step 2: 实体去重 =====

def resolve_entities_against_graph(
    extracted: list[dict],
    existing: list[dict],
    episode_content: str = '',
) -> list[dict]:
    if not existing:
        return extracted

    context = {
        'extracted_entities': extracted,
        'existing_entities': existing,
        'episode_content': episode_content,
    }
    messages = prompt_dedupe_nodes(context)
    content = _call_llm(messages)
    try:
        data = _try_parse_json(content)
    except Exception as exc:
        logger.warning('LLM dedup failed, keeping all extracted: %s', exc)
        return extracted

    if isinstance(data, dict):
        data = data.get('entity_resolutions', data.get('resolutions', []))

    extracted_by_id = {i: e for i, e in enumerate(extracted)}
    existing_by_id = {c.get('candidate_id'): c for c in existing}

    for resolution in (data if isinstance(data, list) else []):
        if not isinstance(resolution, dict):
            continue
        rid = resolution.get('id')
        dup_id = resolution.get('duplicate_candidate_id', -1)
        if rid is None or rid not in extracted_by_id:
            continue
        if dup_id >= 0 and dup_id in existing_by_id:
            extracted_by_id[rid]['_resolved_to'] = existing_by_id[dup_id]
            extracted_by_id[rid]['name'] = resolution.get('name', extracted_by_id[rid]['name'])

    return [e for e in extracted_by_id.values() if '_resolved_to' not in e]


# ===== Step 3: 事实关系抽取 =====

def extract_facts_from_text(
    text: str,
    entities: list[dict],
    reference_time: str = '',
    previous_episodes: list[dict] | None = None,
    stream: bool = False,
) -> list[dict]:
    if len(entities) < 2:
        return []

    context = {
        'episode_content': text,
        'entities': entities,
        'reference_time': reference_time,
        'previous_episodes': previous_episodes or [],
    }
    messages = prompt_extract_facts(context)
    content = _call_llm(messages, stream=stream)
    try:
        data = _try_parse_json(content)
    except Exception as exc:
        logger.error('Failed to parse fact extraction result: %s', exc)
        return []

    if isinstance(data, dict):
        data = data.get('edges', data.get('facts', data.get('relations', [])))

    if not isinstance(data, list):
        return []

    entity_names = {_normalize_string(e.get('name', '')) for e in entities}
    result = []
    for item in data:
        if not isinstance(item, dict):
            continue
        src = _normalize_string(item.get('source_entity_name', ''))
        tgt = _normalize_string(item.get('target_entity_name', ''))
        if src not in entity_names or tgt not in entity_names:
            continue
        if src == tgt:
            continue
        result.append({
            'source_entity_name': item['source_entity_name'],
            'target_entity_name': item['target_entity_name'],
            'relation_type': item.get('relation_type', '关联'),
            'fact': item.get('fact', ''),
            'valid_at': item.get('valid_at', ''),
            'invalid_at': item.get('invalid_at', ''),
            'evidence': item.get('evidence', ''),
            'qualifiers': item.get('qualifiers', []),
            'confidence': item.get('confidence', 0.0),
        })
    return result


# ===== Step 4: 事实去重/矛盾检测 =====

def resolve_facts_against_graph(
    new_fact: dict,
    existing_facts: list[dict],
    invalidation_candidates: list[dict],
) -> dict:
    if not existing_facts:
        return {'is_duplicate': False, 'is_contradicted': False, 'resolved': new_fact}

    context = {
        'new_fact': new_fact.get('fact', ''),
        'existing_facts': existing_facts,
        'invalidation_candidates': invalidation_candidates,
    }
    messages = prompt_dedupe_edges(context)
    content = _call_llm(messages)
    try:
        data = _try_parse_json(content)
    except Exception as exc:
        logger.warning('Fact dedup failed, treating as new: %s', exc)
        return {'is_duplicate': False, 'is_contradicted': False, 'resolved': new_fact}

    if not isinstance(data, dict):
        return {'is_duplicate': False, 'is_contradicted': False, 'resolved': new_fact}
    return {
        'is_duplicate': len(data.get('duplicate_facts', [])) > 0,
        'is_contradicted': len(data.get('contradicted_facts', [])) > 0,
        'resolved': new_fact,
        'duplicate_facts': data.get('duplicate_facts', []),
        'contradicted_facts': data.get('contradicted_facts', []),
    }


# ===== Step 5: 实体摘要 =====

def extract_entity_summary(
    entity_name: str,
    episodes: list[str],
    existing_summary: str = '',
    previous_episodes: list[dict] | None = None,
) -> str:
    context = {
        'entity_name': entity_name,
        'episodes': episodes,
        'existing_summary': existing_summary,
        'previous_episodes': previous_episodes or [],
    }
    messages = prompt_summarize(context)
    content = _call_llm(messages, max_tokens=1024)
    try:
        data = _try_parse_json(content)
    except Exception:
        logger.warning('Failed to parse summary, using empty')
        return ''
    if isinstance(data, dict):
        return data.get('summary', '')
    return ''


# ===== 统一入口（兼容原有接口） =====

EXTRACTION_SYSTEM_PROMPT = """
你是一个专业的会议知识抽取助手。你的任务是从中文会议记录中抽取结构化事实，尤其要抽出更细粒度、更有语义深度的关系。

输出要求：
1. 只输出一个 JSON 对象，不要输出解释文字。
2. 关系抽取不要停留在"部门汇报了工作"这种浅层描述，要尽可能向下细化到：
   - 责任归属
   - 目标值 / 当前值 / 趋势
   - 约束条件
   - 因果 / 影响
   - 时间要求
   - 依赖关系
   - 部署 / 决策 / 要求 / 风险 / 支撑关系
3. 每条关系尽量同时给出：
   - subject / predicate / object
   - fact: 一句自然语言事实表述
   - qualifiers: 限定条件、范围、状态、数值、约束等
   - evidence: 原文中的关键短句或压缩证据
   - confidence: 0 到 1 之间
   - valid_at / invalid_at: 如果文中明确提到时间，可填写；否则留空
4. 如果原文存在多个事实，不要只抽象概括，要拆成多条关系。
5. 避免空泛关系词，优先使用更具体的谓词，例如：
   - 负责 / 汇报 / 目标值 / 当前值 / 低于 / 高于 / 要求 / 督导 / 推进 / 影响 / 支撑 / 依赖 / 计划 / 完成 / 截止于
"""


def extract_meeting_info(text: str, stream: bool = False) -> MeetingExtraction:
    user_prompt = f"""
请从下面会议记录中提取结构化信息，并重点做"深层关系抽取"和"层次结构识别"。

输出 JSON 字段：
- title
- date
- participants
- agenda
- entities: name, entity_type, description
  - entity_type 请使用: Department（部门）、Project（项目）、Metric（指标）、Person（人物）、System（系统）、Document（文档）
- relations:
  - source_entity_name: 源实体名称
  - target_entity_name: 目标实体名称
  - relation_type: 关系类型（如 HAS_PROJECT、HAS_METRIC、负责、汇报、目标值、推进、依赖）
  - fact: 一句自然语言事实描述
  - valid_at（可选）
  - invalid_at（可选）
  - evidence: 原文证据
  - qualifiers: 限定条件列表
  - confidence: 0~1
- action_items: task, assignee, deadline, status, priority
- decisions: content, proposer, status
- metrics: metric_name, value, target, owner, trend, unit
- departments: [{{"name": "部门名称", "description": "", "projects": ["项目名1", "项目名2"]}}]
- summary

层次关系规则：
1. Department 管辖 Project → relation_type 用 HAS_PROJECT
2. Project 拥有 Metric → relation_type 用 HAS_METRIC
3. 其他事实关系（负责、汇报、目标值等）直接用 relation_type 表达

关系抽取规则：
1. 不要只抽"汇报了工作"这种会议动作，要尽量继续下钻出具体事实。
2. 如果一句话里同时包含"主体 + 指标 + 当前值 + 目标值 + 负责人 + 趋势"，应拆成多条关系或在 qualifiers 中保留这些细节。
3. 对于"要求、部署、负责、依赖、影响、约束、目标、风险"类信息优先保留。
4. fact 必须是一句完整、自然、可检索的事实描述。
5. qualifiers 用于补充数值、范围、状态、条件、截止时间、优先级等信息。
6. evidence 用原文中的关键词短句，不要太长。
7. confidence 取值 0 到 1。

会议记录如下：
{text}
"""
    content = _call_llm([
        {'role': 'system', 'content': EXTRACTION_SYSTEM_PROMPT},
        {'role': 'user', 'content': user_prompt},
    ], stream=stream)
    data = _try_parse_json(content)
    data = _normalize_meeting_data(data)
    return MeetingExtraction(**data)


def _normalize_meeting_data(data: dict) -> dict:
    if not isinstance(data, dict):
        return {}
    return {
        'title': _as_str(data.get('title')),
        'date': _as_str(data.get('date')),
        'participants': _as_str_list(data.get('participants')),
        'agenda': _as_str_list(data.get('agenda')),
        'entities': _normalize_entities(data.get('entities')),
        'relations': _normalize_relations(data.get('relations')),
        'action_items': _normalize_action_items(data.get('action_items')),
        'decisions': _normalize_decisions(data.get('decisions')),
        'metrics': _normalize_metrics(data.get('metrics')),
        'departments': _normalize_departments(data.get('departments')),
        'summary': _as_str(data.get('summary')),
    }


def _as_str(value) -> str:
    if value is None:
        return ''
    if isinstance(value, str):
        return value
    return str(value)


def _as_float(value) -> float:
    if value is None or value == '':
        return 0.0
    try:
        numeric = float(value)
        return max(0.0, min(1.0, numeric))
    except (TypeError, ValueError):
        return 0.0


def _as_str_list(value) -> List[str]:
    if isinstance(value, dict):
        items = []
        for key, item in value.items():
            key_text = _as_str(key)
            value_text = _as_str(item)
            if key_text and value_text:
                items.append(f'{key_text}: {value_text}')
            elif key_text:
                items.append(key_text)
            elif value_text:
                items.append(value_text)
        return items
    if not isinstance(value, list):
        return []
    return [_as_str(item) for item in value if item is not None]


def _normalize_entities(value) -> List[dict]:
    if not isinstance(value, list):
        return []
    items = []
    for entity in value:
        if not isinstance(entity, dict):
            continue
        items.append({
            'name': _as_str(entity.get('name')),
            'entity_type': _as_str(entity.get('entity_type')),
            'description': _as_str(entity.get('description')),
        })
    return items


def _normalize_relations(value) -> List[dict]:
    if not isinstance(value, list):
        return []
    items = []
    for relation in value:
        if not isinstance(relation, dict):
            continue
        source = _as_str(relation.get('source_entity_name') or relation.get('subject', ''))
        target = _as_str(relation.get('target_entity_name') or relation.get('object', ''))
        rtype = _as_str(relation.get('relation_type') or relation.get('predicate', ''))
        fact = _as_str(relation.get('fact'))
        if not fact and source and rtype and target:
            fact = f'{source} {rtype} {target}'
        items.append({
            'source_entity_name': source,
            'target_entity_name': target,
            'relation_type': rtype,
            'fact': fact,
            'qualifiers': _as_str_list(relation.get('qualifiers')),
            'evidence': _as_str(relation.get('evidence')),
            'confidence': _as_float(relation.get('confidence')),
            'valid_at': _as_str(relation.get('valid_at')),
            'invalid_at': _as_str(relation.get('invalid_at')),
        })
    return items


def _normalize_action_items(value) -> List[dict]:
    if not isinstance(value, list):
        return []
    items = []
    for action in value:
        if not isinstance(action, dict):
            continue
        items.append({
            'task': _as_str(action.get('task')),
            'assignee': _as_str(action.get('assignee')),
            'deadline': _as_str(action.get('deadline')),
            'status': _as_str(action.get('status')) or '待办',
            'priority': _as_str(action.get('priority')) or '中',
        })
    return items


def _normalize_decisions(value) -> List[dict]:
    if not isinstance(value, list):
        return []
    items = []
    for decision in value:
        if not isinstance(decision, dict):
            continue
        items.append({
            'content': _as_str(decision.get('content')),
            'proposer': _as_str(decision.get('proposer')),
            'status': _as_str(decision.get('status')) or '已决',
        })
    return items


def _normalize_metrics(value) -> List[dict]:
    if not isinstance(value, list):
        return []
    items = []
    for metric in value:
        if not isinstance(metric, dict):
            continue
        items.append({
            'metric_name': _as_str(metric.get('metric_name')),
            'value': _as_str(metric.get('value')),
            'target': _as_str(metric.get('target')),
            'owner': _as_str(metric.get('owner')),
            'trend': _as_str(metric.get('trend')),
            'unit': _as_str(metric.get('unit')),
        })
    return items


def _normalize_departments(value) -> List[dict]:
    if not isinstance(value, list):
        return []
    items = []
    for dept in value:
        if not isinstance(dept, dict):
            continue
        name = _as_str(dept.get('name'))
        if not name:
            continue
        items.append({
            'name': name,
            'description': _as_str(dept.get('description')),
            'projects': _as_str_list(dept.get('projects')),
        })
    return items