""" 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 TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from jinja2 import DictLoader, Environment, select_autoescape from litellm.integrations.custom_prompt_management import CustomPromptManagement if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj else: LiteLLMLoggingObj = Any from litellm.integrations.prompt_management_base import ( PromptManagementBase, PromptManagementClient, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.prompts.init_prompts import PromptSpec 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: Dict[str, Any] = {} 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: Dict[str, Any] = {} 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: Optional[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[BitBucketTemplateManager] = 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: List[AllMessageValues] = parsed_messages else: # If no messages were parsed, prepend the prompt to existing messages final_messages = [ {"role": "user", "content": rendered_prompt} # type: ignore ] + 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(), } # type: ignore ) 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(), } # type: ignore ) 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(), } # type: ignore ) 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()}] # type: ignore return messages # type: ignore 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: Optional[str], prompt_spec: Optional[PromptSpec], 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 prompt_id is not None def _compile_prompt_helper( self, prompt_id: Optional[str], prompt_spec: Optional[PromptSpec], 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 """ if prompt_id is None: raise ValueError("prompt_id is required for BitBucket prompt manager") 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}") async def async_compile_prompt_helper( self, prompt_id: Optional[str], prompt_variables: Optional[dict], dynamic_callback_params: StandardCallbackDynamicParams, prompt_spec: Optional[PromptSpec] = None, prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> PromptManagementClient: """ Async version of compile prompt helper. Since BitBucket operations use sync client, this simply delegates to the sync version. """ if prompt_id is None: raise ValueError("prompt_id is required for BitBucket prompt manager") return self._compile_prompt_helper( prompt_id=prompt_id, prompt_spec=prompt_spec, prompt_variables=prompt_variables, dynamic_callback_params=dynamic_callback_params, prompt_label=prompt_label, prompt_version=prompt_version, ) 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_spec: Optional[PromptSpec] = None, prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ignore_prompt_manager_model: Optional[bool] = False, ignore_prompt_manager_optional_params: Optional[bool] = False, ) -> 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_spec=prompt_spec, prompt_label=prompt_label, prompt_version=prompt_version, ) async def async_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, litellm_logging_obj: LiteLLMLoggingObj, prompt_spec: Optional[PromptSpec] = None, tools: Optional[List[Dict]] = None, prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ignore_prompt_manager_model: Optional[bool] = False, ignore_prompt_manager_optional_params: Optional[bool] = False, ) -> Tuple[str, List[AllMessageValues], dict]: """ Async version - delegates to PromptManagementBase async implementation. """ return await PromptManagementBase.async_get_chat_completion_prompt( self, model, messages, non_default_params, prompt_id=prompt_id, prompt_variables=prompt_variables, litellm_logging_obj=litellm_logging_obj, dynamic_callback_params=dynamic_callback_params, prompt_spec=prompt_spec, tools=tools, prompt_label=prompt_label, prompt_version=prompt_version, ignore_prompt_manager_model=ignore_prompt_manager_model, ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params, )