Adds tests and updates docs for Claude "pre-fill"

docs/my-website/docs/providers/anthropic.md

@@ -1,7 +1,7 @@
 # Anthropic
 LiteLLM supports
 
-- `claude-2.1`
+- `claude-2`
 - `claude-2.1`
 - `claude-instant-1`
 - `claude-instant-1.2`
@@ -14,7 +14,7 @@ import os
 os.environ["ANTHROPIC_API_KEY"] = "your-api-key"
 ```
 
-## Sample Usage
+## Usage
 
 ```python
 import os
@@ -28,7 +28,29 @@ response = completion(model="claude-instant-1", messages=messages)
 print(response)
 ```
 
-## streaming 
+## Usage - "Assistant Pre-fill"
+
+You can "put words in Claude's mouth" by including an `assistant` role message as the last item in the `messages` array.
+
+> [!IMPORTANT]
+> The returned completion will _not_ include your "pre-fill" text, since it is part of the prompt itself. Make sure to prefix Claude's completion with your pre-fill.
+
+```python
+import os
+from litellm import completion
+
+# set env - [OPTIONAL] replace with your anthropic key
+os.environ["ANTHROPIC_API_KEY"] = "your-api-key"
+
+messages = [
+    {"role": "user", "content": "How do you say 'Hello' in German? Return your answer as a JSON object, like this:\n\n{ \"Hello\": \"Hallo\"}"},
+    {"role": "assistant", "content": "{"},
+]
+response = completion(model="claude-2.1", messages=messages)
+print(response)
+```
+
+## Streaming
 Just set `stream=True` when calling completion.
 
 ```python

docs/my-website/docs/providers/bedrock.md

@@ -34,6 +34,58 @@ response = completion(
 )
 ```
 
+## Usage - "Assistant Pre-fill"
+
+If you're using Anthropic's Claude with Bedrock, you can "put words in Claude's mouth" by including an `assistant` role message as the last item in the `messages` array.
+
+> [!IMPORTANT]
+> The returned completion will _**not**_ include your "pre-fill" text, since it is part of the prompt itself. Make sure to prefix Claude's completion with your pre-fill.
+
+```python
+import os
+from litellm import completion
+
+os.environ["AWS_ACCESS_KEY_ID"] = ""
+os.environ["AWS_SECRET_ACCESS_KEY"] = ""
+os.environ["AWS_REGION_NAME"] = ""
+
+messages = [
+    {"role": "user", "content": "How do you say 'Hello' in German? Return your answer as a JSON object, like this:\n\n{ \"Hello\": \"Hallo\"}"},
+    {"role": "assistant", "content": "{"},
+]
+response = completion(model="anthropic.claude-v2", messages=messages)
+```
+
+## Usage - "System" messages
+If you're using Anthropic's Claude 2.1 with Bedrock, `system` role messages are properly formatted for you.
+
+```python
+import os
+from litellm import completion
+
+os.environ["AWS_ACCESS_KEY_ID"] = ""
+os.environ["AWS_SECRET_ACCESS_KEY"] = ""
+os.environ["AWS_REGION_NAME"] = ""
+
+messages = [
+    {"role": "system", "content": "You are a snarky assistant."},
+    {"role": "user", "content": "How do I boil water?"},
+]
+response = completion(model="anthropic.claude-v2:1", messages=messages)
+```
+
+### Example prompt sent to Claude
+
+```
+You are a snarky assistant.
+
+Human: How do I boil water?
+
+Assistant:
+```
+
+
+
 ## Usage - Streaming
 ```python
 import os

litellm/tests/test_prompt_factory.py

@@ -2,27 +2,96 @@
 #    This tests if prompts are being correctly formatted
 import sys
 import os
-import io
 
 sys.path.insert(0, os.path.abspath("../.."))
 
 # from litellm.llms.prompt_templates.factory import prompt_factory
 from litellm import completion
+from litellm.llms.prompt_templates.factory import (
+    anthropic_pt,
+    claude_2_1_pt,
+    llama_2_chat_pt,
+)
 
 
-def codellama_prompt_format():
-    model = "huggingface/codellama/CodeLlama-7b-Instruct-hf"
+def test_codellama_prompt_format():
     messages = [
         {"role": "system", "content": "You are a good bot"},
         {"role": "user", "content": "Hey, how's it going?"},
     ]
-    expected_response = """[INST] <<SYS>>
-You are a good bot
-<</SYS>>
- [/INST]
-[INST] Hey, how's it going? [/INST]"""
-    response = completion(model=model, messages=messages)
-    print(response)
+    expected_prompt = "<s>[INST] <<SYS>>\nYou are a good bot\n<</SYS>>\n [/INST]\n[INST] Hey, how's it going? [/INST]\n"
+    assert llama_2_chat_pt(messages) == expected_prompt
+
+
+def test_claude_2_1_pt_formatting():
+    # Test case: User only, should add Assistant
+    messages = [{"role": "user", "content": "Hello"}]
+    expected_prompt = "\n\nHuman: Hello\n\nAssistant: "
+    assert claude_2_1_pt(messages) == expected_prompt
+
+    # Test case: System, User, and Assistant "pre-fill" sequence,
+    #            Should return pre-fill
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": 'Please return "Hello World" as a JSON object.'},
+        {"role": "assistant", "content": "{"},
+    ]
+    expected_prompt = 'You are a helpful assistant.\n\nHuman: Please return "Hello World" as a JSON object.\n\nAssistant: {'
+    assert claude_2_1_pt(messages) == expected_prompt
+
+    # Test case: System, Assistant sequence, should insert blank Human message
+    #            before Assistant pre-fill
+    messages = [
+        {"role": "system", "content": "You are a storyteller."},
+        {"role": "assistant", "content": "Once upon a time, there "},
+    ]
+    expected_prompt = (
+        "You are a storyteller.\n\nHuman: \n\nAssistant: Once upon a time, there "
+    )
+    assert claude_2_1_pt(messages) == expected_prompt
+
+    # Test case: System, User sequence
+    messages = [
+        {"role": "system", "content": "System reboot"},
+        {"role": "user", "content": "Is everything okay?"},
+    ]
+    expected_prompt = "System reboot\n\nHuman: Is everything okay?\n\nAssistant: "
+    assert claude_2_1_pt(messages) == expected_prompt
+
+
+def test_anthropic_pt_formatting():
+    # Test case: User only, should add Assistant
+    messages = [{"role": "user", "content": "Hello"}]
+    expected_prompt = "\n\nHuman: Hello\n\nAssistant: "
+    assert anthropic_pt(messages) == expected_prompt
+
+    # Test case: System, User, and Assistant "pre-fill" sequence,
+    #            Should return pre-fill
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": 'Please return "Hello World" as a JSON object.'},
+        {"role": "assistant", "content": "{"},
+    ]
+    expected_prompt = '\n\nHuman: <admin>You are a helpful assistant.</admin>\n\nHuman: Please return "Hello World" as a JSON object.\n\nAssistant: {'
+    assert anthropic_pt(messages) == expected_prompt
+
+    # Test case: System, Assistant sequence, should NOT insert blank Human message
+    #            before Assistant pre-fill, because "System" messages are Human
+    #            messages wrapped with <admin></admin>
+    messages = [
+        {"role": "system", "content": "You are a storyteller."},
+        {"role": "assistant", "content": "Once upon a time, there "},
+    ]
+    expected_prompt = "\n\nHuman: <admin>You are a storyteller.</admin>\n\nAssistant: Once upon a time, there "
+    assert anthropic_pt(messages) == expected_prompt
+
+    # Test case: System, User sequence
+    messages = [
+        {"role": "system", "content": "System reboot"},
+        {"role": "user", "content": "Is everything okay?"},
+    ]
+    expected_prompt = "\n\nHuman: <admin>System reboot</admin>\n\nHuman: Is everything okay?\n\nAssistant: "
+    assert anthropic_pt(messages) == expected_prompt
 
 
 # codellama_prompt_format()