mirror of
https://github.com/crewAIInc/crewAI.git
synced 2025-12-16 04:18:35 +00:00
Compare commits
3 Commits
1.2.0
...
devin/1751
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4379ad26d1 | ||
|
|
03ee4c59eb | ||
|
|
ac1080afd8 |
216
docs/oauth2_llm_providers.md
Normal file
216
docs/oauth2_llm_providers.md
Normal file
@@ -0,0 +1,216 @@
|
||||
# OAuth2 LLM Providers
|
||||
|
||||
CrewAI supports OAuth2 authentication for custom LiteLLM providers through configuration files.
|
||||
|
||||
## Features
|
||||
|
||||
- **Automatic Token Management**: Handles OAuth2 token acquisition and refresh automatically
|
||||
- **Multiple Provider Support**: Configure multiple OAuth2 providers in a single configuration file
|
||||
- **Secure Credential Storage**: Keep OAuth2 credentials separate from your code
|
||||
- **Seamless Integration**: Works with existing LiteLLM provider configurations
|
||||
- **Error Handling**: Comprehensive error handling for authentication failures
|
||||
- **Retry Mechanism**: Automatic retry with exponential backoff for token acquisition
|
||||
- **Thread Safety**: Concurrent access protection for token caching
|
||||
- **Configuration Validation**: Automatic validation of URLs and scope formats
|
||||
|
||||
## Error Handling
|
||||
|
||||
The OAuth2 implementation provides specific error classes for different failure scenarios:
|
||||
|
||||
### OAuth2ConfigurationError
|
||||
Raised when there are issues with the OAuth2 configuration:
|
||||
- Invalid configuration file format
|
||||
- Missing required fields
|
||||
- Unknown OAuth2 providers
|
||||
|
||||
### OAuth2AuthenticationError
|
||||
Raised when OAuth2 authentication fails:
|
||||
- Network errors during token acquisition
|
||||
- Invalid credentials
|
||||
- Token endpoint errors
|
||||
|
||||
### OAuth2ValidationError
|
||||
Raised when configuration validation fails:
|
||||
- Invalid URL formats
|
||||
- Malformed scope values
|
||||
|
||||
### Example Error Handling
|
||||
|
||||
```python
|
||||
from crewai import LLM
|
||||
from crewai.llms.oauth2_errors import OAuth2ConfigurationError, OAuth2AuthenticationError
|
||||
|
||||
try:
|
||||
llm = LLM(
|
||||
model="my_custom_provider/my-model",
|
||||
oauth2_config_path="./litellm_config.json"
|
||||
)
|
||||
result = llm.call("test message")
|
||||
except OAuth2ConfigurationError as e:
|
||||
print(f"Configuration error: {e}")
|
||||
except OAuth2AuthenticationError as e:
|
||||
print(f"Authentication failed: {e}")
|
||||
if e.original_error:
|
||||
print(f"Original error: {e.original_error}")
|
||||
```
|
||||
|
||||
## Configuration Validation
|
||||
|
||||
The OAuth2 configuration is automatically validated:
|
||||
|
||||
- **URL Validation**: `token_url` must be a valid HTTP/HTTPS URL
|
||||
- **Scope Validation**: `scope` cannot contain empty values when split by spaces
|
||||
- **Required Fields**: `client_id`, `client_secret`, `token_url`, and `provider_name` are required
|
||||
|
||||
## Retry Mechanism
|
||||
|
||||
Token acquisition includes automatic retry with exponential backoff:
|
||||
- 3 retry attempts by default
|
||||
- Exponential backoff: 1s, 2s, 4s delays
|
||||
- Detailed logging of retry attempts
|
||||
- Thread-safe token caching
|
||||
|
||||
## Configuration
|
||||
|
||||
Create a `litellm_config.json` file in your project directory:
|
||||
|
||||
```json
|
||||
{
|
||||
"oauth2_providers": {
|
||||
"my_custom_provider": {
|
||||
"client_id": "your_client_id",
|
||||
"client_secret": "your_client_secret",
|
||||
"token_url": "https://your-provider.com/oauth/token",
|
||||
"scope": "llm.read llm.write"
|
||||
},
|
||||
"another_provider": {
|
||||
"client_id": "another_client_id",
|
||||
"client_secret": "another_client_secret",
|
||||
"token_url": "https://another-provider.com/token"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from crewai import LLM
|
||||
|
||||
# Initialize LLM with OAuth2 support
|
||||
llm = LLM(
|
||||
model="my_custom_provider/my-model",
|
||||
oauth2_config_path="./litellm_config.json" # Optional, defaults to ./litellm_config.json
|
||||
)
|
||||
|
||||
# Use in CrewAI
|
||||
from crewai import Agent, Task, Crew
|
||||
|
||||
agent = Agent(
|
||||
role="Data Analyst",
|
||||
goal="Analyze data trends",
|
||||
backstory="Expert in data analysis",
|
||||
llm=llm
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description="Analyze the latest sales data",
|
||||
agent=agent
|
||||
)
|
||||
|
||||
crew = Crew(agents=[agent], tasks=[task])
|
||||
result = crew.kickoff()
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
You can also use environment variables in your configuration:
|
||||
|
||||
```json
|
||||
{
|
||||
"oauth2_providers": {
|
||||
"my_provider": {
|
||||
"client_id": "os.environ/MY_CLIENT_ID",
|
||||
"client_secret": "os.environ/MY_CLIENT_SECRET",
|
||||
"token_url": "https://my-provider.com/token"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Supported OAuth2 Flow
|
||||
|
||||
Currently supports the **Client Credentials** OAuth2 flow, which is suitable for server-to-server authentication.
|
||||
|
||||
## Token Management
|
||||
|
||||
- Tokens are automatically cached and refreshed when they expire
|
||||
- A 60-second buffer is used before token expiration to ensure reliability
|
||||
- Failed token acquisition will raise a `RuntimeError` with details
|
||||
|
||||
## Configuration Schema
|
||||
|
||||
The `litellm_config.json` file should follow this schema:
|
||||
|
||||
```json
|
||||
{
|
||||
"oauth2_providers": {
|
||||
"<provider_name>": {
|
||||
"client_id": "string (required)",
|
||||
"client_secret": "string (required)",
|
||||
"token_url": "string (required)",
|
||||
"scope": "string (optional)",
|
||||
"refresh_token": "string (optional)"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Legacy Error Handling
|
||||
|
||||
For backward compatibility, the following error types are still supported:
|
||||
- OAuth2 authentication failures raise `OAuth2AuthenticationError` (previously `RuntimeError`)
|
||||
- Invalid configuration files raise `OAuth2ConfigurationError` (previously `ValueError`)
|
||||
- Network errors during token acquisition are wrapped in `OAuth2AuthenticationError`
|
||||
|
||||
## Examples
|
||||
|
||||
### Basic OAuth2 Provider
|
||||
|
||||
```python
|
||||
from crewai import LLM
|
||||
|
||||
llm = LLM(
|
||||
model="my_provider/gpt-4",
|
||||
oauth2_config_path="./config.json"
|
||||
)
|
||||
|
||||
response = llm.call("Hello, world!")
|
||||
print(response)
|
||||
```
|
||||
|
||||
### Multiple Providers
|
||||
|
||||
```json
|
||||
{
|
||||
"oauth2_providers": {
|
||||
"provider_a": {
|
||||
"client_id": "client_a",
|
||||
"client_secret": "secret_a",
|
||||
"token_url": "https://provider-a.com/token"
|
||||
},
|
||||
"provider_b": {
|
||||
"client_id": "client_b",
|
||||
"client_secret": "secret_b",
|
||||
"token_url": "https://provider-b.com/oauth/token",
|
||||
"scope": "read write"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
```python
|
||||
# Use different providers
|
||||
llm_a = LLM(model="provider_a/model-1", oauth2_config_path="./config.json")
|
||||
llm_b = LLM(model="provider_b/model-2", oauth2_config_path="./config.json")
|
||||
```
|
||||
64
examples/oauth2_llm_example.py
Normal file
64
examples/oauth2_llm_example.py
Normal file
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Example demonstrating OAuth2 authentication with custom LLM providers in CrewAI.
|
||||
|
||||
This example shows how to configure and use OAuth2-authenticated LLM providers.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from crewai import Agent, Task, Crew, LLM
|
||||
|
||||
def create_example_config():
|
||||
"""Create an example OAuth2 configuration file."""
|
||||
config = {
|
||||
"oauth2_providers": {
|
||||
"my_custom_provider": {
|
||||
"client_id": "your_client_id_here",
|
||||
"client_secret": "your_client_secret_here",
|
||||
"token_url": "https://your-provider.com/oauth/token",
|
||||
"scope": "llm.read llm.write"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
config_path = Path("example_oauth2_config.json")
|
||||
with open(config_path, 'w') as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
print(f"Created example config at {config_path}")
|
||||
return config_path
|
||||
|
||||
def main():
|
||||
config_path = create_example_config()
|
||||
|
||||
try:
|
||||
llm = LLM(
|
||||
model="my_custom_provider/my-model",
|
||||
oauth2_config_path=str(config_path)
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
role="Research Assistant",
|
||||
goal="Provide helpful research insights",
|
||||
backstory="An AI assistant specialized in research and analysis",
|
||||
llm=llm
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description="Research the benefits of OAuth2 authentication in AI systems",
|
||||
agent=agent,
|
||||
expected_output="A comprehensive summary of OAuth2 benefits"
|
||||
)
|
||||
|
||||
crew = Crew(agents=[agent], tasks=[task])
|
||||
|
||||
print("Running crew with OAuth2-authenticated LLM...")
|
||||
result = crew.kickoff()
|
||||
print(f"Result: {result}")
|
||||
|
||||
finally:
|
||||
if config_path.exists():
|
||||
config_path.unlink()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -52,6 +52,9 @@ import io
|
||||
from typing import TextIO
|
||||
|
||||
from crewai.llms.base_llm import BaseLLM
|
||||
from crewai.llms.oauth2_config import OAuth2ConfigLoader
|
||||
from crewai.llms.oauth2_token_manager import OAuth2TokenManager
|
||||
from crewai.llms.oauth2_errors import OAuth2ConfigurationError
|
||||
from crewai.utilities.events import crewai_event_bus
|
||||
from crewai.utilities.exceptions.context_window_exceeding_exception import (
|
||||
LLMContextLengthExceededException,
|
||||
@@ -311,6 +314,7 @@ class LLM(BaseLLM):
|
||||
callbacks: List[Any] = [],
|
||||
reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None,
|
||||
stream: bool = False,
|
||||
oauth2_config_path: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.model = model
|
||||
@@ -338,6 +342,10 @@ class LLM(BaseLLM):
|
||||
self.is_anthropic = self._is_anthropic_model(model)
|
||||
self.stream = stream
|
||||
|
||||
self.oauth2_config_loader = OAuth2ConfigLoader(oauth2_config_path)
|
||||
self.oauth2_token_manager = OAuth2TokenManager()
|
||||
self.oauth2_configs = self.oauth2_config_loader.load_config()
|
||||
|
||||
litellm.drop_params = True
|
||||
|
||||
# Normalize self.stop to always be a List[str]
|
||||
@@ -384,7 +392,21 @@ class LLM(BaseLLM):
|
||||
messages = [{"role": "user", "content": messages}]
|
||||
formatted_messages = self._format_messages_for_provider(messages)
|
||||
|
||||
# --- 2) Prepare the parameters for the completion call
|
||||
api_key = self.api_key
|
||||
provider = self._get_custom_llm_provider()
|
||||
|
||||
if provider and provider in self.oauth2_configs:
|
||||
oauth2_config = self.oauth2_configs[provider]
|
||||
try:
|
||||
self._validate_oauth2_provider(provider)
|
||||
access_token = self.oauth2_token_manager.get_access_token(oauth2_config)
|
||||
api_key = access_token
|
||||
logging.debug(f"Using OAuth2 authentication for provider {provider}")
|
||||
except Exception as e:
|
||||
logging.error(f"OAuth2 authentication failed for provider {provider}: {e}")
|
||||
raise
|
||||
|
||||
# --- 3) Prepare the parameters for the completion call
|
||||
params = {
|
||||
"model": self.model,
|
||||
"messages": formatted_messages,
|
||||
@@ -404,7 +426,7 @@ class LLM(BaseLLM):
|
||||
"api_base": self.api_base,
|
||||
"base_url": self.base_url,
|
||||
"api_version": self.api_version,
|
||||
"api_key": self.api_key,
|
||||
"api_key": api_key,
|
||||
"stream": self.stream,
|
||||
"tools": tools,
|
||||
"reasoning_effort": self.reasoning_effort,
|
||||
@@ -1076,6 +1098,12 @@ class LLM(BaseLLM):
|
||||
return self.model.split("/")[0]
|
||||
return None
|
||||
|
||||
def _validate_oauth2_provider(self, provider: str) -> bool:
|
||||
"""Validate that OAuth2 provider exists in configuration"""
|
||||
if provider not in self.oauth2_configs:
|
||||
raise OAuth2ConfigurationError(f"Unknown OAuth2 provider: {provider}")
|
||||
return True
|
||||
|
||||
def _validate_call_params(self) -> None:
|
||||
"""
|
||||
Validate parameters before making a call. Currently this only checks if
|
||||
|
||||
@@ -1 +1,17 @@
|
||||
"""LLM implementations for crewAI."""
|
||||
|
||||
from .base_llm import BaseLLM
|
||||
from .oauth2_config import OAuth2Config, OAuth2ConfigLoader
|
||||
from .oauth2_token_manager import OAuth2TokenManager
|
||||
from .oauth2_errors import OAuth2Error, OAuth2ConfigurationError, OAuth2AuthenticationError, OAuth2ValidationError
|
||||
|
||||
__all__ = [
|
||||
"BaseLLM",
|
||||
"OAuth2Config",
|
||||
"OAuth2ConfigLoader",
|
||||
"OAuth2TokenManager",
|
||||
"OAuth2Error",
|
||||
"OAuth2ConfigurationError",
|
||||
"OAuth2AuthenticationError",
|
||||
"OAuth2ValidationError"
|
||||
]
|
||||
|
||||
65
src/crewai/llms/oauth2_config.py
Normal file
65
src/crewai/llms/oauth2_config.py
Normal file
@@ -0,0 +1,65 @@
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
import json
|
||||
import re
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
from .oauth2_errors import OAuth2ConfigurationError, OAuth2ValidationError
|
||||
|
||||
|
||||
class OAuth2Config(BaseModel):
|
||||
client_id: str = Field(description="OAuth2 client ID")
|
||||
client_secret: str = Field(description="OAuth2 client secret")
|
||||
token_url: str = Field(description="OAuth2 token endpoint URL")
|
||||
scope: Optional[str] = Field(default=None, description="OAuth2 scope")
|
||||
provider_name: str = Field(description="Custom provider name")
|
||||
refresh_token: Optional[str] = Field(default=None, description="OAuth2 refresh token")
|
||||
|
||||
@field_validator('token_url')
|
||||
@classmethod
|
||||
def validate_token_url(cls, v: str) -> str:
|
||||
"""Validate that token_url is a valid HTTP/HTTPS URL."""
|
||||
url_pattern = re.compile(
|
||||
r'^https?://' # http:// or https://
|
||||
r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|' # domain...
|
||||
r'localhost|' # localhost...
|
||||
r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # ...or ip
|
||||
r'(?::\d+)?' # optional port
|
||||
r'(?:/?|[/?]\S+)$', re.IGNORECASE)
|
||||
|
||||
if not url_pattern.match(v):
|
||||
raise OAuth2ValidationError(f"Invalid token URL format: {v}")
|
||||
return v
|
||||
|
||||
@field_validator('scope')
|
||||
@classmethod
|
||||
def validate_scope(cls, v: Optional[str]) -> Optional[str]:
|
||||
"""Validate OAuth2 scope format."""
|
||||
if v:
|
||||
if ' ' in v:
|
||||
raise OAuth2ValidationError("Invalid scope format: scope cannot contain empty values")
|
||||
return v
|
||||
|
||||
|
||||
class OAuth2ConfigLoader:
|
||||
def __init__(self, config_path: Optional[str] = None):
|
||||
self.config_path = Path(config_path) if config_path else Path("litellm_config.json")
|
||||
|
||||
def load_config(self) -> Dict[str, OAuth2Config]:
|
||||
"""Load OAuth2 configurations from litellm_config.json"""
|
||||
if not self.config_path.exists():
|
||||
return {}
|
||||
|
||||
try:
|
||||
with open(self.config_path, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
oauth2_configs = {}
|
||||
for provider_name, config_data in data.get("oauth2_providers", {}).items():
|
||||
oauth2_configs[provider_name] = OAuth2Config(
|
||||
provider_name=provider_name,
|
||||
**config_data
|
||||
)
|
||||
|
||||
return oauth2_configs
|
||||
except (json.JSONDecodeError, KeyError, ValueError) as e:
|
||||
raise OAuth2ConfigurationError(f"Invalid OAuth2 configuration in {self.config_path}: {e}")
|
||||
26
src/crewai/llms/oauth2_errors.py
Normal file
26
src/crewai/llms/oauth2_errors.py
Normal file
@@ -0,0 +1,26 @@
|
||||
"""OAuth2 error classes for CrewAI."""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class OAuth2Error(Exception):
|
||||
"""Base exception class for OAuth2 operation errors."""
|
||||
|
||||
def __init__(self, message: str, original_error: Optional[Exception] = None):
|
||||
super().__init__(message)
|
||||
self.original_error = original_error
|
||||
|
||||
|
||||
class OAuth2ConfigurationError(OAuth2Error):
|
||||
"""Exception raised for OAuth2 configuration errors."""
|
||||
pass
|
||||
|
||||
|
||||
class OAuth2AuthenticationError(OAuth2Error):
|
||||
"""Exception raised for OAuth2 authentication failures."""
|
||||
pass
|
||||
|
||||
|
||||
class OAuth2ValidationError(OAuth2Error):
|
||||
"""Exception raised for OAuth2 validation errors."""
|
||||
pass
|
||||
93
src/crewai/llms/oauth2_token_manager.py
Normal file
93
src/crewai/llms/oauth2_token_manager.py
Normal file
@@ -0,0 +1,93 @@
|
||||
import time
|
||||
import logging
|
||||
import requests
|
||||
from threading import Lock
|
||||
from typing import Dict, Any
|
||||
from .oauth2_config import OAuth2Config
|
||||
from .oauth2_errors import OAuth2AuthenticationError
|
||||
|
||||
|
||||
class OAuth2TokenManager:
|
||||
def __init__(self):
|
||||
self._tokens: Dict[str, Dict[str, Any]] = {}
|
||||
self._lock = Lock()
|
||||
|
||||
def get_access_token(self, config: OAuth2Config) -> str:
|
||||
"""Get valid access token for the provider, refreshing if necessary"""
|
||||
with self._lock:
|
||||
return self._get_access_token_internal(config)
|
||||
|
||||
def _get_access_token_internal(self, config: OAuth2Config) -> str:
|
||||
"""Internal method to get access token (called within lock)"""
|
||||
provider_name = config.provider_name
|
||||
|
||||
if provider_name in self._tokens:
|
||||
token_data = self._tokens[provider_name]
|
||||
if self._is_token_valid(token_data):
|
||||
logging.debug(f"Using cached OAuth2 token for provider {provider_name}")
|
||||
return token_data["access_token"]
|
||||
|
||||
logging.info(f"Acquiring new OAuth2 token for provider {provider_name}")
|
||||
return self._acquire_new_token(config)
|
||||
|
||||
def _is_token_valid(self, token_data: Dict[str, Any]) -> bool:
|
||||
"""Check if token is still valid (not expired)"""
|
||||
if "expires_at" not in token_data:
|
||||
return False
|
||||
|
||||
return time.time() < (token_data["expires_at"] - 60)
|
||||
|
||||
def _acquire_new_token(self, config: OAuth2Config, retry_count: int = 3) -> str:
|
||||
"""Acquire new access token using client credentials flow with retry logic"""
|
||||
for attempt in range(retry_count):
|
||||
try:
|
||||
return self._perform_token_request(config)
|
||||
except requests.RequestException as e:
|
||||
if attempt == retry_count - 1:
|
||||
raise OAuth2AuthenticationError(
|
||||
f"Failed to acquire OAuth2 token for {config.provider_name} after {retry_count} attempts: {e}",
|
||||
original_error=e
|
||||
)
|
||||
wait_time = 2 ** attempt
|
||||
logging.warning(f"OAuth2 token request failed for {config.provider_name}, retrying in {wait_time}s (attempt {attempt + 1}/{retry_count}): {e}")
|
||||
time.sleep(wait_time)
|
||||
|
||||
raise OAuth2AuthenticationError(f"Unexpected error in token acquisition for {config.provider_name}")
|
||||
|
||||
def _perform_token_request(self, config: OAuth2Config) -> str:
|
||||
"""Perform the actual token request"""
|
||||
payload = {
|
||||
"grant_type": "client_credentials",
|
||||
"client_id": config.client_id,
|
||||
"client_secret": config.client_secret,
|
||||
}
|
||||
|
||||
if config.scope:
|
||||
payload["scope"] = config.scope
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
config.token_url,
|
||||
data=payload,
|
||||
timeout=30,
|
||||
headers={"Content-Type": "application/x-www-form-urlencoded"}
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
token_data = response.json()
|
||||
access_token = token_data["access_token"]
|
||||
|
||||
expires_in = token_data.get("expires_in", 3600)
|
||||
self._tokens[config.provider_name] = {
|
||||
"access_token": access_token,
|
||||
"expires_at": time.time() + expires_in,
|
||||
"token_type": token_data.get("token_type", "Bearer")
|
||||
}
|
||||
|
||||
logging.info(f"Successfully acquired OAuth2 token for {config.provider_name}, expires in {expires_in}s")
|
||||
return access_token
|
||||
|
||||
except requests.RequestException:
|
||||
raise
|
||||
except KeyError as e:
|
||||
raise OAuth2AuthenticationError(f"Invalid token response from {config.provider_name}: missing {e}")
|
||||
420
tests/test_oauth2_llm.py
Normal file
420
tests/test_oauth2_llm.py
Normal file
@@ -0,0 +1,420 @@
|
||||
import json
|
||||
import pytest
|
||||
import requests
|
||||
import tempfile
|
||||
import time
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock, patch
|
||||
from crewai.llm import LLM
|
||||
from crewai.llms.oauth2_config import OAuth2Config, OAuth2ConfigLoader
|
||||
from crewai.llms.oauth2_token_manager import OAuth2TokenManager
|
||||
from crewai.llms.oauth2_errors import OAuth2Error, OAuth2ConfigurationError, OAuth2AuthenticationError, OAuth2ValidationError
|
||||
|
||||
|
||||
class TestOAuth2Config:
|
||||
def test_oauth2_config_creation(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider",
|
||||
scope="read write"
|
||||
)
|
||||
assert config.client_id == "test_client"
|
||||
assert config.provider_name == "test_provider"
|
||||
assert config.scope == "read write"
|
||||
|
||||
def test_oauth2_config_loader_missing_file(self):
|
||||
loader = OAuth2ConfigLoader("nonexistent.json")
|
||||
configs = loader.load_config()
|
||||
assert configs == {}
|
||||
|
||||
def test_oauth2_config_loader_valid_config(self):
|
||||
config_data = {
|
||||
"oauth2_providers": {
|
||||
"custom_provider": {
|
||||
"client_id": "test_client",
|
||||
"client_secret": "test_secret",
|
||||
"token_url": "https://example.com/token",
|
||||
"scope": "read"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
|
||||
json.dump(config_data, f)
|
||||
config_path = f.name
|
||||
|
||||
try:
|
||||
loader = OAuth2ConfigLoader(config_path)
|
||||
configs = loader.load_config()
|
||||
|
||||
assert "custom_provider" in configs
|
||||
config = configs["custom_provider"]
|
||||
assert config.client_id == "test_client"
|
||||
assert config.provider_name == "custom_provider"
|
||||
finally:
|
||||
Path(config_path).unlink()
|
||||
|
||||
def test_oauth2_config_loader_invalid_json(self):
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
|
||||
f.write("invalid json")
|
||||
config_path = f.name
|
||||
|
||||
try:
|
||||
loader = OAuth2ConfigLoader(config_path)
|
||||
with pytest.raises(OAuth2ConfigurationError, match="Invalid OAuth2 configuration"):
|
||||
loader.load_config()
|
||||
finally:
|
||||
Path(config_path).unlink()
|
||||
|
||||
|
||||
class TestOAuth2TokenManager:
|
||||
def test_token_acquisition_success(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = {
|
||||
"access_token": "test_token_123",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600
|
||||
}
|
||||
mock_response.raise_for_status.return_value = None
|
||||
|
||||
with patch('requests.post', return_value=mock_response) as mock_post:
|
||||
manager = OAuth2TokenManager()
|
||||
token = manager.get_access_token(config)
|
||||
|
||||
assert token == "test_token_123"
|
||||
mock_post.assert_called_once()
|
||||
|
||||
def test_token_caching(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = {
|
||||
"access_token": "test_token_123",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600
|
||||
}
|
||||
mock_response.raise_for_status.return_value = None
|
||||
|
||||
with patch('requests.post', return_value=mock_response) as mock_post:
|
||||
manager = OAuth2TokenManager()
|
||||
|
||||
token1 = manager.get_access_token(config)
|
||||
assert token1 == "test_token_123"
|
||||
assert mock_post.call_count == 1
|
||||
|
||||
token2 = manager.get_access_token(config)
|
||||
assert token2 == "test_token_123"
|
||||
assert mock_post.call_count == 1
|
||||
|
||||
def test_token_refresh_on_expiry(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = {
|
||||
"access_token": "new_token_456",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600
|
||||
}
|
||||
mock_response.raise_for_status.return_value = None
|
||||
|
||||
with patch('requests.post', return_value=mock_response):
|
||||
manager = OAuth2TokenManager()
|
||||
|
||||
manager._tokens["test_provider"] = {
|
||||
"access_token": "old_token",
|
||||
"expires_at": time.time() - 100
|
||||
}
|
||||
|
||||
token = manager.get_access_token(config)
|
||||
assert token == "new_token_456"
|
||||
|
||||
def test_token_acquisition_failure(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
with patch('requests.post', side_effect=requests.RequestException("Network error")):
|
||||
with patch('time.sleep'):
|
||||
manager = OAuth2TokenManager()
|
||||
|
||||
with pytest.raises(OAuth2AuthenticationError, match="Failed to acquire OAuth2 token"):
|
||||
manager.get_access_token(config)
|
||||
|
||||
def test_invalid_token_response(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = {"invalid": "response"}
|
||||
mock_response.raise_for_status.return_value = None
|
||||
|
||||
with patch('requests.post', return_value=mock_response):
|
||||
manager = OAuth2TokenManager()
|
||||
|
||||
with pytest.raises(OAuth2AuthenticationError, match="Invalid token response"):
|
||||
manager.get_access_token(config)
|
||||
|
||||
|
||||
class TestLLMOAuth2Integration:
|
||||
def test_llm_with_oauth2_config(self):
|
||||
config_data = {
|
||||
"oauth2_providers": {
|
||||
"custom": {
|
||||
"client_id": "test_client",
|
||||
"client_secret": "test_secret",
|
||||
"token_url": "https://example.com/token"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
|
||||
json.dump(config_data, f)
|
||||
config_path = f.name
|
||||
|
||||
try:
|
||||
llm = LLM(
|
||||
model="custom/test-model",
|
||||
oauth2_config_path=config_path
|
||||
)
|
||||
|
||||
assert "custom" in llm.oauth2_configs
|
||||
assert llm.oauth2_configs["custom"].client_id == "test_client"
|
||||
finally:
|
||||
Path(config_path).unlink()
|
||||
|
||||
def test_llm_without_oauth2_config(self):
|
||||
llm = LLM(model="openai/gpt-3.5-turbo")
|
||||
assert llm.oauth2_configs == {}
|
||||
|
||||
@patch('crewai.llm.litellm.completion')
|
||||
def test_llm_oauth2_token_injection(self, mock_completion):
|
||||
config_data = {
|
||||
"oauth2_providers": {
|
||||
"custom": {
|
||||
"client_id": "test_client",
|
||||
"client_secret": "test_secret",
|
||||
"token_url": "https://example.com/token"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mock_completion.return_value = Mock(choices=[Mock(message=Mock(content="test response"))])
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
|
||||
json.dump(config_data, f)
|
||||
config_path = f.name
|
||||
|
||||
try:
|
||||
with patch.object(OAuth2TokenManager, 'get_access_token', return_value="oauth_token_123"):
|
||||
llm = LLM(
|
||||
model="custom/test-model",
|
||||
oauth2_config_path=config_path
|
||||
)
|
||||
|
||||
llm.call("test message")
|
||||
|
||||
call_args = mock_completion.call_args
|
||||
assert call_args[1]['api_key'] == "oauth_token_123"
|
||||
finally:
|
||||
Path(config_path).unlink()
|
||||
|
||||
@patch('crewai.llm.litellm.completion')
|
||||
def test_llm_oauth2_authentication_failure(self, mock_completion):
|
||||
config_data = {
|
||||
"oauth2_providers": {
|
||||
"custom": {
|
||||
"client_id": "test_client",
|
||||
"client_secret": "test_secret",
|
||||
"token_url": "https://example.com/token"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
|
||||
json.dump(config_data, f)
|
||||
config_path = f.name
|
||||
|
||||
try:
|
||||
with patch.object(OAuth2TokenManager, 'get_access_token', side_effect=OAuth2AuthenticationError("Auth failed")):
|
||||
llm = LLM(
|
||||
model="custom/test-model",
|
||||
oauth2_config_path=config_path
|
||||
)
|
||||
|
||||
with pytest.raises(OAuth2AuthenticationError, match="Auth failed"):
|
||||
llm.call("test message")
|
||||
finally:
|
||||
Path(config_path).unlink()
|
||||
|
||||
@patch('crewai.llm.litellm.completion')
|
||||
def test_llm_non_oauth2_provider_unchanged(self, mock_completion):
|
||||
mock_completion.return_value = Mock(choices=[Mock(message=Mock(content="test response"))])
|
||||
|
||||
llm = LLM(
|
||||
model="openai/gpt-3.5-turbo",
|
||||
api_key="original_key"
|
||||
)
|
||||
|
||||
llm.call("test message")
|
||||
|
||||
call_args = mock_completion.call_args
|
||||
assert call_args[1]['api_key'] == "original_key"
|
||||
|
||||
|
||||
class TestOAuth2Validation:
|
||||
def test_valid_url_validation(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
assert config.token_url == "https://example.com/token"
|
||||
|
||||
def test_invalid_url_validation(self):
|
||||
with pytest.raises(OAuth2ValidationError, match="Invalid token URL format"):
|
||||
OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="not-a-url",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
def test_valid_scope_validation(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider",
|
||||
scope="read write"
|
||||
)
|
||||
assert config.scope == "read write"
|
||||
|
||||
def test_invalid_scope_validation(self):
|
||||
with pytest.raises(OAuth2ValidationError, match="Invalid scope format"):
|
||||
OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider",
|
||||
scope="read write"
|
||||
)
|
||||
|
||||
|
||||
class TestOAuth2RetryMechanism:
|
||||
def test_retry_mechanism_success_on_second_attempt(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
responses = [
|
||||
requests.RequestException("Network error"),
|
||||
Mock(json=lambda: {"access_token": "token_123", "expires_in": 3600}, raise_for_status=lambda: None)
|
||||
]
|
||||
|
||||
with patch('requests.post', side_effect=responses):
|
||||
with patch('time.sleep'):
|
||||
manager = OAuth2TokenManager()
|
||||
token = manager.get_access_token(config)
|
||||
assert token == "token_123"
|
||||
|
||||
def test_retry_mechanism_exhausted(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
with patch('requests.post', side_effect=requests.RequestException("Network error")):
|
||||
with patch('time.sleep'):
|
||||
manager = OAuth2TokenManager()
|
||||
|
||||
with pytest.raises(OAuth2AuthenticationError, match="Failed to acquire OAuth2 token.*after 3 attempts"):
|
||||
manager.get_access_token(config)
|
||||
|
||||
|
||||
class TestOAuth2ThreadSafety:
|
||||
def test_concurrent_token_requests(self):
|
||||
config = OAuth2Config(
|
||||
client_id="test_client",
|
||||
client_secret="test_secret",
|
||||
token_url="https://example.com/token",
|
||||
provider_name="test_provider"
|
||||
)
|
||||
|
||||
mock_response = Mock()
|
||||
mock_response.json.return_value = {"access_token": "token_123", "expires_in": 3600}
|
||||
mock_response.raise_for_status.return_value = None
|
||||
|
||||
call_count = 0
|
||||
def mock_post(*args, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
time.sleep(0.1)
|
||||
return mock_response
|
||||
|
||||
with patch('requests.post', side_effect=mock_post):
|
||||
manager = OAuth2TokenManager()
|
||||
|
||||
results = []
|
||||
def get_token():
|
||||
token = manager.get_access_token(config)
|
||||
results.append(token)
|
||||
|
||||
threads = [threading.Thread(target=get_token) for _ in range(5)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
assert all(token == "token_123" for token in results)
|
||||
assert call_count == 1
|
||||
|
||||
|
||||
class TestOAuth2ErrorClasses:
|
||||
def test_oauth2_configuration_error(self):
|
||||
error = OAuth2ConfigurationError("Config error")
|
||||
assert str(error) == "Config error"
|
||||
assert isinstance(error, OAuth2Error)
|
||||
|
||||
def test_oauth2_authentication_error_with_original(self):
|
||||
original = requests.RequestException("Network error")
|
||||
error = OAuth2AuthenticationError("Auth failed", original_error=original)
|
||||
assert str(error) == "Auth failed"
|
||||
assert error.original_error == original
|
||||
|
||||
def test_oauth2_validation_error(self):
|
||||
error = OAuth2ValidationError("Validation failed")
|
||||
assert str(error) == "Validation failed"
|
||||
assert isinstance(error, OAuth2Error)
|
||||
Reference in New Issue
Block a user