feat(anthropic.py): support citations api with new user document mess…

…age format Resolves #7970
BerriAI · Feb 8, 2025 · cda73c9 · cda73c9
1 parent 0c87441
commit cda73c9
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 0 deletions.
diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py
@@ -1421,6 +1421,8 @@ def anthropic_messages_pt(  # noqa: PLR0915
                             )
 
                             user_content.append(_content_element)
+                        elif m.get("type", "") == "document":
+                            user_content.append(cast(AnthropicMessagesDocumentParam, m))
                 elif isinstance(user_message_types_block["content"], str):
                     _anthropic_content_text_element: AnthropicMessagesTextParam = {
                         "type": "text",

diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py
@@ -92,10 +92,17 @@ class AnthropicMessagesImageParam(TypedDict, total=False):
     cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
 
 
+class CitationsObject(TypedDict):
+    enabled: bool
+
+
 class AnthropicMessagesDocumentParam(TypedDict, total=False):
     type: Required[Literal["document"]]
     source: Required[AnthropicContentParamSource]
     cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
+    title: str
+    context: str
+    citations: Optional[CitationsObject]
 
 
 class AnthropicMessagesToolResultContent(TypedDict):

diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py
@@ -382,10 +382,29 @@ class ChatCompletionAudioObject(ChatCompletionContentPartInputAudioParam):
     pass
 
 
+class DocumentObject(TypedDict):
+    type: Literal["text"]
+    media_type: str
+    data: str
+
+
+class CitationsObject(TypedDict):
+    enabled: bool
+
+
+class ChatCompletionDocumentObject(TypedDict):
+    type: Literal["document"]
+    source: DocumentObject
+    title: str
+    context: str
+    citations: Optional[CitationsObject]
+
+
 OpenAIMessageContentListBlock = Union[
     ChatCompletionTextObject,
     ChatCompletionImageObject,
     ChatCompletionAudioObject,
+    ChatCompletionDocumentObject,
 ]
 
 OpenAIMessageContent = Union[
@@ -460,6 +479,7 @@ class ChatCompletionDeveloperMessage(OpenAIChatCompletionDeveloperMessage, total
     "text",
     "image_url",
     "input_audio",
+    "document",
 ]  # used for validating user messages. Prevent users from accidentally sending anthropic messages.
 
 AllMessageValues = Union[

diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py
@@ -1082,3 +1082,55 @@ async def test_anthropic_structured_output():
     assert response is not None
 
     print(response)
+
+
+def test_anthropic_citations_api():
+    """
+    Test the citations API
+    """
+    from litellm import completion
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+    import json
+
+    client = HTTPHandler()
+
+    with patch.object(client, "post") as mock_post:
+        try:
+            resp = completion(
+                model="claude-3-5-sonnet-20241022",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "document",
+                                "source": {
+                                    "type": "text",
+                                    "media_type": "text/plain",
+                                    "data": "The grass is green. The sky is blue.",
+                                },
+                                "title": "My Document",
+                                "context": "This is a trustworthy document.",
+                                "citations": {"enabled": True},
+                            },
+                            {
+                                "type": "text",
+                                "text": "What color is the grass and sky?",
+                            },
+                        ],
+                    }
+                ],
+                client=client,
+            )
+
+            print(resp)
+        except Exception as e:
+            print(e)
+
+        mock_post.assert_called_once()
+
+        print(mock_post.call_args.kwargs)
+
+        json_data = json.loads(mock_post.call_args.kwargs["data"])
+
+        assert json_data["messages"][0]["content"][0]["citations"]["enabled"]