feat: add langfuse
This commit is contained in:
qixinbo
@@ -0,0 +1,15 @@
|
||||
from app.trace.attributes import (
|
||||
build_chat_trace_attributes,
|
||||
build_error_attributes,
|
||||
build_usage_attributes,
|
||||
sanitize_attributes,
|
||||
)
|
||||
from app.trace.service import trace_service
|
||||
|
||||
__all__ = [
|
||||
"trace_service",
|
||||
"sanitize_attributes",
|
||||
"build_chat_trace_attributes",
|
||||
"build_usage_attributes",
|
||||
"build_error_attributes",
|
||||
]
|
||||
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Mapping, Optional
|
||||
|
||||
|
||||
def sanitize_attributes(attributes: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
|
||||
if not attributes:
|
||||
return {}
|
||||
normalized: Dict[str, Any] = {}
|
||||
for key, value in attributes.items():
|
||||
if value is None:
|
||||
continue
|
||||
name = str(key).strip()
|
||||
if not name:
|
||||
continue
|
||||
if isinstance(value, (str, int, float, bool)):
|
||||
normalized[name] = value
|
||||
continue
|
||||
normalized[name] = str(value)
|
||||
return normalized
|
||||
|
||||
|
||||
def build_chat_trace_attributes(
|
||||
*,
|
||||
session_id: str,
|
||||
project_id: Optional[int],
|
||||
model_id: Optional[str],
|
||||
route_mode: str,
|
||||
source: str,
|
||||
knowledge_base_id: Optional[str],
|
||||
) -> Dict[str, Any]:
|
||||
return sanitize_attributes(
|
||||
{
|
||||
"session_id": session_id,
|
||||
"project_id": project_id,
|
||||
"model_id": model_id,
|
||||
"route_mode": route_mode,
|
||||
"source": source,
|
||||
"knowledge_base_id": knowledge_base_id,
|
||||
"component": "chat_stream",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def build_usage_attributes(usage: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
|
||||
if not usage:
|
||||
return {}
|
||||
return sanitize_attributes(
|
||||
{
|
||||
"usage.prompt_tokens": usage.get("prompt_tokens"),
|
||||
"usage.completion_tokens": usage.get("completion_tokens"),
|
||||
"usage.total_tokens": usage.get("total_tokens"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def build_error_attributes(exc: Exception, *, stage: str) -> Dict[str, Any]:
|
||||
return sanitize_attributes(
|
||||
{
|
||||
"error": True,
|
||||
"error.stage": stage,
|
||||
"error.type": exc.__class__.__name__,
|
||||
"error.message": str(exc),
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,187 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, Iterator, Mapping, Optional
|
||||
|
||||
from app.trace.attributes import sanitize_attributes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _NoopSpan:
|
||||
def set_attributes(self, _attributes: Optional[Mapping[str, Any]] = None) -> None:
|
||||
return None
|
||||
|
||||
def update(self, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
def update_trace(self, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
def record_error(self, _exc: Exception, *, stage: str = "unknown") -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _SpanAdapter:
|
||||
def __init__(self, raw_span: Any) -> None:
|
||||
self._raw_span = raw_span
|
||||
|
||||
def set_attributes(self, attributes: Optional[Mapping[str, Any]] = None) -> None:
|
||||
payload = sanitize_attributes(attributes)
|
||||
if not payload:
|
||||
return
|
||||
set_attribute = getattr(self._raw_span, "set_attribute", None)
|
||||
if callable(set_attribute):
|
||||
for key, value in payload.items():
|
||||
set_attribute(key, value)
|
||||
return
|
||||
update = getattr(self._raw_span, "update", None)
|
||||
if callable(update):
|
||||
update(metadata=payload)
|
||||
|
||||
def update(self, **kwargs: Any) -> None:
|
||||
update = getattr(self._raw_span, "update", None)
|
||||
if callable(update):
|
||||
update(**kwargs)
|
||||
|
||||
def update_trace(self, **kwargs: Any) -> None:
|
||||
update_trace = getattr(self._raw_span, "update_trace", None)
|
||||
if callable(update_trace):
|
||||
update_trace(**kwargs)
|
||||
|
||||
def record_error(self, exc: Exception, *, stage: str = "unknown") -> None:
|
||||
self.set_attributes(
|
||||
{
|
||||
"error": True,
|
||||
"error.stage": stage,
|
||||
"error.type": exc.__class__.__name__,
|
||||
"error.message": str(exc),
|
||||
}
|
||||
)
|
||||
self.update(level="ERROR", status_message=str(exc))
|
||||
|
||||
|
||||
class TraceService:
|
||||
def __init__(self) -> None:
|
||||
self._client: Any = None
|
||||
self._enabled = False
|
||||
self._initialized = False
|
||||
self._httpx_instrumented = False
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._enabled
|
||||
|
||||
@property
|
||||
def initialized(self) -> bool:
|
||||
return self._initialized
|
||||
|
||||
def _read_config(self) -> Dict[str, Optional[str]]:
|
||||
return {
|
||||
"public_key": os.getenv("LANGFUSE_PUBLIC_KEY"),
|
||||
"secret_key": os.getenv("LANGFUSE_SECRET_KEY"),
|
||||
"base_url": os.getenv("LANGFUSE_BASE_URL", "http://localhost:3000"),
|
||||
}
|
||||
|
||||
def initialize(self) -> bool:
|
||||
if self._initialized:
|
||||
return self._enabled
|
||||
self._initialized = True
|
||||
cfg = self._read_config()
|
||||
if not cfg["public_key"] or not cfg["secret_key"]:
|
||||
logger.info("Langfuse tracing disabled: missing LANGFUSE_PUBLIC_KEY or LANGFUSE_SECRET_KEY")
|
||||
return False
|
||||
try:
|
||||
from langfuse import Langfuse
|
||||
except Exception as exc:
|
||||
logger.warning("Langfuse tracing disabled: SDK import failed: %s", exc)
|
||||
return False
|
||||
try:
|
||||
self._client = Langfuse(
|
||||
public_key=cfg["public_key"],
|
||||
secret_key=cfg["secret_key"],
|
||||
host=cfg["base_url"],
|
||||
)
|
||||
self._enabled = True
|
||||
logger.info("Langfuse tracing enabled, host=%s", cfg["base_url"])
|
||||
except Exception as exc:
|
||||
logger.warning("Langfuse tracing initialization failed, fallback to no-op: %s", exc)
|
||||
self._client = None
|
||||
self._enabled = False
|
||||
return False
|
||||
|
||||
try:
|
||||
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
|
||||
|
||||
HTTPXClientInstrumentor().instrument()
|
||||
self._httpx_instrumented = True
|
||||
except Exception as exc:
|
||||
logger.warning("HTTPX OTEL instrumentation unavailable: %s", exc)
|
||||
return True
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self._enabled and self._client:
|
||||
flush = getattr(self._client, "flush", None)
|
||||
if callable(flush):
|
||||
try:
|
||||
flush()
|
||||
except Exception:
|
||||
pass
|
||||
close = getattr(self._client, "shutdown", None)
|
||||
if callable(close):
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
pass
|
||||
if self._httpx_instrumented:
|
||||
try:
|
||||
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
|
||||
|
||||
HTTPXClientInstrumentor().uninstrument()
|
||||
except Exception:
|
||||
pass
|
||||
self._client = None
|
||||
self._enabled = False
|
||||
self._initialized = False
|
||||
self._httpx_instrumented = False
|
||||
|
||||
@contextmanager
|
||||
def start_span(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
attributes: Optional[Mapping[str, Any]] = None,
|
||||
input_payload: Optional[Any] = None,
|
||||
) -> Iterator[_SpanAdapter | _NoopSpan]:
|
||||
if not self._enabled or not self._client:
|
||||
yield _NoopSpan()
|
||||
return
|
||||
try:
|
||||
start_observation = getattr(self._client, "start_as_current_observation", None)
|
||||
if callable(start_observation):
|
||||
ctx = start_observation(name=name, as_type="span")
|
||||
else:
|
||||
start_span = getattr(self._client, "start_as_current_span", None)
|
||||
if not callable(start_span):
|
||||
yield _NoopSpan()
|
||||
return
|
||||
ctx = start_span(name=name)
|
||||
except Exception:
|
||||
yield _NoopSpan()
|
||||
return
|
||||
try:
|
||||
with ctx as raw_span:
|
||||
span = _SpanAdapter(raw_span)
|
||||
if attributes:
|
||||
span.set_attributes(attributes)
|
||||
if input_payload is not None:
|
||||
span.update(input=input_payload)
|
||||
yield span
|
||||
except Exception as exc:
|
||||
logger.warning("Langfuse span failure (%s): %s", name, exc)
|
||||
yield _NoopSpan()
|
||||
|
||||
|
||||
trace_service = TraceService()
|
||||
Reference in New Issue
Block a user