From c90ba1c11e2e754161889d8ef89680b8da5404d8 Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 19 Mar 2024 17:35:03 -0700 Subject: [PATCH] [TOOL AGENT] --- README.md | 51 ++++++++---- playground/agents/tool_agent_pydantic.py | 45 ---------- swarms/agents/omni_modal_agent.py | 100 ++++++----------------- swarms/agents/tool_agent.py | 21 +++-- swarms/models/odin.py | 2 +- swarms/tools/tool.py | 7 +- swarms/utils/json_utils.py | 21 +---- tool_agent_pydantic.py | 49 +++++++++++ 8 files changed, 137 insertions(+), 159 deletions(-) delete mode 100644 playground/agents/tool_agent_pydantic.py create mode 100644 tool_agent_pydantic.py diff --git a/README.md b/README.md index 97314c3f..20ce32a3 100644 --- a/README.md +++ b/README.md @@ -68,41 +68,60 @@ agent.run("Generate a 10,000 word blog on health and wellness.") ### `ToolAgent` -ToolAgent is an agent that outputs JSON using any model from huggingface. It takes an example schema and performs a task, outputting JSON. It is versatile, easy to use, and customizable. +ToolAgent is an agent that can use tools through JSON function calling. It intakes any open source model from huggingface and is extremely modular and plug in and play. We need help adding general support to all models soon. ```python -# Import necessary libraries +from pydantic import BaseModel, Field from transformers import AutoModelForCausalLM, AutoTokenizer from swarms import ToolAgent +from swarms.utils.json_utils import base_model_to_json # Load the pre-trained model and tokenizer -model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b") +model = AutoModelForCausalLM.from_pretrained( + "databricks/dolly-v2-12b", + load_in_4bit=True, + device_map="auto", +) tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b") -# Define a JSON schema for person's information -json_schema = { - "type": "object", - "properties": { - "name": {"type": "string"}, - "age": {"type": "number"}, - "is_student": {"type": "boolean"}, - "courses": {"type": "array", "items": {"type": "string"}}, - }, -} + +# Initialize the schema for the person's information +class Schema(BaseModel): + name: str = Field(..., title="Name of the person") + agent: int = Field(..., title="Age of the person") + is_student: bool = Field( + ..., title="Whether the person is a student" + ) + courses: list[str] = Field( + ..., title="List of courses the person is taking" + ) + + +# Convert the schema to a JSON string +tool_schema = base_model_to_json(Schema) # Define the task to generate a person's information -task = "Generate a person's information based on the following schema:" +task = ( + "Generate a person's information based on the following schema:" +) # Create an instance of the ToolAgent class -agent = ToolAgent(model=model, tokenizer=tokenizer, json_schema=json_schema) +agent = ToolAgent( + name="dolly-function-agent", + description="Ana gent to create a child data", + model=model, + tokenizer=tokenizer, + json_schema=tool_schema, +) # Run the agent to generate the person's information generated_data = agent.run(task) # Print the generated data -print(generated_data) +print(f"Generated data: {generated_data}") + ``` diff --git a/playground/agents/tool_agent_pydantic.py b/playground/agents/tool_agent_pydantic.py deleted file mode 100644 index c61fc7b9..00000000 --- a/playground/agents/tool_agent_pydantic.py +++ /dev/null @@ -1,45 +0,0 @@ -# Import necessary libraries -from transformers import AutoModelForCausalLM, AutoTokenizer -from pydantic import BaseModel - -# from swarms import ToolAgent -from swarms.utils.json_utils import base_model_schema_to_json - -# Load the pre-trained model and tokenizer -model = AutoModelForCausalLM.from_pretrained( - "databricks/dolly-v2-12b", - load_in_4bit=True, - device_map="auto", -) -tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b") - - -class Schema(BaseModel): - name: str - agent: int - is_student: bool - courses: list[str] - - -json_schema = str(base_model_schema_to_json(Schema)) -print(json_schema) - -# # Define the task to generate a person's information -# task = ( -# "Generate a person's information based on the following schema:" -# ) - -# # Create an instance of the ToolAgent class -# agent = ToolAgent( -# name="dolly-function-agent", -# description="Ana gent to create a child data", -# model=model, -# tokenizer=tokenizer, -# json_schema=json_schema, -# ) - -# # Run the agent to generate the person's information -# generated_data = agent.run(task) - -# # Print the generated data -# print(f"Generated data: {generated_data}") diff --git a/swarms/agents/omni_modal_agent.py b/swarms/agents/omni_modal_agent.py index 8f2dabc5..4af03906 100644 --- a/swarms/agents/omni_modal_agent.py +++ b/swarms/agents/omni_modal_agent.py @@ -9,11 +9,11 @@ from langchain_experimental.autonomous_agents.hugginggpt.task_planner import ( load_chat_planner, ) from transformers import load_tool +from swarms.utils.loguru_logger import logger +from swarms.structs.agent import Agent -from swarms.structs.message import Message - -class OmniModalAgent: +class OmniModalAgent(Agent): """ OmniModalAgent LLM -> Plans -> Tasks -> Tools -> Response @@ -42,9 +42,13 @@ class OmniModalAgent: def __init__( self, llm: BaseLanguageModel, - # tools: List[BaseTool] + verbose: bool = False, + *args, + **kwargs, ): + super().__init__(llm=llm, *args, **kwargs) self.llm = llm + self.verbose = verbose print("Loading tools...") self.tools = [ @@ -67,79 +71,29 @@ class OmniModalAgent: ] ] + # Load the chat planner and response generator self.chat_planner = load_chat_planner(llm) self.response_generator = load_response_generator(llm) - # self.task_executor = TaskExecutor + self.task_executor = TaskExecutor self.history = [] - def run(self, input: str) -> str: + def run(self, task: str) -> str: """Run the OmniAgent""" - plan = self.chat_planner.plan( - inputs={ - "input": input, - "hf_tools": self.tools, - } - ) - self.task_executor = TaskExecutor(plan) - self.task_executor.run() - - response = self.response_generator.generate( - {"task_execution": self.task_executor} - ) - - return response - - def chat(self, msg: str = None, streaming: bool = False): - """ - Run chat - - Args: - msg (str, optional): Message to send to the agent. Defaults to None. - language (str, optional): Language to use. Defaults to None. - streaming (bool, optional): Whether to stream the response. Defaults to False. - - Returns: - str: Response from the agent - - Usage: - -------------- - agent = MultiModalAgent() - agent.chat("Hello") - - """ - - # add users message to the history - self.history.append(Message("User", msg)) - - # process msg try: - response = self.agent.run(msg) - - # add agent's response to the history - self.history.append(Message("Agent", response)) - - # if streaming is = True - if streaming: - return self._stream_response(response) - else: - response - + plan = self.chat_planner.plan( + inputs={ + "input": task, + "hf_tools": self.tools, + } + ) + self.task_executor = TaskExecutor(plan) + self.task_executor.run() + + response = self.response_generator.generate( + {"task_execution": self.task_executor} + ) + + return response except Exception as error: - error_message = f"Error processing message: {str(error)}" - - # add error to history - self.history.append(Message("Agent", error_message)) - - return error_message - - def _stream_response(self, response: str = None): - """ - Yield the response token by token (word by word) - - Usage: - -------------- - for token in _stream_response(response): - print(token) - - """ - yield from response.split() + logger.error(f"Error running the agent: {error}") + return f"Error running the agent: {error}" diff --git a/swarms/agents/tool_agent.py b/swarms/agents/tool_agent.py index 8e6adf9d..55733215 100644 --- a/swarms/agents/tool_agent.py +++ b/swarms/agents/tool_agent.py @@ -1,10 +1,10 @@ -from typing import Any +from typing import Any, Optional, Callable -from swarms.models.base_llm import AbstractLLM +from swarms.structs.agent import Agent from swarms.tools.format_tools import Jsonformer -class ToolAgent(AbstractLLM): +class ToolAgent(Agent): """ Represents a tool agent that performs a specific task using a model and tokenizer. @@ -67,16 +67,23 @@ class ToolAgent(AbstractLLM): tokenizer: Any = None, json_schema: Any = None, max_number_tokens: int = 500, + parsing_function: Optional[Callable] = None, *args, **kwargs, ): - super().__init__(*args, **kwargs) + super().__init__( + agent_name=name, + agent_description=description, + sop=f"{name} {description} {str(json_schema)}" * args, + **kwargs, + ) self.name = name self.description = description self.model = model self.tokenizer = tokenizer self.json_schema = json_schema self.max_number_tokens = max_number_tokens + self.parsing_function = parsing_function def run(self, task: str, *args, **kwargs): """ @@ -104,7 +111,11 @@ class ToolAgent(AbstractLLM): **kwargs, ) - out = self.toolagent() + if self.parsing_function: + out = self.parsing_function(self.toolagent()) + else: + out = self.toolagent() + return out except Exception as error: print(f"[Error] [ToolAgent] {error}") diff --git a/swarms/models/odin.py b/swarms/models/odin.py index 68bfaffd..288fe0dd 100644 --- a/swarms/models/odin.py +++ b/swarms/models/odin.py @@ -2,7 +2,7 @@ import os import supervision as sv from tqdm import tqdm -from ultralytics_example import YOLO +from ultralytics import YOLO from swarms.models.base_llm import AbstractLLM from swarms.utils.download_weights_from_url import ( diff --git a/swarms/tools/tool.py b/swarms/tools/tool.py index 2a5e3f87..1bb0730f 100644 --- a/swarms/tools/tool.py +++ b/swarms/tools/tool.py @@ -1 +1,6 @@ -from langchain.tools import BaseTool, StructuredTool, Tool, tool # noqa F401 +from langchain.tools import ( + BaseTool, + StructuredTool, + Tool, + tool, +) # noqa F401 diff --git a/swarms/utils/json_utils.py b/swarms/utils/json_utils.py index c1be4dfd..0902d2c7 100644 --- a/swarms/utils/json_utils.py +++ b/swarms/utils/json_utils.py @@ -3,7 +3,7 @@ import json from pydantic import BaseModel -def base_model_schema_to_json(model: BaseModel, indent: int = 3): +def base_model_to_json(model: BaseModel, indent: int = 3): """ Converts the JSON schema of a base model to a formatted JSON string. @@ -13,7 +13,8 @@ def base_model_schema_to_json(model: BaseModel, indent: int = 3): Returns: str: The JSON schema of the base model as a formatted JSON string. """ - return json.dumps(model.model_json_schema(), indent=indent) + out = model.model_json_schema() + return str_to_json(out, indent=indent) def extract_json_from_str(response: str): @@ -34,22 +35,6 @@ def extract_json_from_str(response: str): return json.loads(response[json_start : json_end + 1]) -def base_model_to_json(base_model_instance: BaseModel) -> str: - """ - Convert a Pydantic base model instance to a JSON string. - - Args: - base_model_instance (BaseModel): Instance of the Pydantic base model. - - Returns: - str: JSON string representation of the base model instance. - """ - model_dict = base_model_instance.dict() - json_string = json.dumps(model_dict) - - return json_string - - def str_to_json(response: str, indent: int = 3): """ Converts a string representation of JSON to a JSON object. diff --git a/tool_agent_pydantic.py b/tool_agent_pydantic.py new file mode 100644 index 00000000..da5f4825 --- /dev/null +++ b/tool_agent_pydantic.py @@ -0,0 +1,49 @@ +from pydantic import BaseModel, Field +from transformers import AutoModelForCausalLM, AutoTokenizer + +from swarms import ToolAgent +from swarms.utils.json_utils import base_model_to_json + +# Load the pre-trained model and tokenizer +model = AutoModelForCausalLM.from_pretrained( + "databricks/dolly-v2-12b", + load_in_4bit=True, + device_map="auto", +) +tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b") + + +# Initialize the schema for the person's information +class Schema(BaseModel): + name: str = Field(..., title="Name of the person") + agent: int = Field(..., title="Age of the person") + is_student: bool = Field( + ..., title="Whether the person is a student" + ) + courses: list[str] = Field( + ..., title="List of courses the person is taking" + ) + + +# Convert the schema to a JSON string +tool_schema = base_model_to_json(Schema) + +# Define the task to generate a person's information +task = ( + "Generate a person's information based on the following schema:" +) + +# Create an instance of the ToolAgent class +agent = ToolAgent( + name="dolly-function-agent", + description="Ana gent to create a child data", + model=model, + tokenizer=tokenizer, + json_schema=tool_schema, +) + +# Run the agent to generate the person's information +generated_data = agent.run(task) + +# Print the generated data +print(f"Generated data: {generated_data}")