scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix: move gemini stream field merges to provider

Browse source
This commit is contained in:
Genmin 2026-04-30 15:43:09 -07:00
parent 79e4542703
commit 127d35a92d
3 changed files with 95 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/litellm_core_utils/streaming_provider_specific_fields.py Normal file
View file

@ -0,0 +1,27 @@
from typing import Any, Callable, Dict, Iterable
ProviderSpecificFieldMerger = Callable[[Dict[str, Any], str, Any], bool]
def _get_provider_specific_field_mergers() -> Iterable[ProviderSpecificFieldMerger]:
from litellm.llms.vertex_ai.gemini.streaming_provider_specific_fields import (
merge_gemini_streaming_provider_specific_field,
)
return (merge_gemini_streaming_provider_specific_field,)
def merge_streaming_provider_specific_field(
combined_provider_fields: Dict[str, Any], key: str, value: Any
) -> None:
for merge_field in _get_provider_specific_field_mergers():
if merge_field(combined_provider_fields, key, value):
return
if key not in combined_provider_fields:
combined_provider_fields[key] = value
elif isinstance(value, list) and isinstance(combined_provider_fields[key], list):
# For lists like web_search_results, take the last (most complete) one.
combined_provider_fields[key] = value
else:
combined_provider_fields[key] = value

63
litellm/llms/vertex_ai/gemini/streaming_provider_specific_fields.py Normal file
View file

@ -0,0 +1,63 @@
from typing import Any, Dict, List, Optional
def _merge_server_side_tool_invocations(existing: Optional[Any], incoming: Any) -> Any:
if not isinstance(incoming, list):
return incoming
if not isinstance(existing, list):
existing = []
merged_invocations: List[Any] = []
invocations_by_id: Dict[str, Dict[str, Any]] = {}
for invocation in existing:
if not isinstance(invocation, dict):
merged_invocations.append(invocation)
continue
invocation_copy = dict(invocation)
invocation_id = invocation_copy.get("id")
if isinstance(invocation_id, str) and invocation_id:
invocations_by_id[invocation_id] = invocation_copy
merged_invocations.append(invocation_copy)
for invocation in incoming:
if not isinstance(invocation, dict):
merged_invocations.append(invocation)
continue
invocation_id = invocation.get("id")
if isinstance(invocation_id, str) and invocation_id in invocations_by_id:
existing_invocation = invocations_by_id[invocation_id]
for key, value in invocation.items():
if key not in existing_invocation or existing_invocation[key] is None:
existing_invocation[key] = value
continue
invocation_copy = dict(invocation)
if isinstance(invocation_id, str) and invocation_id:
invocations_by_id[invocation_id] = invocation_copy
merged_invocations.append(invocation_copy)
return merged_invocations
def merge_gemini_streaming_provider_specific_field(
combined_provider_fields: Dict[str, Any], key: str, value: Any
) -> bool:
if key == "server_side_tool_invocations":
combined_provider_fields[key] = _merge_server_side_tool_invocations(
combined_provider_fields.get(key), value
)
return True
if key == "thought_signatures" and isinstance(value, list):
existing = combined_provider_fields.get(key)
if isinstance(existing, list):
existing.extend(value)
else:
combined_provider_fields[key] = list(value)
return True
return False

74
litellm/main.py
View file

@ -7373,74 +7373,6 @@ def stream_chunk_builder_text_completion(
return TextCompletionResponse(**response)
def _merge_server_side_tool_invocations(existing: Optional[Any], incoming: Any) -> Any:
if not isinstance(incoming, list):
return incoming
if not isinstance(existing, list):
existing = []
merged_invocations: List[Any] = []
invocations_by_id: Dict[str, Dict[str, Any]] = {}
for invocation in existing:
if not isinstance(invocation, dict):
merged_invocations.append(invocation)
continue
invocation_copy = dict(invocation)
invocation_id = invocation_copy.get("id")
if isinstance(invocation_id, str) and invocation_id:
invocations_by_id[invocation_id] = invocation_copy
merged_invocations.append(invocation_copy)
for invocation in incoming:
if not isinstance(invocation, dict):
merged_invocations.append(invocation)
continue
invocation_id = invocation.get("id")
if isinstance(invocation_id, str) and invocation_id in invocations_by_id:
existing_invocation = invocations_by_id[invocation_id]
for key, value in invocation.items():
if key not in existing_invocation or existing_invocation[key] is None:
existing_invocation[key] = value
continue
invocation_copy = dict(invocation)
if isinstance(invocation_id, str) and invocation_id:
invocations_by_id[invocation_id] = invocation_copy
merged_invocations.append(invocation_copy)
return merged_invocations
def _merge_streaming_provider_specific_field(
combined_provider_fields: Dict[str, Any], key: str, value: Any
) -> None:
if key == "server_side_tool_invocations":
combined_provider_fields[key] = _merge_server_side_tool_invocations(
combined_provider_fields.get(key), value
)
return
if key == "thought_signatures" and isinstance(value, list):
existing = combined_provider_fields.get(key)
if isinstance(existing, list):
existing.extend(value)
else:
combined_provider_fields[key] = list(value)
return
if key not in combined_provider_fields:
combined_provider_fields[key] = value
elif isinstance(value, list) and isinstance(combined_provider_fields[key], list):
# For lists like web_search_results, take the last (most complete) one
combined_provider_fields[key] = value
else:
combined_provider_fields[key] = value
def stream_chunk_builder( # noqa: PLR0915
chunks: list,
messages: Optional[list] = None,
@ -7680,8 +7612,12 @@ def stream_chunk_builder( # noqa: PLR0915
for chunk in provider_specific_chunks:
fields = chunk["choices"][0]["delta"]["provider_specific_fields"]
if isinstance(fields, dict):
from litellm.litellm_core_utils.streaming_provider_specific_fields import (
merge_streaming_provider_specific_field,
)
for key, value in fields.items():
_merge_streaming_provider_specific_field(
merge_streaming_provider_specific_field(
combined_provider_fields, key, value
)