aient 1.0.29__py3-none-any.whl

aient/models/duckduckgo.py

@@ -0,0 +1,241 @@
+from types import TracebackType
+from collections import defaultdict
+
+import json
+import httpx
+from fake_useragent import UserAgent
+
+class DuckChatException(httpx.HTTPError):
+    """Base exception class for duck_chat."""
+
+
+class RatelimitException(DuckChatException):
+    """Raised for rate limit exceeded errors during API requests."""
+
+
+class ConversationLimitException(DuckChatException):
+    """Raised for conversation limit during API requests to AI endpoint."""
+
+
+from enum import Enum
+class ModelType(Enum):
+    claude = "claude-3-haiku-20240307"
+    llama = "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
+    gpt4omini = "gpt-4o-mini"
+    mixtral = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+
+    @classmethod
+    def _missing_(cls, value):
+        if isinstance(value, str):
+            # 对于完全匹配的情况
+            for member in cls:
+                if member.value == value:
+                    return member
+
+            # 对于部分匹配的情况
+            for member in cls:
+                if value in member.value:
+                    return member
+
+        return None
+
+    def __new__(cls, *args):
+        obj = object.__new__(cls)
+        obj._value_ = args[0]
+        return obj
+
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return f"ModelType({self.value!r})"
+
+class Role(Enum):
+    user = "user"
+    assistant = "assistant"
+
+import msgspec
+class Message(msgspec.Struct):
+    role: Role
+    content: str
+
+    def get(self, key, default=None):
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            return default
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+class History(msgspec.Struct):
+    model: ModelType
+    messages: list[Message]
+
+    def add_to_conversation(self, role: Role, message: str) -> None:
+        self.messages.append(Message(role, message))
+
+    def set_model(self, model_name: str) -> None:
+        self.model = ModelType(model_name)
+
+    def __getitem__(self, index: int) -> list[Message]:
+        return self.messages[index]
+
+    def __len__(self) -> int:
+        return len(self.messages)
+
+class UserHistory(msgspec.Struct):
+    user_history: dict[str, History] = msgspec.field(default_factory=dict)
+
+    def add_to_conversation(self, role: Role, message: str, convo_id: str = "default") -> None:
+        if convo_id not in self.user_history:
+            self.user_history[convo_id] = History(model=ModelType.claude, messages=[])
+        self.user_history[convo_id].add_to_conversation(role, message)
+
+    def get_history(self, convo_id: str = "default") -> History:
+        if convo_id not in self.user_history:
+            self.user_history[convo_id] = History(model=ModelType.claude, messages=[])
+        return self.user_history[convo_id]
+
+    def set_model(self, model_name: str, convo_id: str = "default") -> None:
+        self.get_history(convo_id).set_model(model_name)
+
+    def reset(self, convo_id: str = "default") -> None:
+        self.user_history[convo_id] = History(model=ModelType.claude, messages=[])
+
+    def get_all_convo_ids(self) -> list[str]:
+        return list(self.user_history.keys())
+
+    # 新增方法
+    def __getitem__(self, convo_id: str) -> History:
+        return self.get_history(convo_id)
+
+class DuckChat:
+    def __init__(
+        self,
+        model: ModelType = ModelType.claude,
+        client: httpx.AsyncClient | None = None,
+        user_agent: UserAgent | str = UserAgent(min_version=120.0),
+    ) -> None:
+        if isinstance(user_agent, str):
+            self.user_agent = user_agent
+        else:
+            self.user_agent = user_agent.random
+
+        self._client = client or httpx.AsyncClient(
+            headers={
+                "Host": "duckduckgo.com",
+                "Accept": "text/event-stream",
+                "Accept-Language": "en-US,en;q=0.5",
+                "Accept-Encoding": "gzip, deflate, br",
+                "Referer": "https://duckduckgo.com/",
+                "User-Agent": self.user_agent,
+                "DNT": "1",
+                "Sec-GPC": "1",
+                "Connection": "keep-alive",
+                "Sec-Fetch-Dest": "empty",
+                "Sec-Fetch-Mode": "cors",
+                "Sec-Fetch-Site": "same-origin",
+                "TE": "trailers",
+            }
+        )
+        self.vqd: list[str | None] = []
+        self.history = History(model, [])
+        self.conversation = UserHistory({"default": self.history})
+        self.__encoder = msgspec.json.Encoder()
+        self.__decoder = msgspec.json.Decoder()
+
+        self.tokens_usage = defaultdict(int)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        await self._client.aclose()
+
+    async def add_to_conversation(self, role: Role, message: Message, convo_id: str = "default") -> None:
+        self.conversation.add_to_conversation(role, message, convo_id)
+
+    async def get_vqd(self) -> None:
+        """Get new x-vqd-4 token"""
+        response = await self._client.get(
+            "https://duckduckgo.com/duckchat/v1/status", headers={"x-vqd-accept": "1"}
+        )
+        if response.status_code == 429:
+            try:
+                err_message = self.__decoder.decode(response.content).get("type", "")
+            except Exception:
+                raise DuckChatException(response.text)
+            else:
+                raise RatelimitException(err_message)
+        self.vqd.append(response.headers.get("x-vqd-4"))
+        if not self.vqd:
+            raise DuckChatException("No x-vqd-4")
+
+    async def process_sse_stream(self, convo_id: str = "default"):
+        # print("self.conversation[convo_id]", self.conversation[convo_id])
+        async with self._client.stream(
+            "POST",
+            "https://duckduckgo.com/duckchat/v1/chat",
+            headers={
+                "Content-Type": "application/json",
+                "x-vqd-4": self.vqd[-1],
+            },
+            content=self.__encoder.encode(self.conversation[convo_id]),
+        ) as response:
+            if response.status_code == 400:
+                content = await response.aread()
+                print("response.status_code", response.status_code, content)
+            if response.status_code == 429:
+                raise RatelimitException("Rate limit exceeded")
+
+            async for line in response.aiter_lines():
+                if line.startswith('data: '):
+                    yield line
+
+    async def ask_stream_async(self, query, convo_id, model, **kwargs):
+        """Get answer from chat AI"""
+        if not self.vqd:
+            await self.get_vqd()
+        await self.add_to_conversation(Role.user, query, convo_id)
+        self.conversation.set_model(model, convo_id)
+        full_response = ""
+        async for sse in self.process_sse_stream(convo_id):
+            data = sse.lstrip("data: ")
+            if data == "[DONE]":
+                break
+            resp: dict = json.loads(data)
+            mess = resp.get("message")
+            if mess:
+                yield mess
+                full_response += mess
+        # await self.add_to_conversation(Role.assistant, full_response, convo_id)
+
+    async def reset(self, convo_id: str = "default") -> None:
+        self.conversation.reset(convo_id)
+
+    # async def reask_question(self, num: int) -> str:
+    #     """Get answer from chat AI"""
+
+    #     if num >= len(self.vqd):
+    #         num = len(self.vqd) - 1
+    #     self.vqd = self.vqd[:num]
+
+    #     if not self.history.messages:
+    #         return ""
+
+    #     if not self.vqd:
+    #         await self.get_vqd()
+    #         self.history.messages = [self.history.messages[0]]
+    #     else:
+    #         num = min(num, len(self.vqd))
+    #         self.history.messages = self.history.messages[: (num * 2 - 1)]
+    #     message = await self.get_answer()
+    #     self.add_to_conversation(Role.assistant, message)
+
+    #     return message

aient/models/gemini.py

@@ -0,0 +1,357 @@
+import os
+import re
+import json
+import requests
+
+from .base import BaseLLM
+from ..core.utils import BaseAPI
+
+import copy
+from ..plugins import PLUGINS, get_tools_result_async, function_call_list
+from ..utils.scripts import safe_get
+
+
+class gemini(BaseLLM):
+    def __init__(
+        self,
+        api_key: str = None,
+        engine: str = os.environ.get("GPT_ENGINE") or "gemini-1.5-pro-latest",
+        api_url: str = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}",
+        system_prompt: str = "You are Gemini, a large language model trained by Google. Respond conversationally",
+        temperature: float = 0.5,
+        top_p: float = 0.7,
+        timeout: float = 20,
+        use_plugins: bool = True,
+        print_log: bool = False,
+    ):
+        url = api_url.format(model=engine, stream="streamGenerateContent", api_key=os.environ.get("GOOGLE_AI_API_KEY", api_key))
+        super().__init__(api_key, engine, url, system_prompt=system_prompt, timeout=timeout, temperature=temperature, top_p=top_p, use_plugins=use_plugins, print_log=print_log)
+        self.conversation: dict[str, list[dict]] = {
+            "default": [],
+        }
+
+    def add_to_conversation(
+        self,
+        message: str,
+        role: str,
+        convo_id: str = "default",
+        pass_history: int = 9999,
+        total_tokens: int = 0,
+        function_arguments: str = "",
+    ) -> None:
+        """
+        Add a message to the conversation
+        """
+
+        if convo_id not in self.conversation:
+            self.reset(convo_id=convo_id)
+        # print("message", message)
+
+        if function_arguments:
+            self.conversation[convo_id].append(
+                {
+                    "role": "model",
+                    "parts": [function_arguments]
+                }
+            )
+            function_call_name = function_arguments["functionCall"]["name"]
+            self.conversation[convo_id].append(
+                {
+                    "role": "function",
+                    "parts": [{
+                    "functionResponse": {
+                        "name": function_call_name,
+                        "response": {
+                            "name": function_call_name,
+                            "content": {
+                                "result": message,
+                            }
+                        }
+                    }
+                    }]
+                }
+            )
+
+        else:
+            if isinstance(message, str):
+                message = [{"text": message}]
+            self.conversation[convo_id].append({"role": role, "parts": message})
+
+        history_len = len(self.conversation[convo_id])
+        history = pass_history
+        if pass_history < 2:
+            history = 2
+        while history_len > history:
+            mess_body = self.conversation[convo_id].pop(1)
+            history_len = history_len - 1
+            if mess_body.get("role") == "user":
+                mess_body = self.conversation[convo_id].pop(1)
+                history_len = history_len - 1
+                if safe_get(mess_body, "parts", 0, "functionCall"):
+                    self.conversation[convo_id].pop(1)
+                    history_len = history_len - 1
+
+        if total_tokens:
+            self.tokens_usage[convo_id] += total_tokens
+
+    def reset(self, convo_id: str = "default", system_prompt: str = "You are Gemini, a large language model trained by Google. Respond conversationally") -> None:
+        """
+        Reset the conversation
+        """
+        self.system_prompt = system_prompt or self.system_prompt
+        self.conversation[convo_id] = list()
+
+    def ask_stream(
+        self,
+        prompt: str,
+        role: str = "user",
+        convo_id: str = "default",
+        model: str = "",
+        pass_history: int = 9999,
+        model_max_tokens: int = 4096,
+        system_prompt: str = None,
+        **kwargs,
+    ):
+        self.system_prompt = system_prompt or self.system_prompt
+        if convo_id not in self.conversation or pass_history <= 2:
+            self.reset(convo_id=convo_id, system_prompt=self.system_prompt)
+        self.add_to_conversation(prompt, role, convo_id=convo_id, pass_history=pass_history)
+        # print(self.conversation[convo_id])
+
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        json_post = {
+            "contents": self.conversation[convo_id] if pass_history else [{
+                "role": "user",
+                "content": prompt
+            }],
+            "systemInstruction": {"parts": [{"text": self.system_prompt}]},
+            "safetySettings": [
+                {
+                    "category": "HARM_CATEGORY_HARASSMENT",
+                    "threshold": "BLOCK_NONE"
+                },
+                {
+                    "category": "HARM_CATEGORY_HATE_SPEECH",
+                    "threshold": "BLOCK_NONE"
+                },
+                {
+                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                    "threshold": "BLOCK_NONE"
+                },
+                {
+                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                    "threshold": "BLOCK_NONE"
+                }
+            ],
+        }
+
+        url = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}".format(model=model or self.engine, stream="streamGenerateContent", api_key=os.environ.get("GOOGLE_AI_API_KEY", self.api_key) or kwargs.get("api_key"))
+        self.api_url = BaseAPI(url)
+        url = self.api_url.source_api_url
+
+        if self.print_log:
+            print("url", url)
+            replaced_text = json.loads(re.sub(r';base64,([A-Za-z0-9+/=]+)', ';base64,***', json.dumps(json_post)))
+            print(json.dumps(replaced_text, indent=4, ensure_ascii=False))
+
+        try:
+            response = self.session.post(
+                url,
+                headers=headers,
+                json=json_post,
+                timeout=kwargs.get("timeout", self.timeout),
+                stream=True,
+            )
+        except ConnectionError:
+            print("连接错误，请检查服务器状态或网络连接。")
+            return
+        except requests.exceptions.ReadTimeout:
+            print("请求超时，请检查网络连接或增加超时时间。{e}")
+            return
+        except Exception as e:
+            print(f"发生了未预料的错误: {e}")
+            return
+
+        if response.status_code != 200:
+            print(response.text)
+            raise BaseException(f"{response.status_code} {response.reason} {response.text}")
+        response_role: str = "model"
+        full_response: str = ""
+        try:
+            for line in response.iter_lines():
+                if not line:
+                    continue
+                line = line.decode("utf-8")
+                if line and '\"text\": \"' in line:
+                    content = line.split('\"text\": \"')[1][:-1]
+                    content = "\n".join(content.split("\\n"))
+                    content = content.encode('utf-8').decode('unicode-escape')
+                    full_response += content
+                    yield content
+        except requests.exceptions.ChunkedEncodingError as e:
+            print("Chunked Encoding Error occurred:", e)
+        except Exception as e:
+            print("An error occurred:", e)
+
+        self.add_to_conversation([{"text": full_response}], response_role, convo_id=convo_id, pass_history=pass_history)
+
+    async def ask_stream_async(
+        self,
+        prompt: str,
+        role: str = "user",
+        convo_id: str = "default",
+        model: str = "",
+        pass_history: int = 9999,
+        system_prompt: str = None,
+        language: str = "English",
+        function_arguments: str = "",
+        total_tokens: int = 0,
+        **kwargs,
+    ):
+        self.system_prompt = system_prompt or self.system_prompt
+        if convo_id not in self.conversation or pass_history <= 2:
+            self.reset(convo_id=convo_id, system_prompt=self.system_prompt)
+        self.add_to_conversation(prompt, role, convo_id=convo_id, total_tokens=total_tokens, function_arguments=function_arguments, pass_history=pass_history)
+        # print(self.conversation[convo_id])
+
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        json_post = {
+            "contents": self.conversation[convo_id] if pass_history else [{
+                "role": "user",
+                "content": prompt
+            }],
+            "systemInstruction": {"parts": [{"text": self.system_prompt}]},
+            "safetySettings": [
+                {
+                    "category": "HARM_CATEGORY_HARASSMENT",
+                    "threshold": "BLOCK_NONE"
+                },
+                {
+                    "category": "HARM_CATEGORY_HATE_SPEECH",
+                    "threshold": "BLOCK_NONE"
+                },
+                {
+                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                    "threshold": "BLOCK_NONE"
+                },
+                {
+                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                    "threshold": "BLOCK_NONE"
+                }
+            ],
+        }
+
+        plugins = kwargs.get("plugins", PLUGINS)
+        if all(value == False for value in plugins.values()) == False and self.use_plugins:
+            tools = {
+                "tools": [
+                    {
+                        "function_declarations": [
+
+                        ]
+                    }
+                ],
+                "tool_config": {
+                    "function_calling_config": {
+                        "mode": "AUTO",
+                    },
+                },
+            }
+            json_post.update(copy.deepcopy(tools))
+            for item in plugins.keys():
+                try:
+                    if plugins[item]:
+                        json_post["tools"][0]["function_declarations"].append(function_call_list[item])
+                except:
+                    pass
+
+        url = "https://generativelanguage.googleapis.com/v1beta/models/{model}:{stream}?key={api_key}".format(model=model or self.engine, stream="streamGenerateContent", api_key=os.environ.get("GOOGLE_AI_API_KEY", self.api_key) or kwargs.get("api_key"))
+        self.api_url = BaseAPI(url)
+        url = self.api_url.source_api_url
+
+        if self.print_log:
+            print("url", url)
+            replaced_text = json.loads(re.sub(r';base64,([A-Za-z0-9+/=]+)', ';base64,***', json.dumps(json_post)))
+            print(json.dumps(replaced_text, indent=4, ensure_ascii=False))
+
+        response_role: str = "model"
+        full_response: str = ""
+        function_full_response: str = "{"
+        need_function_call = False
+        revicing_function_call = False
+        total_tokens = 0
+        try:
+            async with self.aclient.stream(
+                "post",
+                url,
+                headers=headers,
+                json=json_post,
+                timeout=kwargs.get("timeout", self.timeout),
+            ) as response:
+                if response.status_code != 200:
+                    error_content = await response.aread()
+                    error_message = error_content.decode('utf-8')
+                    raise BaseException(f"{response.status_code}: {error_message}")
+                try:
+                    async for line in response.aiter_lines():
+                        if not line:
+                            continue
+                        # print(line)
+                        if line and '\"text\": \"' in line:
+                            content = line.split('\"text\": \"')[1][:-1]
+                            content = "\n".join(content.split("\\n"))
+                            full_response += content
+                            yield content
+
+                        if line and '\"totalTokenCount\": ' in line:
+                            content = int(line.split('\"totalTokenCount\": ')[1])
+                            total_tokens = content
+
+                        if line and ('\"functionCall\": {' in line or revicing_function_call):
+                            revicing_function_call = True
+                            need_function_call = True
+                            if ']' in line:
+                                revicing_function_call = False
+                                continue
+
+                            function_full_response += line
+
+                except requests.exceptions.ChunkedEncodingError as e:
+                    print("Chunked Encoding Error occurred:", e)
+                except Exception as e:
+                    print("An error occurred:", e)
+
+        except Exception as e:
+            print(f"发生了未预料的错误: {e}")
+            return
+
+        if response.status_code != 200:
+            await response.aread()
+            print(response.text)
+            raise BaseException(f"{response.status_code} {response.reason} {response.text}")
+        if self.print_log:
+            print("\n\rtotal_tokens", total_tokens)
+        if need_function_call:
+            # print(function_full_response)
+            function_call = json.loads(function_full_response)
+            print(json.dumps(function_call, indent=4, ensure_ascii=False))
+            function_call_name = function_call["functionCall"]["name"]
+            function_full_response = json.dumps(function_call["functionCall"]["args"])
+            function_call_max_tokens = 32000
+            print("\033[32m function_call", function_call_name, "max token:", function_call_max_tokens, "\033[0m")
+            async for chunk in get_tools_result_async(function_call_name, function_full_response, function_call_max_tokens, model or self.engine, gemini, kwargs.get('api_key', self.api_key), self.api_url, use_plugins=False, model=model or self.engine, add_message=self.add_to_conversation, convo_id=convo_id, language=language):
+                if "function_response:" in chunk:
+                    function_response = chunk.replace("function_response:", "")
+                else:
+                    yield chunk
+            response_role = "model"
+            async for chunk in self.ask_stream_async(function_response, response_role, convo_id=convo_id, function_name=function_call_name, total_tokens=total_tokens, model=model or self.engine, function_arguments=function_call, api_key=kwargs.get('api_key', self.api_key), plugins=kwargs.get("plugins", PLUGINS), system_prompt=system_prompt):
+                yield chunk
+        else:
+            self.add_to_conversation([{"text": full_response}], response_role, convo_id=convo_id, total_tokens=total_tokens, pass_history=pass_history)