Adding llm thinking & webscraper tool

alfa2beta3 · May 27, 2023 · 07661a0 · 07661a0
1 parent bf2a0e6
commit 07661a0
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 1 deletion.
diff --git a/superagi/agent/agent_prompt_builder.py b/superagi/agent/agent_prompt_builder.py
@@ -114,7 +114,10 @@ def get_autogpt_prompt(cls, ai_name:str, ai_role: str, goals: List[str], tools:
     prompt_builder.add_constraint("No user assistance")
     prompt_builder.add_constraint("Ensure the command and args are as per current plan and reasoning.")
     prompt_builder.add_constraint(
-      'Exclusively use the commands listed in double quotes e.g. "command name"'
+      'Exclusively use the commands listed in double quotes e.g. "command name".'
+    )
+    prompt_builder.add_constraint(
+      'You should always try to use LlmThinkingTool first instead of directly jumping to google search'
     )
 
     # Add tools to the PromptGenerator object

diff --git a/superagi/tools/llm_thinking/__init__.py b/superagi/tools/llm_thinking/__init__.py
diff --git a/superagi/tools/llm_thinking/tools.py b/superagi/tools/llm_thinking/tools.py
@@ -0,0 +1,34 @@
+from typing import Type
+
+from pydantic import BaseModel, Field
+
+from superagi.llms.base_llm import BaseLlm
+from superagi.tools.base_tool import BaseTool
+
+
+class LlmTaskSchema(BaseModel):
+    task_description: str = Field(
+        ...,
+        description="Text describing the task for which the LLM should generate a response.",
+    )
+
+class LlmThinkingTool(BaseTool):
+    llm: BaseLlm = None
+    name = "LlmThinkingTool"
+    description = (
+        "A tool that interacts with any given LLM "
+        "to generate text given a certain task description."
+    )
+    args_schema: Type[LlmTaskSchema] = LlmTaskSchema
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def _execute(self, task_description: str = "") -> str:
+        try:
+            messages = [{"role": "system", "content": task_description}]
+            result = self.llm.chat_completion(messages)
+            return result["content"]
+        except Exception as e:
+            print(e)
+            return f"Error generating text: {e}"
diff --git a/superagi/tools/web_scraper/__init__.py b/superagi/tools/web_scraper/__init__.py
diff --git a/superagi/tools/web_scraper/tools.py b/superagi/tools/web_scraper/tools.py
@@ -0,0 +1,25 @@
+from typing import Type
+
+from pydantic import BaseModel, Field
+
+from superagi.helper.webpage_extractor import WebpageExtractor
+from superagi.tools.base_tool import BaseTool
+
+
+class WebScraperSchema(BaseModel):
+    website_url: str = Field(
+        ...,
+        description="website url",
+    )
+
+
+class WebScraperTool(BaseTool):
+    name = "WebScraperTool"
+    description = (
+        "Used to scrape website urls and extract text content"
+    )
+    args_schema: Type[WebScraperSchema] = WebScraperSchema
+
+    def _execute(self, website_url: str) -> tuple:
+        content = WebpageExtractor().extract_with_bs4(website_url)
+        return content[:1000]