[FEAT][Postgreqsl] [SQLite]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms-memory"
-version = "0.0.7"
+version = "0.0.8"
 description = "Swarms Memory - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]
@@ -30,6 +30,8 @@ sentence-transformers = "*"
 pinecone = "*"
 faiss-cpu = "*"
 pydantic = "*"
+sqlalchemy = "*"
+sqlite3 = "*"
 
 
 

server/api.py

@@ -58,9 +58,7 @@ def create_collection(request: CreateCollectionRequest):
     Creates a new collection with the specified name.
     """
     try:
-        chroma_client.create_collection(
-            name=request.name
-        )
+        chroma_client.create_collection(name=request.name)
         logger.info(f"Created collection with name: {request.name}")
         return {
             "message": f"Collection '{request.name}' created successfully."

swarms_memory/__init__.py

@@ -1,2 +1,2 @@
-from swarms_memory.vector_dbs import * # noqa: F401, F403
-from swarms_memory.utils import * # noqa: F401, F403
+from swarms_memory.vector_dbs import *  # noqa: F401, F403
+from swarms_memory.utils import *  # noqa: F401, F403

swarms_memory/dbs/__init__.py

@@ -0,0 +1,5 @@
+from swarms_memory.dbs.pg import PostgresDB
+from swarms_memory.dbs.sqlite import SQLiteDB
+
+
+__all__ = ["PostgresDB", "SQLiteDB"]

swarms_memory/dbs/pg.py

@@ -0,0 +1,142 @@
+import uuid
+from typing import Any, List, Optional
+
+from sqlalchemy import JSON, Column, String, create_engine
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import Session
+from swarms_memory import BaseVectorDatabase
+
+
+class PostgresDB(BaseVectorDatabase):
+    """
+    A class representing a Postgres database.
+
+    Args:
+        connection_string (str): The connection string for the Postgres database.
+        table_name (str): The name of the table in the database.
+
+    Attributes:
+        engine: The SQLAlchemy engine for connecting to the database.
+        table_name (str): The name of the table in the database.
+        VectorModel: The SQLAlchemy model representing the vector table.
+
+    """
+
+    def __init__(
+        self, connection_string: str, table_name: str, *args, **kwargs
+    ):
+        """
+        Initializes a new instance of the PostgresDB class.
+
+        Args:
+            connection_string (str): The connection string for the Postgres database.
+            table_name (str): The name of the table in the database.
+
+        """
+        self.engine = create_engine(
+            connection_string, *args, **kwargs
+        )
+        self.table_name = table_name
+        self.VectorModel = self._create_vector_model()
+
+    def _create_vector_model(self):
+        """
+        Creates the SQLAlchemy model for the vector table.
+
+        Returns:
+            The SQLAlchemy model representing the vector table.
+
+        """
+        Base = declarative_base()
+
+        class VectorModel(Base):
+            __tablename__ = self.table_name
+
+            id = Column(
+                UUID(as_uuid=True),
+                primary_key=True,
+                default=uuid.uuid4,
+                unique=True,
+                nullable=False,
+            )
+            vector = Column(
+                String
+            )  # Assuming vector is stored as a string
+            namespace = Column(String)
+            meta = Column(JSON)
+
+        return VectorModel
+
+    def add(
+        self,
+        vector: str,
+        vector_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        meta: Optional[dict] = None,
+    ) -> None:
+        """
+        Adds or updates a vector in the database.
+
+        Args:
+            vector (str): The vector to be added or updated.
+            vector_id (str, optional): The ID of the vector. If not provided, a new ID will be generated.
+            namespace (str, optional): The namespace of the vector.
+            meta (dict, optional): Additional metadata associated with the vector.
+
+        """
+        try:
+            with Session(self.engine) as session:
+                obj = self.VectorModel(
+                    id=vector_id,
+                    vector=vector,
+                    namespace=namespace,
+                    meta=meta,
+                )
+                session.merge(obj)
+                session.commit()
+        except Exception as e:
+            print(f"Error adding or updating vector: {e}")
+
+    def query(
+        self, query: Any, namespace: Optional[str] = None
+    ) -> List[Any]:
+        """
+        Queries vectors from the database based on the given query and namespace.
+
+        Args:
+            query (Any): The query or condition to filter the vectors.
+            namespace (str, optional): The namespace of the vectors to be queried.
+
+        Returns:
+            List[Any]: A list of vectors that match the query and namespace.
+
+        """
+        try:
+            with Session(self.engine) as session:
+                q = session.query(self.VectorModel)
+                if namespace:
+                    q = q.filter_by(namespace=namespace)
+                # Assuming 'query' is a condition or filter
+                q = q.filter(query)
+                return q.all()
+        except Exception as e:
+            print(f"Error querying vectors: {e}")
+            return []
+
+    def delete_vector(self, vector_id):
+        """
+        Deletes a vector from the database based on the given vector ID.
+
+        Args:
+            vector_id: The ID of the vector to be deleted.
+
+        """
+        try:
+            with Session(self.engine) as session:
+                obj = session.get(self.VectorModel, vector_id)
+                if obj:
+                    session.delete(obj)
+                    session.commit()
+        except Exception as e:
+            print(f"Error deleting vector: {e}")