From 0dc4a5fc9d2e55aef54c6a7e55584313eb4dc2cd Mon Sep 17 00:00:00 2001 From: Spyder Rex Date: Sun, 1 Sep 2024 22:19:48 +0000 Subject: [PATCH] Removed Pipeline for now. Added memory. --- .gitignore | 2 + main.py | 95 ++++++------------- requirements.txt | 5 +- squadai/agent.py | 2 +- squadai/cli/cli.py | 6 +- squadai/squad.py | 2 +- squadai/tools/tool_usage.py | 4 +- squadai/utilities/converter.py | 2 +- .../utilities/evaluators/task_evaluator.py | 4 +- squadai/utilities/planning_handler.py | 5 +- system_message.txt | 9 +- 11 files changed, 57 insertions(+), 79 deletions(-) diff --git a/.gitignore b/.gitignore index 6c09bcc..a5e751e 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ __pycache__/ Workspace/ venv/ + +find_os.py diff --git a/main.py b/main.py index 1e4f242..5f23566 100644 --- a/main.py +++ b/main.py @@ -1,8 +1,7 @@ import os import json -import asyncio from typing import List, Dict, Any -from squadai import Agent, Task, Squad, Process, Pipeline +from squadai import Agent, Task, Squad, Process from squadai.squadai_tools import FileWriterTool from langchain_community.tools import DuckDuckGoSearchRun from langchain_community.tools import WikipediaQueryRun @@ -28,8 +27,6 @@ file_search_tool = tool_registry.get("file_search") move_file_tool = tool_registry.get("move_file") - - # Set up Groq API (make sure to set your API key in the environment variables) groq_api_key = os.getenv("GROQ_API_KEY") client = Groq(api_key=groq_api_key) @@ -97,7 +94,7 @@ def create_agent(agent_config: Dict[str, Any]) -> Agent: Create an Agent instance from a configuration dictionary. """ tools = [] - if "search_tool" in agent_config["tools"]: + if "duckduckgo_tool" in agent_config["tools"]: tools.append(duckduckgo_tool) if "wikipedia_tool" in agent_config["tools"]: tools.append(wikipedia_tool) @@ -105,6 +102,18 @@ def create_agent(agent_config: Dict[str, Any]) -> Agent: tools.append(wolframalpha_tool) if "write_file_tool" in agent_config["tools"]: tools.append(write_file_tool) + if "read_file_tool" in agent_config["tools"]: + tools.append(read_file_tool) + if "list_directory_tool" in agent_config["tools"]: + tools.append(list_directory_tool) + if "copy_file_tool" in agent_config["tools"]: + tools.append(copy_file_tool) + if "delete_file_tool" in agent_config["tools"]: + tools.append(delete_file_tool) + if "file_search_tool" in agent_config["tools"]: + tools.append(file_search_tool) + if "move_file_tool" in agent_config["tools"]: + tools.append(move_file_tool) return Agent( role=agent_config["role"], @@ -128,7 +137,7 @@ def create_task(task_config: Dict[str,Any], agents: List[Agent]) -> Task: def create_squad(squad_config: Dict[str, Any], agents: List[Agent], tasks: List[Task]) -> Squad: """ - Create a Squad instace from configuration dictionary, a list of available agents, and a list of tasks. + Create a Squad instance from configuration dictionary, a list of available agents, and a list of tasks. """ squad_agents = [next(agent for agent in agents if agent.role == role) for role in squad_config["agents"]] squad_tasks = [next(task for task in tasks if task.description == desc) for desc in squad_config["tasks"]] @@ -148,89 +157,47 @@ def create_squad(squad_config: Dict[str, Any], agents: List[Agent], tasks: List[ agents=squad_agents, tasks=squad_tasks, process=Process.sequential if squad_config["process"] == "sequential" else Process.hierarchical, + memory=True, + embedder={ + "provider": "cohere", + "config": { + "model": "embed-english-v3.0", "vector_dimension": 1024 + } + }, verbose=squad_config["verbose"], manager_agent=manager ) -def create_pipeline(pipeline_config: List[str], squads: List[Squad]) -> Pipeline: - """ - Create a Pipeline instance from a configuration list and a list of available squads. - """ - pipeline_squads = [next(squad for squad in squads if squad.name == squad_name) for squad_name in pipeline_config] - return Pipeline(stages=pipeline_squads) - -def needs_wolframalpha(config: Dict[str, Any]) -> bool: +def run_squad(config: Dict[str, Any], user_prompt: str) -> str: """ - Check if any agent in the configuration uses the Wolfram Alpha tool. - """ - for agent in config["agents"]: - if "wolframalpha_tool" in agent["tools"]: - return True - return False - -def run_sync_squad(config: Dict[str, Any], user_prompt: str) -> str: - """ - Run the squad synchronously when Wolfram Alpha is needed or there's no pipeline. + Run the squad based on the configuration. """ agents = [create_agent(agent_config) for agent_config in config["agents"]] tasks = [create_task(task_config, agents) for task_config in config["tasks"]] squads = [create_squad(squad_config, agents, tasks) for squad_config in config["squads"]] - if "pipeline" in config: - inputs = [{"initial query": user_prompt}] - pipeline = create_pipeline(config["pipeline"], squads) - result = pipeline.kickoff(inputs=inputs) - return result[0].raw - else: - result = "Running squads individually (sync mode):\n" - for squad in squads: - squad_result = squad.kickoff() - result += f"\n{squad.name}: {squad_result}" - return result - -async def run_async_squad(config: Dict[str, Any], user_prompt: str) -> str: - """ - Run the squad asynchronously when a pipeline is used and Wolfram Alpha is not needed. - """ - agents = [create_agent(agent_config) for agent_config in config["agents"]] - tasks = [create_task(task_config, agents) for task_config in config["tasks"]] - squads = [create_squad(squad_config, agents, tasks) for squad_config in config["squads"]] - - if "pipeline" in config: - inputs = [{"initial query": user_prompt}] - pipeline = create_pipeline(config["pipeline"], squads) - result = await pipeline.kickoff(inputs=inputs) - return result[0].raw - else: - result = "No pipeline defined. Running squads individually (async mode):\n" - for squad in squads: - squad_result = await squad.kickoff() - result += f"\n{squad.name}: {squad_result}" - return result + result = "Running squads:\n" + for squad in squads: + squad_result = squad.kickoff() + result += f"\n{squad.name}: {squad_result}" + return result def run_dynamic_squad(user_prompt: str) -> str: """ Run a dynamically created squadAI based on the user's prompt. """ config = get_squad_config(user_prompt) - - if needs_wolframalpha(config) or "pipeline" not in config: - return run_sync_squad(config, user_prompt) - else: - return asyncio.run(run_async_squad(config, user_prompt)) + return run_squad(config, user_prompt) def main(): user_prompt = input("Enter your goal for squadAI: ") try: result = run_dynamic_squad(user_prompt) - print("-----------------------------") - print("SquadAI Result:") - print(result) - except Exception as e: print(f"An error occurred: {str(e)}") print("Please try again with a different prompt or check your configuration.") + if __name__ == "__main__": main() diff --git a/requirements.txt b/requirements.txt index d8fa4ba..d856dc1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -84,9 +84,10 @@ kubernetes==30.1.0 langchain==0.2.14 langchain-cohere==0.1.9 langchain-community==0.2.12 -langchain-core==0.2.34 +langchain-core==0.2.37 langchain-experimental==0.0.64 langchain-groq==0.1.9 +langchain-ollama==0.1.3 langchain-openai==0.1.22 langchain-text-splitters==0.2.2 langsmith==0.1.104 @@ -106,6 +107,7 @@ multitasking==0.0.11 mypy-extensions==1.0.0 numpy==1.26.4 oauthlib==3.2.2 +ollama==0.3.2 onnxruntime==1.19.0 openai==1.42.0 opentelemetry-api==1.26.0 @@ -137,6 +139,7 @@ pyasn1_modules==0.4.0 pycparser==2.22 pydantic==2.8.2 pydantic_core==2.20.1 +pyee==11.1.0 Pygments==2.18.0 PyJWT==2.9.0 pypdf==4.3.1 diff --git a/squadai/agent.py b/squadai/agent.py index 044f845..2f607d3 100644 --- a/squadai/agent.py +++ b/squadai/agent.py @@ -77,7 +77,7 @@ class Agent(BaseAgent): ) llm: Any = Field( default_factory=lambda: ChatGroq( - model=os.environ.get("GROQ_MODEL_NAME", "llama-3.1-70b-versatile") + model=os.environ.get("GROQ_MODEL_NAME") ), description="Language model that will run the agent.", ) diff --git a/squadai/cli/cli.py b/squadai/cli/cli.py index 84338e4..f031e69 100644 --- a/squadai/cli/cli.py +++ b/squadai/cli/cli.py @@ -1,7 +1,9 @@ +import os from typing import Optional import click import pkg_resources +from dotenv import load_dotenv from squadai.cli.create_squad import create_squad from squadai.cli.create_pipeline import create_pipeline @@ -16,6 +18,8 @@ from .run_squad import run_squad from .train_squad import train_squad +load_dotenv() + @click.group() def squadai(): @@ -142,7 +146,7 @@ def reset_memories(long, short, entities, kickoff_outputs, all): "-m", "--model", type=str, - default="llama-3.1-70b-versatile", + default=os.getenv("GROQ_MODEL_NAME"), help="LLM Model to run the tests on the Squad. For now only accepting only Llama model.", ) def test(n_iterations: int, model: str): diff --git a/squadai/squad.py b/squadai/squad.py index 126244b..fd86de7 100644 --- a/squadai/squad.py +++ b/squadai/squad.py @@ -105,7 +105,7 @@ class Squad(BaseModel): description="Whether the squad should use memory to store memories of it's execution", ) embedder: Optional[dict] = Field( - default={"provider": "spacy"}, + default={"provider": "cohere"}, description="Configuration for the embedder to be used for the squad.", ) usage_metrics: Optional[UsageMetrics] = Field( diff --git a/squadai/tools/tool_usage.py b/squadai/tools/tool_usage.py index 7e041d0..874fd9b 100644 --- a/squadai/tools/tool_usage.py +++ b/squadai/tools/tool_usage.py @@ -6,13 +6,15 @@ from langchain_core.tools import BaseTool from langchain_groq import ChatGroq +from dotenv import load_dotenv from squadai.agents.tools_handler import ToolsHandler from squadai.tools.tool_calling import InstructorToolCalling, ToolCalling from squadai.utilities import I18N, Converter, ConverterError, Printer +load_dotenv() -GROQ_BIGGER_MODELS = ["llama-3.1-70b-versatile"] +GROQ_BIGGER_MODELS = [os.getenv("GROQ_MODEL_NAME")] class ToolUsageErrorException(Exception): diff --git a/squadai/utilities/converter.py b/squadai/utilities/converter.py index 73ac9fa..dba5e87 100644 --- a/squadai/utilities/converter.py +++ b/squadai/utilities/converter.py @@ -195,7 +195,7 @@ def convert_with_instructions( def get_conversion_instructions(model: Type[BaseModel], llm: Any) -> str: instructions = "I'm gonna convert this raw text into valid JSON." - if not is_gpt(llm): + if not is_llama(llm): model_schema = PydanticSchemaParser(model=model).get_schema() instructions = f"{instructions}\n\nThe json should have the following structure, with the following keys:\n{model_schema}" return instructions diff --git a/squadai/utilities/evaluators/task_evaluator.py b/squadai/utilities/evaluators/task_evaluator.py index a9f6047..52692bb 100644 --- a/squadai/utilities/evaluators/task_evaluator.py +++ b/squadai/utilities/evaluators/task_evaluator.py @@ -70,7 +70,7 @@ def evaluate(self, task, output) -> TaskEvaluation: instructions = "Convert all responses into valid JSON output." - if not self._is_gpt(self.llm): + if not self._is_llama(self.llm): model_schema = PydanticSchemaParser(model=TaskEvaluation).get_schema() instructions = f"{instructions}\n\nReturn only valid JSON with the following schema:\n```json\n{model_schema}\n```" @@ -116,7 +116,7 @@ def evaluate_training_data( ) instructions = "I'm gonna convert this raw text into valid JSON." - if not self._is_gpt(self.llm): + if not self._is_llama(self.llm): model_schema = PydanticSchemaParser( model=TrainingTaskEvaluation ).get_schema() diff --git a/squadai/utilities/planning_handler.py b/squadai/utilities/planning_handler.py index 96ae6ad..5b2a576 100644 --- a/squadai/utilities/planning_handler.py +++ b/squadai/utilities/planning_handler.py @@ -1,11 +1,14 @@ +import os from typing import Any, List, Optional from langchain_groq import ChatGroq from pydantic import BaseModel, Field +from dotenv import load_dotenv from squadai.agent import Agent from squadai.task import Task +load_dotenv() class PlanPerTask(BaseModel): task: str = Field(..., description="The task for which the plan is created") @@ -27,7 +30,7 @@ def __init__(self, tasks: List[Task], planning_agent_llm: Optional[Any] = None): self.tasks = tasks if planning_agent_llm is None: - self.planning_agent_llm = ChatGroq(model="llama-3.1-70b-versatile") + self.planning_agent_llm = ChatGroq(model=os.getenv("GROQ_MODEL_NAME")) else: self.planning_agent_llm = planning_agent_llm diff --git a/system_message.txt b/system_message.txt index 502561f..3e3f45f 100644 --- a/system_message.txt +++ b/system_message.txt @@ -1,5 +1,5 @@ You are an AI assistant capable of designing configurations for squadAI, an autonomous agent program. Given a user's goal, - create a JSON configuration for a squadAI setup including agents, tasks, squad details, and optionally a pipeline. Respond only with the JSON and nothing else. Make sure you be specific with the user's query or request amd include this in the configuration. + create a JSON configuration for a squadAI setup including agents, tasks, and squad details. Respond only with the JSON and nothing else. Make sure you be specific with the user's query or request amd include this in the configuration. Follow this structure: { "agents": [ @@ -30,10 +30,8 @@ "verbose": true }, ... - ], - "pipeline": ["squad1", "squad2", ...] + ] } - The "pipeline" field is optional and should only be included if the goal requires a series of interdependent tasks. The list of tools at your disposal along with their descriptions is as follows: duckduckgo_tool: A duckduckgo internet search tool. @@ -46,7 +44,6 @@ delete_file_tool: A tool for deleting files. file_search_tool: A tool for searching files. move_file_tool: A tool for moving files. - - You can access any or all of these tools as needed to complete the task you are given. Keep in mind, however, that the Pipeline option requires that it be run asynchronously, and the wolframalpha_tool cannot be used in async. Therefore if this tool is necessary, you cannot use a Pipeline. But if it isn't necessary, you may use Pipeline to complete the task if you deem it necessary. + You can access any or all of these tools as needed to complete the task you are given.