From 1681bf717520978b35e55f06f4a5d4fc5ee71428 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 26 Sep 2025 04:21:00 +0530 Subject: [PATCH] (Feat) Add BitBucket Integration for Prompt Management (#14882) * 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 --- .../docs/proxy/native_litellm_prompt.md | 90 ++- litellm/__init__.py | 9 + litellm/integrations/bitbucket/README.md | 317 +++++++++++ litellm/integrations/bitbucket/__init__.py | 66 +++ .../bitbucket/bitbucket_client.py | 241 ++++++++ .../bitbucket/bitbucket_prompt_manager.py | 508 +++++++++++++++++ .../custom_logger_registry.py | 3 +- litellm/litellm_core_utils/litellm_logging.py | 39 +- litellm/proxy/proxy_server.py | 50 +- litellm/router.py | 56 +- litellm/types/prompts/init_prompts.py | 3 +- .../bitbucket/test_bitbucket_integration.py | 366 +++++++++++++ .../test_bitbucket_prompt_manager.py | 517 ++++++++++++++++++ 13 files changed, 2214 insertions(+), 51 deletions(-) create mode 100644 litellm/integrations/bitbucket/README.md create mode 100644 litellm/integrations/bitbucket/__init__.py create mode 100644 litellm/integrations/bitbucket/bitbucket_client.py create mode 100644 litellm/integrations/bitbucket/bitbucket_prompt_manager.py create mode 100644 tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py create mode 100644 tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py diff --git a/docs/my-website/docs/proxy/native_litellm_prompt.md b/docs/my-website/docs/proxy/native_litellm_prompt.md index 1e1df999db9..0a4f1ccf35c 100644 --- a/docs/my-website/docs/proxy/native_litellm_prompt.md +++ b/docs/my-website/docs/proxy/native_litellm_prompt.md @@ -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 @@ -41,6 +49,50 @@ response = litellm.completion( ) ``` + + + +**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?"} +) +``` + @@ -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/ # 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/ # 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" + } +) ``` diff --git a/litellm/__init__.py b/litellm/__init__.py index 69e5f78c997..6df9d0780a4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/integrations/bitbucket/README.md b/litellm/integrations/bitbucket/README.md new file mode 100644 index 00000000000..473beeea9e0 --- /dev/null +++ b/litellm/integrations/bitbucket/README.md @@ -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/", # 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. diff --git a/litellm/integrations/bitbucket/__init__.py b/litellm/integrations/bitbucket/__init__.py new file mode 100644 index 00000000000..111d38f78a4 --- /dev/null +++ b/litellm/integrations/bitbucket/__init__.py @@ -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", +] diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py new file mode 100644 index 00000000000..0502422cf8b --- /dev/null +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -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() diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py new file mode 100644 index 00000000000..96cab6d9ccb --- /dev/null +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -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, + ) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index bb8fa580e3c..fbf685c3700 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -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( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d7d82d0440e..5c22bb4ad98 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9fff2478804..23d9acfc1e5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/router.py b/litellm/router.py index 9c87f0bb029..f093096f848 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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) diff --git a/litellm/types/prompts/init_prompts.py b/litellm/types/prompts/init_prompts.py index c168bf31289..3c48cd131c3 100644 --- a/litellm/types/prompts/init_prompts.py +++ b/litellm/types/prompts/init_prompts.py @@ -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] diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py new file mode 100644 index 00000000000..254c468d766 --- /dev/null +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py @@ -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 diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py new file mode 100644 index 00000000000..75d7a94c5e4 --- /dev/null +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -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