[CLEANUP][Feat][Swarms Memory]

pull/529/head
Kye Gomez 6 months ago
parent bbbbc9e3e3
commit b8f31279ff

@ -118,179 +118,7 @@ agent.run("Generate a 10,000 word blog on health and wellness.")
import os
from dotenv import load_dotenv
from swarms import Agent, OpenAIChat
from playground.memory.chromadb_example import ChromaDB
import logging
import os
import uuid
from typing import Optional
import chromadb
from swarms.utils.data_to_text import data_to_text
from swarms.utils.markdown_message import display_markdown_message
from swarms.memory.base_vectordb import BaseVectorDatabase
# Load environment variables
load_dotenv()
# Results storage using local ChromaDB
class ChromaDB(BaseVectorDatabase):
"""
ChromaDB database
Args:
metric (str): The similarity metric to use.
output (str): The name of the collection to store the results in.
limit_tokens (int, optional): The maximum number of tokens to use for the query. Defaults to 1000.
n_results (int, optional): The number of results to retrieve. Defaults to 2.
Methods:
add: _description_
query: _description_
Examples:
>>> chromadb = ChromaDB(
>>> metric="cosine",
>>> output="results",
>>> llm="gpt3",
>>> openai_api_key=OPENAI_API_KEY,
>>> )
>>> chromadb.add(task, result, result_id)
"""
def __init__(
self,
metric: str = "cosine",
output_dir: str = "swarms",
limit_tokens: Optional[int] = 1000,
n_results: int = 3,
docs_folder: str = None,
verbose: bool = False,
*args,
**kwargs,
):
self.metric = metric
self.output_dir = output_dir
self.limit_tokens = limit_tokens
self.n_results = n_results
self.docs_folder = docs_folder
self.verbose = verbose
# Disable ChromaDB logging
if verbose:
logging.getLogger("chromadb").setLevel(logging.INFO)
# Create Chroma collection
chroma_persist_dir = "chroma"
chroma_client = chromadb.PersistentClient(
settings=chromadb.config.Settings(
persist_directory=chroma_persist_dir,
),
*args,
**kwargs,
)
# Create ChromaDB client
self.client = chromadb.Client()
# Create Chroma collection
self.collection = chroma_client.get_or_create_collection(
name=output_dir,
metadata={"hnsw:space": metric},
*args,
**kwargs,
)
display_markdown_message(
"ChromaDB collection created:"
f" {self.collection.name} with metric: {self.metric} and"
f" output directory: {self.output_dir}"
)
# If docs
if docs_folder:
display_markdown_message(
f"Traversing directory: {docs_folder}"
)
self.traverse_directory()
def add(
self,
document: str,
*args,
**kwargs,
):
"""
Add a document to the ChromaDB collection.
Args:
document (str): The document to be added.
condition (bool, optional): The condition to check before adding the document. Defaults to True.
Returns:
str: The ID of the added document.
"""
try:
doc_id = str(uuid.uuid4())
self.collection.add(
ids=[doc_id],
documents=[document],
*args,
**kwargs,
)
print("-----------------")
print("Document added successfully")
print("-----------------")
return doc_id
except Exception as e:
raise Exception(f"Failed to add document: {str(e)}")
def query(
self,
query_text: str,
*args,
**kwargs,
):
"""
Query documents from the ChromaDB collection.
Args:
query (str): The query string.
n_docs (int, optional): The number of documents to retrieve. Defaults to 1.
Returns:
dict: The retrieved documents.
"""
try:
docs = self.collection.query(
query_texts=[query_text],
n_results=self.n_results,
*args,
**kwargs,
)["documents"]
return docs[0]
except Exception as e:
raise Exception(f"Failed to query documents: {str(e)}")
def traverse_directory(self):
"""
Traverse through every file in the given directory and its subdirectories,
and return the paths of all files.
Parameters:
- directory_name (str): The name of the directory to traverse.
Returns:
- list: A list of paths to each file in the directory and its subdirectories.
"""
added_to_db = False
for root, dirs, files in os.walk(self.docs_folder):
for file in files:
file_path = os.path.join(root, file) # Change this line
_, ext = os.path.splitext(file_path)
data = data_to_text(file_path)
added_to_db = self.add(str(data))
print(f"{file_path} added to Database")
return added_to_db
from swarms_memory import ChromaDB
# Get the API key from the environment
api_key = os.environ.get("OPENAI_API_KEY")
@ -334,218 +162,108 @@ An LLM equipped with long term memory and tools, a full stack agent capable of a
```python
import logging
import os
import uuid
from typing import Optional
import chromadb
from dotenv import load_dotenv
from swarms.utils.data_to_text import data_to_text
from swarms.utils.markdown_message import display_markdown_message
from swarms.memory.base_vectordb import BaseVectorDatabase
from swarms import Agent, OpenAIChat
import subprocess
# Making an instance of the ChromaDB class
memory = ChromaDB(
metric="cosine",
n_results=3,
output_dir="results",
docs_folder="docs",
)
# Load environment variables
load_dotenv()
# Results storage using local ChromaDB
class ChromaDB(BaseVectorDatabase):
# Tools
def terminal(
code: str,
):
"""
ChromaDB database
Run code in the terminal.
Args:
metric (str): The similarity metric to use.
output (str): The name of the collection to store the results in.
limit_tokens (int, optional): The maximum number of tokens to use for the query. Defaults to 1000.
n_results (int, optional): The number of results to retrieve. Defaults to 2.
Methods:
add: _description_
query: _description_
Examples:
>>> chromadb = ChromaDB(
>>> metric="cosine",
>>> output="results",
>>> llm="gpt3",
>>> openai_api_key=OPENAI_API_KEY,
>>> )
>>> chromadb.add(task, result, result_id)
code (str): The code to run in the terminal.
Returns:
str: The output of the code.
"""
out = subprocess.run(
code, shell=True, capture_output=True, text=True
).stdout
return str(out)
def __init__(
self,
metric: str = "cosine",
output_dir: str = "swarms",
limit_tokens: Optional[int] = 1000,
n_results: int = 3,
docs_folder: str = None,
verbose: bool = False,
*args,
**kwargs,
):
self.metric = metric
self.output_dir = output_dir
self.limit_tokens = limit_tokens
self.n_results = n_results
self.docs_folder = docs_folder
self.verbose = verbose
# Disable ChromaDB logging
if verbose:
logging.getLogger("chromadb").setLevel(logging.INFO)
# Create Chroma collection
chroma_persist_dir = "chroma"
chroma_client = chromadb.PersistentClient(
settings=chromadb.config.Settings(
persist_directory=chroma_persist_dir,
),
*args,
**kwargs,
)
def browser(query: str):
"""
Search the query in the browser with the `browser` tool.
# Create ChromaDB client
self.client = chromadb.Client()
Args:
query (str): The query to search in the browser.
# Create Chroma collection
self.collection = chroma_client.get_or_create_collection(
name=output_dir,
metadata={"hnsw:space": metric},
*args,
**kwargs,
)
display_markdown_message(
"ChromaDB collection created:"
f" {self.collection.name} with metric: {self.metric} and"
f" output directory: {self.output_dir}"
)
Returns:
str: The search results.
"""
import webbrowser
# If docs
if docs_folder:
display_markdown_message(
f"Traversing directory: {docs_folder}"
)
self.traverse_directory()
url = f"https://www.google.com/search?q={query}"
webbrowser.open(url)
return f"Searching for {query} in the browser."
def add(
self,
document: str,
*args,
**kwargs,
):
def create_file(file_path: str, content: str):
"""
Add a document to the ChromaDB collection.
Create a file using the file editor tool.
Args:
document (str): The document to be added.
condition (bool, optional): The condition to check before adding the document. Defaults to True.
file_path (str): The path to the file.
content (str): The content to write to the file.
Returns:
str: The ID of the added document.
str: The result of the file creation operation.
"""
try:
doc_id = str(uuid.uuid4())
self.collection.add(
ids=[doc_id],
documents=[document],
*args,
**kwargs,
)
print("-----------------")
print("Document added successfully")
print("-----------------")
return doc_id
except Exception as e:
raise Exception(f"Failed to add document: {str(e)}")
def query(
self,
query_text: str,
*args,
**kwargs,
):
with open(file_path, "w") as file:
file.write(content)
return f"File {file_path} created successfully."
def file_editor(file_path: str, mode: str, content: str):
"""
Query documents from the ChromaDB collection.
Edit a file using the file editor tool.
Args:
query (str): The query string.
n_docs (int, optional): The number of documents to retrieve. Defaults to 1.
file_path (str): The path to the file.
mode (str): The mode to open the file in.
content (str): The content to write to the file.
Returns:
dict: The retrieved documents.
"""
try:
docs = self.collection.query(
query_texts=[query_text],
n_results=self.n_results,
*args,
**kwargs,
)["documents"]
return docs[0]
except Exception as e:
raise Exception(f"Failed to query documents: {str(e)}")
def traverse_directory(self):
"""
Traverse through every file in the given directory and its subdirectories,
and return the paths of all files.
Parameters:
- directory_name (str): The name of the directory to traverse.
Returns:
- list: A list of paths to each file in the directory and its subdirectories.
str: The result of the file editing operation.
"""
added_to_db = False
for root, dirs, files in os.walk(self.docs_folder):
for file in files:
file_path = os.path.join(root, file) # Change this line
_, ext = os.path.splitext(file_path)
data = data_to_text(file_path)
added_to_db = self.add(str(data))
print(f"{file_path} added to Database")
return added_to_db
# Making an instance of the ChromaDB class
memory = ChromaDB(
metric="cosine",
n_results=3,
output_dir="results",
docs_folder="docs",
)
with open(file_path, mode) as file:
file.write(content)
return f"File {file_path} edited successfully."
# Initialize a tool
def search_api(query: str):
# Add your logic here
return query
# Initializing the agent with the Gemini instance and other parameters
# Agent
agent = Agent(
agent_name="Covid-19-Chat",
agent_description=(
"This agent provides information about COVID-19 symptoms."
agent_name="Devin",
system_prompt=(
"Autonomous agent that can interact with humans and other"
" agents. Be Helpful and Kind. Use the tools provided to"
" assist the user. Return all code in markdown format."
),
llm=OpenAIChat(),
llm=llm,
max_loops="auto",
autosave=True,
dashboard=False,
streaming_on=True,
verbose=True,
stopping_token="<DONE>",
interactive=True,
tools=[terminal, browser, file_editor, create_file],
code_interpreter=True,
# streaming=True,
long_term_memory=memory,
stopping_condition="finish",
tools=[search_api],
)
# Defining the task and image path
task = ("What are the symptoms of COVID-19?",)
# Running the agent with the specified task and image
out = agent.run(task)
# Run the agent
out = agent("Create a new file for a plan to take over the world.")
print(out)
```

@ -72,7 +72,7 @@ agent.run("Generate a 10,000 word blog on health and wellness.")
```python
from swarms import Agent, OpenAIChat
from playground.memory.chromadb_example import ChromaDB # Copy and paste the code and put it in your own local directory.
from swarms_memory import ChromaDB # Copy and paste the code and put it in your own local directory.
# Making an instance of the ChromaDB class
memory = ChromaDB(

@ -1,5 +1,5 @@
from swarms import Agent, OpenAIChat
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
from swarms.models.tiktoken_wrapper import TikTokenizer
# Initialize the agent

@ -4,7 +4,7 @@ from dotenv import load_dotenv
# Import the OpenAIChat model and the Agent struct
from swarms import Agent, OpenAIChat
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
# Load the environment variables
load_dotenv()

@ -1,6 +1,6 @@
from swarms import Agent
from swarms.models.llama3_hosted import llama3Hosted
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
from swarms.tools.prebuilt.bing_api import fetch_web_articles_bing_api
# Define the research system prompt

@ -10,7 +10,7 @@ $ pip install swarms
"""
from swarms import Agent, OpenAIChat
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
from swarms.tools.prebuilt.bing_api import fetch_web_articles_bing_api
import os
from dotenv import load_dotenv

@ -1,6 +1,6 @@
from swarms import Agent, OpenAIChat
from typing import List
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
memory = ChromaDB(
metric="cosine",

@ -15,7 +15,7 @@ task -> Understanding Agent [understands the problem better] -> Summarize of the
from swarms import Agent, llama3Hosted, AgentRearrange
from pydantic import BaseModel
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
# Initialize the language model agent (e.g., GPT-3)
llm = llama3Hosted(max_tokens=3000)

@ -11,7 +11,7 @@ Todo [Improvements]
from swarms import Agent
from swarms.models.llama3_hosted import llama3Hosted
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
# Model

@ -1,6 +1,6 @@
from swarms import Agent, OpenAIChat
from swarms.structs.mixture_of_agents import MixtureOfAgents
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
SEC_DATA = """

@ -1,6 +1,6 @@
from swarms import Agent, OpenAIChat
from swarms.structs.mixture_of_agents import MixtureOfAgents
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
SEC_DATA = """

@ -6,7 +6,7 @@ from swarms import (
OpenAIChat,
TogetherLLM,
)
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
from dotenv import load_dotenv
# load the environment variables

@ -10,7 +10,7 @@ $ pip install swarms
"""
from swarms import Agent, OpenAIChat
from playground.memory.chromadb_example import ChromaDB
from swarms_memory import ChromaDB
from swarms.tools.prebuilt.bing_api import fetch_web_articles_bing_api
import os
from dotenv import load_dotenv

@ -96,7 +96,7 @@
"outputs": [],
"source": [
"from swarms import Agent, OpenAIChat\n",
"from playground.memory.chromadb_example import ChromaDB\n",
"from swarms_memory import ChromaDB\n",
"\n",
"# Making an instance of the ChromaDB class\n",
"memory = ChromaDB(\n",

@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
[tool.poetry]
name = "swarms"
version = "5.3.2"
version = "5.3.3"
description = "Swarms - Pytorch"
license = "MIT"
authors = ["Kye Gomez <kye@apac.ai>"]
@ -55,6 +55,7 @@ openai = ">=1.30.1,<2.0"
termcolor = "*"
tiktoken = "*"
networkx = "*"
swarms-memory = "*"

@ -29,3 +29,4 @@ termcolor>=2.4.0
pandas>=2.2.2
fastapi>=0.110.1
networkx
swarms-memory

Loading…
Cancel
Save