import json import tempfile from pathlib import Path import pytest from fastapi.testclient import TestClient from unittest.mock import MagicMock, Mock, patch import httpx import litellm from litellm.integrations.dotprompt.dotprompt_manager import DotpromptManager from litellm.integrations.dotprompt.prompt_manager import PromptManager, PromptTemplate from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec def test_prompt_manager_initialization(): """Test basic PromptManager initialization and loading.""" # Test with the existing prompts directory prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) # Should have loaded at least the sample prompts assert len(manager.prompts) >= 3 assert "sample_prompt" in manager.prompts assert "chat_prompt" in manager.prompts assert "coding_assistant" in manager.prompts def test_prompt_template_creation(): """Test PromptTemplate creation and metadata extraction.""" metadata = { "model": "gpt-4", "temperature": 0.7, "input": {"schema": {"text": "string"}}, "output": {"format": "json"}, } template = PromptTemplate( content="Hello {{name}}!", metadata=metadata, 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"} assert template.output_format == "json" def test_render_simple_template(): """Test rendering a simple template with variables.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) # Test sample_prompt rendering rendered = manager.render( "sample_prompt", {"text": "This is a test article about AI."} ) expected_content = "Extract the requested information from the given text. If a piece of information is not present, omit that field from the output.\n\nText: This is a test article about AI." assert rendered == expected_content def test_render_chat_prompt(): """Test rendering the chat prompt with conditional content.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) # Test with system context rendered = manager.render( "chat_prompt", { "user_message": "Hello there!", "system_context": "You are a helpful assistant.", }, ) assert "System: You are a helpful assistant." in rendered assert "User: Hello there!" in rendered # Test without system context rendered_no_system = manager.render("chat_prompt", {"user_message": "Hello there!"}) assert "System:" not in rendered_no_system assert "User: Hello there!" in rendered_no_system def test_render_coding_assistant(): """Test rendering the coding assistant prompt with complex logic.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) rendered = manager.render( "coding_assistant", { "language": "Python", "task": "Create a function to calculate fibonacci numbers", "code": "def fib(n):\n pass", "requirements": ["Use recursion", "Handle edge cases", "Add documentation"], }, ) assert "Focus on Python programming." in rendered assert "Create a function to calculate fibonacci numbers" in rendered assert "def fib(n):" in rendered assert "Use recursion" in rendered assert "Handle edge cases" in rendered assert "Add documentation" in rendered def test_input_validation(): """Test input validation against schema.""" # Create a temporary directory with a test prompt with tempfile.TemporaryDirectory() as temp_dir: prompt_file = Path(temp_dir) / "test_validation.prompt" prompt_file.write_text( """--- input: schema: name: string age: integer active: boolean --- Hello {{name}}, you are {{age}} years old and {'active' if active else 'inactive'}.""" ) manager = PromptManager(prompt_directory=str(temp_dir)) # Valid input should work rendered = manager.render( "test_validation", {"name": "Alice", "age": 30, "active": True} ) assert "Hello Alice, you are 30 years old" in rendered # Invalid type should raise error with pytest.raises(ValueError, match="Invalid type for field 'age'"): manager.render( "test_validation", { "name": "Alice", "age": "thirty", # string instead of int "active": True, }, ) def test_prompt_not_found(): """Test error handling for non-existent prompts.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) with pytest.raises(KeyError, match="Prompt 'nonexistent' not found"): manager.render("nonexistent", {"some": "variable"}) def test_list_prompts(): """Test listing available prompts.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) prompts = manager.list_prompts() assert isinstance(prompts, list) assert "sample_prompt" in prompts assert "chat_prompt" in prompts assert "coding_assistant" in prompts def test_get_prompt_metadata(): """Test retrieving prompt metadata.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) metadata = manager.get_prompt_metadata("sample_prompt") assert metadata is not None assert metadata["model"] == "gemini/gemini-1.5-pro" assert "input" in metadata assert "output" in metadata def test_get_prompt_with_version(): """Test that get_prompt correctly retrieves versioned prompts.""" prompt_dir = Path(__file__).parent manager = PromptManager(prompt_directory=str(prompt_dir)) # Get base prompt (no version) base_prompt = manager.get_prompt(prompt_id="chat_prompt") assert base_prompt is not None assert "User: {{user_message}}" in base_prompt.content # Get version 1 v1_prompt = manager.get_prompt(prompt_id="chat_prompt", version=1) assert v1_prompt is not None assert "Version 1:" in v1_prompt.content assert v1_prompt.model == "gpt-3.5-turbo" # Get version 2 v2_prompt = manager.get_prompt(prompt_id="chat_prompt", version=2) assert v2_prompt is not None assert "Version 2:" in v2_prompt.content assert v2_prompt.model == "gpt-4" def test_add_prompt_programmatically(): """Test adding prompts programmatically.""" prompt_dir = Path( __file__ ).parent # Current directory when running from tests/test_litellm/prompts manager = PromptManager(prompt_directory=str(prompt_dir)) initial_count = len(manager.prompts) manager.add_prompt( "dynamic_prompt", "Hello {{name}}! Welcome to {{place}}.", {"model": "gpt-3.5-turbo", "temperature": 0.5}, ) assert len(manager.prompts) == initial_count + 1 assert "dynamic_prompt" in manager.prompts rendered = manager.render("dynamic_prompt", {"name": "World", "place": "Earth"}) assert rendered == "Hello World! Welcome to Earth." def test_frontmatter_parsing(): """Test YAML frontmatter parsing.""" # Create a temporary directory with a test prompt with tempfile.TemporaryDirectory() as temp_dir: # Test with frontmatter prompt_with_frontmatter = Path(temp_dir) / "with_frontmatter.prompt" prompt_with_frontmatter.write_text( """--- model: gpt-4 temperature: 0.8 input: schema: topic: string --- Write about {{topic}}.""" ) # Test without frontmatter prompt_without_frontmatter = Path(temp_dir) / "without_frontmatter.prompt" prompt_without_frontmatter.write_text("Simple template: {{message}}") manager = PromptManager(prompt_directory=str(temp_dir)) # Check frontmatter was parsed correctly with_meta = manager.get_prompt("with_frontmatter") assert with_meta is not None assert with_meta.model == "gpt-4" assert with_meta.optional_params["temperature"] == 0.8 # Check template without frontmatter still works without_meta = manager.get_prompt("without_frontmatter") assert without_meta is not None assert without_meta.metadata == {} rendered = manager.render("without_frontmatter", {"message": "Hello!"}) assert rendered == "Simple template: Hello!" def test_prompt_manager_json_initialization(): """Test PromptManager initialization with JSON data instead of directory.""" prompt_data = { "json_test_prompt": { "content": "Hello {{name}}! Welcome to {{service}}.", "metadata": {"model": "gpt-4", "temperature": 0.8, "max_tokens": 150}, }, "simple_prompt": { "content": "This is a simple prompt: {{message}}", "metadata": {}, }, } # Initialize PromptManager with JSON data only (no directory) manager = PromptManager(prompt_data=prompt_data) # Should have loaded the JSON prompts assert len(manager.prompts) == 2 assert "json_test_prompt" in manager.prompts assert "simple_prompt" in manager.prompts # Test prompt properties json_prompt = manager.get_prompt("json_test_prompt") assert json_prompt.content == "Hello {{name}}! Welcome to {{service}}." assert json_prompt.model == "gpt-4" assert json_prompt.optional_params["temperature"] == 0.8 assert json_prompt.optional_params["max_tokens"] == 150 def test_prompt_manager_mixed_initialization(): """Test PromptManager with both directory and JSON data.""" # Use existing directory prompt_dir = Path(__file__).parent # Add JSON data json_data = { "json_only_prompt": { "content": "This prompt only exists in JSON: {{data}}", "metadata": {"model": "gpt-3.5-turbo"}, } } manager = PromptManager(prompt_directory=str(prompt_dir), prompt_data=json_data) # Should have prompts from both directory and JSON assert "sample_prompt" in manager.prompts # From directory assert "json_only_prompt" in manager.prompts # From JSON # Test rendering both types json_rendered = manager.render("json_only_prompt", {"data": "test"}) assert json_rendered == "This prompt only exists in JSON: test" def test_load_prompts_from_json_data(): """Test loading additional prompts from JSON data after initialization.""" # Start with directory-based manager prompt_dir = Path(__file__).parent manager = PromptManager(prompt_directory=str(prompt_dir)) initial_count = len(manager.prompts) # Load additional prompts from JSON additional_prompts = { "dynamic_json_prompt": { "content": "Dynamic prompt: {{dynamic_content}}", "metadata": {"model": "claude-3", "temperature": 0.5}, }, "another_json_prompt": { "content": "Another prompt with {{variable}}", "metadata": {"model": "gpt-4"}, }, } manager.load_prompts_from_json_data(additional_prompts) # Should have added the new prompts assert len(manager.prompts) == initial_count + 2 assert "dynamic_json_prompt" in manager.prompts assert "another_json_prompt" in manager.prompts # Test rendering the new prompts rendered = manager.render("dynamic_json_prompt", {"dynamic_content": "test"}) assert rendered == "Dynamic prompt: test" def test_prompt_file_to_json_conversion(): """Test converting .prompt files to JSON format.""" # Create a temporary prompt file with frontmatter with tempfile.TemporaryDirectory() as temp_dir: prompt_file = Path(temp_dir) / "test_conversion.prompt" prompt_file.write_text( """--- model: gpt-4 temperature: 0.7 max_tokens: 200 input: schema: user_input: string context: string output: format: json --- You are an AI assistant. Given the context: {{context}} Please respond to: {{user_input}}""" ) manager = PromptManager() json_data = manager.prompt_file_to_json(prompt_file) # Check the conversion assert "content" in json_data assert "metadata" in json_data expected_content = """You are an AI assistant. Given the context: {{context}} Please respond to: {{user_input}}""" assert json_data["content"] == expected_content metadata = json_data["metadata"] assert metadata["model"] == "gpt-4" assert metadata["temperature"] == 0.7 assert metadata["max_tokens"] == 200 assert metadata["input"]["schema"]["user_input"] == "string" assert metadata["output"]["format"] == "json" def test_json_to_prompt_file_conversion(): """Test converting JSON data back to .prompt file format.""" json_data = { "content": "Hello {{name}}! How can I help you with {{task}}?", "metadata": { "model": "gpt-3.5-turbo", "temperature": 0.8, "max_tokens": 100, "input": {"schema": {"name": "string", "task": "string"}}, }, } manager = PromptManager() prompt_content = manager.json_to_prompt_file(json_data) # Should have YAML frontmatter and content assert prompt_content.startswith("---\n") assert "---\n" in prompt_content[4:] # Second --- delimiter assert "Hello {{name}}! How can I help you with {{task}}?" in prompt_content assert "model: gpt-3.5-turbo" in prompt_content assert "temperature: 0.8" in prompt_content def test_json_to_prompt_file_without_metadata(): """Test converting JSON with no metadata to .prompt format.""" json_data = { "content": "Simple prompt without metadata: {{message}}", "metadata": {}, } manager = PromptManager() prompt_content = manager.json_to_prompt_file(json_data) # Should return just the content without frontmatter assert prompt_content == "Simple prompt without metadata: {{message}}" assert "---" not in prompt_content def test_get_all_prompts_as_json(): """Test exporting all prompts to JSON format.""" prompt_data = { "prompt1": { "content": "First prompt: {{var1}}", "metadata": {"model": "gpt-4"}, }, "prompt2": { "content": "Second prompt: {{var2}}", "metadata": {"model": "claude-3", "temperature": 0.5}, }, } manager = PromptManager(prompt_data=prompt_data) all_prompts_json = manager.get_all_prompts_as_json() assert len(all_prompts_json) == 2 assert "prompt1" in all_prompts_json assert "prompt2" in all_prompts_json # Check structure prompt1_data = all_prompts_json["prompt1"] assert prompt1_data["content"] == "First prompt: {{var1}}" assert prompt1_data["metadata"]["model"] == "gpt-4" prompt2_data = all_prompts_json["prompt2"] assert prompt2_data["content"] == "Second prompt: {{var2}}" assert prompt2_data["metadata"]["model"] == "claude-3" def test_json_prompt_rendering_with_validation(): """Test rendering JSON-based prompts with input validation.""" prompt_data = { "validated_prompt": { "content": "Process {{data}} for user {{user_id}}", "metadata": { "model": "gpt-4", "input": {"schema": {"data": "string", "user_id": "integer"}}, }, } } manager = PromptManager(prompt_data=prompt_data) # Valid input should work rendered = manager.render("validated_prompt", {"data": "test data", "user_id": 123}) assert rendered == "Process test data for user 123" # Invalid input should raise error with pytest.raises(ValueError, match="Invalid type for field 'user_id'"): manager.render( "validated_prompt", {"data": "test data", "user_id": "not_an_int"} ) def test_round_trip_conversion(): """Test converting .prompt file to JSON and back to .prompt file.""" with tempfile.TemporaryDirectory() as temp_dir: # Create original prompt file original_file = Path(temp_dir) / "original.prompt" original_content = """--- model: gpt-4 temperature: 0.6 --- Original prompt content: {{variable}}""" original_file.write_text(original_content) manager = PromptManager() # Convert to JSON json_data = manager.prompt_file_to_json(original_file) # Convert back to prompt file format converted_content = manager.json_to_prompt_file(json_data) # Create new file with converted content converted_file = Path(temp_dir) / "converted.prompt" converted_file.write_text(converted_content) # Load both files and compare manager_original = PromptManager(prompt_directory=str(temp_dir)) original_template = manager_original.get_prompt("original") converted_template = manager_original.get_prompt("converted") # Content should be the same assert original_template.content == converted_template.content assert original_template.model == converted_template.model assert ( original_template.optional_params["temperature"] == converted_template.optional_params["temperature"] ) def test_prompt_main(): """ Integration test placeholder for litellm completion integration. This would be implemented once the PromptManager is integrated with litellm. """ # TODO: Implement once PromptManager is integrated with litellm completion pass @pytest.mark.asyncio async def test_dotprompt_with_prompt_version(): """ Test that dotprompt can load and use specific prompt versions. Versions are stored as separate files with .v{version}.prompt naming convention. """ from litellm.integrations.dotprompt.prompt_manager import PromptManager prompt_dir = Path(__file__).parent prompt_manager = PromptManager(prompt_directory=str(prompt_dir)) # Test version 1 v1_prompt = prompt_manager.get_prompt(prompt_id="chat_prompt", version=1) assert v1_prompt is not None assert v1_prompt.model == "gpt-3.5-turbo" # Verify version 1 content v1_rendered = prompt_manager.render( prompt_id="chat_prompt", prompt_variables={"user_message": "Test v1"}, version=1 ) assert "Version 1:" in v1_rendered assert "Test v1" in v1_rendered # Test version 2 v2_prompt = prompt_manager.get_prompt(prompt_id="chat_prompt", version=2) assert v2_prompt is not None assert v2_prompt.model == "gpt-4" # Verify version 2 content v2_rendered = prompt_manager.render( prompt_id="chat_prompt", prompt_variables={"user_message": "Test v2"}, version=2 ) assert "Version 2:" in v2_rendered assert "Test v2" in v2_rendered def test_keyed_prompt_data_with_prompt_id_keeps_real_content(): prompt_data = { "json_prompt": { "content": "You are a pirate. Begin every reply with AHOY.", "metadata": {"model": "gpt-4o-mini"}, } } manager = PromptManager(prompt_data=prompt_data, prompt_id="agent-prompt") template = manager.get_prompt("json_prompt") assert template is not None assert template.content == "You are a pirate. Begin every reply with AHOY." assert template.model == "gpt-4o-mini" assert "agent-prompt" not in manager.prompts def test_flat_prompt_data_with_prompt_id_registers_under_prompt_id(): manager = PromptManager( prompt_data={"content": "Hello {{name}}", "metadata": {"model": "gpt-4o-mini"}}, prompt_id="flat-prompt", ) template = manager.get_prompt("flat-prompt") assert template is not None assert template.content == "Hello {{name}}" assert manager.render("flat-prompt", {"name": "world"}) == "Hello world" def test_get_prompt_falls_back_to_base_id_for_versioned_id(): manager = PromptManager( prompt_data={"content": "Hi", "metadata": {}}, prompt_id="my-prompt", ) assert manager.get_prompt("my-prompt.v1") is not None assert manager.get_prompt("my-prompt.v12") is not None assert manager.get_prompt("my-prompt.vx") is None assert manager.get_prompt("other-prompt.v1") is None def test_should_run_prompt_management_accepts_versioned_id(): from litellm.integrations.dotprompt import DotpromptManager dotprompt_manager = DotpromptManager( prompt_data={"content": "Hi", "metadata": {}}, prompt_id="versioned-prompt", ) assert dotprompt_manager.should_run_prompt_management("versioned-prompt", None, {}) is True assert dotprompt_manager.should_run_prompt_management("versioned-prompt.v1", None, {}) is True assert dotprompt_manager.should_run_prompt_management("missing-prompt", None, {}) is False def test_prompt_initializer_registers_flat_db_prompt_under_base_id(): from litellm.integrations.dotprompt import DotpromptManager, prompt_initializer from litellm.types.prompts.init_prompts import ( PromptInfo, PromptLiteLLMParams, PromptSpec, ) litellm_params = PromptLiteLLMParams( prompt_integration="dotprompt", prompt_data={"content": "AHOY {{name}}", "metadata": {"model": "gpt-4o-mini"}}, ) prompt_spec = PromptSpec( prompt_id="agent-prompt.v1", litellm_params=litellm_params, prompt_info=PromptInfo(prompt_type="db"), ) dotprompt_manager = prompt_initializer(litellm_params, prompt_spec) assert isinstance(dotprompt_manager, DotpromptManager) template = dotprompt_manager.prompt_manager.get_prompt("agent-prompt") assert template is not None assert template.content == "AHOY {{name}}" def _swap_prompt_manager_and_spec(ignore_prompt_manager_model: bool) -> tuple[DotpromptManager, PromptSpec]: manager = DotpromptManager( prompt_data={"content": "You are a pirate assistant.", "metadata": {"model": "gpt-4o-mini"}}, prompt_id="swap-prompt", ) spec = PromptSpec( prompt_id="swap-prompt", litellm_params=PromptLiteLLMParams( prompt_id="swap-prompt", prompt_integration="dotprompt", ignore_prompt_manager_model=ignore_prompt_manager_model, ), ) return manager, spec @pytest.mark.asyncio async def test_async_prompt_spec_ignore_prompt_manager_model_keeps_requested_model(): from litellm.types.utils import StandardCallbackDynamicParams manager, spec = _swap_prompt_manager_and_spec(ignore_prompt_manager_model=True) model, messages, _ = await manager.async_get_chat_completion_prompt( model="anthropic/claude-haiku-4-5", messages=[{"role": "user", "content": "hi"}], non_default_params={}, prompt_id="swap-prompt", prompt_variables=None, dynamic_callback_params=StandardCallbackDynamicParams(), litellm_logging_obj=MagicMock(), prompt_spec=spec, ) assert model == "anthropic/claude-haiku-4-5" assert len(messages) == 2 assert "pirate" in str(messages[0]["content"]) @pytest.mark.asyncio async def test_async_prompt_spec_without_ignore_flag_swaps_model(): from litellm.types.utils import StandardCallbackDynamicParams manager, spec = _swap_prompt_manager_and_spec(ignore_prompt_manager_model=False) model, _, _ = await manager.async_get_chat_completion_prompt( model="anthropic/claude-haiku-4-5", messages=[{"role": "user", "content": "hi"}], non_default_params={}, prompt_id="swap-prompt", prompt_variables=None, dynamic_callback_params=StandardCallbackDynamicParams(), litellm_logging_obj=MagicMock(), prompt_spec=spec, ) assert model == "gpt-4o-mini" def test_sync_prompt_spec_ignore_prompt_manager_model_keeps_requested_model(): from litellm.types.utils import StandardCallbackDynamicParams manager, spec = _swap_prompt_manager_and_spec(ignore_prompt_manager_model=True) model, _, _ = manager.get_chat_completion_prompt( model="anthropic/claude-haiku-4-5", messages=[{"role": "user", "content": "hi"}], non_default_params={}, prompt_id="swap-prompt", prompt_variables=None, dynamic_callback_params=StandardCallbackDynamicParams(), prompt_spec=spec, ) assert model == "anthropic/claude-haiku-4-5" def test_sync_caller_ignore_flag_survives_missing_prompt_spec(): from litellm.types.utils import StandardCallbackDynamicParams manager, _ = _swap_prompt_manager_and_spec(ignore_prompt_manager_model=False) model, _, _ = manager.get_chat_completion_prompt( model="anthropic/claude-haiku-4-5", messages=[{"role": "user", "content": "hi"}], non_default_params={}, prompt_id="swap-prompt", prompt_variables=None, dynamic_callback_params=StandardCallbackDynamicParams(), prompt_spec=None, ignore_prompt_manager_model=True, ) assert model == "anthropic/claude-haiku-4-5"