diff --git a/README.md b/README.md index 028b998..1e6a01d 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,8 @@ $ pip install swarms-memory ## Usage + + # License MIT diff --git a/pyproject.toml b/pyproject.toml index 6f7c940..c6afe14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,10 @@ classifiers = [ [tool.poetry.dependencies] python = "^3.10" swarms = "*" +chromadb = "*" +loguru = "*" +sentence-transformers = "*" +pinecone = "*" [tool.poetry.group.lint.dependencies] diff --git a/swarms_memory/__init__.py b/swarms_memory/__init__.py index 7b67212..467e68b 100644 --- a/swarms_memory/__init__.py +++ b/swarms_memory/__init__.py @@ -1,3 +1,4 @@ -from swarms_memory.chroma_db_wrapper import ChromaDBWrapper +from swarms_memory.chroma_db_wrapper import ChromaMemory +from swarms_memory.pinecone_wrapper import PineconeMemory -__all__ = ["ChromaDBWrapper"] +__all__ = ["ChromaMemory", "PineconeMemory"] diff --git a/swarms_memory/pinecone_wrapper.py b/swarms_memory/pinecone_wrapper.py index 4e93354..a0f2dbb 100644 --- a/swarms_memory/pinecone_wrapper.py +++ b/swarms_memory/pinecone_wrapper.py @@ -1,10 +1,12 @@ -from typing import List, Dict, Any, Callable, Optional +from typing import Any, Callable, Dict, List, Optional + import pinecone from loguru import logger from sentence_transformers import SentenceTransformer +from swarms.memory.base_vectordb import BaseVectorDatabase -class PineconeMemory: +class PineconeMemory(BaseVectorDatabase): """ A highly customizable wrapper class for Pinecone-based Retrieval-Augmented Generation (RAG) system.