fix citation field name
This commit is contained in:
TomeHirata
parent
38a1dbd13a
commit
ca6d77b479
@ -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")
|
||||
|
||||
@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user