xiaogpt 1.80__tar.gz → 1.81__tar.gz

{xiaogpt-1.80 → xiaogpt-1.81}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xiaogpt
-Version: 1.80
+Version: 1.81
 Summary: Play ChatGPT or other LLM with xiaomi AI speaker
 Author-Email: yihong0618 <zouzou0208@gmail.com>
 License: MIT

{xiaogpt-1.80 → xiaogpt-1.81}/pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
     "numexpr==2.8.6",
 ]
 dynamic = []
-version = "1.80"
+version = "1.81"
 
 [project.license]
 text = "MIT"

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/bot/bard_bot.py

@@ -1,14 +1,15 @@
 """ChatGLM bot"""
 from __future__ import annotations
+
 from typing import Any
 
 from bardapi import BardAsync
 from rich import print
 
-from xiaogpt.bot.base_bot import BaseBot
+from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin
 
 
-class BardBot(BaseBot):
+class BardBot(ChatHistoryMixin, BaseBot):
     def __init__(
         self,
         bard_token: str,

xiaogpt-1.81/xiaogpt/bot/base_bot.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, AsyncGenerator, TypeVar
+
+from xiaogpt.config import Config
+
+T = TypeVar("T", bound="BaseBot")
+
+
+class BaseBot(ABC):
+    @abstractmethod
+    async def ask(self, query: str, **options: Any) -> str:
+        pass
+
+    @abstractmethod
+    async def ask_stream(self, query: str, **options: Any) -> AsyncGenerator[str, None]:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls: type[T], config: Config) -> T:
+        pass
+
+    @abstractmethod
+    def has_history(self) -> bool:
+        pass
+
+    @abstractmethod
+    def change_prompt(self, new_prompt: str) -> None:
+        pass
+
+
+class ChatHistoryMixin:
+    history: list[tuple[str, str]]
+
+    def has_history(self) -> bool:
+        return bool(self.history)
+
+    def change_prompt(self, new_prompt: str) -> None:
+        if self.history:
+            print(self.history)
+            self.history[0][0] = new_prompt
+
+    def get_messages(self) -> list[dict]:
+        ms = []
+        for h in self.history:
+            ms.append({"role": "user", "content": h[0]})
+            ms.append({"role": "assistant", "content": h[1]})
+        return ms
+
+    def add_message(self, query: str, message: str) -> None:
+        self.history.append([f"{query}", message])
+        # only keep 5 history
+        first_history = self.history.pop(0)
+        self.history = [first_history] + self.history[-5:]

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/bot/chatgptapi_bot.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 import openai
 from rich import print
 
-from xiaogpt.bot.base_bot import BaseBot
+from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin
 from xiaogpt.utils import split_sentences
 
 
-class ChatGPTBot(BaseBot):
+class ChatGPTBot(ChatHistoryMixin, BaseBot):
     default_options = {"model": "gpt-3.5-turbo-0613"}
 
     def __init__(
@@ -42,10 +42,7 @@ class ChatGPTBot(BaseBot):
         )
 
     async def ask(self, query, **options):
-        ms = []
-        for h in self.history:
-            ms.append({"role": "user", "content": h[0]})
-            ms.append({"role": "assistant", "content": h[1]})
+        ms = self.get_messages()
         ms.append({"role": "user", "content": f"{query}"})
         kwargs = {**self.default_options, **options}
         try:
@@ -60,18 +57,12 @@ class ChatGPTBot(BaseBot):
             .encode("utf8")
             .decode()
         )
-        self.history.append([f"{query}", message])
-        # only keep 5 history
-        first_history = self.history.pop(0)
-        self.history = [first_history] + self.history[-5:]
+        self.add_message(query, message)
         print(message)
         return message
 
     async def ask_stream(self, query, **options):
-        ms = []
-        for h in self.history:
-            ms.append({"role": "user", "content": h[0]})
-            ms.append({"role": "assistant", "content": h[1]})
+        ms = self.get_messages()
         ms.append({"role": "user", "content": f"{query}"})
         kwargs = {"model": "gpt-3.5-turbo", **options}
         if openai.api_type == "azure":
@@ -99,6 +90,4 @@ class ChatGPTBot(BaseBot):
                 yield sentence
         finally:
             print()
-            self.history.append([f"{query}", message])
-            first_history = self.history.pop(0)
-            self.history = [first_history] + self.history[-5:]
+            self.add_message(query, message)

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/bot/glm_bot.py

@@ -1,21 +1,18 @@
 """ChatGLM bot"""
 from __future__ import annotations
 
-from typing import Any, AsyncGenerator
+from typing import Any
 
 import zhipuai
 from rich import print
 
-from xiaogpt.bot.base_bot import BaseBot
+from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin
 
 
-class GLMBot(BaseBot):
+class GLMBot(ChatHistoryMixin, BaseBot):
     default_options = {"model": "chatglm_130b"}
 
-    def __init__(
-        self,
-        glm_key: str,
-    ) -> None:
+    def __init__(self, glm_key: str) -> None:
         self.history = []
         zhipuai.api_key = glm_key
 
@@ -24,10 +21,7 @@ class GLMBot(BaseBot):
         return cls(glm_key=config.glm_key)
 
     def ask(self, query, **options):
-        ms = []
-        for h in self.history:
-            ms.append({"role": "user", "content": h[0]})
-            ms.append({"role": "assistant", "content": h[1]})
+        ms = self.get_messages()
         kwargs = {**self.default_options, **options}
         kwargs["prompt"] = ms
         ms.append({"role": "user", "content": f"{query}"})
@@ -40,10 +34,7 @@ class GLMBot(BaseBot):
         for i in r.events():
             message += str(i.data)
 
-        self.history.append([f"{query}", message])
-        # only keep 5 history
-        first_history = self.history.pop(0)
-        self.history = [first_history] + self.history[-5:]
+        self.add_message(query, message)
         print(message)
         return message
 

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/bot/gpt3_bot.py

@@ -1,11 +1,11 @@
 import openai
 from rich import print
 
-from xiaogpt.bot.base_bot import BaseBot
+from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin
 from xiaogpt.utils import split_sentences
 
 
-class GPT3Bot(BaseBot):
+class GPT3Bot(ChatHistoryMixin, BaseBot):
     def __init__(self, openai_key, api_base=None, proxy=None):
         openai.api_key = openai_key
         if api_base:

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/bot/langchain_bot.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import os
 
+from langchain.memory import ConversationBufferWindowMemory
 from rich import print
 
 from xiaogpt.bot.base_bot import BaseBot
@@ -26,8 +27,14 @@ class LangChainBot(BaseBot):
             os.environ["OPENAI_API_BASE"] = api_base
         if proxy:
             os.environ["OPENAI_PROXY"] = proxy
-        # Todo，Plan to implement within langchain
-        self.history = []
+        self.memory = ConversationBufferWindowMemory(return_messages=True)
+
+    def has_history(self) -> bool:
+        return len(self.memory.chat_memory.messages) > 0
+
+    def change_prompt(self, new_prompt: str) -> None:
+        self.memory.clear()
+        self.memory.chat_memory.add_user_message(new_prompt)
 
     @classmethod
     def from_config(cls, config):
@@ -39,14 +46,11 @@ class LangChainBot(BaseBot):
         )
 
     async def ask(self, query, **options):
-        # Todo，Currently only supports stream
-        raise Exception(
-            "The bot does not support it. Please use 'ask_stream, add: --stream'"
-        )
+        return await agent_search(query, self.memory)
 
     async def ask_stream(self, query, **options):
         callback = AsyncIteratorCallbackHandler()
-        task = asyncio.create_task(agent_search(query, callback))
+        task = asyncio.create_task(agent_search(query, self.memory, callback))
         try:
             async for message in split_sentences(callback.aiter()):
                 yield message

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/bot/newbing_bot.py

@@ -4,13 +4,13 @@ import re
 
 from EdgeGPT import Chatbot, ConversationStyle
 
-from xiaogpt.bot.base_bot import BaseBot
+from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin
 from xiaogpt.utils import split_sentences
 
 _reference_link_re = re.compile(r"\[\d+\]: .+?\n+")
 
 
-class NewBingBot(BaseBot):
+class NewBingBot(ChatHistoryMixin, BaseBot):
     def __init__(
         self,
         bing_cookie_path: str = "",
@@ -52,7 +52,7 @@ class NewBingBot(BaseBot):
         kwargs = {"conversation_style": ConversationStyle.balanced, **options}
         try:
             completion = self._bot.ask_stream(prompt=query, **kwargs)
-        except Exception as e:
+        except Exception:
             return
 
         async def text_gen():

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/langchain/callbacks.py

@@ -33,7 +33,7 @@ class AsyncIteratorCallbackHandler(AsyncCallbackHandler):
 
     async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
         if token is not None and token != "":
-            print(token, end="")
+            print(token, end="", flush=True)
             self.queue.put_nowait(token)
 
     async def on_chain_end(

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/langchain/chain.py

@@ -1,11 +1,16 @@
+from __future__ import annotations
+
 from langchain.agents import AgentType, Tool, initialize_agent
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import LLMMathChain
 from langchain.chat_models import ChatOpenAI
+from langchain.schema.memory import BaseMemory
 from langchain.utilities import SerpAPIWrapper
 
 
-async def agent_search(query: str, callback: BaseCallbackHandler) -> None:
+async def agent_search(
+    query: str, memeory: BaseMemory, callback: BaseCallbackHandler | None = None
+) -> str:
     llm = ChatOpenAI(
         streaming=True,
         temperature=0,
@@ -23,8 +28,8 @@ async def agent_search(query: str, callback: BaseCallbackHandler) -> None:
     ]
 
     agent = initialize_agent(
-        tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=False
+        tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=False, memory=memeory
     )
-
+    callbacks = [callback] if callback else None
     # query eg：'杭州亚运会中国队获得了多少枚金牌？' // '计算3的2次方'
-    await agent.arun(query, callbacks=[callback])
+    return await agent.arun(query, callbacks=callbacks)

{xiaogpt-1.80 → xiaogpt-1.81}/xiaogpt/xiaogpt.py

@@ -231,9 +231,7 @@ class MiGPT:
         new_prompt = "以下都" + new_prompt
         print(f"Prompt from {self.config.prompt} change to {new_prompt}")
         self.config.prompt = new_prompt
-        if self.chatbot.history:
-            print(self.chatbot.history)
-            self.chatbot.history[0][0] = new_prompt
+        self.chatbot.change_prompt(new_prompt)
 
     async def get_latest_ask_from_xiaoai(self, session: ClientSession) -> dict | None:
         retries = 3
@@ -489,7 +487,7 @@ class MiGPT:
 
                 print("-" * 20)
                 print("问题：" + query + "？")
-                if not self.chatbot.history:
+                if not self.chatbot.has_history():
                     query = f"{query}，{self.config.prompt}"
                 if self.config.mute_xiaoai:
                     await self.stop_if_xiaoai_is_playing()

xiaogpt-1.80/xiaogpt/bot/base_bot.py

@@ -1,25 +0,0 @@
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from typing import Any, AsyncGenerator, TypeVar
-
-from xiaogpt.config import Config
-
-T = TypeVar("T", bound="BaseBot")
-
-
-class BaseBot(ABC):
-    history: list
-
-    @abstractmethod
-    async def ask(self, query: str, **options: Any) -> str:
-        pass
-
-    @abstractmethod
-    async def ask_stream(self, query: str, **options: Any) -> AsyncGenerator[str, None]:
-        pass
-
-    @classmethod
-    @abstractmethod
-    def from_config(cls: type[T], config: Config) -> T:
-        pass