Update services/utils.py

services/utils.py

@@ -19,7 +19,11 @@ import json
 from typing import Optional
 import tiktoken
 import time
-
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 #Setting the openai api key
 api_key=os.getenv('OPENAI_API_KEY')
@@ -54,6 +58,11 @@ class Candidate(BaseModel):
     payment_method: Optional[str] = Field(default=None, description="INSERT PAYMENT METHOD FROM THE RECEIPT OCR TEXT. LOOK FOR KEYWORDS LIKE 'CASH', 'CARD', 'CREDIT', 'DEBIT', 'VISA', 'MASTERCARD', 'AMEX', 'PAYPAL', ETC. IF NOT FOUND RETURN null.")
 
 
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+def openai_response(model:OpenAI,input:str):
+    result = model.invoke(input)
+    return result
+
     
 def strcuture_document_data(raw_text:str)->dict:
     
@@ -79,16 +88,12 @@ def strcuture_document_data(raw_text:str)->dict:
             input_variables=["query"],
             partial_variables={"format_instructions": parser.get_format_instructions(),"raw_text":raw_text},
         )
-        # print("parser.get_format_instructions()")
-        # print(parser.get_format_instructions())
+
         input = prompt.format_prompt(query=doc_query)
-        with get_openai_callback() as cb:
-            result = model.invoke(input.to_string())
-            
-        time.sleep(0.25)
+
+        result = openai_response(model,input.to_string())   
         print(f"GPT Response {result}")
-        # result = extract_json_from_string(result)
-        # print(f"Formatted Response : {result}")
+
         
         class_object= parser.parse(result)
         dict_object=class_object.__dict__