PMFlow/backend/app/services/semantic_fact_service/shared.py

from __future__ import annotations

import hashlib
import json
import re
from datetime import UTC, datetime
from typing import Any, Optional

from sqlalchemy import and_, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from ...config.ai_prompts import (
    PROMPT_SEMANTIC_FACT_EXTRACT_DOCUMENT,
    PROMPT_SEMANTIC_FACT_EXTRACT_PROJECT,
    PROMPT_SEMANTIC_FACT_EXTRACT_RELEASE,
    PROMPT_SEMANTIC_FACT_EXTRACT_REQUIREMENT,
    PROMPT_SEMANTIC_FACT_EXTRACT_REVIEW,
    PROMPT_SEMANTIC_FACT_EXTRACT_SIGNAL,
    PROMPT_SEMANTIC_FACT_EXTRACT_TASK,
    PROMPT_SEMANTIC_FACT_EXTRACT_USER,
)
from ...models.ingest_event import IngestEvent
from ...models.knowledge_base_chunk import KnowledgeBaseChunk
from ...models.knowledge_base_doc import KnowledgeBaseDoc
from ...models.project import Project
from ...models.semantic_fact import SemanticFact
from ...models.semantic_fact_extraction_run import SemanticFactExtractionRun
from ...models.signal import Signal
from ...models.task import Task
from ...models.truth_node import TruthNode
from ...models.user import User

SEMANTIC_FACT_EXTRACT_VERSION = "semantic-facts-v1"
SEMANTIC_FACT_EXTRACT_CHANNEL = "semantic_fact_extract"
SEMANTIC_FACT_SUPPORTED_ENTITY_TYPES = {
    "project",
    "task",
    "signal",
    "document",
    "user",
    "requirement",
    "release",
    "review",
}
SEMANTIC_FACT_INACTIVE_STATUSES = {
    "archived",
    "closed",
    "deleted",
    "disabled",
    "inactive",
    "offline",
    "已关闭",
    "已归档",
}
TRUTH_NODE_TYPE_TO_ENTITY_TYPE = {
    "charter": "requirement",
    "release": "release",
    "review": "review",
}
ENTITY_TYPE_TO_TRUTH_NODE_TYPE = {
    value: key for key, value in TRUTH_NODE_TYPE_TO_ENTITY_TYPE.items()
}
SEMANTIC_FACT_PROMPTS = {
    "project": PROMPT_SEMANTIC_FACT_EXTRACT_PROJECT,
    "task": PROMPT_SEMANTIC_FACT_EXTRACT_TASK,
    "signal": PROMPT_SEMANTIC_FACT_EXTRACT_SIGNAL,
    "document": PROMPT_SEMANTIC_FACT_EXTRACT_DOCUMENT,
    "user": PROMPT_SEMANTIC_FACT_EXTRACT_USER,
    "requirement": PROMPT_SEMANTIC_FACT_EXTRACT_REQUIREMENT,
    "release": PROMPT_SEMANTIC_FACT_EXTRACT_RELEASE,
    "review": PROMPT_SEMANTIC_FACT_EXTRACT_REVIEW,
}
SEMANTIC_FACT_FUNCTION_NAMES = {
    entity_type: f"semantic_fact_extract_{entity_type}"
    for entity_type in SEMANTIC_FACT_SUPPORTED_ENTITY_TYPES
}


def _utc_now_iso() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat()


def _normalize_fact_value(value: Any) -> str:
    text = str(value or "").strip().lower()
    if not text:
        return ""
    text = text.replace("a股", "a股")
    text = re.sub(r"\s+", "_", text)
    text = re.sub(r"[^a-z0-9\u4e00-\u9fff_]+", "_", text)
    text = re.sub(r"_+", "_", text).strip("_")
    return text


def _normalized_entity_type(entity_type: Any) -> str:
    value = str(entity_type or "").strip().lower()
    if value == "doc":
        return "document"
    return value


def _normalized_status(value: Any) -> str:
    return str(value or "").strip().lower()


def _status_is_inactive(value: Any) -> bool:
    return _normalized_status(value) in SEMANTIC_FACT_INACTIVE_STATUSES


def _parse_iso_datetime(value: Any) -> Optional[datetime]:
    text = str(value or "").strip()
    if not text:
        return None
    try:
        parsed = datetime.fromisoformat(text.replace("Z", "+00:00"))
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=UTC)
    return parsed.astimezone(UTC)


__all__ = [
    "AsyncSession",
    "ENTITY_TYPE_TO_TRUTH_NODE_TYPE",
    "IngestEvent",
    "KnowledgeBaseChunk",
    "KnowledgeBaseDoc",
    "Project",
    "PROMPT_SEMANTIC_FACT_EXTRACT_DOCUMENT",
    "PROMPT_SEMANTIC_FACT_EXTRACT_PROJECT",
    "PROMPT_SEMANTIC_FACT_EXTRACT_RELEASE",
    "PROMPT_SEMANTIC_FACT_EXTRACT_REQUIREMENT",
    "PROMPT_SEMANTIC_FACT_EXTRACT_REVIEW",
    "PROMPT_SEMANTIC_FACT_EXTRACT_SIGNAL",
    "PROMPT_SEMANTIC_FACT_EXTRACT_TASK",
    "PROMPT_SEMANTIC_FACT_EXTRACT_USER",
    "SEMANTIC_FACT_EXTRACT_VERSION",
    "SEMANTIC_FACT_EXTRACT_CHANNEL",
    "SEMANTIC_FACT_FUNCTION_NAMES",
    "SEMANTIC_FACT_INACTIVE_STATUSES",
    "SEMANTIC_FACT_PROMPTS",
    "SEMANTIC_FACT_SUPPORTED_ENTITY_TYPES",
    "SemanticFact",
    "SemanticFactExtractionRun",
    "Signal",
    "Task",
    "TRUTH_NODE_TYPE_TO_ENTITY_TYPE",
    "TruthNode",
    "User",
    "_normalize_fact_value",
    "_normalized_entity_type",
    "_normalized_status",
    "_parse_iso_datetime",
    "_status_is_inactive",
    "_utc_now_iso",
    "and_",
    "func",
    "hashlib",
    "json",
    "or_",
    "select",
]