[FEAT][BlockList]

pull/339/head
Kye 1 year ago committed by evelynmitchell
parent 39d61cd4de
commit f17ee89a3c

@ -0,0 +1,97 @@
import os
from dotenv import load_dotenv
from transformers import AutoModelForCausalLM, AutoTokenizer
# Import the models, structs, and telemetry modules
from swarms import (
Gemini,
GPT4VisionAPI,
Mixtral,
OpenAI,
ToolAgent,
BlocksList,
)
# Load the environment variables
load_dotenv()
# Get the environment variables
openai_api_key = os.getenv("OPENAI_API_KEY")
gemini_api_key = os.getenv("GEMINI_API_KEY")
# Tool Agent
model = AutoModelForCausalLM.from_pretrained(
"databricks/dolly-v2-12b"
)
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
json_schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "number"},
"is_student": {"type": "boolean"},
"courses": {"type": "array", "items": {"type": "string"}},
},
}
toolagent = ToolAgent(
model=model, tokenizer=tokenizer, json_schema=json_schema
)
# Blocks List which enables you to build custom swarms by adding classes or functions
swarm = BlocksList(
"SocialMediaSwarm",
"A swarm of social media agents",
[
OpenAI(openai_api_key=openai_api_key),
Mixtral(),
GPT4VisionAPI(openai_api_key=openai_api_key),
Gemini(gemini_api_key=gemini_api_key),
],
)
# Add the new block to the swarm
swarm.add(toolagent)
# Remove a block from the swarm
swarm.remove(toolagent)
# Update a block in the swarm
swarm.update(toolagent)
# Get a block at a specific index
block_at_index = swarm.get(0)
# Get all blocks in the swarm
all_blocks = swarm.get_all()
# Get blocks by name
openai_blocks = swarm.get_by_name("OpenAI")
# Get blocks by type
gpt4_blocks = swarm.get_by_type("GPT4VisionAPI")
# Get blocks by ID
block_by_id = swarm.get_by_id(toolagent.id)
# Get blocks by parent
blocks_by_parent = swarm.get_by_parent(swarm)
# Get blocks by parent ID
blocks_by_parent_id = swarm.get_by_parent_id(swarm.id)
# Get blocks by parent name
blocks_by_parent_name = swarm.get_by_parent_name(swarm.name)
# Get blocks by parent type
blocks_by_parent_type = swarm.get_by_parent_type(type(swarm).__name__)
# Get blocks by parent description
blocks_by_parent_description = swarm.get_by_parent_description(
swarm.description
)
# Run the block in the swarm
inference = swarm.run_block(toolagent, "Hello World")
print(inference)

@ -170,6 +170,7 @@ class HuggingfaceLLM(AbstractLLM):
"bnb_4bit_use_double_quant": True,
"bnb_4bit_quant_type": "nf4",
"bnb_4bit_compute_dtype": dtype,
"bnb_4bit_compute_dtype": dtype,
}
bnb_config = BitsAndBytesConfig(**quantization_config)
@ -189,6 +190,7 @@ class HuggingfaceLLM(AbstractLLM):
self.model_id, *args, **kwargs
).to(self.device)
def print_error(self, error: str):
"""Print error"""
print(colored(f"Error: {error}", "red"))
@ -263,7 +265,7 @@ class HuggingfaceLLM(AbstractLLM):
*args,
**kwargs,
)
return self.tokenizer.decode(
outputs[0], skip_special_tokens=True
)

@ -1,5 +1,6 @@
from typing import Callable, Any
import logging
from functools import wraps
from typing import Any, Callable
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@ -20,6 +21,7 @@ def block(
Callable[..., Any]: The transformed function.
"""
@wraps(function)
def wrapper(*args, **kwargs):
# Here you can add code to execute the function on various hardwares
# For now, we'll just call the function normally

@ -36,6 +36,15 @@ class BlocksList(BaseStructure):
get_by_parent_name(parent_name: str): Get blocks by parent name.
get_by_parent_type(parent_type: str): Get blocks by parent type.
get_by_parent_description(parent_description: str): Get blocks by parent description.
Examples:
>>> from swarms.structs.block import Block
>>> from swarms.structs.blockslist import BlocksList
>>> block = Block("block", "A block")
>>> blockslist = BlocksList("blockslist", "A list of blocks", [block])
>>> blockslist
"""
def __init__(
@ -47,8 +56,10 @@ class BlocksList(BaseStructure):
**kwargs,
):
super().__init__(name=name, description=description, **kwargs)
self.parent = parent
self.name = name
self.description = description
self.blocks = blocks
self.parent = parent
def add(self, block: Any):
self.blocks.append(block)
@ -64,6 +75,26 @@ class BlocksList(BaseStructure):
def get_all(self):
return self.blocks
def run_block(self, block: Any, task: str, *args, **kwargs):
"""Run the block for the specified task.
Args:
task (str): The task to be performed by the block.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
The output of the block.
Raises:
Exception: If an error occurs during the execution of the block.
"""
try:
return block.run(task, *args, **kwargs)
except Exception as error:
print(f"[Error] [Block] {error}")
raise error
def get_by_name(self, name: str):
return [block for block in self.blocks if block.name == name]

Loading…
Cancel
Save