forked from phoenix/litellm-mirror
fix databricks streaming test
This commit is contained in:
parent
5d1fcc545b
commit
51da6ab64e
2 changed files with 21 additions and 2 deletions
16
litellm/litellm_core_utils/streaming_utils.py
Normal file
16
litellm/litellm_core_utils/streaming_utils.py
Normal file
|
@ -0,0 +1,16 @@
|
||||||
|
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||||
|
|
||||||
|
|
||||||
|
def generic_chunk_has_all_required_fields(chunk: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Checks if the provided chunk dictionary contains all required fields for GenericStreamingChunk.
|
||||||
|
|
||||||
|
:param chunk: The dictionary to check.
|
||||||
|
:return: True if all required fields are present, False otherwise.
|
||||||
|
"""
|
||||||
|
_all_fields = GChunk.__annotations__
|
||||||
|
|
||||||
|
# this is an optional field in GenericStreamingChunk, it's not required to be present
|
||||||
|
_all_fields.pop("provider_specific_fields", None)
|
||||||
|
|
||||||
|
return all(key in chunk for key in _all_fields)
|
|
@ -9565,12 +9565,15 @@ class CustomStreamWrapper:
|
||||||
try:
|
try:
|
||||||
# return this for all models
|
# return this for all models
|
||||||
completion_obj = {"content": ""}
|
completion_obj = {"content": ""}
|
||||||
|
from litellm.litellm_core_utils.streaming_utils import (
|
||||||
|
generic_chunk_has_all_required_fields,
|
||||||
|
)
|
||||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||||
|
|
||||||
if (
|
if (
|
||||||
isinstance(chunk, dict)
|
isinstance(chunk, dict)
|
||||||
and all(
|
and generic_chunk_has_all_required_fields(
|
||||||
key in chunk for key in GChunk.__annotations__
|
chunk=chunk
|
||||||
) # check if chunk is a generic streaming chunk
|
) # check if chunk is a generic streaming chunk
|
||||||
) or (
|
) or (
|
||||||
self.custom_llm_provider
|
self.custom_llm_provider
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue