74 lines
1.9 KiB
Python
74 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
import json
|
|
import os
|
|
from typing import Any, Iterator, Mapping, Optional
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
|
|
load_dotenv()
|
|
|
|
try:
|
|
from langfuse import propagate_attributes
|
|
except Exception: # pragma: no cover - optional dependency fallback
|
|
propagate_attributes = None # type: ignore[assignment]
|
|
|
|
|
|
def _langfuse_is_configured() -> bool:
|
|
return bool(os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"))
|
|
|
|
|
|
def _stringify_metadata(
|
|
metadata: Optional[Mapping[str, Any]],
|
|
) -> Optional[dict[str, str]]:
|
|
if not metadata:
|
|
return None
|
|
|
|
result: dict[str, str] = {}
|
|
for key, value in metadata.items():
|
|
if value is None:
|
|
continue
|
|
if isinstance(value, str):
|
|
result[key] = value
|
|
continue
|
|
try:
|
|
result[key] = json.dumps(value, ensure_ascii=False)
|
|
except Exception:
|
|
result[key] = str(value)
|
|
|
|
return result or None
|
|
|
|
|
|
@contextmanager
|
|
def propagate_conversation_session(
|
|
*,
|
|
conversation_id: str,
|
|
metadata: Optional[Mapping[str, Any]] = None,
|
|
tags: Optional[list[str]] = None,
|
|
) -> Iterator[None]:
|
|
"""Attach `session_id=conversation_id` to existing spans in the current execution scope.
|
|
|
|
SimpleLLMFunc already creates the turn/message spans for us. We only need to make sure
|
|
all spans emitted while handling the same conversation inherit the same Langfuse session.
|
|
"""
|
|
|
|
if not _langfuse_is_configured() or propagate_attributes is None:
|
|
yield
|
|
return
|
|
|
|
propagated_metadata = _stringify_metadata(
|
|
{
|
|
"conversation_id": conversation_id,
|
|
**(dict(metadata) if metadata else {}),
|
|
}
|
|
)
|
|
|
|
with propagate_attributes(
|
|
session_id=conversation_id,
|
|
metadata=propagated_metadata,
|
|
tags=tags,
|
|
):
|
|
yield
|