mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
Add AgentOps Integration to LiteLLM (#9685)
* feat(sidebars): add new item for agentops integration in Logging & Observability category * Update agentops_integration.md to enhance title formatting and remove redundant section * Enhance AgentOps integration in documentation and codebase by removing LiteLLMCallbackHandler references, adding environment variable configurations, and updating logging initialization for AgentOps support. * Update AgentOps integration documentation to include instructions for obtaining API keys and clarify environment variable setup. * Add unit tests for AgentOps integration and improve error handling in token fetching * Add unit tests for AgentOps configuration and token fetching functionality * Corrected agentops test directory * Linting fix * chore: add OpenTelemetry dependencies to pyproject.toml * chore: update OpenTelemetry dependencies and add new packages in pyproject.toml and poetry.lock
This commit is contained in:
parent
ebfff975d4
commit
b2955a2bdd
9 changed files with 608 additions and 67 deletions
|
@ -113,6 +113,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"pagerduty",
|
||||
"humanloop",
|
||||
"gcs_pubsub",
|
||||
"agentops",
|
||||
"anthropic_cache_control_hook",
|
||||
]
|
||||
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
|
||||
|
|
3
litellm/integrations/agentops/__init__.py
Normal file
3
litellm/integrations/agentops/__init__.py
Normal file
|
@ -0,0 +1,3 @@
|
|||
from .agentops import AgentOps
|
||||
|
||||
__all__ = ["AgentOps"]
|
118
litellm/integrations/agentops/agentops.py
Normal file
118
litellm/integrations/agentops/agentops.py
Normal file
|
@ -0,0 +1,118 @@
|
|||
"""
|
||||
AgentOps integration for LiteLLM - Provides OpenTelemetry tracing for LLM calls
|
||||
"""
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Dict, Any
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
|
||||
@dataclass
|
||||
class AgentOpsConfig:
|
||||
endpoint: str = "https://otlp.agentops.cloud/v1/traces"
|
||||
api_key: Optional[str] = None
|
||||
service_name: Optional[str] = None
|
||||
deployment_environment: Optional[str] = None
|
||||
auth_endpoint: str = "https://api.agentops.ai/v3/auth/token"
|
||||
|
||||
@classmethod
|
||||
def from_env(cls):
|
||||
return cls(
|
||||
endpoint="https://otlp.agentops.cloud/v1/traces",
|
||||
api_key=os.getenv("AGENTOPS_API_KEY"),
|
||||
service_name=os.getenv("AGENTOPS_SERVICE_NAME", "agentops"),
|
||||
deployment_environment=os.getenv("AGENTOPS_ENVIRONMENT", "production"),
|
||||
auth_endpoint="https://api.agentops.ai/v3/auth/token"
|
||||
)
|
||||
|
||||
class AgentOps(OpenTelemetry):
|
||||
"""
|
||||
AgentOps integration - built on top of OpenTelemetry
|
||||
|
||||
Example usage:
|
||||
```python
|
||||
import litellm
|
||||
|
||||
litellm.success_callback = ["agentops"]
|
||||
|
||||
response = litellm.completion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
)
|
||||
```
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[AgentOpsConfig] = None,
|
||||
):
|
||||
if config is None:
|
||||
config = AgentOpsConfig.from_env()
|
||||
|
||||
# Prefetch JWT token for authentication
|
||||
jwt_token = None
|
||||
project_id = None
|
||||
if config.api_key:
|
||||
try:
|
||||
response = self._fetch_auth_token(config.api_key, config.auth_endpoint)
|
||||
jwt_token = response.get("token")
|
||||
project_id = response.get("project_id")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
headers = f"Authorization=Bearer {jwt_token}" if jwt_token else None
|
||||
|
||||
otel_config = OpenTelemetryConfig(
|
||||
exporter="otlp_http",
|
||||
endpoint=config.endpoint,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
# Initialize OpenTelemetry with our config
|
||||
super().__init__(
|
||||
config=otel_config,
|
||||
callback_name="agentops"
|
||||
)
|
||||
|
||||
# Set AgentOps-specific resource attributes
|
||||
resource_attrs = {
|
||||
"service.name": config.service_name or "litellm",
|
||||
"deployment.environment": config.deployment_environment or "production",
|
||||
"telemetry.sdk.name": "agentops",
|
||||
}
|
||||
|
||||
if project_id:
|
||||
resource_attrs["project.id"] = project_id
|
||||
|
||||
self.resource_attributes = resource_attrs
|
||||
|
||||
def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Fetch JWT authentication token from AgentOps API
|
||||
|
||||
Args:
|
||||
api_key: AgentOps API key
|
||||
auth_endpoint: Authentication endpoint
|
||||
|
||||
Returns:
|
||||
Dict containing JWT token and project ID
|
||||
"""
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
|
||||
client = _get_httpx_client()
|
||||
try:
|
||||
response = client.post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
|
||||
return response.json()
|
||||
finally:
|
||||
client.close()
|
|
@ -28,6 +28,7 @@ from litellm._logging import _is_debugging_on, verbose_logger
|
|||
from litellm.batches.batch_utils import _handle_completed_batch
|
||||
from litellm.caching.caching import DualCache, InMemoryCache
|
||||
from litellm.caching.caching_handler import LLMCachingHandler
|
||||
|
||||
from litellm.constants import (
|
||||
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
|
||||
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
|
||||
|
@ -36,6 +37,7 @@ from litellm.cost_calculator import (
|
|||
RealtimeAPITokenUsageProcessor,
|
||||
_select_model_name_for_cost_calc,
|
||||
)
|
||||
from litellm.integrations.agentops import AgentOps
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
from litellm.integrations.arize.arize import ArizeLogger
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
@ -2685,7 +2687,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
"""
|
||||
try:
|
||||
custom_logger_init_args = custom_logger_init_args or {}
|
||||
if logging_integration == "lago":
|
||||
if logging_integration == "agentops": # Add AgentOps initialization
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, AgentOps):
|
||||
return callback # type: ignore
|
||||
|
||||
agentops_logger = AgentOps()
|
||||
_in_memory_loggers.append(agentops_logger)
|
||||
return agentops_logger # type: ignore
|
||||
elif logging_integration == "lago":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, LagoLogger):
|
||||
return callback # type: ignore
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue