phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Clarifai - Added streaming and async completion support

Browse source
This commit is contained in:
mogith-pn 2024-05-03 14:03:38 +00:00
parent d770df2259
commit 723ef9963e
6 changed files with 259 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

28
litellm/utils.py
View file

@ -2807,6 +2807,7 @@ def client(original_function):
)
else:
return result
return result
# Prints Exactly what was passed to litellm function - don't execute any logic here - it should just print
@ -2910,6 +2911,7 @@ def client(original_function):
model_response_object=ModelResponse(),
stream=kwargs.get("stream", False),
)
if kwargs.get("stream", False) == True:
cached_result = CustomStreamWrapper(
completion_stream=cached_result,
@ -9905,6 +9907,27 @@ class CustomStreamWrapper:
return {"text": "", "is_finished": False}
except Exception as e:
raise e
def handle_clarifai_completion_chunk(self, chunk):
try:
if isinstance(chunk, dict):
parsed_response = chunk
if isinstance(chunk, (str, bytes)):
if isinstance(chunk, bytes):
parsed_response = chunk.decode("utf-8")
else:
parsed_response = chunk
data_json = json.loads(parsed_response)
text = data_json.get("outputs", "")[0].get("data", "").get("text", "").get("raw","")
prompt_tokens = len(encoding.encode(data_json.get("outputs", "")[0].get("input","").get("data", "").get("text", "").get("raw","")))
completion_tokens = len(encoding.encode(text))
return {
"text": text,
"is_finished": True,
}
except:
traceback.print_exc()
return ""
def model_response_creator(self):
model_response = ModelResponse(stream=True, model=self.model)
@ -9949,6 +9972,11 @@ class CustomStreamWrapper:
completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
elif (
self.custom_llm_provider and self.custom_llm_provider == "clarifai"
):
response_obj = self.handle_clarifai_completion_chunk(chunk)
completion_obj["content"] = response_obj["text"]
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
response_obj = self.handle_replicate_chunk(chunk)
completion_obj["content"] = response_obj["text"]