From 29147cf37c0760d99527f2cbd5d0606e30cce4df Mon Sep 17 00:00:00 2001
From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
Date: Wed, 14 May 2025 15:11:51 -0400
Subject: [PATCH] fix: Retry and timeout logic

---
 src/crud.py               | 24 ++++++++++++------------
 src/utils/model_client.py |  4 +++-
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/crud.py b/src/crud.py
index b5ac4f60..eecc28cd 100644
--- a/src/crud.py
+++ b/src/crud.py
@@ -1346,9 +1346,9 @@ async def query_documents(
     top_k: int = 5,
 ) -> Sequence[models.Document]:
     # Using async client with await
-    response = await openai_client.embeddings.create(
-        model="text-embedding-3-small", input=query
-    )
+    response = await openai_client.with_options(
+        timeout=10, max_retries=3
+    ).embeddings.create(model="text-embedding-3-small", input=query)
     embedding_query = response.data[0].embedding
     stmt = (
         select(models.Document)
@@ -1402,9 +1402,9 @@ async def create_document(
     )
 
     # Using async client with await
-    response = await openai_client.embeddings.create(
-        input=document.content, model="text-embedding-3-small"
-    )
+    response = await openai_client.with_options(
+        timeout=10, max_retries=3
+    ).embeddings.create(input=document.content, model="text-embedding-3-small")
 
     embedding = response.data[0].embedding
 
@@ -1460,9 +1460,9 @@ async def update_document(
     if document.content is not None:
         honcho_document.content = document.content
         # Using async client with await
-        response = await openai_client.embeddings.create(
-            input=document.content, model="text-embedding-3-small"
-        )
+        response = await openai_client.with_options(
+            timeout=10, max_retries=3
+        ).embeddings.create(input=document.content, model="text-embedding-3-small")
         embedding = response.data[0].embedding
         honcho_document.embedding = embedding
         honcho_document.created_at = func.now()
@@ -1531,9 +1531,9 @@ async def get_duplicate_documents(
     """
     # Get embedding for the content
     # Using async client with await
-    response = await openai_client.embeddings.create(
-        input=content, model="text-embedding-3-small"
-    )
+    response = await openai_client.with_options(
+        timeout=10, max_retries=3
+    ).embeddings.create(input=content, model="text-embedding-3-small")
     embedding = response.data[0].embedding
 
     # Find documents with similar embeddings
diff --git a/src/utils/model_client.py b/src/utils/model_client.py
index 8c98a800..a802548d 100644
--- a/src/utils/model_client.py
+++ b/src/utils/model_client.py
@@ -246,7 +246,9 @@ class ModelClient:
         formatted_messages.extend(messages)
 
         # Make the API call with the OpenAI client
-        response = await self.openai_client.chat.completions.create(
+        response = await self.openai_client.with_options(
+            timeout=10.0, max_retries=2
+        ).chat.completions.create(
             model=self.model,
             messages=formatted_messages,
             max_tokens=max_tokens,