add base url for OpenAI (#166)

rag/llm/chat_model.py

@@ -31,8 +31,9 @@ class Base(ABC):
 
 
 class GptTurbo(Base):
-    def __init__(self, key, model_name="gpt-3.5-turbo"):
-        self.client = OpenAI(api_key=key)
+    def __init__(self, key, model_name="gpt-3.5-turbo", base_url="https://api.openai.com/v1"):
+        if not base_url: base_url="https://api.openai.com/v1"
+        self.client = OpenAI(api_key=key, base_url=base_url)
         self.model_name = model_name
 
     def chat(self, system, history, gen_conf):
@@ -53,9 +54,10 @@ class GptTurbo(Base):
 
 
 class MoonshotChat(GptTurbo):
-    def __init__(self, key, model_name="moonshot-v1-8k"):
+    def __init__(self, key, model_name="moonshot-v1-8k", base_url="https://api.moonshot.cn/v1"):
+        if not base_url: base_url="https://api.moonshot.cn/v1"
         self.client = OpenAI(
-            api_key=key, base_url="https://api.moonshot.cn/v1",)
+            api_key=key, base_url=base_url)
         self.model_name = model_name
 
     def chat(self, system, history, gen_conf):
@@ -76,7 +78,7 @@ class MoonshotChat(GptTurbo):
 
 
 class QWenChat(Base):
-    def __init__(self, key, model_name=Generation.Models.qwen_turbo):
+    def __init__(self, key, model_name=Generation.Models.qwen_turbo, **kwargs):
         import dashscope
         dashscope.api_key = key
         self.model_name = model_name
@@ -105,7 +107,7 @@ class QWenChat(Base):
 
 
 class ZhipuChat(Base):
-    def __init__(self, key, model_name="glm-3-turbo"):
+    def __init__(self, key, model_name="glm-3-turbo", **kwargs):
         self.client = ZhipuAI(api_key=key)
         self.model_name = model_name
 
@@ -154,7 +156,7 @@ class LocalLLM(Base):
 
             return do_rpc
 
-    def __init__(self, key, model_name="glm-3-turbo"):
+    def __init__(self, **kwargs):
         self.client = LocalLLM.RPCProxy("127.0.0.1", 7860)
 
     def chat(self, system, history, gen_conf):

rag/llm/cv_model.py

@@ -67,8 +67,9 @@ class Base(ABC):
 
 
 class GptV4(Base):
-    def __init__(self, key, model_name="gpt-4-vision-preview", lang="Chinese"):
-        self.client = OpenAI(api_key=key)
+    def __init__(self, key, model_name="gpt-4-vision-preview", lang="Chinese", base_url="https://api.openai.com/v1"):
+        if not base_url: base_url="https://api.openai.com/v1"
+        self.client = OpenAI(api_key=key, base_url=base_url)
         self.model_name = model_name
         self.lang = lang
 
@@ -84,7 +85,7 @@ class GptV4(Base):
 
 
 class QWenCV(Base):
-    def __init__(self, key, model_name="qwen-vl-chat-v1", lang="Chinese"):
+    def __init__(self, key, model_name="qwen-vl-chat-v1", lang="Chinese", **kwargs):
         import dashscope
         dashscope.api_key = key
         self.model_name = model_name
@@ -123,7 +124,7 @@ class QWenCV(Base):
 
 
 class Zhipu4V(Base):
-    def __init__(self, key, model_name="glm-4v", lang="Chinese"):
+    def __init__(self, key, model_name="glm-4v", lang="Chinese", **kwargs):
         self.client = ZhipuAI(api_key=key)
         self.model_name = model_name
         self.lang = lang
@@ -140,7 +141,7 @@ class Zhipu4V(Base):
 
 
 class LocalCV(Base):
-    def __init__(self, key, model_name="glm-4v", lang="Chinese"):
+    def __init__(self, key, model_name="glm-4v", lang="Chinese", **kwargs):
         pass
 
     def describe(self, image, max_tokens=1024):

rag/llm/embedding_model.py

@@ -51,7 +51,7 @@ class Base(ABC):
 
 
 class HuEmbedding(Base):
-    def __init__(self, key="", model_name=""):
+    def __init__(self, **kwargs):
         """
         If you have trouble downloading HuggingFace models, -_^ this might help!!
 
@@ -81,8 +81,9 @@ class HuEmbedding(Base):
 
 
 class OpenAIEmbed(Base):
-    def __init__(self, key, model_name="text-embedding-ada-002"):
-        self.client = OpenAI(api_key=key)
+    def __init__(self, key, model_name="text-embedding-ada-002", base_url="https://api.openai.com/v1"):
+        if not base_url: base_url="https://api.openai.com/v1"
+        self.client = OpenAI(api_key=key, base_url=base_url)
         self.model_name = model_name
 
     def encode(self, texts: list, batch_size=32):
@@ -98,7 +99,7 @@ class OpenAIEmbed(Base):
 
 
 class QWenEmbed(Base):
-    def __init__(self, key, model_name="text_embedding_v2"):
+    def __init__(self, key, model_name="text_embedding_v2", **kwargs):
         dashscope.api_key = key
         self.model_name = model_name
 
@@ -131,7 +132,7 @@ class QWenEmbed(Base):
 
 
 class ZhipuEmbed(Base):
-    def __init__(self, key, model_name="embedding-2"):
+    def __init__(self, key, model_name="embedding-2", **kwargs):
         self.client = ZhipuAI(api_key=key)
         self.model_name = model_name
 

rag/svr/task_executor.py

@@ -280,4 +280,5 @@ if __name__ == "__main__":
     from mpi4py import MPI
 
     comm = MPI.COMM_WORLD
-    main(int(sys.argv[2]), int(sys.argv[1]))
+    while True:
+        main(int(sys.argv[2]), int(sys.argv[1]))