Add retry logic for OpenAI chat completions

Author: chatrtham

src/fastapi_app/rag_advanced.py

@@ -1,15 +1,13 @@
 import copy
+import logging
 import pathlib
 from collections.abc import AsyncGenerator
-from typing import (
-    Any,
-)
+from typing import Any
 
 from openai import AsyncOpenAI
-from openai.types.chat import (
-    ChatCompletion,
-)
+from openai.types.chat import ChatCompletion
 from openai_messages_token_helper import get_token_limit
+from tenacity import before_sleep_log, retry, stop_after_attempt, wait_random_exponential
 
 from .api_models import ThoughtStep
 from .postgres_searcher import PostgresSearcher
@@ -20,6 +18,9 @@
     handle_specify_package_function_call,
 )
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class AdvancedRAGChat:
     def __init__(
@@ -40,13 +41,17 @@ def __init__(
         self.query_prompt_template = open(current_dir / "prompts/query.txt").read()
         self.answer_prompt_template = open(current_dir / "prompts/answer.txt").read()
 
+    @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6), before_sleep=before_sleep_log(logger, logging.WARNING))
+    async def openai_chat_completion(self, *args, **kwargs) -> ChatCompletion:
+        return await self.openai_chat_client.chat.completions.create(*args, **kwargs)
+
     async def hybrid_search(self, messages, top, vector_search, text_search):
         # Generate an optimized keyword search query based on the chat history and the last question
         query_messages = copy.deepcopy(messages)
         query_messages.insert(0, {"role": "system", "content": self.query_prompt_template})
         query_response_token_limit = 500
 
-        query_chat_completion: ChatCompletion = await self.openai_chat_client.chat.completions.create(
+        query_chat_completion: ChatCompletion = await self.openai_chat_completion(
             messages=query_messages,
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
             temperature=0.0,
@@ -110,7 +115,7 @@ async def run(
         specify_package_messages.insert(0, {"role": "system", "content": self.specify_package_prompt_template})
         specify_package_token_limit = 300
 
-        specify_package_chat_completion: ChatCompletion = await self.openai_chat_client.chat.completions.create(
+        specify_package_chat_completion: ChatCompletion = await self.openai_chat_completion(
             messages=specify_package_messages,
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
             temperature=0.0,
@@ -155,9 +160,9 @@ async def run(
         # Build messages for the final chat completion
         messages.insert(0, {"role": "system", "content": self.answer_prompt_template})
         messages[-1]["content"].append({"type": "text", "text": "\n\nSources:\n" + content})
-        response_token_limit = 1024
+        response_token_limit = 4096
 
-        chat_completion_response = await self.openai_chat_client.chat.completions.create(
+        chat_completion_response = await self.openai_chat_completion(
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
             messages=messages,
             temperature=overrides.get("temperature", 0.3),

src/requirements.txt

@@ -21,7 +21,7 @@ email_validator==2.1.1
 environs==11.0.0
 fastapi==0.111.0
 fastapi-cli==0.0.4
--e git+https://github.com/azure-samples/rag-postgres-openai-python@1b189c6a227119d31a3947afc36d229cc0b2ac58#egg=fastapi_app&subdirectory=src
+-e git+https://github.com/chatrtham/rag-postgres-openai-python.git@7a91d2ab7d3814bb2ed6286a8b89255405309e94#egg=fastapi_app&subdirectory=src
 filelock==3.14.0
 frozenlist==1.4.1
 gitdb==4.0.11
@@ -76,6 +76,7 @@ smmap==5.0.1
 sniffio==1.3.1
 SQLAlchemy==2.0.30
 starlette==0.37.2
+tenacity==8.4.1
 tiktoken==0.7.0
 tqdm==4.66.4
 typer==0.12.3