mseep-agentops 0.4.18__py3-none-any.whl

agentops/client/api/base.py

@@ -0,0 +1,162 @@
+"""
+Base API client classes for making HTTP requests.
+
+This module provides the foundation for all API clients in the AgentOps SDK.
+"""
+
+from typing import Any, Dict, Optional, Protocol
+
+import requests
+
+from agentops.client.http.http_client import HttpClient
+from agentops.helpers.version import get_agentops_version
+
+
+class TokenFetcher(Protocol):
+    """Protocol for token fetching functions"""
+
+    def __call__(self, api_key: str) -> str:
+        ...
+
+
+class BaseApiClient:
+    """
+    Base class for API communication with connection pooling.
+
+    This class provides the core HTTP functionality without authentication.
+    It should be used for APIs that don't require authentication.
+    """
+
+    def __init__(self, endpoint: str):
+        """
+        Initialize the base API client.
+
+        Args:
+            endpoint: The base URL for the API
+        """
+        self.endpoint = endpoint
+        self.http_client = HttpClient()
+        self.last_response: Optional[requests.Response] = None
+
+    def prepare_headers(self, custom_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+        """
+        Prepare headers for API requests.
+
+        Args:
+            custom_headers: Additional headers to include
+
+        Returns:
+            Headers dictionary with standard headers and any custom headers
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "Connection": "keep-alive",
+            "Keep-Alive": "timeout=10, max=1000",
+            "User-Agent": f"agentops-python/{get_agentops_version() or 'unknown'}",
+        }
+
+        if custom_headers:
+            headers.update(custom_headers)
+
+        return headers
+
+    def _get_full_url(self, path: str) -> str:
+        """
+        Get the full URL for a path.
+
+        Args:
+            path: The API endpoint path
+
+        Returns:
+            The full URL
+        """
+        return f"{self.endpoint}{path}"
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        data: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+        timeout: int = 30,
+    ) -> requests.Response:
+        """
+        Make a generic HTTP request
+
+        Args:
+            method: HTTP method (e.g., 'get', 'post', 'put', 'delete')
+            path: API endpoint path
+            data: Request payload (for POST, PUT methods)
+            headers: Request headers
+            timeout: Request timeout in seconds
+
+        Returns:
+            Response from the API
+
+        Raises:
+            Exception: If the request fails
+        """
+        url = self._get_full_url(path)
+
+        try:
+            response = self.http_client.request(method=method, url=url, data=data, headers=headers, timeout=timeout)
+
+            self.last_response = response
+            return response
+        except requests.RequestException as e:
+            self.last_response = None
+            raise Exception(f"{method.upper()} request failed: {str(e)}") from e
+
+    def post(self, path: str, data: Dict[str, Any], headers: Dict[str, str]) -> requests.Response:
+        """
+        Make POST request
+
+        Args:
+            path: API endpoint path
+            data: Request payload
+            headers: Request headers
+
+        Returns:
+            Response from the API
+        """
+        return self.request("post", path, data=data, headers=headers)
+
+    def get(self, path: str, headers: Dict[str, str]) -> requests.Response:
+        """
+        Make GET request
+
+        Args:
+            path: API endpoint path
+            headers: Request headers
+
+        Returns:
+            Response from the API
+        """
+        return self.request("get", path, headers=headers)
+
+    def put(self, path: str, data: Dict[str, Any], headers: Dict[str, str]) -> requests.Response:
+        """
+        Make PUT request
+
+        Args:
+            path: API endpoint path
+            data: Request payload
+            headers: Request headers
+
+        Returns:
+            Response from the API
+        """
+        return self.request("put", path, data=data, headers=headers)
+
+    def delete(self, path: str, headers: Dict[str, str]) -> requests.Response:
+        """
+        Make DELETE request
+
+        Args:
+            path: API endpoint path
+            headers: Request headers
+
+        Returns:
+            Response from the API
+        """
+        return self.request("delete", path, headers=headers)

agentops/client/api/types.py

@@ -0,0 +1,21 @@
+"""
+Common types used across API client modules.
+
+This module contains type definitions used by multiple API client modules.
+"""
+
+from typing import TypedDict
+
+
+class AuthTokenResponse(TypedDict):
+    """Response from the auth/token endpoint"""
+
+    token: str
+    project_id: str
+
+
+class UploadedObjectResponse(TypedDict):
+    """Response from the v4/objects/upload endpoint"""
+
+    url: str
+    size: int

agentops/client/api/versions/__init__.py

@@ -0,0 +1,10 @@
+"""
+API client versions package.
+
+This package contains client implementations for different API versions.
+"""
+
+from agentops.client.api.versions.v3 import V3Client
+from agentops.client.api.versions.v4 import V4Client
+
+__all__ = ["V3Client", "V4Client"]

agentops/client/api/versions/v3.py

@@ -0,0 +1,65 @@
+"""
+V3 API client for the AgentOps API.
+
+This module provides the client for the V3 version of the AgentOps API.
+"""
+
+from agentops.client.api.base import BaseApiClient
+from agentops.client.api.types import AuthTokenResponse
+from agentops.exceptions import ApiServerException
+from agentops.logging import logger
+from termcolor import colored
+
+
+class V3Client(BaseApiClient):
+    """Client for the AgentOps V3 API"""
+
+    def __init__(self, endpoint: str):
+        """
+        Initialize the V3 API client.
+
+        Args:
+            endpoint: The base URL for the API
+        """
+        # Set up with V3-specific auth endpoint
+        super().__init__(endpoint)
+
+    def fetch_auth_token(self, api_key: str) -> AuthTokenResponse:
+        path = "/v3/auth/token"
+        data = {"api_key": api_key}
+        headers = self.prepare_headers()
+
+        r = self.post(path, data, headers)
+
+        if r.status_code != 200:
+            error_msg = f"Authentication failed: {r.status_code}"
+            try:
+                error_data = r.json()
+                if "error" in error_data:
+                    error_msg = f"{error_data['error']}"
+            except Exception:
+                pass
+            logger.error(f"{error_msg} - Perhaps an invalid API key?")
+            raise ApiServerException(error_msg)
+
+        try:
+            jr = r.json()
+            token = jr.get("token")
+            if not token:
+                raise ApiServerException("No token in authentication response")
+
+            # Check project premium status
+            if jr.get("project_prem_status") != "pro":
+                logger.info(
+                    colored(
+                        "\x1b[34mYou're on the agentops free plan 🤔\x1b[0m",
+                        "blue",
+                    )
+                )
+
+            return jr
+        except Exception as e:
+            logger.error(f"Failed to process authentication response: {str(e)}")
+            raise ApiServerException(f"Failed to process authentication response: {str(e)}")
+
+    # Add V3-specific API methods here

agentops/client/api/versions/v4.py

@@ -0,0 +1,104 @@
+"""
+V4 API client for the AgentOps API.
+
+This module provides the client for the V4 version of the AgentOps API.
+"""
+
+from typing import Optional, Union, Dict
+
+from agentops.client.api.base import BaseApiClient
+from agentops.exceptions import ApiServerException
+from agentops.client.api.types import UploadedObjectResponse
+from agentops.helpers.version import get_agentops_version
+
+
+class V4Client(BaseApiClient):
+    """Client for the AgentOps V4 API"""
+
+    auth_token: str
+
+    def set_auth_token(self, token: str):
+        """
+        Set the authentication token for API requests.
+
+        Args:
+            token: The authentication token to set
+        """
+        self.auth_token = token
+
+    def prepare_headers(self, custom_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+        """
+        Prepare headers for API requests.
+
+        Args:
+            custom_headers: Additional headers to include
+        Returns:
+            Headers dictionary with standard headers and any custom headers
+        """
+        headers = {
+            "Authorization": f"Bearer {self.auth_token}",
+            "User-Agent": f"agentops-python/{get_agentops_version() or 'unknown'}",
+        }
+        if custom_headers:
+            headers.update(custom_headers)
+        return headers
+
+    def upload_object(self, body: Union[str, bytes]) -> UploadedObjectResponse:
+        """
+        Upload an object to the API and return the response.
+
+        Args:
+            body: The object to upload, either as a string or bytes.
+        Returns:
+            UploadedObjectResponse: The response from the API after upload.
+        """
+        if isinstance(body, bytes):
+            body = body.decode("utf-8")
+
+        response = self.post("/v4/objects/upload/", body, self.prepare_headers())
+
+        if response.status_code != 200:
+            error_msg = f"Upload failed: {response.status_code}"
+            try:
+                error_data = response.json()
+                if "error" in error_data:
+                    error_msg = error_data["error"]
+            except Exception:
+                pass
+            raise ApiServerException(error_msg)
+
+        try:
+            response_data = response.json()
+            return UploadedObjectResponse(**response_data)
+        except Exception as e:
+            raise ApiServerException(f"Failed to process upload response: {str(e)}")
+
+    def upload_logfile(self, body: Union[str, bytes], trace_id: int) -> UploadedObjectResponse:
+        """
+        Upload an log file to the API and return the response.
+
+        Args:
+            body: The log file to upload, either as a string or bytes.
+        Returns:
+            UploadedObjectResponse: The response from the API after upload.
+        """
+        if isinstance(body, bytes):
+            body = body.decode("utf-8")
+
+        response = self.post("/v4/logs/upload/", body, {**self.prepare_headers(), "Trace-Id": str(trace_id)})
+
+        if response.status_code != 200:
+            error_msg = f"Upload failed: {response.status_code}"
+            try:
+                error_data = response.json()
+                if "error" in error_data:
+                    error_msg = error_data["error"]
+            except Exception:
+                pass
+            raise ApiServerException(error_msg)
+
+        try:
+            response_data = response.json()
+            return UploadedObjectResponse(**response_data)
+        except Exception as e:
+            raise ApiServerException(f"Failed to process upload response: {str(e)}")

agentops/client/client.py

@@ -0,0 +1,211 @@
+import atexit
+from typing import Optional, Any
+
+from agentops.client.api import ApiClient
+from agentops.config import Config
+from agentops.exceptions import NoApiKeyException
+from agentops.instrumentation import instrument_all
+from agentops.logging import logger
+from agentops.logging.config import configure_logging, intercept_opentelemetry_logging
+from agentops.sdk.core import TraceContext, tracer
+from agentops.legacy import Session
+
+# Global variables to hold the client's auto-started trace and its legacy session wrapper
+_client_init_trace_context: Optional[TraceContext] = None
+_client_legacy_session_for_init_trace: Optional[Session] = None
+
+# Single atexit handler registered flag
+_atexit_registered = False
+
+
+def _end_init_trace_atexit():
+    """Global atexit handler to end the client's auto-initialized trace during shutdown."""
+    global _client_init_trace_context, _client_legacy_session_for_init_trace
+    if _client_init_trace_context is not None:
+        logger.debug("Auto-ending client's init trace during shutdown.")
+        try:
+            # Use global tracer to end the trace directly
+            if tracer.initialized and _client_init_trace_context.span.is_recording():
+                tracer.end_trace(_client_init_trace_context, end_state="Shutdown")
+        except Exception as e:
+            logger.warning(f"Error ending client's init trace during shutdown: {e}")
+        finally:
+            _client_init_trace_context = None
+            _client_legacy_session_for_init_trace = None  # Clear its legacy wrapper too
+
+
+class Client:
+    """Singleton client for AgentOps service"""
+
+    config: Config
+    _initialized: bool
+    _init_trace_context: Optional[TraceContext] = None  # Stores the context of the auto-started trace
+    _legacy_session_for_init_trace: Optional[
+        Session
+    ] = None  # Stores the legacy Session wrapper for the auto-started trace
+
+    __instance = None  # Class variable for singleton pattern
+
+    api: ApiClient
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> "Client":
+        if cls.__instance is None:
+            cls.__instance = super(Client, cls).__new__(cls)
+            # Initialize instance variables that should only be set once per instance
+            cls.__instance._init_trace_context = None
+            cls.__instance._legacy_session_for_init_trace = None
+        return cls.__instance
+
+    def __init__(self):
+        # Initialization of attributes like config, _initialized should happen here if they are instance-specific
+        # and not shared via __new__ for a true singleton that can be re-configured.
+        # However, the current pattern re-initializes config in init().
+        if (
+            not hasattr(self, "_initialized") or not self._initialized
+        ):  # Ensure init logic runs only once per actual initialization intent
+            self.config = Config()  # Initialize config here for the instance
+            self._initialized = False
+            # self._init_trace_context = None # Already done in __new__
+            # self._legacy_session_for_init_trace = None # Already done in __new__
+
+    def init(self, **kwargs: Any) -> None:  # Return type updated to None
+        # Recreate the Config object to parse environment variables at the time of initialization
+        # This allows re-init with new env vars if needed, though true singletons usually init once.
+        self.config = Config()
+        self.configure(**kwargs)
+
+        # Only treat as re-initialization if a different non-None API key is explicitly provided
+        provided_api_key = kwargs.get("api_key")
+        if self.initialized and provided_api_key is not None and provided_api_key != self.config.api_key:
+            logger.warning("AgentOps Client being re-initialized with a different API key. This is unusual.")
+            # Reset initialization status to allow re-init with new key/config
+            self._initialized = False
+            if self._init_trace_context and self._init_trace_context.span.is_recording():
+                logger.warning("Ending previously auto-started trace due to re-initialization.")
+                tracer.end_trace(self._init_trace_context, "Reinitialized")
+            self._init_trace_context = None
+            self._legacy_session_for_init_trace = None
+
+        if self.initialized:
+            logger.debug("AgentOps Client already initialized.")
+            # If auto_start_session was true, return the existing legacy session wrapper
+            if self.config.auto_start_session:
+                return self._legacy_session_for_init_trace
+            return None  # If not auto-starting, and already initialized, return None
+
+        if not self.config.api_key:
+            raise NoApiKeyException
+
+        configure_logging(self.config)
+        intercept_opentelemetry_logging()
+
+        self.api = ApiClient(self.config.endpoint)
+
+        try:
+            response = self.api.v3.fetch_auth_token(self.config.api_key)
+            if response is None:
+                # If auth fails, we cannot proceed with tracer initialization that depends on project_id
+                logger.error("Failed to fetch auth token. AgentOps SDK will not be initialized.")
+                return None  # Explicitly return None if auth fails
+        except Exception as e:
+            # Re-raise authentication exceptions so they can be caught by tests and calling code
+            logger.error(f"Authentication failed: {e}")
+            raise
+
+        self.api.v4.set_auth_token(response["token"])
+
+        tracing_config = self.config.dict()
+        tracing_config["project_id"] = response["project_id"]
+
+        tracer.initialize_from_config(tracing_config, jwt=response["token"])
+
+        if self.config.instrument_llm_calls:
+            instrument_all()
+
+        # self._initialized = True # Set initialized to True here - MOVED to after trace start attempt
+
+        global _atexit_registered
+        if not _atexit_registered:
+            atexit.register(_end_init_trace_atexit)  # Register new atexit handler
+            _atexit_registered = True
+
+        # Auto-start trace if configured
+        if self.config.auto_start_session:
+            if self._init_trace_context is None or not self._init_trace_context.span.is_recording():
+                logger.debug("Auto-starting init trace.")
+                trace_name = self.config.trace_name or "default"
+                self._init_trace_context = tracer.start_trace(
+                    trace_name=trace_name,
+                    tags=list(self.config.default_tags) if self.config.default_tags else None,
+                    is_init_trace=True,
+                )
+                if self._init_trace_context:
+                    self._legacy_session_for_init_trace = Session(self._init_trace_context)
+
+                    # For backward compatibility, also update the global references in legacy and client modules
+                    # These globals are what old code might have been using via agentops.legacy.get_session() or similar indirect access.
+                    global _client_init_trace_context, _client_legacy_session_for_init_trace
+                    _client_init_trace_context = self._init_trace_context
+                    _client_legacy_session_for_init_trace = self._legacy_session_for_init_trace
+
+                    # Update legacy module's _current_session and _current_trace_context
+                    # This is tricky; direct access to another module's globals is not ideal.
+                    # Prefer explicit calls if possible, but for maximum BC:
+                    try:
+                        import agentops.legacy
+
+                        agentops.legacy._current_session = self._legacy_session_for_init_trace
+                        agentops.legacy._current_trace_context = self._init_trace_context
+                    except ImportError:
+                        pass  # Should not happen
+
+                else:
+                    logger.error("Failed to start the auto-init trace.")
+                    # Even if auto-start fails, core services up to the tracer might be initialized.
+                    # Set self.initialized to True if tracer is up, but return None.
+                    self._initialized = tracer.initialized
+                    return None  # Failed to start trace
+
+            self._initialized = True  # Successfully initialized and auto-trace started (if configured)
+            # For backward compatibility, return the legacy session wrapper when auto_start_session=True
+            return self._legacy_session_for_init_trace
+        else:
+            logger.debug("Auto-start session is disabled. No init trace started by client.")
+            self._initialized = True  # Successfully initialized, just no auto-trace
+            return None  # No auto-session, so return None
+
+    def configure(self, **kwargs: Any) -> None:
+        """Update client configuration"""
+        self.config.configure(**kwargs)
+
+    @property
+    def initialized(self) -> bool:
+        return self._initialized
+
+    @initialized.setter
+    def initialized(self, value: bool) -> None:
+        if self._initialized and self._initialized != value:
+            # Allow re-setting to False if we are intentionally re-initializing
+            # This logic is now partly in init() to handle re-init cases
+            pass
+        self._initialized = value
+
+    # ------------------------------------------------------------
+    # Remove the old __instance = None at the end of the class definition if it's a repeat
+    # __instance = None # This was a class variable, should be defined once
+
+    # Make _init_trace_context and _legacy_session_for_init_trace accessible
+    # to the atexit handler if it becomes a static/class method or needs access
+    # For now, the atexit handler is global and uses global vars copied from these.
+
+    # Deprecate and remove the old global _active_session from this module.
+    # Consumers should use agentops.start_trace() or rely on the auto-init trace.
+    # For a transition, the auto-init trace's legacy wrapper is set to legacy module's globals.
+
+
+# Ensure the global _active_session (if needed for some very old compatibility) points to the client's legacy session for init trace.
+# This specific global _active_session in client.py is problematic and should be phased out.
+# For now, _client_legacy_session_for_init_trace is the primary global for the auto-init trace's legacy Session.
+
+# Remove the old global _active_session defined at the top of this file if it's no longer the primary mechanism.
+# The new globals _client_init_trace_context and _client_legacy_session_for_init_trace handle the auto-init trace.