Update models/gpt4_model.py

models/gpt4_model.py

@@ -1,7 +1,11 @@
 from .base_model import BaseModel
 
 import openai
+from openai import AsyncOpenAI, OpenAI
 from tqdm import tqdm
+import asyncio
+import os
+
 class GPT4Model(BaseModel):
     def __init__(self, 
                  generation_model="gpt-4-vision-preview", 
@@ -12,46 +16,52 @@ class GPT4Model(BaseModel):
         self.embedding_model = embedding_model
         self.temperature = temperature
     
-    def respond(self, messages: list) -> str:
-        try:
-            response = openai.ChatCompletion.create(
-                messages=messages,
-                model=self.generation_model,
-                temperature=self.temperature,
-                max_tokens=1000,
-            ).choices[0]['message']['content']
-        except:
-            try:
-                response = openai.ChatCompletion.create(
-                    messages=messages,
-                    model=self.generation_model,
-                    temperature=self.temperature,
-                    max_tokens=1000,
-                ).choices[0]['message']['content']
-            except:
-                try:
-                    response = openai.ChatCompletion.create(
-                        messages=messages,
-                        model=self.generation_model,
-                        temperature=self.temperature,
-                        max_tokens=1000,
-                    ).choices[0]['message']['content']
-                except:
-                    response = "No answer was provided."
+    async def respond_async(self, messages: list[dict]) -> str:
+        client = AsyncOpenAI(
+            api_key=os.environ["OPENAI_API_KEY"],
+            base_url=os.environ["OPENAI_API_BASE"]
+        )
+        print("start api call")
+        output = await client.chat.completions.create(
+            messages=messages,
+            model=self.generation_model,
+            temperature=self.temperature,
+            max_tokens=1000,
+        )
+        print("end api call")
+        response = output.choices[0].message.content
         # content = response.choices[0]['message']['content']
         
         return response
     
-    def embedding(self, texts: list) -> list:
+    def respond(self, messages: list[dict]) -> str:
+        client = OpenAI(
+            api_key=os.environ["OPENAI_API_KEY"],
+            base_url=os.environ["OPENAI_API_BASE"]
+        )
+        # OpenAI.api_key=os.environ["OPENAI_API_KEY"]
+        # OpenAI.api_base=os.environ["OPENAI_API_BASE"]
+        response = client.chat.completions.create(
+            messages=messages,
+            model=self.generation_model,
+            temperature=self.temperature,
+            max_tokens=1000,
+        ).choices[0].message.content
+        return response
+    
+    def embedding(self, texts: list[str]) -> list[float]:
+        client = OpenAI(
+            api_key=os.environ["OPENAI_API_KEY"],
+            base_url=os.environ["OPENAI_API_BASE"]
+        )
         data = []
         # print(f"{self.embedding_model} Embedding:")
         for i in range(0, len(texts), 2048):
             lower = i
             upper = min(i+2048, len(texts))
-            data += openai.Embedding.create(input=texts[lower:upper],
+            data += client.embeddings.create(input=texts[lower:upper],
                                             model=self.embedding_model
-                                            )["data"]
-            
-        embeddings = [d["embedding"] for d in data]
+                                            ).data
+        embeddings = [d.embedding for d in data]
         
         return embeddings