Unverified Commit de584807 authored by Chenhe Gu's avatar Chenhe Gu Committed by GitHub

fix streaming (#1944)

parent a1285cbf
......@@ -337,9 +337,9 @@ class OAIAPICompatLargeLanguageModel(_CommonOAI_API_Compat, LargeLanguageModel):
)
)
for chunk in response.iter_content(chunk_size=2048):
for chunk in response.iter_lines(decode_unicode=True, delimiter='\n\n'):
if chunk:
decoded_chunk = chunk.decode('utf-8').strip().lstrip('data: ').lstrip()
decoded_chunk = chunk.strip().lstrip('data: ').lstrip()
chunk_json = None
try:
......@@ -356,7 +356,7 @@ class OAIAPICompatLargeLanguageModel(_CommonOAI_API_Compat, LargeLanguageModel):
continue
choice = chunk_json['choices'][0]
chunk_index = choice['index'] if 'index' in choice else chunk_index
chunk_index += 1
if 'delta' in choice:
delta = choice['delta']
......@@ -408,12 +408,6 @@ class OAIAPICompatLargeLanguageModel(_CommonOAI_API_Compat, LargeLanguageModel):
message=assistant_prompt_message,
)
)
else:
yield create_final_llm_result_chunk(
index=chunk_index + 1,
message=AssistantPromptMessage(content=""),
finish_reason="End of stream."
)
chunk_index += 1
......
......@@ -22,7 +22,7 @@ def test_validate_credentials():
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
'api_key': 'invalid_key',
'endpoint_url': 'https://api.together.xyz/v1/chat/completions',
'endpoint_url': 'https://api.together.xyz/v1/',
'mode': 'chat'
}
)
......@@ -31,7 +31,7 @@ def test_validate_credentials():
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'),
'endpoint_url': 'https://api.together.xyz/v1/chat/completions',
'endpoint_url': 'https://api.together.xyz/v1/',
'mode': 'chat'
}
)
......@@ -43,7 +43,7 @@ def test_invoke_model():
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'),
'endpoint_url': 'https://api.together.xyz/v1/completions',
'endpoint_url': 'https://api.together.xyz/v1/',
'mode': 'completion'
},
prompt_messages=[
......@@ -74,7 +74,7 @@ def test_invoke_stream_model():
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'),
'endpoint_url': 'https://api.together.xyz/v1/chat/completions',
'endpoint_url': 'https://api.together.xyz/v1/',
'mode': 'chat'
},
prompt_messages=[
......@@ -110,7 +110,7 @@ def test_invoke_chat_model_with_tools():
model='gpt-3.5-turbo',
credentials={
'api_key': os.environ.get('OPENAI_API_KEY'),
'endpoint_url': 'https://api.openai.com/v1/chat/completions',
'endpoint_url': 'https://api.openai.com/v1/',
'mode': 'chat'
},
prompt_messages=[
......@@ -165,7 +165,7 @@ def test_get_num_tokens():
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
'api_key': os.environ.get('OPENAI_API_KEY'),
'endpoint_url': 'https://api.openai.com/v1/chat/completions'
'endpoint_url': 'https://api.openai.com/v1/'
},
prompt_messages=[
SystemPromptMessage(
......
......@@ -18,9 +18,8 @@ def test_validate_credentials():
model='text-embedding-ada-002',
credentials={
'api_key': 'invalid_key',
'endpoint_url': 'https://api.openai.com/v1/embeddings',
'context_size': 8184,
'max_chunks': 32
'endpoint_url': 'https://api.openai.com/v1/',
'context_size': 8184
}
)
......@@ -29,9 +28,8 @@ def test_validate_credentials():
model='text-embedding-ada-002',
credentials={
'api_key': os.environ.get('OPENAI_API_KEY'),
'endpoint_url': 'https://api.openai.com/v1/embeddings',
'context_size': 8184,
'max_chunks': 32
'endpoint_url': 'https://api.openai.com/v1/',
'context_size': 8184
}
)
......@@ -43,9 +41,8 @@ def test_invoke_model():
model='text-embedding-ada-002',
credentials={
'api_key': os.environ.get('OPENAI_API_KEY'),
'endpoint_url': 'https://api.openai.com/v1/embeddings',
'context_size': 8184,
'max_chunks': 32
'endpoint_url': 'https://api.openai.com/v1/',
'context_size': 8184
},
texts=[
"hello",
......@@ -67,8 +64,7 @@ def test_get_num_tokens():
credentials={
'api_key': os.environ.get('OPENAI_API_KEY'),
'endpoint_url': 'https://api.openai.com/v1/embeddings',
'context_size': 8184,
'max_chunks': 32
'context_size': 8184
},
texts=[
"hello",
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment