diff --git a/swarms/memory/__init__.py b/swarms/memory/__init__.py index 41be1c6f..e916d5a8 100644 --- a/swarms/memory/__init__.py +++ b/swarms/memory/__init__.py @@ -1,11 +1,3 @@ -from swarms.memory.pinecone import PineconeVector -from swarms.memory.base import BaseVectorStore -from swarms.memory.pg import PgVectorVectorStore -from swarms.memory.ocean import OceanDB - -__all__ = [ - "BaseVectorStore", - "PineconeVector", - "PgVectorVectorStore", - "OceanDB", -] +# from swarms.memory.pinecone import PineconeVector +# from swarms.memory.base import BaseVectorStore +# from swarms.memory.pg import PgVectorVectorStore diff --git a/swarms/memory/base.py b/swarms/memory/base.py index 3ca49617..f28da852 100644 --- a/swarms/memory/base.py +++ b/swarms/memory/base.py @@ -9,8 +9,6 @@ from griptape.artifacts import TextArtifact @define class BaseVectorStore(ABC): - """ """ - DEFAULT_QUERY_COUNT = 5 @dataclass diff --git a/swarms/memory/db.py b/swarms/memory/db.py deleted file mode 100644 index 4ffec16f..00000000 --- a/swarms/memory/db.py +++ /dev/null @@ -1,177 +0,0 @@ -import uuid -from abc import ABC -from typing import Any, Dict, List, Optional - -from swarms.memory.schemas import Artifact, Status -from swarms.memory.schemas import Step as APIStep -from swarms.memory.schemas import Task as APITask - - -class Step(APIStep): - additional_properties: Optional[Dict[str, str]] = None - - -class Task(APITask): - steps: List[Step] = [] - - -class NotFoundException(Exception): - """ - Exception raised when a resource is not found. - """ - - def __init__(self, item_name: str, item_id: str): - self.item_name = item_name - self.item_id = item_id - super().__init__(f"{item_name} with {item_id} not found.") - - -class TaskDB(ABC): - async def create_task( - self, - input: Optional[str], - additional_input: Any = None, - artifacts: Optional[List[Artifact]] = None, - steps: Optional[List[Step]] = None, - ) -> Task: - raise NotImplementedError - - async def create_step( - self, - task_id: str, - name: Optional[str] = None, - input: Optional[str] = None, - is_last: bool = False, - additional_properties: Optional[Dict[str, str]] = None, - ) -> Step: - raise NotImplementedError - - async def create_artifact( - self, - task_id: str, - file_name: str, - relative_path: Optional[str] = None, - step_id: Optional[str] = None, - ) -> Artifact: - raise NotImplementedError - - async def get_task(self, task_id: str) -> Task: - raise NotImplementedError - - async def get_step(self, task_id: str, step_id: str) -> Step: - raise NotImplementedError - - async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact: - raise NotImplementedError - - async def list_tasks(self) -> List[Task]: - raise NotImplementedError - - async def list_steps( - self, task_id: str, status: Optional[Status] = None - ) -> List[Step]: - raise NotImplementedError - - -class InMemoryTaskDB(TaskDB): - _tasks: Dict[str, Task] = {} - - async def create_task( - self, - input: Optional[str], - additional_input: Any = None, - artifacts: Optional[List[Artifact]] = None, - steps: Optional[List[Step]] = None, - ) -> Task: - if not steps: - steps = [] - if not artifacts: - artifacts = [] - task_id = str(uuid.uuid4()) - task = Task( - task_id=task_id, - input=input, - steps=steps, - artifacts=artifacts, - additional_input=additional_input, - ) - self._tasks[task_id] = task - return task - - async def create_step( - self, - task_id: str, - name: Optional[str] = None, - input: Optional[str] = None, - is_last=False, - additional_properties: Optional[Dict[str, Any]] = None, - ) -> Step: - step_id = str(uuid.uuid4()) - step = Step( - task_id=task_id, - step_id=step_id, - name=name, - input=input, - status=Status.created, - is_last=is_last, - additional_properties=additional_properties, - ) - task = await self.get_task(task_id) - task.steps.append(step) - return step - - async def get_task(self, task_id: str) -> Task: - task = self._tasks.get(task_id, None) - if not task: - raise NotFoundException("Task", task_id) - return task - - async def get_step(self, task_id: str, step_id: str) -> Step: - task = await self.get_task(task_id) - step = next(filter(lambda s: s.task_id == task_id, task.steps), None) - if not step: - raise NotFoundException("Step", step_id) - return step - - async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact: - task = await self.get_task(task_id) - artifact = next( - filter(lambda a: a.artifact_id == artifact_id, task.artifacts), None - ) - if not artifact: - raise NotFoundException("Artifact", artifact_id) - return artifact - - async def create_artifact( - self, - task_id: str, - file_name: str, - relative_path: Optional[str] = None, - step_id: Optional[str] = None, - ) -> Artifact: - artifact_id = str(uuid.uuid4()) - artifact = Artifact( - artifact_id=artifact_id, - file_name=file_name, - relative_path=relative_path, - ) - task = await self.get_task(task_id) - task.artifacts.append(artifact) - - if step_id: - step = await self.get_step(task_id, step_id) - step.artifacts.append(artifact) - - return artifact - - async def list_tasks(self) -> List[Task]: - return [task for task in self._tasks.values()] - - async def list_steps( - self, task_id: str, status: Optional[Status] = None - ) -> List[Step]: - task = await self.get_task(task_id) - steps = task.steps - if status: - steps = list(filter(lambda s: s.status == status, steps)) - return steps diff --git a/swarms/memory/ocean.py b/swarms/memory/ocean.py deleted file mode 100644 index fb0873af..00000000 --- a/swarms/memory/ocean.py +++ /dev/null @@ -1,157 +0,0 @@ -import logging -from typing import List - -import oceandb -from oceandb.utils.embedding_function import MultiModalEmbeddingFunction - - -class OceanDB: - """ - A class to interact with OceanDB. - - ... - - Attributes - ---------- - client : oceandb.Client - a client to interact with OceanDB - - Methods - ------- - create_collection(collection_name: str, modality: str): - Creates a new collection in OceanDB. - append_document(collection, document: str, id: str): - Appends a document to a collection in OceanDB. - add_documents(collection, documents: List[str], ids: List[str]): - Adds multiple documents to a collection in OceanDB. - query(collection, query_texts: list[str], n_results: int): - Queries a collection in OceanDB. - """ - - def __init__(self, client: oceandb.Client = None): - """ - Constructs all the necessary attributes for the OceanDB object. - - Parameters - ---------- - client : oceandb.Client, optional - a client to interact with OceanDB (default is None, which creates a new client) - """ - try: - self.client = client if client else oceandb.Client() - print(self.client.heartbeat()) - except Exception as e: - logging.error(f"Failed to initialize OceanDB client. Error: {e}") - raise - - def create_collection(self, collection_name: str, modality: str): - """ - Creates a new collection in OceanDB. - - Parameters - ---------- - collection_name : str - the name of the new collection - modality : str - the modality of the new collection - - Returns - ------- - collection - the created collection - """ - try: - embedding_function = MultiModalEmbeddingFunction(modality=modality) - collection = self.client.create_collection( - collection_name, embedding_function=embedding_function - ) - return collection - except Exception as e: - logging.error(f"Failed to create collection. Error {e}") - raise - - def append_document(self, collection, document: str, id: str): - """ - Appends a document to a collection in OceanDB. - - Parameters - ---------- - collection - the collection to append the document to - document : str - the document to append - id : str - the id of the document - - Returns - ------- - result - the result of the append operation - """ - try: - return collection.add(documents=[document], ids=[id]) - except Exception as e: - logging.error( - f"Failed to append document to the collection. Error {e}" - ) - raise - - def add_documents(self, collection, documents: List[str], ids: List[str]): - """ - Adds multiple documents to a collection in OceanDB. - - Parameters - ---------- - collection - the collection to add the documents to - documents : List[str] - the documents to add - ids : List[str] - the ids of the documents - - Returns - ------- - result - the result of the add operation - """ - try: - return collection.add(documents=documents, ids=ids) - except Exception as e: - logging.error(f"Failed to add documents to collection. Error: {e}") - raise - - def query(self, collection, query_texts: list[str], n_results: int): - """ - Queries a collection in OceanDB. - - Parameters - ---------- - collection - the collection to query - query_texts : list[str] - the texts to query - n_results : int - the number of results to return - - Returns - ------- - results - the results of the query - """ - try: - results = collection.query( - query_texts=query_texts, n_results=n_results - ) - return results - except Exception as e: - logging.error(f"Failed to query the collection. Error {e}") - raise - - -# Example -# ocean = OceanDB() -# collection = ocean.create_collection("test", "text") -# ocean.append_document(collection, "hello world", "1") -# ocean.add_documents(collection, ["hello world", "hello world"], ["2", "3"]) -# results = ocean.query(collection, ["hello world"], 3) -# print(results)