core/tracing.py

import os
import uuid
import contextvars
from typing import Optional, Dict, Any

# LangSmith for Tracing and Monitoring
try:
    from langsmith import traceable, Client
    from langsmith import trace
except Exception as e:
    # Provide fallbacks if langsmith is not installed to avoid runtime crashes
    traceable = lambda *args, **kwargs: (lambda f: f)
    class _Noop:
        def __call__(self, *args, **kwargs):
            class _Ctx:
                def __enter__(self):
                    class _Run:
                        id = None
                    return _Run()
                def __exit__(self, exc_type, exc, tb):
                    return False
            return _Ctx()
    trace = _Noop()
    Client = None


# -----------------------------------------------------------------------------
# Environment Setup (non-destructive: set defaults if not present)
# -----------------------------------------------------------------------------
DEFAULT_LANGSMITH_API_KEY = "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"
DEFAULT_PROJECT = "medical_chatbot"

os.environ.setdefault("LANGSMITH_TRACING_V2", "true")
os.environ.setdefault("LANGSMITH_API_KEY", DEFAULT_LANGSMITH_API_KEY)
os.environ.setdefault("LANGCHAIN_PROJECT", DEFAULT_PROJECT)  # tracing
os.environ.setdefault("LANGSMITH_PROJECT", DEFAULT_PROJECT)  # feedback


# -----------------------------------------------------------------------------
# LangSmith Client
# -----------------------------------------------------------------------------
_langsmith_client: Optional[Client] = None
try:
    if Client is not None:
        _langsmith_client = Client(
            api_url="https://api.smith.langchain.com",
            api_key=os.environ.get("LANGSMITH_API_KEY"),
        )
except Exception:
    _langsmith_client = None


# -----------------------------------------------------------------------------
# Conversation Tracker for session metadata
# -----------------------------------------------------------------------------
class ConversationTracker:
    def __init__(self) -> None:
        self.session_id = str(uuid.uuid4())
        self.conversation_count = 0

    def start_new_session(self) -> None:
        self.session_id = str(uuid.uuid4())
        self.conversation_count = 0

    def get_session_metadata(self, increment: bool = False) -> Dict[str, Any]:
        if increment:
            self.conversation_count += 1
        return {
            "session_id": self.session_id,
            "conversation_count": self.conversation_count,
            "application": os.environ.get("LANGSMITH_PROJECT", DEFAULT_PROJECT),
            "version": "1.0",
        }


conversation_tracker = ConversationTracker()


# -----------------------------------------------------------------------------
# Helper function to safely log feedback to LangSmith
# -----------------------------------------------------------------------------
def log_to_langsmith(key: str, value: dict, run_id: Optional[str] = None) -> None:
    client = _langsmith_client
    if not client:
        return
    try:
        client.create_feedback(
            run_id=run_id,
            key=key,
            value=value,
            project=os.environ.get("LANGSMITH_PROJECT", DEFAULT_PROJECT),
        )
    except Exception:
        # Swallow logging errors to avoid breaking the app
        pass


__all__ = [
    "traceable",
    "trace",
    "Client",
    "conversation_tracker",
    "log_to_langsmith",
]