Merge pull request #79 from ruanrongman/main

fix iot tool ModuleNotFoundError

promptulate/client/chat.py

@@ -37,6 +37,7 @@
     HumanFeedBackTool
 )
 from promptulate.tools.shell import ShellTool
+
 from promptulate.utils import set_proxy_mode, print_text
 
 MODEL_MAPPING = {"OpenAI": ChatOpenAI, "ErnieBot": ErnieBot}

promptulate/config.py

@@ -53,6 +53,7 @@ def __init__(self):
         self.ernie_bot_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions"
         self.ernie_bot_turbo_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant"
         self.ernie_bot_token_url = "https://aip.baidubce.com/oauth/2.0/token"
+        self.ernie_embedding_v1_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1"
         self.key_default_retry_times = 5
         """If llm(like OpenAI) unable to obtain data, retry request until the data is obtained."""
         self.enable_stdout_hook = True

promptulate/tools/__init__.py

@@ -17,6 +17,8 @@
     SemanticScholarReferenceTool,
     SemanticScholarCitationTool,
 )
+from promptulate.tools.shell import ShellTool
+
 from promptulate.tools.sleep import SleepTool
 
 __all__ = [
@@ -33,10 +35,11 @@
     "ArxivQueryTool",
     "SemanticScholarReferenceTool",
     "SemanticScholarCitationTool",
+    "IotSwitchTool",
     "SemanticScholarQueryTool",
     "PythonREPLTool",
+    "ShellTool",
     "Calculator",
     "SleepTool",
-    "IotSwitchTool",
     "HumanFeedBackTool",
 ]

promptulate/tools/iot_swith_mqtt/api_wrapper.py

@@ -1,7 +1,5 @@
-import paho.mqtt.client as mqtt
-
-
 class IotSwitchAPIWrapper:
-    def run(self, client: mqtt, topic: str, command: str) -> str:
+
+    def run(self, client, topic: str, command: str) -> str:
         client.publish(topic, command)
         return "ok"

promptulate/tools/iot_swith_mqtt/tools.py

@@ -1,8 +1,6 @@
 import re
 from typing import Dict, List
 
-import paho.mqtt.client as mqtt
-
 from promptulate.llms import ChatOpenAI, BaseLLM
 from promptulate.tools import Tool
 from promptulate.tools.iot_swith_mqtt.api_wrapper import IotSwitchAPIWrapper
@@ -13,7 +11,7 @@
 
 
 class IotSwitchTool(Tool):
-    """A tool for running python code in a REPL."""
+    """A tool for switching operations on various devices"""
 
     name: str = "Iot_Switch_Mqtt"
     description: str = (
@@ -23,17 +21,25 @@ class IotSwitchTool(Tool):
         "If the operation of the device is successful, an OK will be returned, otherwise a failure will be returned."
     )
     llm_prompt_template: StringTemplate = prompt_template
-    client: mqtt.Client
     rule_table: List[Dict]
-    api_wrapper: IotSwitchAPIWrapper = IotSwitchAPIWrapper()
 
     def __init__(
         self,
+        client,
         llm: BaseLLM = None,
-        client: mqtt.Client = None,
         rule_table: List[Dict] = None,
+        api_wrapper: IotSwitchAPIWrapper = IotSwitchAPIWrapper(),
         **kwargs
     ):
+        """
+        Args:
+            llm: BaseLLM
+            client: mqtt.Client
+             rule_table: List[Dict]
+             api_wrapper: IotSwitchAPIWrapper
+             **kwargs
+        """
+        self.api_wrapper = api_wrapper
         self.llm: BaseLLM = llm or ChatOpenAI(
             temperature=0.1, enable_preset_description=False
         )
@@ -43,6 +49,14 @@ def __init__(
         super().__init__(**kwargs)
 
     def _run(self, question: str, *args, **kwargs) -> str:
+        try:
+            import paho.mqtt.client as mqtt
+        except ImportError:
+            raise ImportError(
+                "Could not import paho python package. "
+                "This is needed in order to for IotSwitchTool. "
+                "Please install it with `pip install paho-mqtt`."
+            )
         if len(self.rule_table) == 0:
             raise Exception("rule_table is empty")
         else:
@@ -52,7 +66,6 @@ def _run(self, question: str, *args, **kwargs) -> str:
                 key = key + str(index) + "." + s["content"] + "\n"
                 index = index + 1
             prompt = self.llm_prompt_template.format(question=question, rule_key=key)
-            logger.debug(prompt)
             llm_output = self.llm(prompt)
             return self._process_llm_result(llm_output)