ForcePilot/backend/package/yuxi/services/langfuse_service.py

from __future__ import annotations

import asyncio
import os
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any

from yuxi.utils.logging_config import logger

try:
    from langfuse import Langfuse
    from langfuse.langchain import CallbackHandler
except Exception:  # pragma: no cover - optional dependency during local test collection
    Langfuse = None  # type: ignore[assignment]
    CallbackHandler = None  # type: ignore[assignment]


_FALSE_VALUES = {"0", "false", "no", "off"}


@dataclass(slots=True)
class LangfuseRunContext:
    callbacks: list[Any] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    tags: list[str] = field(default_factory=list)
    trace_id: str | None = None


def is_langfuse_enabled() -> bool:
    enabled_raw = (os.getenv("LANGFUSE_ENABLED") or "true").strip().lower()
    if enabled_raw in _FALSE_VALUES:
        return False

    if Langfuse is None or CallbackHandler is None:
        return False

    return bool(os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"))


@lru_cache(maxsize=1)
def get_langfuse_client() -> Langfuse | None:
    if not is_langfuse_enabled():
        return None

    if Langfuse is None:
        return None

    kwargs: dict[str, Any] = {
        "public_key": os.getenv("LANGFUSE_PUBLIC_KEY"),
        "secret_key": os.getenv("LANGFUSE_SECRET_KEY"),
    }
    host = os.getenv("LANGFUSE_BASE_URL")
    if host:
        kwargs["host"] = host

    try:
        return Langfuse(**kwargs)
    except Exception as exc:
        logger.warning(f"初始化 Langfuse 客户端失败，将跳过 tracing: {exc}")
        return None


def build_trace_metadata(
    *,
    user_id: str,
    thread_id: str,
    agent_id: str,
    request_id: str,
    operation: str,
    agent_config_id: int | None = None,
    message_type: str | None = None,
    username: str | None = None,
    login_user_id: str | None = None,
    department_id: int | str | None = None,
) -> dict[str, Any]:
    metadata: dict[str, Any] = {
        "langfuse_user_id": user_id,
        "langfuse_session_id": thread_id,
        "request_id": request_id,
        "thread_id": thread_id,
        "agent_id": agent_id,
        "operation": operation,
        "source": "yuxi",
        "feature": "chat",
    }

    if agent_config_id is not None:
        metadata["agent_config_id"] = str(agent_config_id)
    if message_type:
        metadata["message_type"] = message_type
    if username:
        metadata["username"] = username
    if login_user_id:
        metadata["login_user_id"] = login_user_id
    if department_id is not None:
        metadata["department_id"] = str(department_id)

    return metadata


def build_trace_tags(*, agent_id: str, operation: str, message_type: str | None = None) -> list[str]:
    tags = ["yuxi", "chat", operation, f"agent:{agent_id}"]
    if message_type:
        tags.append(f"message_type:{message_type}")
    return tags


def build_run_context(
    *,
    user_id: str,
    thread_id: str,
    agent_id: str,
    request_id: str,
    operation: str,
    agent_config_id: int | None = None,
    message_type: str | None = None,
    username: str | None = None,
    login_user_id: str | None = None,
    department_id: int | str | None = None,
) -> LangfuseRunContext:
    metadata = build_trace_metadata(
        user_id=user_id,
        thread_id=thread_id,
        agent_id=agent_id,
        request_id=request_id,
        operation=operation,
        agent_config_id=agent_config_id,
        message_type=message_type,
        username=username,
        login_user_id=login_user_id,
        department_id=department_id,
    )
    tags = build_trace_tags(agent_id=agent_id, operation=operation, message_type=message_type)

    client = get_langfuse_client()
    if client is None or CallbackHandler is None:
        return LangfuseRunContext(metadata=metadata, tags=tags)

    trace_id = client.create_trace_id(seed=request_id)
    handler = CallbackHandler(trace_context={"trace_id": trace_id})
    return LangfuseRunContext(callbacks=[handler], metadata=metadata, tags=tags, trace_id=trace_id)


def get_trace_info(run_context: LangfuseRunContext | None) -> dict[str, Any]:
    if run_context is None:
        return {}

    metadata = run_context.metadata or {}
    trace_id = run_context.trace_id
    if run_context.callbacks:
        last_trace_id = getattr(run_context.callbacks[0], "last_trace_id", None)
        if last_trace_id:
            trace_id = last_trace_id

    if not trace_id:
        return {}

    trace_info = {
        "langfuse_trace_id": trace_id,
        "langfuse_user_id": metadata.get("langfuse_user_id"),
        "langfuse_session_id": metadata.get("langfuse_session_id"),
    }

    # Do not fetch trace_url on the request critical path. Langfuse resolves the
    # project id via a remote API call, which can add noticeable latency when the
    # base URL is slow or unreachable. If a trace URL is still needed, fetch it
    # later via get_trace_url_async() and patch message metadata asynchronously.
    return trace_info


async def get_trace_url_async(
    run_context: LangfuseRunContext | None,
    *,
    timeout: float = 5.0,
) -> str | None:
    if run_context is None:
        return None

    trace_id = run_context.trace_id
    if run_context.callbacks:
        last_trace_id = getattr(run_context.callbacks[0], "last_trace_id", None)
        if last_trace_id:
            trace_id = last_trace_id

    if not trace_id:
        return None

    client = get_langfuse_client()
    if client is None:
        return None

    try:
        return await asyncio.wait_for(
            asyncio.to_thread(client.get_trace_url, trace_id=trace_id),
            timeout=timeout,
        )
    except Exception:
        return None


def flush_langfuse() -> None:
    client = get_langfuse_client()
    if client is None:
        return

    try:
        client.flush()
    except Exception as exc:
        logger.warning(f"刷新 Langfuse 事件失败: {exc}")