anthropic and together ai streaming no longer waits till completion to return response

This commit is contained in:
Krrish Dholakia 2023-08-26 19:00:53 -07:00
parent 95ce560f6d
commit 320fe89a24
6 changed files with 78 additions and 114 deletions

View file

@ -35,8 +35,6 @@ from litellm.utils import (
) )
from litellm.utils import ( from litellm.utils import (
get_ollama_response_stream, get_ollama_response_stream,
stream_to_string,
together_ai_completion_streaming,
) )
####### ENVIRONMENT VARIABLES ################### ####### ENVIRONMENT VARIABLES ###################
@ -542,46 +540,56 @@ def completion(
logging.pre_call(input=prompt, api_key=TOGETHER_AI_TOKEN) logging.pre_call(input=prompt, api_key=TOGETHER_AI_TOKEN)
print(f"TOGETHER_AI_TOKEN: {TOGETHER_AI_TOKEN}") print(f"TOGETHER_AI_TOKEN: {TOGETHER_AI_TOKEN}")
res = requests.post(
endpoint,
json={
"model": model,
"prompt": prompt,
"request_type": "language-model-inference",
**optional_params,
},
headers=headers,
)
if "stream_tokens" in optional_params and optional_params["stream_tokens"] == True: if "stream_tokens" in optional_params and optional_params["stream_tokens"] == True:
res = requests.post(
endpoint,
json={
"model": model,
"prompt": prompt,
"request_type": "language-model-inference",
**optional_params,
},
stream=optional_params["stream_tokens"],
headers=headers,
)
response = CustomStreamWrapper( response = CustomStreamWrapper(
res.iter_lines(), model, custom_llm_provider="together_ai" res.iter_lines(), model, custom_llm_provider="together_ai"
) )
return response return response
## LOGGING else:
logging.post_call( res = requests.post(
input=prompt, api_key=TOGETHER_AI_TOKEN, original_response=res.text endpoint,
) json={
# make this safe for reading, if output does not exist raise an error "model": model,
json_response = res.json() "prompt": prompt,
if "output" not in json_response: "request_type": "language-model-inference",
raise Exception( **optional_params,
f"liteLLM: Error Making TogetherAI request, JSON Response {json_response}" },
headers=headers,
) )
completion_response = json_response["output"]["choices"][0]["text"] ## LOGGING
prompt_tokens = len(encoding.encode(prompt)) logging.post_call(
completion_tokens = len(encoding.encode(completion_response)) input=prompt, api_key=TOGETHER_AI_TOKEN, original_response=res.text
## RESPONSE OBJECT )
model_response["choices"][0]["message"]["content"] = completion_response # make this safe for reading, if output does not exist raise an error
model_response["created"] = time.time() json_response = res.json()
model_response["model"] = model if "output" not in json_response:
model_response["usage"] = { raise Exception(
"prompt_tokens": prompt_tokens, f"liteLLM: Error Making TogetherAI request, JSON Response {json_response}"
"completion_tokens": completion_tokens, )
"total_tokens": prompt_tokens + completion_tokens, completion_response = json_response["output"]["choices"][0]["text"]
} prompt_tokens = len(encoding.encode(prompt))
response = model_response completion_tokens = len(encoding.encode(completion_response))
## RESPONSE OBJECT
model_response["choices"][0]["message"]["content"] = completion_response
model_response["created"] = time.time()
model_response["model"] = model
model_response["usage"] = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
response = model_response
elif model in litellm.vertex_chat_models: elif model in litellm.vertex_chat_models:
# import vertexai/if it fails then pip install vertexai# import cohere/if it fails then pip install cohere # import vertexai/if it fails then pip install vertexai# import cohere/if it fails then pip install cohere
install_and_import("vertexai") install_and_import("vertexai")

View file

@ -62,27 +62,22 @@ messages = [{"content": user_message, "role": "user"}]
# test on anthropic completion call # test on anthropic completion call
try: # try:
response = completion( # response = completion(
model="claude-instant-1", messages=messages, stream=True, logger_fn=logger_fn # model="claude-instant-1", messages=messages, stream=True, logger_fn=logger_fn
) # )
complete_response = "" # complete_response = ""
start_time = time.time() # start_time = time.time()
time_since_initial_request = [] # for chunk in response:
for chunk in response: # chunk_time = time.time()
chunk_time = time.time() # print(f"time since initial request: {chunk_time - start_time:.5f}")
time_since_initial_request.append(chunk_time - start_time) # print(chunk["choices"][0]["delta"])
print(f"time since initial request: {chunk_time - start_time:.5f}") # complete_response += chunk["choices"][0]["delta"]["content"]
print(chunk["choices"][0]["delta"]) # if complete_response == "":
complete_response += chunk["choices"][0]["delta"]["content"] # raise Exception("Empty response received")
if complete_response == "": # except:
raise Exception("Empty response received") # print(f"error occurred: {traceback.format_exc()}")
print(f"set(time_since_initial_request): {set(time_since_initial_request)}") # pass
if len(set(time_since_initial_request)) == 1:
raise Exception("All time since initial request is the same")
except:
print(f"error occurred: {traceback.format_exc()}")
pass
# # test on huggingface completion call # # test on huggingface completion call
@ -104,23 +99,23 @@ except:
# pass # pass
# test on together ai completion call # test on together ai completion call
# try: try:
# start_time = time.time() start_time = time.time()
# response = completion( response = completion(
# model="Replit-Code-3B", messages=messages, logger_fn=logger_fn, stream= True model="Replit-Code-3B", messages=messages, logger_fn=logger_fn, stream= True
# ) )
# complete_response = "" complete_response = ""
# print(f"returned response object: {response}") print(f"returned response object: {response}")
# for chunk in response: for chunk in response:
# chunk_time = time.time() chunk_time = time.time()
# print(f"time since initial request: {chunk_time - start_time:.2f}") print(f"time since initial request: {chunk_time - start_time:.2f}")
# print(chunk["choices"][0]["delta"]) print(chunk["choices"][0]["delta"])
# complete_response += chunk["choices"][0]["delta"]["content"] if len(chunk["choices"][0]["delta"].keys()) > 0 else "" complete_response += chunk["choices"][0]["delta"]["content"] if len(chunk["choices"][0]["delta"].keys()) > 0 else ""
# if complete_response == "": if complete_response == "":
# raise Exception("Empty response received") raise Exception("Empty response received")
# except: except:
# print(f"error occurred: {traceback.format_exc()}") print(f"error occurred: {traceback.format_exc()}")
# pass pass
# # test on azure completion call # # test on azure completion call

View file

@ -1593,45 +1593,6 @@ async def stream_to_string(generator):
return response return response
########## Together AI streaming ############################# [TODO] move together ai to it's own llm class
async def together_ai_completion_streaming(json_data, headers):
session = aiohttp.ClientSession()
url = "https://api.together.xyz/inference"
# headers = {
# 'Authorization': f'Bearer {together_ai_token}',
# 'Content-Type': 'application/json'
# }
# data = {
# "model": "togethercomputer/llama-2-70b-chat",
# "prompt": "write 1 page on the topic of the history of the united state",
# "max_tokens": 1000,
# "temperature": 0.7,
# "top_p": 0.7,
# "top_k": 50,
# "repetition_penalty": 1,
# "stream_tokens": True
# }
try:
async with session.post(url, json=json_data, headers=headers) as resp:
async for line in resp.content.iter_any():
# print(line)
if line:
try:
json_chunk = line.decode("utf-8")
json_string = json_chunk.split("data: ")[1]
# Convert the JSON string to a dictionary
data_dict = json.loads(json_string)
completion_response = data_dict["choices"][0]["text"]
completion_obj = {"role": "assistant", "content": ""}
completion_obj["content"] = completion_response
yield {"choices": [{"delta": completion_obj}]}
except:
pass
finally:
await session.close()
def completion_with_fallbacks(**kwargs): def completion_with_fallbacks(**kwargs):
response = None response = None
rate_limited_models = set() rate_limited_models = set()

View file

@ -1,6 +1,6 @@
[tool.poetry] [tool.poetry]
name = "litellm" name = "litellm"
version = "0.1.489" version = "0.1.490"
description = "Library to easily interface with LLM API providers" description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"] authors = ["BerriAI"]
license = "MIT License" license = "MIT License"