Restructing tools adding FunctionTool and BaseTool

chat19 · May 18, 2023 · 72f9c8e · 72f9c8e
1 parent 6c2a856
commit 72f9c8e
Show file tree

Hide file tree

Showing 8 changed files with 141 additions and 51 deletions.
diff --git a/superagi/agent/agent_prompt_builder.py b/superagi/agent/agent_prompt_builder.py
@@ -3,7 +3,7 @@
 from pydantic.types import List
 
 from superagi.agent.agent_prompt import AgentPrompt
-from superagi.tools.base_tool import Tool
+from superagi.tools.base_tool import BaseTool
 
 FINISH_NAME = "finish"
 
@@ -78,14 +78,14 @@ def add_tools_to_prompt(self, final_string):
     final_string = final_string + finish_string + "\n\n"
     return final_string
 
-  def _generate_command_string(self, tool: Tool) -> str:
+  def _generate_command_string(self, tool: BaseTool) -> str:
     output = f"{tool.name}: {tool.description}"
     print(tool.args)
     output += f", args json schema: {json.dumps(tool.args)}"
     return output
 
   @classmethod
-  def get_autogpt_prompt(cls, ai_name:str, ai_role: str, goals: List[str], tools: List[Tool]) -> str:
+  def get_autogpt_prompt(cls, ai_name:str, ai_role: str, goals: List[str], tools: List[BaseTool]) -> str:
     # Initialize the PromptGenerator object
     prompt_builder = AgentPromptBuilder()
     prompt_builder.set_ai_name(ai_name)

diff --git a/superagi/agent/super_agi.py b/superagi/agent/super_agi.py
@@ -10,7 +10,7 @@
 from superagi.agent.output_parser import BaseOutputParser, AgentOutputParser
 from superagi.types.common import BaseMessage, HumanMessage, AIMessage, SystemMessage
 from superagi.llms.base_llm import BaseLlm
-from superagi.tools.base_tool import Tool
+from superagi.tools.base_tool import BaseTool
 from superagi.vector_store.base import VectorStore
 from superagi.vector_store.document import Document
 
@@ -24,7 +24,7 @@ def __init__(self,
                  llm: BaseLlm,
                  memory: VectorStore,
                  output_parser: BaseOutputParser,
-                 tools: List[Tool],
+                 tools: List[BaseTool],
                  ):
         self.ai_name = ai_name
         self.ai_role = ai_role
@@ -40,7 +40,7 @@ def from_llm_and_tools(
             ai_name: str,
             ai_role: str,
             memory: VectorStore,
-            tools: List[Tool],
+            tools: List[BaseTool],
             llm: BaseLlm
     ) -> SuperAgi:
         return cls(

diff --git a/superagi/tools/base_tool.py b/superagi/tools/base_tool.py
@@ -1,85 +1,117 @@
+from abc import abstractmethod
 from functools import wraps
 from typing import Optional, Type, Callable, Any, Union
 
 from pydantic import BaseModel, Field, create_model, validate_arguments, Extra
 from inspect import signature
 
-class _SchemaConfig:
+
+class SchemaSettings:
     """Configuration for the pydantic model."""
     extra = Extra.forbid
     arbitrary_types_allowed = True
 
-def get_filtered_args(
-    inferred_model: Type[BaseModel],
-    func: Callable,
+
+def extract_valid_parameters(
+        inferred_type: Type[BaseModel],
+        function: Callable,
 ) -> dict:
     """Get the arguments from a function's signature."""
-    schema = inferred_model.schema()["properties"]
-    valid_keys = signature(func).parameters
-    return {k: schema[k] for k in valid_keys if k != "run_manager"}
+    schema = inferred_type.schema()["properties"]
+    valid_params = signature(function).parameters
+    return {param: schema[param] for param in valid_params if param != "run_manager"}
 
-def _create_subset_model(
-    name: str, model: BaseModel, field_names: list
+
+def _construct_model_subset(
+        model_name: str, original_model: BaseModel, required_fields: list
 ) -> Type[BaseModel]:
     """Create a pydantic model with only a subset of model's fields."""
     fields = {
-        field_name: (
-            model.__fields__[field_name].type_,
-            model.__fields__[field_name].default,
+        field: (
+            original_model.__fields__[field].type_,
+            original_model.__fields__[field].default,
         )
-        for field_name in field_names
-        if field_name in model.__fields__
+        for field in required_fields
+        if field in original_model.__fields__
     }
-    return create_model(name, **fields)  # type: ignore
+    return create_model(model_name, **fields)  # type: ignore
+
 
-def create_schema_from_function(
-    model_name: str,
-    func: Callable,
+def create_function_schema(
+        schema_name: str,
+        function: Callable,
 ) -> Type[BaseModel]:
     """Create a pydantic schema from a function's signature."""
-    validated = validate_arguments(func, config=_SchemaConfig)  # type: ignore
-    inferred_model = validated.model  # type: ignore
-    if "run_manager" in inferred_model.__fields__:
-        del inferred_model.__fields__["run_manager"]
-    # Pydantic adds placeholder virtual fields we need to strip
-    filtered_args = get_filtered_args(inferred_model, func)
-    return _create_subset_model(
-        f"{model_name}Schema", inferred_model, list(filtered_args)
+    validated = validate_arguments(function, config=SchemaSettings)  # type: ignore
+    inferred_type = validated.model  # type: ignore
+    if "run_manager" in inferred_type.__fields__:
+        del inferred_type.__fields__["run_manager"]
+    valid_parameters = extract_valid_parameters(inferred_type, function)
+    return _construct_model_subset(
+        f"{schema_name}Schema", inferred_type, list(valid_parameters)
     )
 
-class Tool(BaseModel):
+
+class BaseTool(BaseModel):
     name: str = None
     description: str
-    func: Callable
     args_schema: Type[BaseModel] = None
-    coroutine: Optional[Callable] = None
 
     @property
     def args(self):
         # print("args_schema", self.args_schema)
         if self.args_schema is not None:
-            return self.args_schema.__fields__ or self.args_schema.schema()["properties"]
+            return self.args_schema.schema()["properties"]
         else:
-            name = self.name or self.func.__name__
-            args_schema = create_schema_from_function(f"{name}Schema", self.func)
+            name = self.name
+            args_schema = create_function_schema(f"{name}Schema", self.execute)
             # print("args:", args_schema.schema()["properties"])
             return args_schema.schema()["properties"]
 
+    @abstractmethod
     def execute(self, *tool_input):
-        return self.func(*tool_input)
+        pass
+
+    @classmethod
+    def from_function(cls, func: Callable, args_schema: Type[BaseModel] = None):
+        if args_schema:
+            return cls(description=func.__doc__, args_schema=args_schema)
+        else:
+            return cls(description=func.__doc__)
+
+
+class FunctionalTool(BaseTool):
+    name: str = None
+    description: str
+    func: Callable
+    args_schema: Type[BaseModel] = None
+
+    @property
+    def args(self):
+        if self.args_schema is not None:
+            return self.args_schema.schema()["properties"]
+        else:
+            name = self.name
+            args_schema = create_function_schema(f"{name}Schema", self.execute)
+            # print("args:", args_schema.schema()["properties"])
+            return args_schema.schema()["properties"]
 
+    def execute(self, *tool_input):
+        return self.func(*tool_input)
     @classmethod
     def from_function(cls, func: Callable, args_schema: Type[BaseModel] = None):
         if args_schema:
-            return cls(description=func.__doc__, args_schema=args_schema, func=func)
+            return cls(description=func.__doc__, args_schema=args_schema)
         else:
-            return cls(description=func.__doc__, func=func)
+            return cls(description=func.__doc__)
+
 
-def tool(*args: Union[str, Callable], return_direct: bool = False, args_schema: Optional[Type[BaseModel]] = None) -> Callable:
+def tool(*args: Union[str, Callable], return_direct: bool = False,
+         args_schema: Optional[Type[BaseModel]] = None) -> Callable:
     def decorator(func: Callable) -> Callable:
         nonlocal args_schema
 
-        tool_instance = Tool.from_function(func, args_schema)
+        tool_instance = FunctionalTool.from_function(func, args_schema)
 
         @wraps(func)
         def wrapper(*tool_args, **tool_kwargs):

diff --git a/superagi/tools/jira/__init__.py b/superagi/tools/jira/__init__.py
diff --git a/superagi/tools/jira/tool.py b/superagi/tools/jira/tool.py
@@ -0,0 +1,13 @@
+from superagi.tools.base_tool import BaseTool
+
+
+class JiraTool(BaseTool):
+    def __init__(self):
+        super().__init__("Jira", "Helps to create Jira tickets", self.create_jira_ticket)
+
+    def execute(self):
+        print("Jira tool")
+
+    def create_jira_ticket(self, name: str):
+        print("hello ramram", name)
+        return
diff --git a/superagi/tools/twitter/__init__.py b/superagi/tools/twitter/__init__.py
diff --git a/superagi/tools/twitter/send_tweet.py b/superagi/tools/twitter/send_tweet.py
@@ -0,0 +1,43 @@
+import os
+from typing import Type
+
+from pydantic import BaseModel, Field
+
+from superagi.tools.base_tool import BaseTool
+
+import tweepy
+
+class SentTweetSchema(BaseModel):
+    tweet_text: str = Field(
+        ...,
+        description="Text post you want to write on twitter.",
+    )
+
+
+class SendTweetTool(BaseTool):
+    name = "SendTweet"
+    description = (
+        "A wrapper around Twitter. "
+        "Useful to send/write tweet on twitter "
+        "Input should be a search query."
+    )
+    args_schema: Type[SentTweetSchema] = SentTweetSchema
+
+    def execute(self, tweet_text: str):
+        consumer_key = os.environ.get("TW_CONSUMER_KEY")
+        consumer_secret = os.environ.get("TW_CONSUMER_SECRET")
+        access_token = os.environ.get("TW_ACCESS_TOKEN")
+        access_token_secret = os.environ.get("TW_ACCESS_TOKEN_SECRET")
+        bearer_token = os.environ.get("TW_BEARER_SECRET")
+        auth = tweepy.OAuthHandler(consumer_key, consumer_secret)
+        auth.set_access_token(access_token, access_token_secret)
+
+        try:
+            api = tweepy.API(auth)
+            api.update_status(tweet_text)
+            # response = client.create_tweet(text=tweet_text, user_auth=False)
+            # print(response)
+            return "Tweet sent successfully!"
+        except tweepy.TweepyException as e:
+            print(e)
+            return f"Error sending tweet: {e}"
diff --git a/test.py b/test.py
@@ -1,13 +1,13 @@
 from superagi.agent.super_agi import SuperAgi
 from superagi.llms.openai import OpenAi
-from superagi.tools.base_tool import Tool
-from superagi.vector_store.document import Document
+from superagi.tools.base_tool import FunctionalTool
+from superagi.tools.twitter.send_tweet import SendTweetTool
 from superagi.vector_store.embedding.openai import OpenAiEmbedding
 from superagi.vector_store.vector_factory import VectorFactory
 
 memory = VectorFactory.get_vector_storage("PineCone", "super-agent-index1", OpenAiEmbedding())
-memory.add_documents([Document("Hello world")])
-memory.get_matching_text("Hello world")
+# memory.add_documents([Document("Hello world")])
+# memory.get_matching_text("Hello world")
 
 def test_function(name: str):
     print("hello ramram", name)
@@ -19,9 +19,11 @@ def create_campaign(campaign_name: str):
 
 
 tools = [
-    Tool(name="Search", description="Helps to search google", func=test_function),
-    Tool(name="Campaign Create", description="Creates campaign", func=create_campaign),
+    FunctionalTool(name="Search", description="Helps to search google", func=test_function),
+    FunctionalTool(name="Campaign Create", description="Creates campaign", func=create_campaign)
 ]
 
-superagi = SuperAgi.from_llm_and_tools("Super AGI", "Super AGI", memory, tools, OpenAi())
-superagi.execute(["I want to send campaign"])
+send_tool = SendTweetTool()
+send_tool.execute("Innovation isn't a one-time event; it's a culture. It's about daring to question the status quo, nurturing a curiosity that stretches horizons, and constantly seeking new ways to add value #Innovation #ChangeTheWorld")
+# superagi = SuperAgi.from_llm_and_tools("Super AGI", "Super AGI", memory, tools, OpenAi())
+# superagi.execute(["I want to send campaign"])