fix: openai compatable with llama and gemini

application/llm/docsgpt_provider.py

@@ -1,34 +1,132 @@
-from application.llm.base import BaseLLM
 import json
-import requests
+import sys
+
+from application.core.settings import settings
+from application.llm.base import BaseLLM
 
 
 class DocsGPTAPILLM(BaseLLM):
 
     def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
+        from openai import OpenAI
+
         super().__init__(*args, **kwargs)
-        self.api_key = api_key
+        self.client = OpenAI(api_key="sk-docsgpt-public", base_url="https://oai.arc53.com")
         self.user_api_key = user_api_key
-        self.endpoint = "https://llm.arc53.com"
+        self.api_key = api_key
 
-    def _raw_gen(self, baseself, model, messages, stream=False, *args, **kwargs):
-        response = requests.post(
-            f"{self.endpoint}/answer", json={"messages": messages, "max_new_tokens": 30}
-        )
-        response_clean = response.json()["a"].replace("###", "")
+    def _clean_messages_openai(self, messages):
+        cleaned_messages = []
+        for message in messages:
+            role = message.get("role")
+            content = message.get("content")
 
-        return response_clean
+            if role == "model":
+                role = "assistant"
 
-    def _raw_gen_stream(self, baseself, model, messages, stream=True, *args, **kwargs):
-        response = requests.post(
-            f"{self.endpoint}/stream",
-            json={"messages": messages, "max_new_tokens": 256},
-            stream=True,
-        )
+            if role and content is not None:
+                if isinstance(content, str):
+                    cleaned_messages.append({"role": role, "content": content})
+                elif isinstance(content, list):
+                    for item in content:
+                        if "text" in item:
+                            cleaned_messages.append(
+                                {"role": role, "content": item["text"]}
+                            )
+                        elif "function_call" in item:
+                            tool_call = {
+                                "id": item["function_call"]["call_id"],
+                                "type": "function",
+                                "function": {
+                                    "name": item["function_call"]["name"],
+                                    "arguments": json.dumps(
+                                        item["function_call"]["args"]
+                                    ),
+                                },
+                            }
+                            cleaned_messages.append(
+                                {
+                                    "role": "assistant",
+                                    "content": None,
+                                    "tool_calls": [tool_call],
+                                }
+                            )
+                        elif "function_response" in item:
+                            cleaned_messages.append(
+                                {
+                                    "role": "tool",
+                                    "tool_call_id": item["function_response"][
+                                        "call_id"
+                                    ],
+                                    "content": json.dumps(
+                                        item["function_response"]["response"]["result"]
+                                    ),
+                                }
+                            )
+                        else:
+                            raise ValueError(
+                                f"Unexpected content dictionary format: {item}"
+                            )
+                else:
+                    raise ValueError(f"Unexpected content type: {type(content)}")
 
-        for line in response.iter_lines():
-            if line:
-                data_str = line.decode("utf-8")
-                if data_str.startswith("data: "):
-                    data = json.loads(data_str[6:])
-                    yield data["a"]
+        return cleaned_messages
+
+    def _raw_gen(
+        self,
+        baseself,
+        model,
+        messages,
+        stream=False,
+        tools=None,
+        engine=settings.AZURE_DEPLOYMENT_NAME,
+        **kwargs,
+    ):
+        messages = self._clean_messages_openai(messages)
+        if tools:
+            response = self.client.chat.completions.create(
+                model="docsgpt",
+                messages=messages,
+                stream=stream,
+                tools=tools,
+                **kwargs,
+            )
+            return response.choices[0]
+        else:
+            response = self.client.chat.completions.create(
+                model="docsgpt", messages=messages, stream=stream, **kwargs
+            )
+            return response.choices[0].message.content
+
+    def _raw_gen_stream(
+        self,
+        baseself,
+        model,
+        messages,
+        stream=True,
+        tools=None,
+        engine=settings.AZURE_DEPLOYMENT_NAME,
+        **kwargs,
+    ):
+        messages = self._clean_messages_openai(messages)
+        if tools:
+            response = self.client.chat.completions.create(
+                model="docsgpt",
+                messages=messages,
+                stream=stream,
+                tools=tools,
+                **kwargs,
+            )
+        else:
+            response = self.client.chat.completions.create(
+                model="docsgpt", messages=messages, stream=stream, **kwargs
+            )
+
+        for line in response:
+            if len(line.choices) > 0 and line.choices[0].delta.content is not None and len(line.choices[0].delta.content) > 0:
+                yield line.choices[0].delta.content
+            elif len(line.choices) > 0:
+                yield line.choices[0]
+
+    def _supports_tools(self):
+        return True