From ca6d77b479b771b53ea0cf0107719db06b06165c Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Mon, 1 Sep 2025 16:41:35 +0900 Subject: [PATCH] fix citation field name --- .../llms/databricks/chat/transformation.py | 28 +++++---- .../test_databricks_chat_transformation.py | 57 +++++++++++++++++-- 2 files changed, 69 insertions(+), 16 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 5600d5c642..9330b01923 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -26,7 +26,6 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( - handle_messages_with_content_list_to_str_conversion, strip_name_from_messages, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -301,7 +300,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: """ Databricks does not support: - - content in list format. - 'name' in user message. """ new_messages = [] @@ -311,7 +309,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): else: _message = message new_messages.append(_message) - new_messages = handle_messages_with_content_list_to_str_conversion(new_messages) new_messages = strip_name_from_messages(new_messages) if is_async: @@ -388,10 +385,16 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): citations: Optional[List[Any]] = None if isinstance(content, list): for item in content: + text = item.get("text", None) if item.get("citations") is not None: if citations is None: citations = [] - citations.append(item["citations"]) + citations.append( + [ + {**citation, "supported_text": text} + for citation in item["citations"] + ] + ) return citations def _transform_dbrx_choices( @@ -583,12 +586,17 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): for _tc in tool_calls: if _tc.get("function", {}).get("arguments") == "{}": _tc["function"]["arguments"] = "" # avoid invalid json - citation = choice["delta"].get("citation") - if citation is not None: - choice["delta"].setdefault("provider_specific_fields", {})[ - "citation" - ] = citation - choice["delta"].pop("citation", None) + if isinstance(choice["delta"]["content"], list) and ( + content := choice["delta"]["content"] + ): + if citations := content[0].get("citations"): + # TODO: Databricks delta does not include supported text or chunk type. + # Add either here once Databricks supports it to enable citation linkage. + choice["delta"].setdefault("provider_specific_fields", {})[ + "citation" + ] = citations[ + 0 + ] # Databricks Content item always has citation as a list of list # extract the content str content_str = DatabricksConfig.extract_content_str( choice["delta"].get("content") diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index d61f826e89..51a2e971c0 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -88,7 +88,7 @@ def test_transform_choices_without_signature(): assert choices[0].message.reasoning_content == "i'm thinking without signature." assert choices[0].message.thinking_blocks is not None assert len(choices[0].message.thinking_blocks) == 1 - + # Verify the thinking block was created successfully without signature thinking_block = choices[0].message.thinking_blocks[0] assert thinking_block["type"] == "thinking" @@ -104,8 +104,17 @@ def test_transform_choices_with_citations(): "content": [ { "type": "text", - "text": "Paris", - "citations": [{"source": "wiki"}], + "text": "Blue", + "citations": [ + { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + } + ], } ], }, @@ -117,7 +126,19 @@ def test_transform_choices_with_citations(): choices = config._transform_dbrx_choices(choices=databricks_choices) assert choices[0].message.provider_specific_fields == { - "citations": [[{"source": "wiki"}]] + "citations": [ + [ + { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + "supported_text": "Blue", + } + ] + ] } @@ -130,7 +151,24 @@ def test_chunk_parser_with_citation(): "model": "test", "choices": [ { - "delta": {"citation": {"source": "wiki"}}, + "delta": { + "content": [ + { + "type": "text", + "text": "", + "citations": [ + { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + } + ], + } + ], + }, "index": 0, "finish_reason": None, } @@ -139,5 +177,12 @@ def test_chunk_parser_with_citation(): parsed = iterator.chunk_parser(chunk) assert parsed.choices[0].delta.provider_specific_fields == { - "citation": {"source": "wiki"} + "citation": { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + } }