mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
b2955a2bdd
* 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
118 lines
No EOL
3.7 KiB
Python
118 lines
No EOL
3.7 KiB
Python
"""
|
|
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() |