'load llm infomation from a json file and add support for OpenRouter' (#1533)

### What problem does this PR solve?

#1467 

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: Zhedong Cen <cenzhedong2@126.com>

rag/llm/__init__.py

@@ -45,7 +45,8 @@ CvModel = {
     "Tongyi-Qianwen": QWenCV,
     "ZHIPU-AI": Zhipu4V,
     "Moonshot": LocalCV,
-    'Gemini':GeminiCV
+    'Gemini':GeminiCV,
+    'OpenRouter':OpenRouterCV
 }
 
 
@@ -65,7 +66,8 @@ ChatModel = {
     "Mistral": MistralChat,
     'Gemini' : GeminiChat,
     "Bedrock": BedrockChat,
-    "Groq": GroqChat
+    "Groq": GroqChat,
+    'OpenRouter':OpenRouterChat
 }
 
 

rag/llm/chat_model.py

@@ -685,7 +685,6 @@ class GeminiChat(Base):
         yield  response._chunks[-1].usage_metadata.total_token_count
 
 
-
 class GroqChat:
     def __init__(self, key, model_name,base_url=''):
         self.client = Groq(api_key=key)
@@ -697,7 +696,6 @@ class GroqChat:
         for k in list(gen_conf.keys()):
             if k not in ["temperature", "top_p", "max_tokens"]:
                 del gen_conf[k]
-
         ans = ""
         try:
             response = self.client.chat.completions.create(
@@ -707,7 +705,7 @@ class GroqChat:
             )
             ans = response.choices[0].message.content
             if response.choices[0].finish_reason == "length":
-                ans += "...\nFor the content length reason, it stopped, continue?" if self.is_english(
+                ans += "...\nFor the content length reason, it stopped, continue?" if is_english(
                     [ans]) else "······\n由于长度的原因，回答被截断了，要继续吗？"
             return ans, response.usage.total_tokens
         except Exception as e:
@@ -734,11 +732,20 @@ class GroqChat:
                 ans += resp.choices[0].delta.content
                 total_tokens += 1
                 if resp.choices[0].finish_reason == "length":
-                    ans += "...\nFor the content length reason, it stopped, continue?" if self.is_english(
+                    ans += "...\nFor the content length reason, it stopped, continue?" if is_english(
                         [ans]) else "······\n由于长度的原因，回答被截断了，要继续吗？"
                 yield ans
 
         except Exception as e:
             yield ans + "\n**ERROR**: " + str(e)
 
-        yield total_tokens
+        yield total_tokens
+
+
+## openrouter
+class OpenRouterChat(Base):
+    def __init__(self, key, model_name, base_url="https://openrouter.ai/api/v1"):
+        self.base_url = "https://openrouter.ai/api/v1"
+        self.client = OpenAI(base_url=self.base_url, api_key=key)
+        self.model_name = model_name
+

rag/llm/cv_model.py

@@ -23,6 +23,8 @@ from openai import OpenAI
 import os
 import base64
 from io import BytesIO
+import json
+import requests
 
 from api.utils import get_uuid
 from api.utils.file_utils import get_project_base_directory
@@ -212,7 +214,7 @@ class GeminiCV(Base):
         self.model = GenerativeModel(model_name=self.model_name)
         self.model._client = _client
         self.lang = lang 
-    
+
     def describe(self, image, max_tokens=2048):
         from PIL.Image import open
         gen_config = {'max_output_tokens':max_tokens}
@@ -227,6 +229,63 @@ class GeminiCV(Base):
         )
         return res.text,res.usage_metadata.total_token_count
 
+
+class OpenRouterCV(Base):
+    def __init__(
+        self,
+        key,
+        model_name,
+        lang="Chinese",
+        base_url="https://openrouter.ai/api/v1/chat/completions",
+    ):
+        self.model_name = model_name
+        self.lang = lang
+        self.base_url = "https://openrouter.ai/api/v1/chat/completions"
+        self.key = key
+
+    def describe(self, image, max_tokens=300):
+        b64 = self.image2base64(image)
+        response = requests.post(
+            url=self.base_url,
+            headers={
+                "Authorization": f"Bearer {self.key}",
+            },
+            data=json.dumps(
+                {
+                    "model": self.model_name,
+                    "messages": self.prompt(b64),
+                    "max_tokens": max_tokens,
+                }
+            ),
+        )
+        response = response.json()
+        return (
+            response["choices"][0]["message"]["content"].strip(),
+            response["usage"]["total_tokens"],
+        )
+
+    def prompt(self, b64):
+        return [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
+                    },
+                    {
+                        "type": "text",
+                        "text": (
+                            "请用中文详细描述一下图中的内容，比如时间，地点，人物，事情，人物心情等，如果有数据请提取出数据。"
+                            if self.lang.lower() == "chinese"
+                            else "Please describe the content of this picture, like where, when, who, what happen. If it has number data, please extract them out."
+                        ),
+                    },
+                ],
+            }
+        ]
+
+
 class LocalCV(Base):
     def __init__(self, key, model_name="glm-4v", lang="Chinese", **kwargs):
         pass