implementing chroma

memory
Sashin 1 year ago
parent 1911d8db94
commit e09531131d

@ -18,13 +18,13 @@ class VectorDatabase(ABC):
@abstractmethod @abstractmethod
def query( def query(
self, vector: Dict[str, Any], num_results: int self, text: str, num_results: int
) -> Dict[str, Any]: ) -> Dict[str, Any]:
""" """
Query the database for vectors similar to the given vector. Query the database for vectors similar to the given vector.
Args: Args:
vector (Dict[str, Any]): The vector to compare against. text (Dict[str, Any]): The vector to compare against.
num_results (int): The number of similar vectors to return. num_results (int): The number of similar vectors to return.
Returns: Returns:

@ -1,10 +1,13 @@
from typing import List from abc import ABC, abstractmethod
from typing import Any, Dict, List
from chromadb.utils import embedding_functions from chromadb.utils import embedding_functions
from httpx import RequestError from httpx import RequestError
import chromadb import chromadb
from swarms.memory.base_vector_db import VectorDatabase
class ChromaClient:
class ChromaClient(VectorDatabase):
def __init__( def __init__(
self, self,
collection_name: str = "chromadb-collection", collection_name: str = "chromadb-collection",
@ -90,23 +93,14 @@ class ChromaClient:
print(f"Error searching vectors: {e}") print(f"Error searching vectors: {e}")
return None return None
def search_vectors_formatted(self, query: str, limit: int = 2): def add(self, vector: Dict[str, Any], metadata: Dict[str, Any]) -> None:
""" pass
Searches the collection for vectors similar to the query vector.
Args: def query(self, vector: Dict[str, Any], num_results: int) -> Dict[str, Any]:
query (str): The query string to be converted into a vector and used for searching. pass
limit (int): The number of search results to return. Defaults to 3.
Returns: def delete(self, vector_id: str) -> None:
SearchResult or None: Returns the search results if successful, otherwise None. pass
"""
try: def update(self, vector_id: str, vector: Dict[str, Any], metadata: Dict[str, Any]) -> None:
search_result = self.collection.query( pass
query_texts=query,
n_results=limit,
)
return search_result
except Exception as e:
print(f"Error searching vectors: {e}")
return None

@ -635,6 +635,10 @@ class Agent:
AGENT_SYSTEM_PROMPT_3, response AGENT_SYSTEM_PROMPT_3, response
) )
# Retreiving long term memory
if self.memory:
task = self.agent_memory_prompt(response, task)
attempt = 0 attempt = 0
while attempt < self.retry_attempts: while attempt < self.retry_attempts:
try: try:
@ -755,6 +759,33 @@ class Agent:
""" """
return agent_history_prompt return agent_history_prompt
def agent_memory_prompt(
self,
query,
prompt
):
"""
Generate the agent long term memory prompt
Args:
system_prompt (str): The system prompt
history (List[str]): The history of the conversation
Returns:
str: The agent history prompt
"""
context_injected_prompt = prompt
if self.memory:
ltr = self.memory.query(query)
context_injected_prompt = f"""{prompt}
################ CONTEXT ####################
{ltr}
"""
return context_injected_prompt
async def run_concurrent(self, tasks: List[str], **kwargs): async def run_concurrent(self, tasks: List[str], **kwargs):
""" """
Run a batch of tasks concurrently and handle an infinite level of task inputs. Run a batch of tasks concurrently and handle an infinite level of task inputs.

Loading…
Cancel
Save