Shepherd MCP

"""Langfuse provider client for Shepherd MCP.""" from __future__ import annotations import base64 import os from datetime import datetime from typing import Any import httpx from shepherd_mcp.models.langfuse import ( LangfuseObservation, LangfuseObservationsResponse, LangfuseScore, LangfuseScoresResponse, LangfuseSession, LangfuseSessionsResponse, LangfuseTrace, LangfuseTracesResponse, ) from shepherd_mcp.providers.base import ( AuthenticationError, BaseProvider, NotFoundError, ProviderError, RateLimitError, ) class LangfuseClient(BaseProvider): """Client for Langfuse API. Uses Basic Auth with public_key:secret_key. API Reference: https://api.reference.langfuse.com/ """ DEFAULT_HOST = "https://cloud.langfuse.com" def __init__( self, public_key: str | None = None, secret_key: str | None = None, host: str | None = None, ) -> None: """Initialize the client. Args: public_key: Langfuse public API key. If not provided, reads from LANGFUSE_PUBLIC_KEY env var. secret_key: Langfuse secret API key. If not provided, reads from LANGFUSE_SECRET_KEY env var. host: Langfuse host URL. If not provided, reads from LANGFUSE_HOST env var or defaults to cloud.langfuse.com. """ self.public_key = public_key or os.environ.get("LANGFUSE_PUBLIC_KEY") self.secret_key = secret_key or os.environ.get("LANGFUSE_SECRET_KEY") if not self.public_key or not self.secret_key: raise AuthenticationError( "No API keys provided. Set LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY " "environment variables." ) self.host = (host or os.environ.get("LANGFUSE_HOST", self.DEFAULT_HOST)).rstrip("/") # Create Basic Auth header credentials = f"{self.public_key}:{self.secret_key}" encoded = base64.b64encode(credentials.encode()).decode() self._auth_header = f"Basic {encoded}" self._client = httpx.Client( timeout=30.0, headers={ "Authorization": self._auth_header, "Content-Type": "application/json", }, ) @property def name(self) -> str: """Return the provider name.""" return "langfuse" def _handle_error_response(self, response: httpx.Response) -> None: """Handle error responses from the API.""" if response.status_code == 401: raise AuthenticationError("Authentication failed. Check your API keys.") if response.status_code == 404: try: detail = response.json().get("message", "Resource not found") except Exception: detail = "Resource not found" raise NotFoundError(detail) if response.status_code == 429: raise RateLimitError("Rate limit exceeded. Please try again later.") if response.status_code >= 400: try: detail = response.json().get("message", f"HTTP {response.status_code}") except Exception: detail = f"HTTP {response.status_code}" raise ProviderError(detail) def _parse_timestamp(self, timestamp: str | None) -> str | None: """Parse timestamp to ISO 8601 format.""" if not timestamp: return None # If already in ISO format, return as-is if "T" in timestamp: return timestamp # Try common date formats formats = [ "%Y-%m-%d", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", ] for fmt in formats: try: dt = datetime.strptime(timestamp, fmt) return dt.isoformat() + "Z" except ValueError: continue # Return as-is if no format matches return timestamp def _get(self, path: str, params: dict[str, Any] | None = None) -> dict[str, Any]: """Make a GET request.""" response = self._client.get(f"{self.host}{path}", params=params) self._handle_error_response(response) return response.json() def _post(self, path: str, json: dict[str, Any] | None = None) -> dict[str, Any]: """Make a POST request.""" response = self._client.post(f"{self.host}{path}", json=json) self._handle_error_response(response) return response.json() # ======================================================================== # Traces API # ======================================================================== def list_traces( self, limit: int = 50, page: int = 1, user_id: str | None = None, name: str | None = None, session_id: str | None = None, tags: list[str] | None = None, from_timestamp: str | None = None, to_timestamp: str | None = None, ) -> LangfuseTracesResponse: """List traces with pagination and filters. Args: limit: Maximum number of results per page. page: Page number (1-indexed). user_id: Filter by user ID. name: Filter by trace name. session_id: Filter by session ID. tags: Filter by tags. from_timestamp: Filter by start timestamp. to_timestamp: Filter by end timestamp. Returns: LangfuseTracesResponse with traces data and pagination meta. """ params: dict[str, Any] = { "limit": limit, "page": page, } if user_id: params["userId"] = user_id if name: params["name"] = name if session_id: params["sessionId"] = session_id if tags: params["tags"] = tags if from_timestamp: params["fromTimestamp"] = self._parse_timestamp(from_timestamp) if to_timestamp: params["toTimestamp"] = self._parse_timestamp(to_timestamp) data = self._get("/api/public/traces", params) return LangfuseTracesResponse(**data) def get_trace(self, trace_id: str) -> LangfuseTrace: """Get a specific trace with its observations. Args: trace_id: The trace ID to fetch. Returns: LangfuseTrace with full trace data including observations. """ data = self._get(f"/api/public/traces/{trace_id}") return LangfuseTrace(**data) # ======================================================================== # Sessions API # ======================================================================== def list_sessions( self, limit: int = 50, page: int = 1, from_timestamp: str | None = None, to_timestamp: str | None = None, ) -> LangfuseSessionsResponse: """List sessions with pagination. Args: limit: Maximum number of results per page. page: Page number (1-indexed). from_timestamp: Filter by start timestamp. to_timestamp: Filter by end timestamp. Returns: LangfuseSessionsResponse with sessions data and pagination meta. """ params: dict[str, Any] = { "limit": limit, "page": page, } if from_timestamp: params["fromTimestamp"] = self._parse_timestamp(from_timestamp) if to_timestamp: params["toTimestamp"] = self._parse_timestamp(to_timestamp) data = self._get("/api/public/sessions", params) return LangfuseSessionsResponse(**data) def get_session(self, session_id: str) -> LangfuseSession: """Get a specific session. Args: session_id: The session ID to fetch. Returns: LangfuseSession with session data. """ data = self._get(f"/api/public/sessions/{session_id}") return LangfuseSession(**data) # ======================================================================== # Observations API # ======================================================================== def list_observations( self, limit: int = 50, page: int = 1, name: str | None = None, user_id: str | None = None, trace_id: str | None = None, obs_type: str | None = None, from_timestamp: str | None = None, to_timestamp: str | None = None, ) -> LangfuseObservationsResponse: """List observations with pagination and filters. Args: limit: Maximum number of results per page. page: Page number (1-indexed). name: Filter by observation name. user_id: Filter by user ID. trace_id: Filter by trace ID. obs_type: Filter by type (GENERATION, SPAN, EVENT). from_timestamp: Filter by start timestamp. to_timestamp: Filter by end timestamp. Returns: LangfuseObservationsResponse with observations data. """ params: dict[str, Any] = { "limit": limit, "page": page, } if name: params["name"] = name if user_id: params["userId"] = user_id if trace_id: params["traceId"] = trace_id if obs_type: params["type"] = obs_type if from_timestamp: params["fromTimestamp"] = self._parse_timestamp(from_timestamp) if to_timestamp: params["toTimestamp"] = self._parse_timestamp(to_timestamp) data = self._get("/api/public/observations", params) return LangfuseObservationsResponse(**data) def get_observation(self, observation_id: str) -> LangfuseObservation: """Get a specific observation. Args: observation_id: The observation ID to fetch. Returns: LangfuseObservation with observation data. """ data = self._get(f"/api/public/observations/{observation_id}") return LangfuseObservation(**data) # ======================================================================== # Scores API # ======================================================================== def list_scores( self, limit: int = 50, page: int = 1, name: str | None = None, user_id: str | None = None, trace_id: str | None = None, from_timestamp: str | None = None, to_timestamp: str | None = None, ) -> LangfuseScoresResponse: """List scores with pagination and filters. Args: limit: Maximum number of results per page. page: Page number (1-indexed). name: Filter by score name. user_id: Filter by user ID. trace_id: Filter by trace ID. from_timestamp: Filter by start timestamp. to_timestamp: Filter by end timestamp. Returns: LangfuseScoresResponse with scores data. """ params: dict[str, Any] = { "limit": limit, "page": page, } if name: params["name"] = name if user_id: params["userId"] = user_id if trace_id: params["traceId"] = trace_id if from_timestamp: params["fromTimestamp"] = self._parse_timestamp(from_timestamp) if to_timestamp: params["toTimestamp"] = self._parse_timestamp(to_timestamp) data = self._get("/api/public/scores", params) return LangfuseScoresResponse(**data) def get_score(self, score_id: str) -> LangfuseScore: """Get a specific score. Args: score_id: The score ID to fetch. Returns: LangfuseScore with score data. """ data = self._get(f"/api/public/scores/{score_id}") return LangfuseScore(**data) def close(self) -> None: """Close the HTTP client.""" self._client.close() def __enter__(self) -> LangfuseClient: return self def __exit__(self, *args) -> None: self.close()