fix(anthropic/chat/transformation.py): return citations as a provider…

…-specific-field Resolves #7970
BerriAI · Feb 8, 2025 · fa68e3b · fa68e3b
1 parent cda73c9
commit fa68e3b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 38 deletions.
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
@@ -628,6 +628,7 @@ def transform_response(
             )
         else:
             text_content = ""
+            citations: List[Any] = []
             tool_calls: List[ChatCompletionToolCallChunk] = []
             for idx, content in enumerate(completion_response["content"]):
                 if content["type"] == "text":
@@ -645,10 +646,14 @@ def transform_response(
                             index=idx,
                         )
                     )
+                ## CITATIONS
+                if content.get("citations", None) is not None:
+                    citations.append(content["citations"])
 
             _message = litellm.Message(
                 tool_calls=tool_calls,
                 content=text_content or None,
+                provider_specific_fields={"citations": citations},
             )
 
             ## HANDLE JSON MODE - anthropic returns single function call

diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py
@@ -1092,45 +1092,32 @@ def test_anthropic_citations_api():
     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=[
+    resp = completion(
+        model="claude-3-5-sonnet-20241022",
+        messages=[
+            {
+                "role": "user",
+                "content": [
                     {
-                        "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?",
-                            },
-                        ],
-                    }
+                        "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"])
+    citations = resp.choices[0].message.provider_specific_fields["citations"]
 
-        assert json_data["messages"][0]["content"][0]["citations"]["enabled"]
+    assert citations is not None