|
|
|
@ -1,610 +1,176 @@
|
|
|
|
|
# Building Custom Vector Memory Databases with the BaseVectorDatabase Class
|
|
|
|
|
# Integrating the Agent Class with Memory Systems in the Swarms Memory Framework
|
|
|
|
|
|
|
|
|
|
In the age of large language models (LLMs) and AI-powered applications, efficient memory management has become a crucial component. Vector databases, which store and retrieve data in high-dimensional vector spaces, have emerged as powerful tools for handling the vast amounts of data generated and consumed by AI systems. However, integrating vector databases into your applications can be a daunting task, requiring in-depth knowledge of their underlying architectures and APIs.
|
|
|
|
|
## Introduction
|
|
|
|
|
|
|
|
|
|
Enter the `BaseVectorDatabase` class, a powerful abstraction layer designed to simplify the process of creating and integrating custom vector memory databases into your AI applications. By inheriting from this class, developers can build tailored vector database solutions that seamlessly integrate with their existing systems, enabling efficient storage, retrieval, and manipulation of high-dimensional data.
|
|
|
|
|
In this guide, we will cover how to integrate various memory systems from the Swarms Memory framework into an agent class. The Swarms Memory framework allows for the integration of different database-backed memory systems, enabling agents to retain and query long-term knowledge effectively. We'll walk through examples of integrating with Pinecone, ChromaDB, and Faiss, showcasing how to configure custom functions and embed memory functionality into an agent class.
|
|
|
|
|
|
|
|
|
|
In this comprehensive guide, we'll explore the `BaseVectorDatabase` class in detail, covering its core functionality and diving deep into the process of creating custom vector memory databases using popular solutions like PostgreSQL, Pinecone, Chroma, FAISS, and more. Whether you're a seasoned AI developer or just starting to explore the world of vector databases, this guide will provide you with the knowledge and tools necessary to build robust, scalable, and efficient memory solutions for your AI applications.
|
|
|
|
|
## Installation
|
|
|
|
|
|
|
|
|
|
## Understanding the BaseVectorDatabase Class
|
|
|
|
|
|
|
|
|
|
Before we dive into the implementation details, let's take a closer look at the `BaseVectorDatabase` class and its core functionality.
|
|
|
|
|
|
|
|
|
|
The `BaseVectorDatabase` class is an abstract base class that defines the interface for interacting with a vector database. It serves as a blueprint for creating concrete implementations of vector databases, ensuring a consistent and standardized approach to database operations across different systems.
|
|
|
|
|
|
|
|
|
|
The class provides a set of abstract methods that define the essential functionality required for working with vector databases, such as connecting to the database, executing queries, and performing CRUD (Create, Read, Update, Delete) operations.
|
|
|
|
|
|
|
|
|
|
Here's a breakdown of the abstract methods defined in the `BaseVectorDatabase` class:
|
|
|
|
|
|
|
|
|
|
1\. `connect()`: This method establishes a connection to the vector database.
|
|
|
|
|
|
|
|
|
|
2\. `close()`: This method closes the connection to the vector database.
|
|
|
|
|
|
|
|
|
|
3\. `query(query: str)`: This method executes a given query on the vector database.
|
|
|
|
|
|
|
|
|
|
4\. `fetch_all()`: This method retrieves all rows from the result set of a query.
|
|
|
|
|
|
|
|
|
|
5\. `fetch_one()`: This method retrieves a single row from the result set of a query.
|
|
|
|
|
|
|
|
|
|
6\. `add(doc: str)`: This method adds a new record to the vector database.
|
|
|
|
|
|
|
|
|
|
7\. `get(query: str)`: This method retrieves a record from the vector database based on a given query.
|
|
|
|
|
|
|
|
|
|
8\. `update(doc)`: This method updates a record in the vector database.
|
|
|
|
|
|
|
|
|
|
9\. `delete(message)`: This method deletes a record from the vector database.
|
|
|
|
|
|
|
|
|
|
By inheriting from the `BaseVectorDatabase` class and implementing these abstract methods, developers can create concrete vector database implementations tailored to their specific needs and requirements.
|
|
|
|
|
|
|
|
|
|
## Creating a Custom Vector Memory Database
|
|
|
|
|
|
|
|
|
|
Now that we have a solid understanding of the `BaseVectorDatabase` class, let's dive into the process of creating a custom vector memory database by inheriting from this class. Throughout this guide, we'll explore various vector database solutions, including PostgreSQL, Pinecone, Chroma, FAISS, and more, showcasing how to integrate them seamlessly into your AI applications.
|
|
|
|
|
|
|
|
|
|
### Step 1: Inherit from the BaseVectorDatabase Class
|
|
|
|
|
|
|
|
|
|
The first step in creating a custom vector memory database is to inherit from the `BaseVectorDatabase` class. This will provide your custom implementation with the foundational structure and interface defined by the abstract class.
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
|
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
|
from swarms import BaseVectorDatabase
|
|
|
|
|
|
|
|
|
|
class MyCustomVectorDatabase(BaseVectorDatabase):
|
|
|
|
|
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
|
|
|
|
|
# Custom initialization logic
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
First, you need to install the Swarms Memory package:
|
|
|
|
|
|
|
|
|
|
```bash
|
|
|
|
|
$ pip install swarms-memory
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
In the example above, we define a new class `MyCustomVectorDatabase` that inherits from the `BaseVectorDatabase` class. Within the `__init__` method, you can add any custom initialization logic specific to your vector database implementation.
|
|
|
|
|
## Usage Examples
|
|
|
|
|
|
|
|
|
|
### Step 2: Implement the Abstract Methods
|
|
|
|
|
### Integrating Pinecone with the Agent Class
|
|
|
|
|
|
|
|
|
|
The next step is to implement the abstract methods defined in the `BaseVectorDatabase` class. These methods provide the core functionality for interacting with your vector database, such as connecting, querying, and performing CRUD operations.
|
|
|
|
|
Pinecone is a vector database for high-speed querying of large-scale datasets. Here's how you can integrate it with your agent:
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
from swarms import BaseVectorDatabase
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class MyCustomVectorDatabase(BaseVectorDatabase):
|
|
|
|
|
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
|
|
|
|
|
# Custom initialization logic
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def connect(self):
|
|
|
|
|
|
|
|
|
|
# Implementation for connecting to the vector database
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def close(self):
|
|
|
|
|
|
|
|
|
|
# Implementation for closing the vector database connection
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def query(self, query: str):
|
|
|
|
|
|
|
|
|
|
# Implementation for executing a query on the vector database
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def fetch_all(self):
|
|
|
|
|
|
|
|
|
|
# Implementation for fetching all rows from the result set
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def fetch_one(self):
|
|
|
|
|
|
|
|
|
|
# Implementation for fetching a single row from the result set
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def add(self, doc: str):
|
|
|
|
|
|
|
|
|
|
# Implementation for adding a new record to the vector database
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def get(self, query: str):
|
|
|
|
|
|
|
|
|
|
# Implementation for retrieving a record from the vector database
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def update(self, doc):
|
|
|
|
|
|
|
|
|
|
# Implementation for updating a record in the vector database
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def delete(self, message):
|
|
|
|
|
|
|
|
|
|
# Implementation for deleting a record from the vector database
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
In this example, we define placeholders for each of the abstract methods within the `MyCustomVectorDatabase` class. These placeholders will be replaced with the actual implementation logic specific to your chosen vector database solution.
|
|
|
|
|
|
|
|
|
|
### Step 3: Choose and Integrate Your Vector Database Solution
|
|
|
|
|
|
|
|
|
|
With the foundational structure in place, it's time to choose a specific vector database solution and integrate it into your custom implementation. In this guide, we'll explore several popular vector database solutions, including PostgreSQL, Pinecone, Chroma, FAISS, and more, providing examples and guidance on how to integrate them seamlessly.
|
|
|
|
|
|
|
|
|
|
### PostgreSQL Integration
|
|
|
|
|
|
|
|
|
|
PostgreSQL is a powerful open-source relational database management system that supports vector data types and operations, making it a viable choice for building custom vector memory databases.
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
|
|
|
|
|
import psycopg2
|
|
|
|
|
from swarms import BaseVectorDatabase
|
|
|
|
|
|
|
|
|
|
class PostgreSQLVectorDatabase(MyCustomVectorDatabase):
|
|
|
|
|
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
|
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
# PostgreSQL connection details
|
|
|
|
|
|
|
|
|
|
self.conn = psycopg2.connect(
|
|
|
|
|
|
|
|
|
|
host="localhost",
|
|
|
|
|
|
|
|
|
|
database="vector_db",
|
|
|
|
|
|
|
|
|
|
user="postgres",
|
|
|
|
|
|
|
|
|
|
password="your_password"
|
|
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
self.cur = self.conn.cursor()
|
|
|
|
|
|
|
|
|
|
def connect(self):
|
|
|
|
|
|
|
|
|
|
# PostgreSQL connection logic
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def close(self):
|
|
|
|
|
|
|
|
|
|
# Close PostgreSQL connection
|
|
|
|
|
|
|
|
|
|
self.cur.close()
|
|
|
|
|
|
|
|
|
|
self.conn.close()
|
|
|
|
|
|
|
|
|
|
def query(self, query: str):
|
|
|
|
|
|
|
|
|
|
# Execute PostgreSQL query
|
|
|
|
|
|
|
|
|
|
self.cur.execute(query)
|
|
|
|
|
|
|
|
|
|
def fetch_all(self):
|
|
|
|
|
|
|
|
|
|
# Fetch all rows from PostgreSQL result set
|
|
|
|
|
|
|
|
|
|
return self.cur.fetchall()
|
|
|
|
|
|
|
|
|
|
# Implement other abstract methods
|
|
|
|
|
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
In this example, we define a `PostgreSQLVectorDatabase` class that inherits from `MyCustomVectorDatabase`. Within the `__init__` method, we establish a connection to a PostgreSQL database using the `psycopg2` library. We then implement the `connect()`, `close()`, `query()`, and `fetch_all()` methods specific to PostgreSQL.
|
|
|
|
|
|
|
|
|
|
### Pinecone Integration
|
|
|
|
|
|
|
|
|
|
Pinecone is a managed vector database service that provides efficient storage, retrieval, and manipulation of high-dimensional vector data.
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
|
|
|
|
|
import pinecone
|
|
|
|
|
from swarms import BaseVectorDatabase
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class PineconeVectorDatabase(MyCustomVectorDatabase):
|
|
|
|
|
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
|
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
# Pinecone initialization
|
|
|
|
|
|
|
|
|
|
pinecone.init(api_key="your_api_key", environment="your_environment")
|
|
|
|
|
|
|
|
|
|
self.index = pinecone.Index("your_index_name")
|
|
|
|
|
|
|
|
|
|
def connect(self):
|
|
|
|
|
|
|
|
|
|
# Pinecone connection logic
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def close(self):
|
|
|
|
|
|
|
|
|
|
# Close Pinecone connection
|
|
|
|
|
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
def query(self, query: str):
|
|
|
|
|
|
|
|
|
|
# Execute Pinecone query
|
|
|
|
|
|
|
|
|
|
results = self.index.query(query)
|
|
|
|
|
|
|
|
|
|
return results
|
|
|
|
|
|
|
|
|
|
def add(self, doc: str):
|
|
|
|
|
from typing import List, Dict, Any
|
|
|
|
|
from swarms_memory import PineconeMemory
|
|
|
|
|
from swarms import Agent, Anthropic
|
|
|
|
|
from transformers import AutoTokenizer, AutoModel
|
|
|
|
|
import torch
|
|
|
|
|
import os
|
|
|
|
|
|
|
|
|
|
# Add document to Pinecone index
|
|
|
|
|
# Custom embedding function using a HuggingFace model
|
|
|
|
|
def custom_embedding_function(text: str) -> List[float]:
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
|
|
|
|
|
model = AutoModel.from_pretrained("bert-base-uncased")
|
|
|
|
|
inputs = tokenizer(
|
|
|
|
|
text,
|
|
|
|
|
return_tensors="pt",
|
|
|
|
|
padding=True,
|
|
|
|
|
truncation=True,
|
|
|
|
|
max_length=512,
|
|
|
|
|
)
|
|
|
|
|
with torch.no_grad():
|
|
|
|
|
outputs = model(**inputs)
|
|
|
|
|
embeddings = (
|
|
|
|
|
outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
|
|
|
|
|
)
|
|
|
|
|
return embeddings
|
|
|
|
|
|
|
|
|
|
# Initialize the Pinecone memory wrapper
|
|
|
|
|
pinecone_memory = PineconeMemory(
|
|
|
|
|
api_key="your-api-key",
|
|
|
|
|
environment="your-environment",
|
|
|
|
|
index_name="your-index-name",
|
|
|
|
|
embedding_function=custom_embedding_function,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
self.index.upsert([("id", doc)])
|
|
|
|
|
# Model
|
|
|
|
|
model = Anthropic(anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))
|
|
|
|
|
|
|
|
|
|
# Implement other abstract methods
|
|
|
|
|
# Initialize the agent with Pinecone memory
|
|
|
|
|
agent = Agent(
|
|
|
|
|
agent_name="Financial-Analysis-Agent",
|
|
|
|
|
system_prompt="Agent system prompt here",
|
|
|
|
|
agent_description="Agent performs financial analysis.",
|
|
|
|
|
llm=model,
|
|
|
|
|
long_term_memory=pinecone_memory,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# Run a query
|
|
|
|
|
agent.run("What is the current market trend in AI investments?")
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
In this example, we define a `PineconeVectorDatabase` class that inherits from `MyCustomVectorDatabase`. Within the `__init__` method, we initialize the Pinecone client and create an index. We then implement the `query()` and `add()` methods specific to the Pinecone API.
|
|
|
|
|
### Integrating ChromaDB with the Agent Class
|
|
|
|
|
|
|
|
|
|
### Chroma Integration
|
|
|
|
|
|
|
|
|
|
Chroma is an open-source vector database library that provides efficient storage, retrieval, and manipulation of vector data using various backends, including DuckDB, Chromadb, and more.
|
|
|
|
|
ChromaDB is a simple, high-performance vector store for use with embeddings. Here's how you can integrate ChromaDB:
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
import logging
|
|
|
|
|
from swarms_memory import ChromaDB
|
|
|
|
|
from swarms import Agent, Anthropic
|
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
# 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 = os.path.join(self.docs_folder, file)
|
|
|
|
|
_, ext = os.path.splitext(file)
|
|
|
|
|
data = data_to_text(file)
|
|
|
|
|
added_to_db = self.add(str(data))
|
|
|
|
|
print(f"{file} added to Database")
|
|
|
|
|
|
|
|
|
|
return added_to_db
|
|
|
|
|
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
In this example, we define a `ChromaVectorDatabase` class that inherits from `MyCustomVectorDatabase`. Within the `__init__` method, we create a Chroma client and get or create a collection. We then implement the `query()` and `add()` methods specific to the Chroma API.
|
|
|
|
|
|
|
|
|
|
### FAISS Integration
|
|
|
|
|
|
|
|
|
|
FAISS (Facebook AI Similarity Search) is a library for efficient similarity search and clustering of dense vectors, developed by Meta AI.
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
|
|
|
|
|
import faiss
|
|
|
|
|
|
|
|
|
|
class FAISSVectorDatabase(MyCustomVectorDatabase):
|
|
|
|
|
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
|
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
# FAISS initialization
|
|
|
|
|
|
|
|
|
|
self.index = faiss.IndexFlatL2(64) # Assuming 64-dimensional vectors
|
|
|
|
|
|
|
|
|
|
self.index_path = "faiss_index.index"
|
|
|
|
|
|
|
|
|
|
def connect(self):
|
|
|
|
|
|
|
|
|
|
# FAISS connection logic
|
|
|
|
|
|
|
|
|
|
self.index = faiss.read_index(self.index_path)
|
|
|
|
|
|
|
|
|
|
def close(self):
|
|
|
|
|
|
|
|
|
|
# Close FAISS connection
|
|
|
|
|
|
|
|
|
|
faiss.write_index(self.index, self.index_path)
|
|
|
|
|
|
|
|
|
|
def query(self, query: str):
|
|
|
|
|
|
|
|
|
|
# Execute FAISS query
|
|
|
|
|
|
|
|
|
|
query_vector = # Convert query to vector
|
|
|
|
|
|
|
|
|
|
distances, indices = self.index.search(query_vector, k=10)
|
|
|
|
|
|
|
|
|
|
return [(self.index.reconstruct(i), d) for i, d in zip(indices, distances)]
|
|
|
|
|
|
|
|
|
|
def add(self, doc: str):
|
|
|
|
|
|
|
|
|
|
# Add document to FAISS index
|
|
|
|
|
|
|
|
|
|
doc_vector = # Convert doc to vector
|
|
|
|
|
# Initialize the ChromaDB client
|
|
|
|
|
chromadb_memory = ChromaDB(
|
|
|
|
|
metric="cosine",
|
|
|
|
|
output_dir="finance_agent_rag",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
self.index.add(doc_vector)
|
|
|
|
|
# Model
|
|
|
|
|
model = Anthropic(anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))
|
|
|
|
|
|
|
|
|
|
# Implement other abstract methods
|
|
|
|
|
# Initialize the agent with ChromaDB memory
|
|
|
|
|
agent = Agent(
|
|
|
|
|
agent_name="Financial-Analysis-Agent",
|
|
|
|
|
system_prompt="Agent system prompt here",
|
|
|
|
|
agent_description="Agent performs financial analysis.",
|
|
|
|
|
llm=model,
|
|
|
|
|
long_term_memory=chromadb_memory,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# Run a query
|
|
|
|
|
agent.run("What are the components of a startup's stock incentive equity plan?")
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
Now, how do you integrate a vector datbase with an agent? This is how:
|
|
|
|
|
### Integrating Faiss with the Agent Class
|
|
|
|
|
|
|
|
|
|
## Integrate Memory with `Agent`
|
|
|
|
|
Faiss is a library for efficient similarity search and clustering of dense vectors. Here's how you can integrate Faiss:
|
|
|
|
|
|
|
|
|
|
```python
|
|
|
|
|
from typing import List, Dict, Any
|
|
|
|
|
from swarms_memory.faiss_wrapper import FAISSDB
|
|
|
|
|
from swarms import Agent, Anthropic
|
|
|
|
|
from transformers import AutoTokenizer, AutoModel
|
|
|
|
|
import torch
|
|
|
|
|
import os
|
|
|
|
|
|
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
|
|
|
|
|
|
# Import the OpenAIChat model and the Agent struct
|
|
|
|
|
from swarms import Agent, OpenAIChat
|
|
|
|
|
|
|
|
|
|
# Load the environment variables
|
|
|
|
|
load_dotenv()
|
|
|
|
|
|
|
|
|
|
# Get the API key from the environment
|
|
|
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# Initilaize the chromadb client
|
|
|
|
|
faiss = FAISSVectorDatabase()
|
|
|
|
|
|
|
|
|
|
# Initialize the language model
|
|
|
|
|
llm = OpenAIChat(
|
|
|
|
|
temperature=0.5,
|
|
|
|
|
model_name="gpt-4",
|
|
|
|
|
openai_api_key=api_key,
|
|
|
|
|
max_tokens=1000,
|
|
|
|
|
# Custom embedding function using a HuggingFace model
|
|
|
|
|
def custom_embedding_function(text: str) -> List[float]:
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
|
|
|
|
|
model = AutoModel.from_pretrained("bert-base-uncased")
|
|
|
|
|
inputs = tokenizer(
|
|
|
|
|
text,
|
|
|
|
|
return_tensors="pt",
|
|
|
|
|
padding=True,
|
|
|
|
|
truncation=True,
|
|
|
|
|
max_length=512,
|
|
|
|
|
)
|
|
|
|
|
with torch.no_grad():
|
|
|
|
|
outputs = model(**inputs)
|
|
|
|
|
embeddings = (
|
|
|
|
|
outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
|
|
|
|
|
)
|
|
|
|
|
return embeddings
|
|
|
|
|
|
|
|
|
|
# Initialize the FAISS memory wrapper
|
|
|
|
|
faiss_memory = FAISSDB(
|
|
|
|
|
dimension=768,
|
|
|
|
|
index_type="Flat",
|
|
|
|
|
embedding_function=custom_embedding_function,
|
|
|
|
|
metric="cosine",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
## Initialize the workflow
|
|
|
|
|
# Model
|
|
|
|
|
model = Anthropic(anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))
|
|
|
|
|
|
|
|
|
|
# Initialize the agent with Faiss memory
|
|
|
|
|
agent = Agent(
|
|
|
|
|
llm=llm,
|
|
|
|
|
max_loops=4,
|
|
|
|
|
autosave=True,
|
|
|
|
|
dashboard=True,
|
|
|
|
|
long_term_memory=faiss,
|
|
|
|
|
agent_name="Financial-Analysis-Agent",
|
|
|
|
|
system_prompt="Agent system prompt here",
|
|
|
|
|
agent_description="Agent performs financial analysis.",
|
|
|
|
|
llm=model,
|
|
|
|
|
long_term_memory=faiss_memory,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# Run the workflow on a task
|
|
|
|
|
out = agent.run("Generate a 10,000 word blog on health and wellness.")
|
|
|
|
|
print(out)
|
|
|
|
|
# Run a query
|
|
|
|
|
agent.run("Explain the differences between various types of financial instruments.")
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
In this example, we define a `FAISSVectorDatabase` class that inherits from `MyCustomVectorDatabase`. Within the `__init__` method, we create a FAISS index and set the index path. We then implement the `connect()`, `close()`, `query()`, and `add()` methods specific to the FAISS library, assuming 64-dimensional vectors for simplicity.
|
|
|
|
|
|
|
|
|
|
These examples provide a starting point for integrating various vector database solutions into your custom implementation. Each solution has its own strengths, weaknesses, and trade-offs, so it's essential to carefully evaluate your requirements and choose the solution that best fits your needs.
|
|
|
|
|
|
|
|
|
|
### Step 4: Add Custom Functionality and Optimizations
|
|
|
|
|
|
|
|
|
|
Once you've integrated your chosen vector database solution, you can further extend and optimize your custom implementation by adding custom functionality and performance optimizations.
|
|
|
|
|
|
|
|
|
|
#### Custom Functionality:
|
|
|
|
|
|
|
|
|
|
- **Indexing Strategies**: Implement custom indexing strategies to optimize search performance and memory usage.
|
|
|
|
|
|
|
|
|
|
- **Data Preprocessing**: Add data preprocessing logic to handle different data formats, perform embedding, and prepare data for storage in the vector database.
|
|
|
|
|
## Mermaid Graphs for Visualizing Integration
|
|
|
|
|
|
|
|
|
|
- **Query Optimization**: Introduce query optimization techniques, such as query caching, result filtering, or query rewriting, to improve query performance.
|
|
|
|
|
To help visualize the integration process, here's a Mermaid graph illustrating how an agent interacts with the memory systems:
|
|
|
|
|
|
|
|
|
|
- **Data Partitioning**: Implement data partitioning strategies to distribute data across multiple nodes or shards for better scalability and performance.
|
|
|
|
|
|
|
|
|
|
- **Metadata Management**: Introduce metadata management capabilities to store and retrieve additional information associated with the vector data.
|
|
|
|
|
|
|
|
|
|
Performance Optimizations:
|
|
|
|
|
|
|
|
|
|
- **Caching**: Implement caching mechanisms to reduce redundant computations and improve response times.
|
|
|
|
|
|
|
|
|
|
- **Asynchronous Operations**: Utilize asynchronous programming techniques to improve concurrency and responsiveness.
|
|
|
|
|
|
|
|
|
|
- **Multithreading and Parallelization**: Leverage multithreading and parallelization to distribute computationally intensive tasks across multiple cores or processors.
|
|
|
|
|
|
|
|
|
|
- **Load Balancing**: Implement load balancing strategies to distribute workloads evenly across multiple nodes or instances.
|
|
|
|
|
|
|
|
|
|
- **Monitoring and Profiling**: Introduce monitoring and profiling tools to identify performance bottlenecks and optimize critical sections of your code.
|
|
|
|
|
|
|
|
|
|
By adding custom functionality and performance optimizations, you can tailor your custom vector memory database to meet the specific requirements of your AI applications, ensuring efficient and scalable data management.
|
|
|
|
|
|
|
|
|
|
### Best Practices and Considerations
|
|
|
|
|
|
|
|
|
|
Building custom vector memory databases is a powerful but complex endeavor. To ensure the success and longevity of your implementation, it's essential to follow best practices and consider potential challenges and considerations.
|
|
|
|
|
|
|
|
|
|
1\. **Scalability and Performance Testing**: Vector databases can quickly grow in size and complexity as your AI applications handle increasing amounts of data. Thoroughly test your implementation for scalability and performance under various load conditions, and optimize accordingly.
|
|
|
|
|
|
|
|
|
|
2\. **Data Quality and Integrity**: Ensure that the data stored in your vector database is accurate, consistent, and free from duplicates or errors. Implement data validation and cleansing mechanisms to maintain data quality and integrity.
|
|
|
|
|
|
|
|
|
|
3\. **Security and Access Control**: Vector databases may store sensitive or proprietary data. Implement robust security measures, such as encryption, access controls, and auditing mechanisms, to protect your data from unauthorized access or breaches.
|
|
|
|
|
|
|
|
|
|
4\. **Distributed Architectures**: As your data and workloads grow, consider implementing distributed architectures to distribute the storage and computational load across multiple nodes or clusters. This can improve scalability, fault tolerance, and overall performance.
|
|
|
|
|
|
|
|
|
|
5\. **Data Versioning and Backup**: Implement data versioning and backup strategies to ensure data integrity and enable recovery in case of errors or system failures.
|
|
|
|
|
|
|
|
|
|
6\. **Documentation and Maintainability**: Well-documented code and comprehensive documentation are essential for ensuring the long-term maintainability and extensibility of your custom vector memory database implementation.
|
|
|
|
|
|
|
|
|
|
7\. **Continuous Integration and Deployment**: Adopt continuous integration and deployment practices to streamline the development, testing, and deployment processes, ensuring that changes are thoroughly tested and deployed efficiently.
|
|
|
|
|
|
|
|
|
|
8\. **Compliance and Regulatory Requirements**: Depending on your industry and use case, ensure that your custom vector memory database implementation complies with relevant regulations and standards, such as data privacy laws or industry-specific guidelines.
|
|
|
|
|
|
|
|
|
|
9\. **Community Engagement and Collaboration**: Stay engaged with the vector database community, participate in discussions, and collaborate with other developers to share knowledge, best practices, and insights.
|
|
|
|
|
|
|
|
|
|
By following these best practices and considering potential challenges, you can build robust, scalable, and efficient custom vector memory databases that meet the demanding requirements of modern AI applications.
|
|
|
|
|
|
|
|
|
|
# Conclusion
|
|
|
|
|
|
|
|
|
|
In this comprehensive guide, we've explored the `BaseVectorDatabase` class and its role in simplifying the process of creating custom vector memory databases. We've covered the core functionality of the class, walked through the step-by-step process of inheriting and extending its functionality, and provided examples of integrating popular vector database solutions like PostgreSQL, Pinecone, Chroma, and FAISS.
|
|
|
|
|
```mermaid
|
|
|
|
|
graph TD;
|
|
|
|
|
A[Agent] -->|Queries| B[Memory System]
|
|
|
|
|
B --> C{Pinecone / ChromaDB / Faiss}
|
|
|
|
|
C --> D[Embedding Function]
|
|
|
|
|
D --> E[LLM Model]
|
|
|
|
|
E --> F[Query Results]
|
|
|
|
|
F -->|Returns| A
|
|
|
|
|
```
|
|
|
|
|
|
|
|
|
|
Building custom vector memory databases empowers developers to create tailored and efficient data management solutions that seamlessly integrate with their AI applications. By leveraging the power of vector databases, you can unlock new possibilities in data storage, retrieval, and manipulation, enabling your AI systems to handle vast amounts of high-dimensional data with ease.
|
|
|
|
|
This graph shows the flow from the agent sending queries to the memory system, which processes them using the embedding function and LLM model, and finally returns the results back to the agent.
|
|
|
|
|
|
|
|
|
|
Remember, the journey of building custom vector memory databases is an iterative and collaborative process that requires continuous learning, adaptation, and refinement. Embrace the challenges, stay up-to-date with the latest developments in vector databases and AI, and continuously strive to optimize and enhance your implementations.
|
|
|
|
|
## Conclusion
|
|
|
|
|
|
|
|
|
|
As you embark on this journey, keep in mind the importance of scalability, performance, data quality, security, and compliance. Foster an environment of collaboration, knowledge sharing, and community engagement to ensure that your custom vector memory databases are robust, reliable, and capable of meeting the ever-evolving demands of the AI landscape.
|
|
|
|
|
Integrating various memory systems from the Swarms Memory framework into the agent class enables the creation of powerful, memory-augmented agents capable of retaining and recalling information over time. Whether you're using Pinecone, ChromaDB, or Faiss, the process involves initializing the memory system, embedding functions, and then passing this memory system to the agent class. The examples and visualizations provided should help you get started with building your own memory-augmented agents.
|
|
|
|
|
|
|
|
|
|
So, dive in, leverage the power of the `BaseVectorDatabase` class, and create the custom vector memory databases that will drive the future of AI-powered applications.
|
|
|
|
|
Happy coding!
|
|
|
|
|