[SWARMS]Memory]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms-memory"
-version = "0.0.3"
+version = "0.0.5"
 description = "Swarms Memory - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]

server/Dockerfile

@@ -0,0 +1,32 @@
+# Start with the official Python 3.11 slim image
+FROM python:3.11-slim
+
+# Set environment variables to prevent Python from writing .pyc files to disk
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+ENV API_HOST="0.0.0.0"
+ENV API_PORT=8000
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies and update package list
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    gcc \
+    libpq-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Copy the FastAPI application code
+COPY . .
+
+# Expose the port FastAPI will run on
+EXPOSE 8000
+
+# Command to run the application with Uvicorn and Gunicorn
+CMD ["gunicorn", "main:api", "-w", "4", "-k", "uvicorn.workers.UvicornWorker", "--bind", "0.0.0.0:8000"]

server/api.py

@@ -0,0 +1,259 @@
+from fastapi import FastAPI, HTTPException, Path, Body
+from pydantic import BaseModel, Field
+from typing import List, Optional
+from loguru import logger
+import chromadb
+import uuid
+import os
+import uvicorn
+
+app = FastAPI()
+
+# Initialize the ChromaDB client
+chroma_client = chromadb.Client()
+
+# Logger configuration
+logger.add("api_logs.log", rotation="500 MB")
+
+
+# Pydantic models
+class CreateCollectionRequest(BaseModel):
+    name: str = Field(..., description="The name of the collection")
+
+
+class AddDocumentsRequest(BaseModel):
+    documents: List[str] = Field(
+        ...,
+        description="List of documents to be added to the collection",
+    )
+    ids: Optional[List[str]] = Field(
+        None,
+        description="Optional list of document IDs. If not provided, IDs will be autogenerated",
+    )
+
+
+class QueryDocumentsRequest(BaseModel):
+    query_texts: List[str] = Field(
+        ...,
+        description="List of query texts to search for similar documents",
+    )
+    n_results: int = Field(
+        1, description="Number of similar results to return"
+    )
+
+
+class UpdateDocumentRequest(BaseModel):
+    document: str = Field(
+        ..., description="The updated content of the document"
+    )
+
+
+# Routes
+
+
+@app.post("/collections", response_model=dict)
+def create_collection(request: CreateCollectionRequest):
+    """
+    Creates a new collection with the specified name.
+    """
+    try:
+        collection = chroma_client.create_collection(
+            name=request.name
+        )
+        logger.info(f"Created collection with name: {request.name}")
+        return {
+            "message": f"Collection '{request.name}' created successfully."
+        }
+    except Exception as e:
+        logger.error(f"Error creating collection: {e}")
+        raise HTTPException(
+            status_code=500, detail="Failed to create collection."
+        )
+
+
+# @app.get("/collections/{collection_id}/get_all_docs", response_model=dict)
+# def get_all_docs(
+#     collection_id: str = Path(
+#         ...,
+#         description="The ID of the collection"
+#     ),
+# ):
+#     collection = chroma_client.get_collection(collection_id)
+
+#     if not collection:
+#         raise HTTPException(
+#             status_code=404, detail="Collection not found."
+#         )
+
+#     return collection.get_all_docs()
+
+
+@app.post(
+    "/collections/{collection_id}/documents", response_model=dict
+)
+def add_documents(
+    collection_id: str = Path(
+        ..., description="The ID of the collection"
+    ),
+    request: AddDocumentsRequest = Body(...),
+):
+    """
+    Adds one or more documents to the specified collection.
+    """
+    try:
+        collection = chroma_client.get_collection(collection_id)
+        if not collection:
+            raise HTTPException(
+                status_code=404, detail="Collection not found."
+            )
+
+        ids = request.ids or [
+            str(uuid.uuid4()) for _ in range(len(request.documents))
+        ]
+        collection.add(documents=request.documents, ids=ids)
+        logger.info(
+            f"Added {len(request.documents)} documents to collection {collection_id}."
+        )
+        return {
+            "message": f"Documents added successfully to collection {collection_id}.",
+            "ids": ids,
+        }
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        logger.error(
+            f"Error adding documents to collection {collection_id}: {e}"
+        )
+        raise HTTPException(
+            status_code=500,
+            detail="Failed to add documents to collection.",
+        )
+
+
+@app.get(
+    "/collections/{collection_id}/documents", response_model=dict
+)
+def query_documents(
+    collection_id: str = Path(
+        ..., description="The ID of the collection"
+    ),
+    query: QueryDocumentsRequest = Body(...),
+):
+    """
+    Queries the collection for the most similar documents based on the provided query texts.
+    """
+    try:
+        collection = chroma_client.get_collection(collection_id)
+        if not collection:
+            raise HTTPException(
+                status_code=404, detail="Collection not found."
+            )
+
+        results = collection.query(
+            query_texts=query.query_texts, n_results=query.n_results
+        )
+        logger.info(
+            f"Queried collection {collection_id} with {query.query_texts}."
+        )
+        return results
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        logger.error(
+            f"Error querying documents in collection {collection_id}: {e}"
+        )
+        raise HTTPException(
+            status_code=500, detail="Failed to query documents."
+        )
+
+
+@app.delete(
+    "/collections/{collection_id}/documents/{document_id}",
+    response_model=dict,
+)
+def delete_document(
+    collection_id: str = Path(
+        ..., description="The ID of the collection"
+    ),
+    document_id: str = Path(
+        ..., description="The ID of the document to delete"
+    ),
+):
+    """
+    Deletes a specific document from the collection.
+    """
+    try:
+        collection = chroma_client.get_collection(collection_id)
+        if not collection:
+            raise HTTPException(
+                status_code=404, detail="Collection not found."
+            )
+
+        collection.delete(ids=[document_id])
+        logger.info(
+            f"Deleted document {document_id} from collection {collection_id}."
+        )
+        return {
+            "message": f"Document {document_id} deleted successfully from collection {collection_id}."
+        }
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        logger.error(
+            f"Error deleting document {document_id} from collection {collection_id}: {e}"
+        )
+        raise HTTPException(
+            status_code=500, detail="Failed to delete document."
+        )
+
+
+@app.put(
+    "/collections/{collection_id}/documents/{document_id}",
+    response_model=dict,
+)
+def update_document(
+    collection_id: str = Path(
+        ..., description="The ID of the collection"
+    ),
+    document_id: str = Path(
+        ..., description="The ID of the document to update"
+    ),
+    request: UpdateDocumentRequest = Body(...),
+):
+    """
+    Updates the content of a specific document within a collection.
+    """
+    try:
+        collection = chroma_client.get_collection(collection_id)
+        if not collection:
+            raise HTTPException(
+                status_code=404, detail="Collection not found."
+            )
+
+        collection.update(
+            documents=[request.document], ids=[document_id]
+        )
+        logger.info(
+            f"Updated document {document_id} in collection {collection_id}."
+        )
+        return {
+            "message": f"Document {document_id} updated successfully in collection {collection_id}."
+        }
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        logger.error(
+            f"Error updating document {document_id} in collection {collection_id}: {e}"
+        )
+        raise HTTPException(
+            status_code=500, detail="Failed to update document."
+        )
+
+
+if __name__ == "__main__":
+
+    uvicorn.run(
+        app,
+        host=os.getenv("API_HOST"),
+        port=os.getenv("API_PORT", 8000),
+    )