scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

(Feat) Add BitBucket Integration for Prompt Management (#14882)

Browse source 
* Add bitbucket integration for prompt management

* remove not needed file

* remove not needed file

* Add the correct readme

* Add the correct readme

* fix test for bitbucket

* fix test for bitbucket
This commit is contained in:
Sameer Kankute 2025-09-26 04:21:00 +05:30 • committed by GitHub
parent dcbccd1fea
commit 1681bf7175
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 2214 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
docs/my-website/docs/proxy/native_litellm_prompt.md
View file

@ -1,11 +1,19 @@
import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
import '@theme/IdealImage'
import '@theme/TabItem'
import '@theme/Tabs'
import Image
import TabItem
import Tabs
# LiteLLM Prompt Management (GitOps)
Store prompts as `.prompt` files in your repository and use them directly with LiteLLM. No external services required.
## Supported Integrations
- **File System**: Store `.prompt` files locally
- **BitBucket**: Store `.prompt` files in BitBucket repositories with team-based access control
## Quick Start
<Tabs>
@ -41,6 +49,50 @@ response = litellm.completion(
)
```
</TabItem>
<TabItem value="bitbucket" label="BITBUCKET">
**1. Create a .prompt file in BitBucket**
Create `prompts/hello.prompt` in your BitBucket repository:
```yaml
---
model: gpt-4
temperature: 0.7
---
System: You are a helpful assistant.
User: {{user_message}}
```
**2. Configure BitBucket access**
```python
import litellm
# Configure BitBucket access
bitbucket_config = {
"workspace": "your-workspace",
"repository": "your-repo",
"access_token": "your-access-token",
"branch": "main"
}
# Set global BitBucket configuration
litellm.set_global_bitbucket_config(bitbucket_config)
```
**3. Use with LiteLLM**
```python
response = litellm.completion(
model="bitbucket/gpt-4",
prompt_id="hello",
prompt_variables={"user_message": "What is the capital of France?"}
)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
@ -70,6 +122,12 @@ model_list:
litellm_settings:
global_prompt_directory: "./prompts"
# Or use BitBucket for team-based prompt management
global_bitbucket_config:
workspace: "your-workspace"
repository: "your-repo"
access_token: "your-access-token"
branch: "main"
```
**3. Start the proxy**
@ -142,21 +200,43 @@ User: {{user_message}}
### API Reference
For dotprompt integration, use these parameters:
For prompt integrations, use these parameters:
**File System (dotprompt):**
```
model: dotprompt/<base_model> # required (e.g., dotprompt/gpt-4)
prompt_id: str # required - the .prompt filename without extension
prompt_variables: Optional[dict] # optional - variables for template rendering
```
**Example API call:**
**BitBucket:**
```
model: bitbucket/<base_model> # required (e.g., bitbucket/gpt-4)
prompt_id: str # required - the .prompt filename without extension
prompt_variables: Optional[dict] # optional - variables for template rendering
bitbucket_config: Optional[dict] # optional - BitBucket configuration (if not set globally)
```
**Example API calls:**
```python
# File system integration
response = litellm.completion(
model="dotprompt/gpt-4",
prompt_id="hello",
prompt_variables={"user_message": "Hello world"},
messages=[{"role": "user", "content": "This will be ignored"}]
)
# BitBucket integration
response = litellm.completion(
model="bitbucket/gpt-4",
prompt_id="hello",
prompt_variables={"user_message": "Hello world"},
bitbucket_config={
"workspace": "your-workspace",
"repository": "your-repo",
"access_token": "your-token"
}
)
```

9
litellm/__init__.py
View file

@ -151,6 +151,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"aws_sqs",
"vector_store_pre_call_hook",
"dotprompt",
"bitbucket",
"cloudzero",
"posthog",
]
@ -1352,3 +1353,11 @@ from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_k
### PASSTHROUGH ###
from .passthrough import allm_passthrough_route, llm_passthrough_route
### GLOBAL CONFIG ###
global_bitbucket_config: Optional[Dict[str, Any]] = None
def set_global_bitbucket_config(config: Dict[str, Any]) -> None:
"""Set global BitBucket configuration for prompt management."""
global global_bitbucket_config
global_bitbucket_config = config

317
litellm/integrations/bitbucket/README.md Normal file
View file

@ -0,0 +1,317 @@
# LiteLLM BitBucket Prompt Management
A powerful prompt management system for LiteLLM that fetches `.prompt` files from BitBucket repositories. This enables team-based prompt management with BitBucket's built-in access control and version control capabilities.
## Features
- **🏢 Team-based access control**: Leverage BitBucket's workspace and repository permissions
- **📁 Repository-based prompt storage**: Store prompts in BitBucket repositories
- **🔐 Multiple authentication methods**: Support for access tokens and basic auth
- **🎯 YAML frontmatter**: Define model, parameters, and schemas in file headers
- **🔧 Handlebars templating**: Use `{{variable}}` syntax with Jinja2 backend
- **✅ Input validation**: Automatic validation against defined schemas
- **🔗 LiteLLM integration**: Works seamlessly with `litellm.completion()`
- **💬 Smart message parsing**: Converts prompts to proper chat messages
- **⚙️ Parameter extraction**: Automatically applies model settings from prompts
## Quick Start
### 1. Set up BitBucket Repository
Create a repository in your BitBucket workspace and add `.prompt` files:
```
your-repo/
├── prompts/
│ ├── chat_assistant.prompt
│ ├── code_reviewer.prompt
│ └── data_analyst.prompt
```
### 2. Create a `.prompt` file
Create a file called `prompts/chat_assistant.prompt`:
```yaml
---
model: gpt-4
temperature: 0.7
max_tokens: 150
input:
schema:
user_message: string
system_context?: string
---
{% if system_context %}System: {{system_context}}
{% endif %}User: {{user_message}}
```
### 3. Configure BitBucket Access
#### Option A: Access Token (Recommended)
```python
import litellm
# Configure BitBucket access
bitbucket_config = {
"workspace": "your-workspace",
"repository": "your-repo",
"access_token": "your-access-token",
"branch": "main" # optional, defaults to main
}
# Set global BitBucket configuration
litellm.set_global_bitbucket_config(bitbucket_config)
```
#### Option B: Basic Authentication
```python
import litellm
# Configure BitBucket access with basic auth
bitbucket_config = {
"workspace": "your-workspace",
"repository": "your-repo",
"username": "your-username",
"access_token": "your-app-password", # Use app password for basic auth
"auth_method": "basic",
"branch": "main"
}
litellm.set_global_bitbucket_config(bitbucket_config)
```
### 4. Use with LiteLLM
```python
# Use with completion - the model prefix 'bitbucket/' tells LiteLLM to use BitBucket prompt management
response = litellm.completion(
model="bitbucket/gpt-4", # The actual model comes from the .prompt file
prompt_id="prompts/chat_assistant", # Location of the prompt file
prompt_variables={
"user_message": "What is machine learning?",
"system_context": "You are a helpful AI tutor."
},
# Any additional messages will be appended after the prompt
messages=[{"role": "user", "content": "Please explain it simply."}]
)
print(response.choices[0].message.content)
```
## Proxy Server Configuration
### 1. Create a `.prompt` file
Create `prompts/hello.prompt`:
```yaml
---
model: gpt-4
temperature: 0.7
---
System: You are a helpful assistant.
User: {{user_message}}
```
### 2. Setup config.yaml
```yaml
model_list:
- model_name: my-bitbucket-model
litellm_params:
model: bitbucket/gpt-4
prompt_id: "prompts/hello"
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
global_bitbucket_config:
workspace: "your-workspace"
repository: "your-repo"
access_token: "your-access-token"
branch: "main"
```
### 3. Start the proxy
```bash
litellm --config config.yaml --detailed_debug
```
### 4. Test it!
```bash
curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "my-bitbucket-model",
"messages": [{"role": "user", "content": "IGNORED"}],
"prompt_variables": {
"user_message": "What is the capital of France?"
}
}'
```
## Prompt File Format
### Basic Structure
```yaml
---
# Model configuration
model: gpt-4
temperature: 0.7
max_tokens: 500
# Input schema (optional)
input:
schema:
user_message: string
system_context?: string
---
System: You are a helpful {{role}} assistant.
User: {{user_message}}
```
### Advanced Features
**Multi-role conversations:**
```yaml
---
model: gpt-4
temperature: 0.3
---
System: You are a helpful coding assistant.
User: {{user_question}}
```
**Dynamic model selection:**
```yaml
---
model: "{{preferred_model}}" # Model can be a variable
temperature: 0.7
---
System: You are a helpful assistant specialized in {{domain}}.
User: {{user_message}}
```
## Team-Based Access Control
BitBucket's built-in permission system provides team-based access control:
1. **Workspace-level permissions**: Control access to entire workspaces
2. **Repository-level permissions**: Control access to specific repositories
3. **Branch-level permissions**: Control access to specific branches
4. **User and group management**: Manage team members and their access levels
### Setting up Team Access
1. **Create workspaces for each team**:
```
team-a-prompts/
team-b-prompts/
team-c-prompts/
```
2. **Configure repository permissions**:
- Grant read access to team members
- Grant write access to prompt maintainers
- Use branch protection rules for production prompts
3. **Use different access tokens**:
- Each team can have their own access token
- Tokens can be scoped to specific repositories
- Use app passwords for additional security
## API Reference
### BitBucket Configuration
```python
bitbucket_config = {
"workspace": str, # Required: BitBucket workspace name
"repository": str, # Required: Repository name
"access_token": str, # Required: BitBucket access token or app password
"branch": str, # Optional: Branch to fetch from (default: "main")
"base_url": str, # Optional: Custom BitBucket API URL
"auth_method": str, # Optional: "token" or "basic" (default: "token")
"username": str, # Optional: Username for basic auth
"base_url" : str # Optional: Incase where the base url is not https://api.bitbucket.org/2.0
}
```
### LiteLLM Integration
```python
response = litellm.completion(
model="bitbucket/<base_model>", # required (e.g., bitbucket/gpt-4)
prompt_id=str, # required - the .prompt filename without extension
prompt_variables=dict, # optional - variables for template rendering
bitbucket_config=dict, # optional - BitBucket configuration (if not set globally)
messages=list, # optional - additional messages
)
```
## Error Handling
The BitBucket integration provides detailed error messages for common issues:
- **Authentication errors**: Invalid access tokens or credentials
- **Permission errors**: Insufficient access to workspace/repository
- **File not found**: Missing .prompt files
- **Network errors**: Connection issues with BitBucket API
## Security Considerations
1. **Access Token Security**: Store access tokens securely using environment variables or secret management systems
2. **Repository Permissions**: Use BitBucket's permission system to control access
3. **Branch Protection**: Protect main branches from unauthorized changes
4. **Audit Logging**: BitBucket provides audit logs for all repository access
## Troubleshooting
### Common Issues
1. **"Access denied" errors**: Check your BitBucket permissions for the workspace and repository
2. **"Authentication failed" errors**: Verify your access token or credentials
3. **"File not found" errors**: Ensure the .prompt file exists in the specified branch
4. **Template rendering errors**: Check your Handlebars syntax in the .prompt file
### Debug Mode
Enable debug logging to troubleshoot issues:
```python
import litellm
litellm.set_verbose = True
# Your BitBucket prompt calls will now show detailed logs
response = litellm.completion(
model="bitbucket/gpt-4",
prompt_id="your_prompt",
prompt_variables={"key": "value"}
)
```
## Migration from File-Based Prompts
If you're currently using file-based prompts with the dotprompt integration, you can easily migrate to BitBucket:
1. **Upload your .prompt files** to a BitBucket repository
2. **Update your configuration** to use BitBucket instead of local files
3. **Set up team access** using BitBucket's permission system
4. **Update your code** to use `bitbucket/` model prefix instead of `dotprompt/`
This provides better collaboration, version control, and team-based access control for your prompts.

66
litellm/integrations/bitbucket/__init__.py Normal file
View file

@ -0,0 +1,66 @@
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from .bitbucket_prompt_manager import BitBucketPromptManager
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.types.prompts.init_prompts import SupportedPromptIntegrations
from .bitbucket_prompt_manager import BitBucketPromptManager
# Global instances
global_bitbucket_config: Optional[dict] = None
def set_global_bitbucket_config(config: dict) -> None:
"""
Set the global BitBucket configuration for prompt management.
Args:
config: Dictionary containing BitBucket configuration
- workspace: BitBucket workspace name
- repository: Repository name
- access_token: BitBucket access token
- branch: Branch to fetch prompts from (default: main)
"""
import litellm
litellm.global_bitbucket_config = config # type: ignore
def prompt_initializer(
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
) -> "CustomPromptManagement":
"""
Initialize a prompt from a BitBucket repository.
"""
bitbucket_config = getattr(litellm_params, "bitbucket_config", None)
prompt_id = getattr(litellm_params, "prompt_id", None)
if not bitbucket_config:
raise ValueError(
"bitbucket_config is required for BitBucket prompt integration"
)
try:
bitbucket_prompt_manager = BitBucketPromptManager(
bitbucket_config=bitbucket_config,
prompt_id=prompt_id,
)
return bitbucket_prompt_manager
except Exception as e:
raise e
prompt_initializer_registry = {
SupportedPromptIntegrations.BITBUCKET.value: prompt_initializer,
}
# Export public API
__all__ = [
"BitBucketPromptManager",
"set_global_bitbucket_config",
"global_bitbucket_config",
]

241
litellm/integrations/bitbucket/bitbucket_client.py Normal file
View file

@ -0,0 +1,241 @@
"""
BitBucket API client for fetching .prompt files from BitBucket repositories.
"""
import base64
from typing import Any, Dict, List, Optional
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class BitBucketClient:
"""
Client for interacting with BitBucket API to fetch .prompt files.
Supports:
- Authentication with access tokens
- Fetching file contents from repositories
- Team-based access control through BitBucket permissions
- Branch-specific file fetching
"""
def __init__(self, config: Dict[str, Any]):
"""
Initialize the BitBucket client.
Args:
config: Dictionary containing:
- workspace: BitBucket workspace name
- repository: Repository name
- access_token: BitBucket access token (or app password)
- branch: Branch to fetch from (default: main)
- base_url: Custom BitBucket API base URL (optional)
- auth_method: Authentication method ('token' or 'basic', default: 'token')
- username: Username for basic auth (optional)
"""
self.workspace = config.get("workspace")
self.repository = config.get("repository")
self.access_token = config.get("access_token")
self.branch = config.get("branch", "main")
self.base_url = config.get("", "https://api.bitbucket.org/2.0")
self.auth_method = config.get("auth_method", "token")
self.username = config.get("username")
if not all([self.workspace, self.repository, self.access_token]):
raise ValueError("workspace, repository, and access_token are required")
# Set up authentication headers
self.headers = {
"Accept": "application/json",
"Content-Type": "application/json",
}
if self.auth_method == "basic" and self.username:
# Use basic auth with username and app password
credentials = f"{self.username}:{self.access_token}"
encoded_credentials = base64.b64encode(credentials.encode()).decode()
self.headers["Authorization"] = f"Basic {encoded_credentials}"
else:
# Use token-based authentication (default)
self.headers["Authorization"] = f"Bearer {self.access_token}"
# Initialize HTTPHandler
self.http_handler = HTTPHandler()
def get_file_content(self, file_path: str) -> Optional[str]:
"""
Fetch the content of a file from the BitBucket repository.
Args:
file_path: Path to the file in the repository
Returns:
File content as string, or None if file not found
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
try:
response = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
# BitBucket returns file content as base64 encoded
if response.headers.get("content-type", "").startswith("text/"):
return response.text
else:
# For binary files or when content-type is not text, try to decode as base64
try:
return base64.b64decode(response.content).decode("utf-8")
except Exception:
return response.text
except Exception as e:
# Check if it's an HTTP error
if hasattr(e, "response") and hasattr(e.response, "status_code"):
if e.response.status_code == 404:
return None
elif e.response.status_code == 403:
raise Exception(
f"Access denied to file '{file_path}'. Check your BitBucket permissions for workspace '{self.workspace}' and repository '{self.repository}'."
)
elif e.response.status_code == 401:
raise Exception(
"Authentication failed. Check your BitBucket access token and permissions."
)
else:
raise Exception(f"Failed to fetch file '{file_path}': {e}")
else:
raise Exception(f"Error fetching file '{file_path}': {e}")
def list_files(
self, directory_path: str = "", file_extension: str = ".prompt"
) -> List[str]:
"""
List files in a directory with a specific extension.
Args:
directory_path: Directory path in the repository (empty for root)
file_extension: File extension to filter by (default: .prompt)
Returns:
List of file paths
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{directory_path}"
try:
response = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data = response.json()
files = []
for item in data.get("values", []):
if item.get("type") == "commit_file":
file_path = item.get("path", "")
if file_path.endswith(file_extension):
files.append(file_path)
return files
except Exception as e:
# Check if it's an HTTP error
if hasattr(e, "response") and hasattr(e.response, "status_code"):
if e.response.status_code == 404:
return []
elif e.response.status_code == 403:
raise Exception(
f"Access denied to directory '{directory_path}'. Check your BitBucket permissions for workspace '{self.workspace}' and repository '{self.repository}'."
)
elif e.response.status_code == 401:
raise Exception(
"Authentication failed. Check your BitBucket access token and permissions."
)
else:
raise Exception(f"Failed to list files in '{directory_path}': {e}")
else:
raise Exception(f"Error listing files in '{directory_path}': {e}")
def get_repository_info(self) -> Dict[str, Any]:
"""
Get information about the repository.
Returns:
Dictionary containing repository information
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}"
try:
response = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
return response.json()
except Exception as e:
raise Exception(f"Failed to get repository info: {e}")
def test_connection(self) -> bool:
"""
Test the connection to the BitBucket repository.
Returns:
True if connection is successful, False otherwise
"""
try:
self.get_repository_info()
return True
except Exception:
return False
def get_branches(self) -> List[Dict[str, Any]]:
"""
Get list of branches in the repository.
Returns:
List of branch information dictionaries
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/refs/branches"
try:
response = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data = response.json()
return data.get("values", [])
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str) -> Optional[Dict[str, Any]]:
"""
Get metadata about a file (size, last modified, etc.).
Args:
file_path: Path to the file in the repository
Returns:
Dictionary containing file metadata, or None if file not found
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
try:
# Use GET with Range header to get just the headers (HEAD equivalent)
headers = self.headers.copy()
headers["Range"] = "bytes=0-0" # Request only first byte to get headers
response = self.http_handler.get(url, headers=headers)
response.raise_for_status()
return {
"content_type": response.headers.get("content-type"),
"content_length": response.headers.get("content-length"),
"last_modified": response.headers.get("last-modified"),
}
except Exception as e:
# Check if it's an HTTP error
if hasattr(e, "response") and hasattr(e.response, "status_code"):
if e.response.status_code == 404:
return None
raise Exception(f"Failed to get file metadata for '{file_path}': {e}")
else:
raise Exception(f"Error getting file metadata for '{file_path}': {e}")
def close(self):
"""Close the HTTP handler to free resources."""
if hasattr(self, "http_handler"):
self.http_handler.close()

508
litellm/integrations/bitbucket/bitbucket_prompt_manager.py Normal file
View file

@ -0,0 +1,508 @@
"""
BitBucket prompt manager that integrates with LiteLLM's prompt management system.
Fetches .prompt files from BitBucket repositories and provides team-based access control.
"""
from typing import Any, Dict, List, Optional, Tuple, Union
from jinja2 import DictLoader, Environment, select_autoescape
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.integrations.prompt_management_base import (
PromptManagementBase,
PromptManagementClient,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import StandardCallbackDynamicParams
from .bitbucket_client import BitBucketClient
class BitBucketPromptTemplate:
"""
Represents a prompt template loaded from BitBucket.
"""
def __init__(
self,
template_id: str,
content: str,
metadata: Dict[str, Any],
model: Optional[str] = None,
):
self.template_id = template_id
self.content = content
self.metadata = metadata
self.model = model or metadata.get("model")
self.temperature = metadata.get("temperature")
self.max_tokens = metadata.get("max_tokens")
self.input_schema = metadata.get("input", {}).get("schema", {})
self.optional_params = {
k: v for k, v in metadata.items() if k not in ["model", "input", "content"]
}
def __repr__(self):
return f"BitBucketPromptTemplate(id='{self.template_id}', model='{self.model}')"
class BitBucketTemplateManager:
"""
Manager for loading and rendering .prompt files from BitBucket repositories.
Supports:
- Fetching .prompt files from BitBucket repositories
- Team-based access control through BitBucket permissions
- YAML frontmatter for metadata
- Handlebars-style templating (using Jinja2)
- Input/output schema validation
- Model configuration
"""
def __init__(
self,
bitbucket_config: Dict[str, Any],
prompt_id: Optional[str] = None,
):
self.bitbucket_config = bitbucket_config
self.prompt_id = prompt_id
self.prompts: Dict[str, BitBucketPromptTemplate] = {}
self.bitbucket_client = BitBucketClient(bitbucket_config)
self.jinja_env = Environment(
loader=DictLoader({}),
autoescape=select_autoescape(["html", "xml"]),
# Use Handlebars-style delimiters to match Dotprompt spec
variable_start_string="{{",
variable_end_string="}}",
block_start_string="{%",
block_end_string="%}",
comment_start_string="{#",
comment_end_string="#}",
)
# Load prompts from BitBucket if prompt_id is provided
if self.prompt_id:
self._load_prompt_from_bitbucket(self.prompt_id)
def _load_prompt_from_bitbucket(self, prompt_id: str) -> None:
"""Load a specific .prompt file from BitBucket."""
try:
# Fetch the .prompt file from BitBucket
prompt_content = self.bitbucket_client.get_file_content(
f"{prompt_id}.prompt"
)
if prompt_content:
template = self._parse_prompt_file(prompt_content, prompt_id)
self.prompts[prompt_id] = template
except Exception as e:
raise Exception(f"Failed to load prompt '{prompt_id}' from BitBucket: {e}")
def _parse_prompt_file(
self, content: str, prompt_id: str
) -> BitBucketPromptTemplate:
"""Parse a .prompt file content and extract metadata and template."""
# Split frontmatter and content
if content.startswith("---"):
parts = content.split("---", 2)
if len(parts) >= 3:
frontmatter_str = parts[1].strip()
template_content = parts[2].strip()
else:
frontmatter_str = ""
template_content = content
else:
frontmatter_str = ""
template_content = content
# Parse YAML frontmatter
metadata = {}
if frontmatter_str:
try:
import yaml
metadata = yaml.safe_load(frontmatter_str) or {}
except ImportError:
# Fallback to basic parsing if PyYAML is not available
metadata = self._parse_yaml_basic(frontmatter_str)
except Exception:
metadata = {}
return BitBucketPromptTemplate(
template_id=prompt_id,
content=template_content,
metadata=metadata,
)
def _parse_yaml_basic(self, yaml_str: str) -> Dict[str, Any]:
"""Basic YAML parser for simple cases when PyYAML is not available."""
result = {}
for line in yaml_str.split("\n"):
line = line.strip()
if ":" in line and not line.startswith("#"):
key, value = line.split(":", 1)
key = key.strip()
value = value.strip()
# Try to parse value as appropriate type
if value.lower() in ["true", "false"]:
result[key] = value.lower() == "true"
elif value.isdigit():
result[key] = int(value)
elif value.replace(".", "").isdigit():
result[key] = float(value)
else:
result[key] = value.strip("\"'")
return result
def render_template(
self, template_id: str, variables: Dict[str, Any] = None
) -> str:
"""Render a template with the given variables."""
if template_id not in self.prompts:
raise ValueError(f"Template '{template_id}' not found")
template = self.prompts[template_id]
jinja_template = self.jinja_env.from_string(template.content)
return jinja_template.render(**(variables or {}))
def get_template(self, template_id: str) -> Optional[BitBucketPromptTemplate]:
"""Get a template by ID."""
return self.prompts.get(template_id)
def list_templates(self) -> List[str]:
"""List all available template IDs."""
return list(self.prompts.keys())
class BitBucketPromptManager(CustomPromptManagement):
"""
BitBucket prompt manager that integrates with LiteLLM's prompt management system.
This class enables using .prompt files from BitBucket repositories with the
litellm completion() function by implementing the PromptManagementBase interface.
Usage:
# Configure BitBucket access
bitbucket_config = {
"workspace": "your-workspace",
"repository": "your-repo",
"access_token": "your-token",
"branch": "main" # optional, defaults to main
}
# Use with completion
response = litellm.completion(
model="bitbucket/gpt-4",
prompt_id="my_prompt",
prompt_variables={"variable": "value"},
bitbucket_config=bitbucket_config,
messages=[{"role": "user", "content": "This will be combined with the prompt"}]
)
"""
def __init__(
self,
bitbucket_config: Dict[str, Any],
prompt_id: Optional[str] = None,
):
self.bitbucket_config = bitbucket_config
self.prompt_id = prompt_id
self._prompt_manager: Optional[BitBucketPromptManager] = None
@property
def integration_name(self) -> str:
"""Integration name used in model names like 'bitbucket/gpt-4'."""
return "bitbucket"
@property
def prompt_manager(self) -> BitBucketTemplateManager:
"""Get or create the prompt manager instance."""
if self._prompt_manager is None:
self._prompt_manager = BitBucketTemplateManager(
bitbucket_config=self.bitbucket_config,
prompt_id=self.prompt_id,
)
return self._prompt_manager
def get_prompt_template(
self,
prompt_id: str,
prompt_variables: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict[str, Any]]:
"""
Get a prompt template and render it with variables.
Args:
prompt_id: The ID of the prompt template
prompt_variables: Variables to substitute in the template
Returns:
Tuple of (rendered_prompt, metadata)
"""
template = self.prompt_manager.get_template(prompt_id)
if not template:
raise ValueError(f"Prompt template '{prompt_id}' not found")
# Render the template
rendered_prompt = self.prompt_manager.render_template(
prompt_id, prompt_variables or {}
)
# Extract metadata
metadata = {
"model": template.model,
"temperature": template.temperature,
"max_tokens": template.max_tokens,
**template.optional_params,
}
return rendered_prompt, metadata
def pre_call_hook(
self,
user_id: Optional[str],
messages: List[AllMessageValues],
function_call: Optional[Union[Dict[str, Any], str]] = None,
litellm_params: Optional[Dict[str, Any]] = None,
prompt_id: Optional[str] = None,
prompt_variables: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]:
"""
Pre-call hook that processes the prompt template before making the LLM call.
"""
if not prompt_id:
return messages, litellm_params
try:
# Get the rendered prompt and metadata
rendered_prompt, prompt_metadata = self.get_prompt_template(
prompt_id, prompt_variables
)
# Parse the rendered prompt into messages
parsed_messages = self._parse_prompt_to_messages(rendered_prompt)
# Merge with existing messages
if parsed_messages:
# If we have parsed messages, use them instead of the original messages
final_messages = parsed_messages
else:
# If no messages were parsed, prepend the prompt to existing messages
final_messages = [
{"role": "user", "content": rendered_prompt}
] + messages
# Update litellm_params with prompt metadata
if litellm_params is None:
litellm_params = {}
# Apply model and parameters from prompt metadata
if prompt_metadata.get("model"):
litellm_params["model"] = prompt_metadata["model"]
for param in [
"temperature",
"max_tokens",
"top_p",
"frequency_penalty",
"presence_penalty",
]:
if param in prompt_metadata:
litellm_params[param] = prompt_metadata[param]
return final_messages, litellm_params
except Exception as e:
# Log error but don't fail the call
import litellm
litellm._logging.verbose_proxy_logger.error(
f"Error in BitBucket prompt pre_call_hook: {e}"
)
return messages, litellm_params
def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]:
"""
Parse prompt content into a list of messages.
Handles both simple prompts and multi-role conversations.
"""
messages = []
lines = prompt_content.strip().split("\n")
current_role = None
current_content = []
for line in lines:
line = line.strip()
if not line:
continue
# Check for role indicators
if line.lower().startswith("system:"):
if current_role and current_content:
messages.append(
{
"role": current_role,
"content": "\n".join(current_content).strip(),
}
)
current_role = "system"
current_content = [line[7:].strip()] # Remove "System:" prefix
elif line.lower().startswith("user:"):
if current_role and current_content:
messages.append(
{
"role": current_role,
"content": "\n".join(current_content).strip(),
}
)
current_role = "user"
current_content = [line[5:].strip()] # Remove "User:" prefix
elif line.lower().startswith("assistant:"):
if current_role and current_content:
messages.append(
{
"role": current_role,
"content": "\n".join(current_content).strip(),
}
)
current_role = "assistant"
current_content = [line[10:].strip()] # Remove "Assistant:" prefix
else:
# Continue building current message
current_content.append(line)
# Add the last message
if current_role and current_content:
messages.append(
{"role": current_role, "content": "\n".join(current_content).strip()}
)
# If no role indicators found, treat as a single user message
if not messages and prompt_content.strip():
messages = [{"role": "user", "content": prompt_content.strip()}]
return messages
def post_call_hook(
self,
user_id: Optional[str],
response: Any,
input_messages: List[AllMessageValues],
function_call: Optional[Union[Dict[str, Any], str]] = None,
litellm_params: Optional[Dict[str, Any]] = None,
prompt_id: Optional[str] = None,
prompt_variables: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Any:
"""
Post-call hook for any post-processing after the LLM call.
"""
return response
def get_available_prompts(self) -> List[str]:
"""Get list of available prompt IDs."""
return self.prompt_manager.list_templates()
def reload_prompts(self) -> None:
"""Reload prompts from BitBucket."""
if self.prompt_id:
self._prompt_manager = None # Reset to force reload
self.prompt_manager # This will trigger reload
def should_run_prompt_management(
self,
prompt_id: str,
dynamic_callback_params: StandardCallbackDynamicParams,
) -> bool:
"""
Determine if prompt management should run based on the prompt_id.
For BitBucket, we always return True and handle the prompt loading
in the _compile_prompt_helper method.
"""
return True
def _compile_prompt_helper(
self,
prompt_id: str,
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> PromptManagementClient:
"""
Compile a BitBucket prompt template into a PromptManagementClient structure.
This method:
1. Loads the prompt template from BitBucket
2. Renders it with the provided variables
3. Converts the rendered text into chat messages
4. Extracts model and optional parameters from metadata
"""
try:
# Load the prompt from BitBucket if not already loaded
if prompt_id not in self.prompt_manager.prompts:
self.prompt_manager._load_prompt_from_bitbucket(prompt_id)
# Get the rendered prompt and metadata
rendered_prompt, prompt_metadata = self.get_prompt_template(
prompt_id, prompt_variables
)
# Convert rendered content to chat messages
messages = self._parse_prompt_to_messages(rendered_prompt)
# Extract model from metadata (if specified)
template_model = prompt_metadata.get("model")
# Extract optional parameters from metadata
optional_params = {}
for param in [
"temperature",
"max_tokens",
"top_p",
"frequency_penalty",
"presence_penalty",
]:
if param in prompt_metadata:
optional_params[param] = prompt_metadata[param]
return PromptManagementClient(
prompt_id=prompt_id,
prompt_template=messages,
prompt_template_model=template_model,
prompt_template_optional_params=optional_params,
completed_messages=None,
)
except Exception as e:
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
def get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Get chat completion prompt from BitBucket and return processed model, messages, and parameters.
"""
return PromptManagementBase.get_chat_completion_prompt(
self,
model,
messages,
non_default_params,
prompt_id,
prompt_variables,
dynamic_callback_params,
prompt_label,
prompt_version,
)

3
litellm/litellm_core_utils/custom_logger_registry.py
View file

@ -39,6 +39,7 @@ try:
from litellm_enterprise.integrations.prometheus import PrometheusLogger
except Exception:
PrometheusLogger = None
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
from litellm.integrations.dotprompt import DotpromptManager
from litellm.integrations.s3_v2 import S3Logger
@ -90,6 +91,7 @@ class CustomLoggerRegistry:
"dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3,
"vector_store_pre_call_hook": VectorStorePreCallHook,
"dotprompt": DotpromptManager,
"bitbucket": BitBucketPromptManager,
"cloudzero": CloudZeroLogger,
"posthog": PostHogLogger,
}
@ -157,7 +159,6 @@ class CustomLoggerRegistry:
if callback_class == class_type:
callback_strs.append(callback_str)
return callback_strs
@classmethod
def get_class_type_for_custom_logger_name(

39
litellm/litellm_core_utils/litellm_logging.py
View file

@ -1259,7 +1259,9 @@ class Logging(LiteLLMLoggingBaseClass):
"standard_built_in_tools_params": self.standard_built_in_tools_params,
"router_model_id": router_model_id,
"litellm_logging_obj": self,
"service_tier": self.optional_params.get("service_tier") if self.optional_params else None,
"service_tier": self.optional_params.get("service_tier")
if self.optional_params
else None,
}
except Exception as e: # error creating kwargs for cost calculation
debug_info = StandardLoggingModelCostFailureDebugInformation(
@ -3637,6 +3639,27 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
dotprompt_logger = DotpromptManager()
_in_memory_loggers.append(dotprompt_logger)
return dotprompt_logger # type: ignore
elif logging_integration == "bitbucket":
from litellm.integrations.bitbucket.bitbucket_prompt_manager import (
BitBucketPromptManager,
)
for callback in _in_memory_loggers:
if isinstance(callback, BitBucketPromptManager):
return callback
# Get global BitBucket config
import litellm
bitbucket_config = getattr(litellm, "global_bitbucket_config", None)
if bitbucket_config is None:
raise ValueError(
"BitBucket configuration not found. Please set litellm.global_bitbucket_config first."
)
bitbucket_logger = BitBucketPromptManager(bitbucket_config=bitbucket_config)
_in_memory_loggers.append(bitbucket_logger)
return bitbucket_logger # type: ignore
return None
except Exception as e:
verbose_logger.exception(
@ -4222,16 +4245,22 @@ class StandardLoggingPayloadSetup:
# Get the actual s3_path from the configured cold storage logger instance
s3_path = "" # default value
# Try to get the actual logger instance from the logger name
try:
custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(configured_cold_storage_logger)
if custom_logger and hasattr(custom_logger, 's3_path') and custom_logger.s3_path:
custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(
configured_cold_storage_logger
)
if (
custom_logger
and hasattr(custom_logger, "s3_path")
and custom_logger.s3_path
):
s3_path = custom_logger.s3_path
except Exception:
# If any error occurs in getting the logger instance, use default empty s3_path
pass
s3_object_key = get_s3_object_key(
s3_path=s3_path, # Use actual s3_path from logger configuration
team_alias_prefix="", # Don't split by team alias for cold storage

50
litellm/proxy/proxy_server.py
View file

@ -458,9 +458,9 @@ except ImportError:
server_root_path = os.getenv("SERVER_ROOT_PATH", "")
_license_check = LicenseCheck()
premium_user: bool = _license_check.is_premium()
premium_user_data: Optional["EnterpriseLicenseData"] = (
_license_check.airgapped_license_data
)
premium_user_data: Optional[
"EnterpriseLicenseData"
] = _license_check.airgapped_license_data
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
)
@ -954,9 +954,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
dual_cache=user_api_key_cache
)
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
redis_usage_cache: Optional[RedisCache] = (
None # redis cache used for tracking spend, tpm/rpm limits
)
redis_usage_cache: Optional[
RedisCache
] = None # redis cache used for tracking spend, tpm/rpm limits
user_custom_auth = None
user_custom_key_generate = None
user_custom_sso = None
@ -1287,9 +1287,9 @@ async def update_cache( # noqa: PLR0915
_id = "team_id:{}".format(team_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
await user_api_key_cache.async_get_cache(key=_id)
)
existing_spend_obj: Optional[
LiteLLM_TeamTable
] = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
# do nothing if team not in api key cache
return
@ -1856,6 +1856,15 @@ class ProxyConfig:
verbose_proxy_logger.info(
f"{blue_color_code}Set Global Prompt Directory on LiteLLM Proxy{reset_color_code}"
)
elif key == "global_bitbucket_config":
from litellm.integrations.bitbucket import (
set_global_bitbucket_config,
)
set_global_bitbucket_config(value)
verbose_proxy_logger.info(
f"{blue_color_code}Set Global BitBucket Config on LiteLLM Proxy{reset_color_code}"
)
elif key == "callbacks":
initialize_callbacks_on_proxy(
value=value,
@ -3102,10 +3111,10 @@ class ProxyConfig:
)
try:
guardrails_in_db: List[Guardrail] = (
await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
guardrails_in_db: List[
Guardrail
] = await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
verbose_proxy_logger.debug(
"guardrails from the DB %s", str(guardrails_in_db)
@ -3335,9 +3344,9 @@ async def initialize( # noqa: PLR0915
user_api_base = api_base
dynamic_config[user_model]["api_base"] = api_base
if api_version:
os.environ["AZURE_API_VERSION"] = (
api_version # set this for azure - litellm can read this from the env
)
os.environ[
"AZURE_API_VERSION"
] = api_version # set this for azure - litellm can read this from the env
if max_tokens: # model-specific param
dynamic_config[user_model]["max_tokens"] = max_tokens
if temperature: # model-specific param
@ -6109,12 +6118,10 @@ def _add_team_models_to_all_models(
team_models: Dict[str, Set[str]] = {}
for team_object in team_db_objects_typed:
if (
len(team_object.models) == 0 # empty list = all model access
or SpecialModelNames.all_proxy_models.value in team_object.models
):
model_list = llm_router.get_model_list()
if model_list is not None:
for model in model_list:
@ -6265,7 +6272,6 @@ async def get_all_team_and_direct_access_models(
for _model in all_models:
model_id = _model.get("model_info", {}).get("id", None)
if model_id is not None and model_id in direct_access_models:
_model["model_info"]["direct_access"] = True
## FILTER OUT MODELS THAT ARE NOT IN DIRECT_ACCESS_MODELS OR ACCESS_VIA_TEAM_IDS - only show user models they can call
@ -8625,9 +8631,9 @@ async def get_config_list(
hasattr(sub_field_info, "description")
and sub_field_info.description is not None
):
nested_fields[idx].field_description = (
sub_field_info.description
)
nested_fields[
idx
].field_description = sub_field_info.description
idx += 1
_stored_in_db = None

56
litellm/router.py
View file

@ -2003,7 +2003,14 @@ class Router:
prompt_label=prompt_label,
)
kwargs = {**data, **kwargs, **optional_params}
# Filter out prompt management specific parameters from data before merging
prompt_management_params = {
"bitbucket_config", "dotprompt_config", "prompt_id",
"prompt_variables", "prompt_label", "prompt_version"
}
filtered_data = {k: v for k, v in data.items() if k not in prompt_management_params}
kwargs = {**filtered_data, **kwargs, **optional_params}
kwargs["model"] = model
kwargs["messages"] = messages
kwargs["litellm_logging_obj"] = litellm_logging_object
@ -5124,21 +5131,38 @@ class Router:
import os
#### VALIDATE MODEL ########
# check if model provider in supported providers
(
_model,
custom_llm_provider,
dynamic_api_key,
api_base,
) = litellm.get_llm_provider(
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.get(
"custom_llm_provider", None
),
)
# done reading model["litellm_params"]
if custom_llm_provider not in litellm.provider_list:
raise Exception(f"Unsupported provider - {custom_llm_provider}")
# Check if this is a prompt management model before validating as LLM provider
litellm_model = deployment.litellm_params.model
is_prompt_management_model = False
if "/" in litellm_model:
split_litellm_model = litellm_model.split("/")[0]
if split_litellm_model in litellm._known_custom_logger_compatible_callbacks:
is_prompt_management_model = True
if is_prompt_management_model:
# For prompt management models, skip LLM provider validation
# The actual model will be resolved at runtime from the prompt file
_model = litellm_model
custom_llm_provider = None
dynamic_api_key = None
api_base = None
else:
# check if model provider in supported providers
(
_model,
custom_llm_provider,
dynamic_api_key,
api_base,
) = litellm.get_llm_provider(
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.get(
"custom_llm_provider", None
),
)
# done reading model["litellm_params"]
if custom_llm_provider not in litellm.provider_list:
raise Exception(f"Unsupported provider - {custom_llm_provider}")
#### DEPLOYMENT NAMES INIT ########
self.deployment_names.append(deployment.litellm_params.model)

3
litellm/types/prompts/init_prompts.py
View file

@ -3,13 +3,13 @@ from enum import Enum
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, ConfigDict
from typing_extensions import Required, TypedDict
class SupportedPromptIntegrations(str, Enum):
DOT_PROMPT = "dotprompt"
LANGFUSE = "langfuse"
CUSTOM = "custom"
BITBUCKET = "bitbucket"
class PromptInfo(BaseModel):
@ -55,6 +55,5 @@ class PromptInfoResponse(BaseModel):
raw_prompt_template: Optional[PromptTemplateBase] = None
class ListPromptsResponse(BaseModel):
prompts: List[PromptSpec]

366
tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py Normal file
View file

@ -0,0 +1,366 @@
import json
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import litellm
from litellm.integrations.bitbucket import BitBucketPromptManager
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_integration_with_litellm(mock_client_class):
"""Test BitBucket prompt integration with LiteLLM completion."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
temperature: 0.7
max_tokens: 150
---
System: You are a helpful assistant.
User: {{user_message}}"""
mock_client_class.return_value = mock_client
# Configure BitBucket
bitbucket_config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
# Set global BitBucket configuration
litellm.set_global_bitbucket_config(bitbucket_config)
# Test that the configuration was set
assert litellm.global_bitbucket_config == bitbucket_config
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_initialization(mock_client_class):
"""Test BitBucketPromptManager initialization."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
temperature: 0.7
---
Hello {{name}}!"""
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Should have loaded the prompt
assert "test_prompt" in manager.prompt_manager.prompts
template = manager.prompt_manager.prompts["test_prompt"]
assert template.model == "gpt-4"
assert template.temperature == 0.7
# Test rendering
rendered = manager.prompt_manager.render_template("test_prompt", {"name": "World"})
assert rendered == "Hello World!"
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_error_handling(mock_client_class):
"""Test BitBucketPromptManager error handling."""
# Mock the BitBucket client to raise an error
mock_client = MagicMock()
mock_client.get_file_content.side_effect = Exception("BitBucket API error")
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
with pytest.raises(Exception, match="Failed to load prompt 'test_prompt' from BitBucket"):
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
_ = manager.prompt_manager # This triggers the error
def test_bitbucket_prompt_manager_config_validation():
"""Test BitBucketPromptManager configuration validation."""
# Test missing required fields - validation happens when prompt_manager is accessed
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
manager = BitBucketPromptManager({})
_ = manager.prompt_manager # This triggers validation
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
manager = BitBucketPromptManager({"workspace": "test"})
_ = manager.prompt_manager # This triggers validation
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
manager = BitBucketPromptManager({"repository": "test"})
_ = manager.prompt_manager # This triggers validation
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
manager = BitBucketPromptManager({"access_token": "test"})
_ = manager.prompt_manager # This triggers validation
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_complex_prompt(mock_client_class):
"""Test BitBucketPromptManager with complex prompt structure."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
temperature: 0.3
max_tokens: 500
input:
schema:
user_question: string
context?: string
language: string
---
System: You are a helpful {{language}} programming assistant.
{% if context %}Context: {{context}}
{% endif %}User: {{user_question}}
Please provide a detailed response in {{language}}."""
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="complex_prompt")
# Should have loaded the prompt
assert "complex_prompt" in manager.prompt_manager.prompts
template = manager.prompt_manager.prompts["complex_prompt"]
assert template.model == "gpt-4"
assert template.temperature == 0.3
assert template.max_tokens == 500
assert template.input_schema == {
"user_question": "string",
"context?": "string",
"language": "string"
}
# Test rendering with all variables
rendered = manager.prompt_manager.render_template(
"complex_prompt",
{
"user_question": "How do I create a class?",
"context": "Python programming",
"language": "Python"
}
)
assert "You are a helpful Python programming assistant." in rendered
assert "Context: Python programming" in rendered
assert "How do I create a class?" in rendered
assert "Please provide a detailed response in Python." in rendered
# Test rendering without optional context
rendered_no_context = manager.prompt_manager.render_template(
"complex_prompt",
{
"user_question": "What is inheritance?",
"language": "Java"
}
)
assert "You are a helpful Java programming assistant." in rendered_no_context
assert "Context:" not in rendered_no_context
assert "What is inheritance?" in rendered_no_context
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_message_parsing(mock_client_class):
"""Test BitBucketPromptManager message parsing for different prompt formats."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
---
System: You are a helpful assistant.
User: {{user_message}}
Assistant: I'll help you with that."""
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="conversation_prompt")
# Test message parsing
messages = manager._parse_prompt_to_messages(
"System: You are a helpful assistant.\n\nUser: Hello!\n\nAssistant: Hi there!"
)
assert len(messages) == 3
assert messages[0]["role"] == "system"
assert messages[0]["content"] == "You are a helpful assistant."
assert messages[1]["role"] == "user"
assert messages[1]["content"] == "Hello!"
assert messages[2]["role"] == "assistant"
assert messages[2]["content"] == "Hi there!"
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_pre_call_hook_integration(mock_client_class):
"""Test BitBucketPromptManager pre_call_hook integration."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
temperature: 0.8
max_tokens: 200
---
System: You are a helpful assistant.
User: {{user_message}}"""
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Test pre_call_hook
original_messages = [{"role": "user", "content": "This will be ignored"}]
litellm_params = {"api_key": "test-key"}
result_messages, result_params = manager.pre_call_hook(
user_id="test_user",
messages=original_messages,
litellm_params=litellm_params,
prompt_id="test_prompt",
prompt_variables={"user_message": "What is AI?"}
)
# Should have parsed the prompt into messages
assert len(result_messages) == 2
assert result_messages[0]["role"] == "system"
assert result_messages[0]["content"] == "You are a helpful assistant."
assert result_messages[1]["role"] == "user"
assert result_messages[1]["content"] == "What is AI?"
# Should have updated litellm_params
assert result_params["model"] == "gpt-4"
assert result_params["temperature"] == 0.8
assert result_params["max_tokens"] == 200
assert result_params["api_key"] == "test-key" # Original params preserved
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_post_call_hook(mock_client_class):
"""Test BitBucketPromptManager post_call_hook."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = "Simple prompt: {{message}}"
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Mock response
mock_response = MagicMock()
mock_response.choices = [MagicMock()]
mock_response.choices[0].message = MagicMock()
mock_response.choices[0].message.content = "Test response"
# Test post_call_hook
result = manager.post_call_hook(
user_id="test_user",
response=mock_response,
input_messages=[{"role": "user", "content": "test"}],
litellm_params={},
prompt_id="test_prompt"
)
# Should return the response unchanged
assert result == mock_response
def test_bitbucket_prompt_manager_integration_name():
"""Test BitBucketPromptManager integration name."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
assert manager.integration_name == "bitbucket"
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_get_template(mock_client_class):
"""Test BitBucketPromptManager get_template method."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = "Test content"
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Test getting existing template
template = manager.prompt_manager.get_template("test_prompt")
assert template is not None
assert template.template_id == "test_prompt"
# Test getting non-existing template
template = manager.prompt_manager.get_template("nonexistent")
assert template is None
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_list_templates(mock_client_class):
"""Test BitBucketPromptManager list_templates method."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = "Test content"
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Test listing templates
templates = manager.prompt_manager.list_templates()
assert isinstance(templates, list)
assert "test_prompt" in templates

517
tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py Normal file
View file

@ -0,0 +1,517 @@
import json
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.integrations.bitbucket.bitbucket_client import BitBucketClient
from litellm.integrations.bitbucket.bitbucket_prompt_manager import (
BitBucketPromptManager,
BitBucketPromptTemplate,
)
def test_bitbucket_prompt_template_creation():
"""Test BitBucketPromptTemplate creation and metadata extraction."""
metadata = {
"model": "gpt-4",
"temperature": 0.7,
"input": {"schema": {"text": "string"}},
"output": {"format": "json"},
}
template = BitBucketPromptTemplate(
template_id="test_template",
content="Hello {{name}}!",
metadata=metadata,
)
assert template.template_id == "test_template"
assert template.content == "Hello {{name}}!"
assert template.model == "gpt-4"
assert template.optional_params["temperature"] == 0.7
assert template.input_schema == {"text": "string"}
def test_bitbucket_client_initialization():
"""Test BitBucketClient initialization with different auth methods."""
# Test token-based auth
config_token = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
"branch": "main",
}
client = BitBucketClient(config_token)
assert client.workspace == "test-workspace"
assert client.repository == "test-repo"
assert client.access_token == "test-token"
assert client.branch == "main"
assert client.auth_method == "token"
# Test basic auth
config_basic = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-password",
"username": "test-user",
"auth_method": "basic",
}
client_basic = BitBucketClient(config_basic)
assert client_basic.auth_method == "basic"
assert client_basic.username == "test-user"
def test_bitbucket_client_missing_required_fields():
"""Test BitBucketClient initialization with missing required fields."""
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
BitBucketClient({"workspace": "test"})
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
BitBucketClient({"repository": "test"})
with pytest.raises(ValueError, match="workspace, repository, and access_token are required"):
BitBucketClient({"access_token": "test"})
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get")
def test_bitbucket_client_get_file_content_success(mock_get):
"""Test successful file content retrieval from BitBucket."""
# Mock successful response
mock_response = MagicMock()
mock_response.text = "file content"
mock_response.headers = {"content-type": "text/plain"}
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
client = BitBucketClient(config)
content = client.get_file_content("test.prompt")
assert content == "file content"
mock_get.assert_called_once()
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get")
def test_bitbucket_client_get_file_content_not_found(mock_get):
"""Test file content retrieval when file doesn't exist."""
# Mock 404 response
import httpx
mock_response = MagicMock()
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError("404 Not Found", request=MagicMock(), response=mock_response)
mock_response.status_code = 404
mock_response.response = mock_response
mock_get.return_value = mock_response
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
client = BitBucketClient(config)
content = client.get_file_content("nonexistent.prompt")
assert content is None
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get")
def test_bitbucket_client_get_file_content_access_denied(mock_get):
"""Test file content retrieval with access denied error."""
# Mock 403 response
import httpx
mock_response = MagicMock()
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError("403 Forbidden", request=MagicMock(), response=mock_response)
mock_response.status_code = 403
mock_response.response = mock_response
mock_get.return_value = mock_response
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
client = BitBucketClient(config)
with pytest.raises(Exception, match="Access denied to file 'test.prompt'"):
client.get_file_content("test.prompt")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get")
def test_bitbucket_client_get_file_content_auth_failed(mock_get):
"""Test file content retrieval with authentication failure."""
# Mock 401 response
import httpx
mock_response = MagicMock()
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError("401 Unauthorized", request=MagicMock(), response=mock_response)
mock_response.status_code = 401
mock_response.response = mock_response
mock_get.return_value = mock_response
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
client = BitBucketClient(config)
with pytest.raises(Exception, match="Authentication failed"):
client.get_file_content("test.prompt")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get")
def test_bitbucket_client_list_files_success(mock_get):
"""Test successful file listing from BitBucket."""
# Mock successful response
mock_response = MagicMock()
mock_response.json.return_value = {
"values": [
{"type": "commit_file", "path": "prompts/test1.prompt"},
{"type": "commit_file", "path": "prompts/test2.prompt"},
{"type": "commit_file", "path": "prompts/other.txt"},
]
}
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
client = BitBucketClient(config)
files = client.list_files("prompts", ".prompt")
assert len(files) == 2
assert "prompts/test1.prompt" in files
assert "prompts/test2.prompt" in files
def test_bitbucket_prompt_manager_parse_prompt_file():
"""Test parsing .prompt file content with YAML frontmatter."""
prompt_content = """---
model: gpt-4
temperature: 0.7
max_tokens: 150
input:
schema:
user_message: string
system_context?: string
---
{% if system_context %}System: {{system_context}}
{% endif %}User: {{user_message}}"""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
template = manager.prompt_manager._parse_prompt_file(prompt_content, "test_prompt")
assert template.template_id == "test_prompt"
assert template.model == "gpt-4"
assert template.temperature == 0.7
assert template.max_tokens == 150
assert template.input_schema == {"user_message": "string", "system_context?": "string"}
assert "{% if system_context %}" in template.content
def test_bitbucket_prompt_manager_parse_prompt_file_no_frontmatter():
"""Test parsing .prompt file content without YAML frontmatter."""
prompt_content = "Simple prompt: {{message}}"
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
template = manager.prompt_manager._parse_prompt_file(prompt_content, "simple_prompt")
assert template.template_id == "simple_prompt"
assert template.content == "Simple prompt: {{message}}"
assert template.metadata == {}
def test_bitbucket_prompt_manager_render_template():
"""Test template rendering with variables."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
# Add a test template
template = BitBucketPromptTemplate(
template_id="test_template",
content="Hello {{name}}! Welcome to {{place}}.",
metadata={"model": "gpt-4"},
)
manager.prompt_manager.prompts["test_template"] = template
rendered = manager.prompt_manager.render_template("test_template", {"name": "World", "place": "Earth"})
assert rendered == "Hello World! Welcome to Earth."
def test_bitbucket_prompt_manager_render_template_not_found():
"""Test template rendering when template doesn't exist."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
with pytest.raises(ValueError, match="Template 'nonexistent' not found"):
manager.prompt_manager.render_template("nonexistent", {"some": "variable"})
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_integration(mock_client_class):
"""Test BitBucketPromptManager integration with BitBucketClient."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
temperature: 0.7
---
Hello {{name}}!"""
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Should have loaded the prompt
assert "test_prompt" in manager.prompt_manager.prompts
template = manager.prompt_manager.prompts["test_prompt"]
assert template.model == "gpt-4"
assert template.temperature == 0.7
# Test rendering
rendered = manager.prompt_manager.render_template("test_prompt", {"name": "World"})
assert rendered == "Hello World!"
def test_bitbucket_prompt_manager_parse_prompt_to_messages():
"""Test parsing prompt content into messages."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
# Test simple user message
simple_prompt = "Hello, how can I help you?"
messages = manager._parse_prompt_to_messages(simple_prompt)
assert len(messages) == 1
assert messages[0]["role"] == "user"
assert messages[0]["content"] == "Hello, how can I help you?"
# Test multi-role conversation
multi_role_prompt = """System: You are a helpful assistant.
User: What is the capital of France?
Assistant: The capital of France is Paris."""
messages = manager._parse_prompt_to_messages(multi_role_prompt)
assert len(messages) == 3
assert messages[0]["role"] == "system"
assert messages[0]["content"] == "You are a helpful assistant."
assert messages[1]["role"] == "user"
assert messages[1]["content"] == "What is the capital of France?"
assert messages[2]["role"] == "assistant"
assert messages[2]["content"] == "The capital of France is Paris."
def test_bitbucket_prompt_manager_pre_call_hook():
"""Test the pre_call_hook method."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
# Add a test template
template = BitBucketPromptTemplate(
template_id="test_prompt",
content="System: You are a helpful assistant.\n\nUser: {{user_message}}",
metadata={"model": "gpt-4", "temperature": 0.7},
)
manager.prompt_manager.prompts["test_prompt"] = template
# Test pre_call_hook
messages = [{"role": "user", "content": "This will be ignored"}]
litellm_params = {}
result_messages, result_params = manager.pre_call_hook(
user_id="test_user",
messages=messages,
litellm_params=litellm_params,
prompt_id="test_prompt",
prompt_variables={"user_message": "Hello!"}
)
# Should have parsed the prompt into messages
assert len(result_messages) == 2
assert result_messages[0]["role"] == "system"
assert result_messages[0]["content"] == "You are a helpful assistant."
assert result_messages[1]["role"] == "user"
assert result_messages[1]["content"] == "Hello!"
# Should have updated litellm_params
assert result_params["model"] == "gpt-4"
assert result_params["temperature"] == 0.7
def test_bitbucket_prompt_manager_pre_call_hook_no_prompt_id():
"""Test pre_call_hook when no prompt_id is provided."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
messages = [{"role": "user", "content": "Hello"}]
litellm_params = {}
result_messages, result_params = manager.pre_call_hook(
user_id="test_user",
messages=messages,
litellm_params=litellm_params,
prompt_id=None,
)
# Should return original messages and params unchanged
assert result_messages == messages
assert result_params == litellm_params
def test_bitbucket_prompt_manager_get_available_prompts():
"""Test getting list of available prompts."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
# Add some test templates
template1 = BitBucketPromptTemplate("prompt1", "content1", {})
template2 = BitBucketPromptTemplate("prompt2", "content2", {})
manager.prompt_manager.prompts["prompt1"] = template1
manager.prompt_manager.prompts["prompt2"] = template2
available_prompts = manager.get_available_prompts()
assert set(available_prompts) == {"prompt1", "prompt2"}
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
def test_bitbucket_prompt_manager_reload_prompts(mock_client_class):
"""Test reloading prompts from BitBucket."""
# Mock the BitBucket client
mock_client = MagicMock()
mock_client.get_file_content.return_value = """---
model: gpt-4
temperature: 0.7
---
Hello {{name}}!"""
mock_client_class.return_value = mock_client
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config, prompt_id="test_prompt")
# Mock the prompt manager to test reload
with patch.object(manager, '_prompt_manager', None):
manager.reload_prompts()
# Should trigger reload by accessing prompt_manager property
_ = manager.prompt_manager
def test_bitbucket_prompt_manager_yaml_parsing_fallback():
"""Test YAML parsing fallback when PyYAML is not available."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
# Test basic YAML parsing fallback
yaml_content = """model: gpt-4
temperature: 0.7
max_tokens: 150"""
parsed = manager.prompt_manager._parse_yaml_basic(yaml_content)
assert parsed["model"] == "gpt-4"
assert parsed["temperature"] == 0.7
assert parsed["max_tokens"] == 150
def test_bitbucket_prompt_manager_yaml_parsing_with_types():
"""Test YAML parsing with different data types."""
config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
}
manager = BitBucketPromptManager(config)
yaml_content = """model: gpt-4
temperature: 0.7
max_tokens: 150
enabled: true
disabled: false
count: 42
rate: 0.5"""
parsed = manager.prompt_manager._parse_yaml_basic(yaml_content)
assert parsed["model"] == "gpt-4"
assert parsed["temperature"] == 0.7
assert parsed["max_tokens"] == 150
assert parsed["enabled"] is True
assert parsed["disabled"] is False
assert parsed["count"] == 42
assert parsed["rate"] == 0.5