klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-02 21:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Gracefully close thread when there's an exception in the openai llm thread. Closes #894.

Browse Source
This commit is contained in:
sabaimran
2024-09-03 13:14:55 -07:00
parent 6ed68b574b
commit 17901406aa
1 changed files with 32 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
src/khoj/processor/conversation/openai/utils.py
View File

@@ -100,34 +100,37 @@ def chat_completion_with_backoff(
def llm_thread(g, messages, model_name, temperature, openai_api_key=None, api_base_url=None, model_kwargs=None):
client_key = f"{openai_api_key}--{api_base_url}"
if client_key not in openai_clients:
client: openai.OpenAI = openai.OpenAI(
api_key=openai_api_key,
base_url=api_base_url,
try:
client_key = f"{openai_api_key}--{api_base_url}"
if client_key not in openai_clients:
client: openai.OpenAI = openai.OpenAI(
api_key=openai_api_key,
base_url=api_base_url,
)
openai_clients[client_key] = client
else:
client: openai.OpenAI = openai_clients[client_key]
formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
chat = client.chat.completions.create(
stream=True,
messages=formatted_messages,
model=model_name, # type: ignore
temperature=temperature,
timeout=20,
**(model_kwargs or dict()),
)
openai_clients[client_key] = client
else:
client: openai.OpenAI = openai_clients[client_key]
formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
chat = client.chat.completions.create(
stream=True,
messages=formatted_messages,
model=model_name, # type: ignore
temperature=temperature,
timeout=20,
**(model_kwargs or dict()),
)
for chunk in chat:
if len(chunk.choices) == 0:
continue
delta_chunk = chunk.choices[0].delta
if isinstance(delta_chunk, str):
g.send(delta_chunk)
elif delta_chunk.content:
g.send(delta_chunk.content)
g.close()
for chunk in chat:
if len(chunk.choices) == 0:
continue
delta_chunk = chunk.choices[0].delta
if isinstance(delta_chunk, str):
g.send(delta_chunk)
elif delta_chunk.content:
g.send(delta_chunk.content)
except Exception as e:
logger.error(f"Error in llm_thread: {e}")
finally:
g.close()