src/hermes/observability/integrations.py

"""Third-party observability integrations: LangSmith, W&B, Arize."""

from __future__ import annotations

import logging
from typing import Any

from hermes.config.settings import get_settings

logger = logging.getLogger(__name__)


class LangSmithIntegration:
    """LangSmith tracing integration for LLM observability."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._client: Any = None
        self._enabled = False

    async def initialize(self) -> None:
        """Initialize LangSmith client if configured."""
        if not self.settings.observability.langsmith_api_key:
            logger.info("LangSmith not configured (no API key)")
            return

        try:
            from langsmith import Client

            self._client = Client(
                api_key=self.settings.observability.langsmith_api_key,
                api_url=self.settings.observability.langsmith_endpoint,
            )
            self._enabled = True
            logger.info("LangSmith initialized")
        except ImportError:
            logger.warning("langsmith package not installed. Run: pip install langsmith")
        except Exception as e:
            logger.warning(f"Failed to initialize LangSmith: {e}")

    async def create_run(
        self,
        name: str,
        run_type: str = "chain",
        inputs: dict[str, Any] | None = None,
        project_name: str | None = None,
    ) -> str | None:
        """Create a trace run."""
        if not self._enabled or not self._client:
            return None

        try:
            project = project_name or self.settings.observability.langsmith_project
            run = self._client.create_run(
                name=name,
                run_type=run_type,
                inputs=inputs or {},
                project_name=project,
            )
            return str(run.id) if hasattr(run, "id") else None
        except Exception as e:
            logger.debug(f"LangSmith create_run failed: {e}")
            return None

    async def update_run(
        self,
        run_id: str,
        outputs: dict[str, Any] | None = None,
        error: str | None = None,
        end_time: float | None = None,
    ) -> None:
        """Update a trace run with outputs or error."""
        if not self._enabled or not self._client:
            return

        try:
            self._client.update_run(
                run_id=run_id,
                outputs=outputs,
                error=error,
                end_time=end_time,
            )
        except Exception as e:
            logger.debug(f"LangSmith update_run failed: {e}")

    async def trace_llm_call(
        self,
        model: str,
        prompt: str,
        response: str,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        duration_ms: float = 0.0,
        tags: list[str] | None = None,
    ) -> None:
        """Trace an LLM call to LangSmith."""
        if not self._enabled:
            return

        run_id = await self.create_run(
            name=f"llm_call_{model}",
            run_type="llm",
            inputs={"prompt": prompt, "model": model},
        )
        if run_id:
            await self.update_run(
                run_id=run_id,
                outputs={
                    "response": response,
                    "token_usage": {
                        "prompt_tokens": prompt_tokens,
                        "completion_tokens": completion_tokens,
                    },
                    "duration_ms": duration_ms,
                },
            )

    @property
    def enabled(self) -> bool:
        """Check if LangSmith is enabled."""
        return self._enabled


class WeightsAndBiasesIntegration:
    """Weights & Biases integration for experiment tracking."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._run: Any = None
        self._enabled = False

    async def initialize(self, project: str = "hermes-platform") -> None:
        """Initialize W&B run."""
        try:
            import wandb

            self._run = wandb.init(
                project=project,
                config={
                    "observability_enabled": self.settings.observability.enabled,
                    "model_provider": self.settings.model.provider,
                    "model_name": self.settings.model.name,
                },
                reinit=True,
            )
            self._enabled = True
            logger.info("W&B initialized")
        except ImportError:
            logger.warning("wandb package not installed. Run: pip install wandb")
        except Exception as e:
            logger.warning(f"Failed to initialize W&B: {e}")

    async def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
        """Log metrics to W&B."""
        if not self._enabled or not self._run:
            return
        try:
            self._run.log(metrics, step=step)
        except Exception as e:
            logger.debug(f"W&B log_metrics failed: {e}")

    async def log_llm_call(
        self,
        model: str,
        prompt_tokens: int,
        completion_tokens: int,
        duration_ms: float,
        success: bool = True,
    ) -> None:
        """Log LLM call metrics to W&B."""
        await self.log_metrics({
            f"llm/{model}/prompt_tokens": prompt_tokens,
            f"llm/{model}/completion_tokens": completion_tokens,
            f"llm/{model}/total_tokens": prompt_tokens + completion_tokens,
            f"llm/{model}/duration_ms": duration_ms,
            f"llm/{model}/success": 1.0 if success else 0.0,
        })

    async def finish(self) -> None:
        """Finish the W&B run."""
        if self._enabled and self._run:
            try:
                self._run.finish()
            except Exception as e:
                logger.debug(f"W&B finish failed: {e}")

    @property
    def enabled(self) -> bool:
        """Check if W&B is enabled."""
        return self._enabled


class ArizeIntegration:
    """Arize AI integration for LLM observability."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._client: Any = None
        self._enabled = False

    async def initialize(self, api_key: str = "", space_key: str = "") -> None:
        """Initialize Arize client."""
        if not api_key and not self.settings.observability.langsmith_api_key:
            logger.info("Arize not configured")
            return

        try:
            from arize.api import Client

            self._client = Client(
                api_key=api_key or self.settings.observability.langsmith_api_key,
                space_key=space_key or "hermes-platform",
            )
            self._enabled = True
            logger.info("Arize initialized")
        except ImportError:
            logger.warning("arize package not installed. Run: pip install arize")
        except Exception as e:
            logger.warning(f"Failed to initialize Arize: {e}")

    async def log_llm_event(
        self,
        model: str,
        prompt: str,
        response: str,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        latency_ms: float = 0.0,
        tags: dict[str, str] | None = None,
    ) -> None:
        """Log an LLM event to Arize."""
        if not self._enabled or not self._client:
            return
        try:
            self._client.log_llm_record(
                model=model,
                prompt=prompt,
                response=response,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                latency_ms=latency_ms,
                tags=tags or {},
            )
        except Exception as e:
            logger.debug(f"Arize log_llm_event failed: {e}")

    @property
    def enabled(self) -> bool:
        """Check if Arize is enabled."""
        return self._enabled


class ObservabilityManager:
    """Central manager for all observability integrations."""

    def __init__(self) -> None:
        self.langsmith = LangSmithIntegration()
        self.wandb = WeightsAndBiasesIntegration()
        self.arize = ArizeIntegration()

    async def initialize_all(self) -> None:
        """Initialize all configured integrations."""
        await self.langsmith.initialize()
        await self.wandb.initialize()
        await self.arize.initialize()

    async def trace_llm_call(
        self,
        model: str,
        prompt: str,
        response: str,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        duration_ms: float = 0.0,
    ) -> None:
        """Trace an LLM call across all configured providers."""
        await self.langsmith.trace_llm_call(
            model=model,
            prompt=prompt,
            response=response,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            duration_ms=duration_ms,
        )
        await self.wandb.log_llm_call(
            model=model,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            duration_ms=duration_ms,
        )
        await self.arize.log_llm_event(
            model=model,
            prompt=prompt,
            response=response,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            latency_ms=duration_ms,
        )

    async def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
        """Log metrics across all configured providers."""
        await self.wandb.log_metrics(metrics, step=step)

    async def shutdown(self) -> None:
        """Shutdown all integrations gracefully."""
        await self.wandb.finish()


obs_manager = ObservabilityManager()