Skip to content

Commit

Permalink
Removed Pipeline for now. Added memory.
Browse files Browse the repository at this point in the history
  • Loading branch information
SpyderRex committed Sep 1, 2024
1 parent ffd303b commit 0dc4a5f
Show file tree
Hide file tree
Showing 11 changed files with 57 additions and 79 deletions.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,5 @@ __pycache__/
Workspace/

venv/

find_os.py
95 changes: 31 additions & 64 deletions main.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -97,14 +94,26 @@ 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)
if "wolframalpha_tool" in agent_config["tools"]:
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"],
Expand All @@ -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"]]
Expand All @@ -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()
5 changes: 4 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion squadai/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
)
Expand Down
6 changes: 5 additions & 1 deletion squadai/cli/cli.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -16,6 +18,8 @@
from .run_squad import run_squad
from .train_squad import train_squad

load_dotenv()


@click.group()
def squadai():
Expand Down Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion squadai/squad.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
4 changes: 3 additions & 1 deletion squadai/tools/tool_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion squadai/utilities/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions squadai/utilities/evaluators/task_evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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```"

Expand Down Expand Up @@ -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()
Expand Down
5 changes: 4 additions & 1 deletion squadai/utilities/planning_handler.py
Original file line number Diff line number Diff line change
@@ -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")
Expand All @@ -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

Expand Down
9 changes: 3 additions & 6 deletions system_message.txt
Original file line number Diff line number Diff line change
@@ -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": [
Expand Down Expand Up @@ -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.
Expand All @@ -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.


0 comments on commit 0dc4a5f

Please sign in to comment.