Merge pull request #41 from small-thinking/create-abstract-llm

Create abstract llm

.gitignore

@@ -167,4 +167,5 @@ scripts
 
 # Polymind learned facts and tools
 knowledge/facts/**
-knowledge/tools/**
+knowledge/tools/**
+use_cases

polymind/__init__.py

@@ -6,4 +6,4 @@
 from .core.tool import BaseTool
 
 # Expose the tools
-from .tools.oai_tool import OpenAIChatTool
+from .tools.llm_tool import OpenAIChatTool

polymind/tools/oai_tool.py → polymind/tools/llm_tool.py

@@ -4,6 +4,7 @@
 """
 
 import os
+from abc import ABC, abstractmethod
 from typing import List
 
 import numpy as np
@@ -16,7 +17,67 @@
 from polymind.tools.rest_api_tool import RestAPITool
 
 
-class OpenAIChatTool(BaseTool):
+class LLMTool(BaseTool, ABC):
+    """LLM tool defines the basic properties of the language model tools."""
+
+    max_tokens: int = Field(..., description="The maximum number of tokens for the chat.")
+    temperature: float = Field(default=1.0, description="The temperature for the chat.")
+    top_p: float = Field(
+        default=0.1,
+        description="The top p for the chat. Top p is used to prevent the model from generating unlikely words.",
+    )
+    stop: str = Field(default=None, description="The stop sequence for the chat.")
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._set_client()
+
+    @abstractmethod
+    def _set_client(self):
+        """Set the client for the language model."""
+        pass
+
+    @abstractmethod
+    async def _invoke(self, input: Message) -> Message:
+        """Invoke the language model with the input message and return the response message.
+
+        Args:
+            input (Message): The input message to the language model. The message should contain the below keys:
+                - prompt: The prompt for the chat.
+                - system_prompt: The system prompt for the chat.
+                - max_tokens: The maximum number of tokens for the chat.
+                - temperature: The temperature for the chat.
+                - top_p: The top p for the chat.
+                - stop: The stop sequence for the chat.
+        """
+        pass
+
+    async def _execute(self, input: Message) -> Message:
+        """Execute the tool and return the result.
+        The input message should contain a "prompt" and optionally a "system_prompt".
+        """
+
+        # Validate the input message.
+        prompt = input.get("prompt", "")
+        system_prompt = input.get("system_prompt", self.system_prompt)
+        if not prompt:
+            raise ValueError("Prompt cannot be empty.")
+        input.content.update(
+            {
+                "max_tokens": self.max_tokens,
+                "temperature": self.temperature,
+                "top_p": self.top_p,
+                "system_prompt": system_prompt,
+            }
+        )
+        if self.stop:
+            input.content["stop"] = self.stop
+
+        response_message = await self._invoke(input)
+        return response_message
+
+
+class OpenAIChatTool(LLMTool):
     """OpenAITool is a bridge to OpenAI APIs.
     The tool can be initialized with llm_name, system_prompt, max_tokens, and temperature.
     The input message of this tool should contain a "prompt", and optionally a "system_prompt".
@@ -40,9 +101,10 @@ class OpenAIChatTool(BaseTool):
     system_prompt: str = Field(default="You are a helpful AI assistant.")
     max_tokens: int = Field(default=1500)
     temperature: float = Field(default=0.7)
+    stop: str = Field(default=None)
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    def _set_client(self):
+        """Set the client for the language model."""
         self.client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
     def input_spec(self) -> List[Param]:
@@ -65,6 +127,24 @@ def input_spec(self) -> List[Param]:
                 example="hello, how are you?",
                 description="The prompt for the chat.",
             ),
+            Param(
+                name="max_tokens",
+                type="int",
+                example="1500",
+                description="The maximum number of tokens for the chat.",
+            ),
+            Param(
+                name="temperature",
+                type="float",
+                example="0.7",
+                description="The temperature for the chat.",
+            ),
+            Param(
+                name="top_p",
+                type="float",
+                example="0.1",
+                description="The top p for the chat.",
+            ),
         ]
 
     def output_spec(self) -> List[Param]:
@@ -83,27 +163,39 @@ def output_spec(self) -> List[Param]:
             ),
         ]
 
-    async def _execute(self, input: Message) -> Message:
+    async def _invoke(self, input: Message) -> Message:
         """Execute the tool and return the result.
         The derived class must implement this method to define the behavior of the tool.
 
         Args:
-            input (Message): The input to the tool carried in a message.
+            input (Message): The input to the tool carried in a message. The message should contain the below keys:
+                - prompt: The prompt for the chat.
+                - system_prompt: The system prompt for the chat.
+                - max_tokens: The maximum number of tokens for the chat.
+                - temperature: The temperature for the chat.
+                - top_p: The top p for the chat.
+                - stop: The stop sequence for the chat.
 
         Returns:
             Message: The result of the tool carried in a message.
         """
         prompt = input.get("prompt", "")
         system_prompt = input.get("system_prompt", self.system_prompt)
-        if not prompt:
-            raise ValueError("Prompt cannot be empty.")
+        temperature = input.get("temperature", self.temperature)
+        max_tokens = input.get("max_tokens", self.max_tokens)
+        top_p = input.get("top_p", self.top_p)
+        stop = input.get("stop", self.stop)
 
         response = await self.client.chat.completions.create(
             model=self.llm_name,
             messages=[
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": prompt},
             ],
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            stop=stop,
         )
         content = response.choices[0].message.content
         response_message = Message(content={"response": content})

polymind/tools/tool_management_tool.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 from typing import Any, Dict, List
 
 import faiss
@@ -10,7 +11,8 @@
 from polymind.core.indexer import Indexer
 from polymind.core.message import Message
 from polymind.core.tool import BaseTool, Param
-from polymind.tools.oai_tool import OpenAIEmbeddingTool
+from polymind.tools.llm_tool import (LLMTool, OpenAIChatTool,
+                                     OpenAIEmbeddingTool)
 
 
 class ToolIndexer(Indexer):
@@ -36,7 +38,7 @@ def _extra_input_spec(self) -> List[Param]:
                 example="rest-api-tool",
             ),
             Param(
-                name="desscriptions",
+                name="descriptions",
                 type="List[str]",
                 description="The descriptions of the tool to be indexed.",
                 example="""[
@@ -193,3 +195,188 @@ async def _execute(self, input_message: Message) -> Message:
         embedding = await self.embedder._embedding([requirement])
         candidates = self._find_top_k_candidates(embedding)
         return Message(content={"candidates": candidates})
+
+
+class ToolCreator(BaseTool):
+    tool_name: str = "tool-creator"
+
+    llm_tool: LLMTool = Field(default=None, description="The LLM tool to generate the tool code.")
+
+    descriptions: List[str] = [
+        "ToolCreator is a tool to generate the tool code based on the requirement",
+        "ToolCreator is a tool to create the tool.",
+        "ToolCreator is a codegen tool to generate the tool code based on the requirement.",
+    ]
+
+    learned_tool_folder: str = Field(default="./knowledge/tools", description="The folder to store the learned tools.")
+
+    system_prompt: str = Field(
+        default="""
+        You are a code generator. You are given a requirement to create a tool.
+        The generated tool should inherit the below BaseTool class.
+        Please put your answer into a ```python``` code block.
+
+        --- python below ---
+        # File: {{tool_name}}.py
+
+        class BaseTool(BaseModel, ABC):
+            '''The base class of the tool.
+            In an agent system, a tool is an object that can be used to perform a task.
+            For example, search for information from the internet, query a database,
+            or perform a calculation.
+            '''
+
+            tool_name: str = Field(..., description="The name of the tool.")
+            descriptions: List[str] = Field(
+                ...,
+                min_length=3,
+                description='''The descriptions of the tool. The descriptions will be
+                converted to embeddings and used to index the tool. One good practice is to
+                describe the tools with the following aspects: what the tool does, and describe
+                the tools from different perspectives.
+                ''',
+            )
+
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                load_dotenv(override=True)
+
+            def __str__(self):
+                return self.tool_name
+
+            @field_validator("tool_name")
+            def check_tool_name(cls, v: str) -> str:
+                if not v:
+                    raise ValueError("The tool_name must not be empty.")
+                return v
+
+            @field_validator("descriptions")
+            def check_descriptions(cls, v: List[str]) -> List[str]:
+                if len(v) < 3:
+                    raise ValueError("The descriptions must have at least 3 items. The more the better.")
+                return v
+
+            def get_descriptions(self) -> List[str]:
+                return self.descriptions
+
+            async def __call__(self, input: Message) -> Message:
+                '''Makes the instance callable, delegating to the execute method.
+                This allows the instance to be used as a callable object, simplifying the execution of the tool.
+
+                Args:
+                    input (Message): The input message to the tool.
+
+                Returns:
+                    Message: The output message from the tool.
+                '''
+                return await self._execute(input)
+
+            def get_spec(self) -> str:
+                '''Return the input and output specification of the tool.
+
+                Returns:
+                    Tuple[List[Param], List[Param]]: The input and output specification of the tool.
+                '''
+                input_json_obj = []
+                for param in self.input_spec():
+                    input_json_obj.append(param.to_json_obj())
+                output_json_obj = []
+                for param in self.output_spec():
+                    output_json_obj.append(param.to_json_obj())
+                spec_json_obj = {
+                    "input_message": input_json_obj,
+                    "output_message": output_json_obj,
+                }
+                return json.dumps(spec_json_obj, indent=4)
+
+            @abstractmethod
+            def input_spec(self) -> List[Param]:
+                '''Return the specification of the input parameters.'''
+                pass
+
+            @abstractmethod
+            def output_spec(self) -> List[Param]:
+                '''Return the specification of the output parameters.'''
+                pass
+
+            @abstractmethod
+            async def _execute(self, input: Message) -> Message:
+                '''Execute the tool and return the result.
+                The derived class must implement this method to define the behavior of the tool.
+
+                Args:
+                    input (Message): The input to the tool carried in a message.
+
+                Returns:
+                    Message: The result of the tool carried in a message.
+                '''
+                pass
+        --- python above ---
+
+        In addition to any needed package, the implementation should also import the below packages:
+
+        from polymind.core.message import Message
+        from polymind.core.tool import BaseTool, Param
+        """,
+        description="The system prompt to generate the tool code.",
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        if not self.llm_tool:
+            self.llm_tool = OpenAIChatTool(llm_name="gpt-3.5-turbo", system_prompt=self.system_prompt)
+
+    def input_spec(self) -> List[Param]:
+        return [
+            Param(
+                name="requirement",
+                type="str",
+                description="Text-based requirement description for tool creation.",
+                example="I need a tool to call REST API.",
+            )
+        ]
+
+    def output_spec(self) -> List[Param]:
+        return [
+            Param(
+                name="status",
+                type="str",
+                description="Whether the tool is successfully created and indexed.",
+                example="success",
+            )
+        ]
+
+    async def _execute(self, input_message: Message) -> Message:
+        """Leverage the LLM tool to generate the tool code based on the requirement.
+
+        Args:
+            input_message (Message): The input message containing the requirement.
+                The message should contain the "requirement" key.
+        """
+        requirement = input_message.content.get("requirement", "")
+        if not requirement:
+            raise ValueError("Requirement not provided.")
+        input_message = Message(content={"prompt": requirement, "system_prompt": self.system_prompt})
+        tool_code = await self.llm_tool(input=input_message)
+
+        # Check the ```python``` code block in the tool_code and extract the source code.
+        tool_code = tool_code.content["response"]
+        tool_code = re.search(r"```python(.*?)```", tool_code, re.DOTALL)
+        if not tool_code:
+            raise ValueError("Tool code not found.")
+        tool_code = tool_code.group(1).strip()
+        # Save the tool code to a file.
+        tool_name = re.search(r"class\s+(.*?)\(", tool_code).group(1)
+        if not tool_name:
+            raise ValueError("Tool name not found.")
+        tool_name = tool_name.strip().lower()
+        tool_file_name = f"{tool_name}.py"
+        tool_file_path = os.path.join(self.learned_tool_folder, "test_learned", tool_file_name)
+
+        if not os.path.exists(os.path.dirname(tool_file_path)):
+            os.makedirs(os.path.dirname(tool_file_path), exist_ok=True)
+        with open(tool_file_path, "w") as f:
+            f.write(tool_code)
+
+        # Output message
+        return Message(content={"status": "success"})