Adding all other methods

main.py

@@ -1,13 +1,15 @@
 import fastapi
 import json
 import uvicorn
-from fastapi import HTTPException
-from fastapi.responses import HTMLResponse
+from fastapi import HTTPException , status
+from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FASTAPI as _FastAPI, Response
 from sse_starlette.sse import EventSourceResponse
 from starlette.responses import StreamingResponse
+from starlette.requests import Request
 from pydantic import BaseModel
-from typing import List, Dict, Any, Generator, Optional, cast
+from typing import List, Dict, Any, Generator, Optional, cast, Callable
 from server import client
 from chromadb.api.types import (
     Documents,
@@ -50,8 +52,29 @@ from uuid import UUID
 from chromadb.telemetry import Telemetry
 from overrides import override
 
+async def catch_exceptions_middleware(
+    request: Request, call_next: Callable[[Request], Any]
+) -> Response:
+    try:
+        return await call_next(request)
+    except ChromaError as e:
+        return JSONResponse(
+            content={"error": e.name(), "message": e.message()}, status_code=e.code()
+        )
+    except Exception as e:
+        return JSONResponse(content={"error": repr(e)}, status_code=500)
+
+
+
+def _uuid(uuid_str: str) -> UUID:
+    try:
+        return UUID(uuid_str)
+    except ValueError:
+        raise InvalidUUIDError(f"Could not parse {uuid_str} as a UUID")
+    
 
 app = fastapi.FastAPI(title="ChromaDB")
+app.middleware("http")(catch_exceptions_middleware)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -62,6 +85,7 @@ app.add_middleware(
 api_base="/api/v1"
 bkend=client()
 
+    
 @app.get(api_base+"")
 def heartbeat():
     print("Received heartbeat request")
@@ -98,18 +122,77 @@ def list_collections():
     return bkend.list_collections()
 
 @app.post(api_base+"/collections")
-def create_collection(
-        collection: CreateCollection 
-    ) -> Collection:
+def create_collection( collection: CreateCollection ) -> Collection:
     print("Received request to create_collection")
     return bkend.create_collection(name=collection.name,metadata=collection.metadata,get_or_create=collection.get_or_create)
 
 @app.get(api_base+"/collections/{collection_name}")
-def get_collection(
-        collection_name: str,
-    ) -> Collection:
+def get_collection( collection_name: str) -> Collection:
+    print("Received get_collection request")
     return bkend.get_collection(collection_name)
 
+@app.post(api_base+"/collections/{collection_id}/add")
+def add(collection_id:str , add:AddEmbedding) -> None:
+    print("Received add request")
+    try:
+        result=bkend.add(collection_id=_uuid(collection_id),embeddings=add.embeddings,metadatas=add.metadatas,documents=add.documents,ids=add.ids,increment_index=add.increment_index)
+    except InvalidDimensionException as e:
+            raise HTTPException(status_code=500, detail=str(e))
+    return result   
+
+@app.post(api_base+"/collections/{collection_id}/update")
+def update(collection_id:str , update:UpdateEmbedding) -> None:
+    print("Received update request")
+    return bkend.update(ids=update.ids, collection_id=_uuid(collection_id), embeddings=update.embeddings, documents=update.documents, metadatas=update.metadatas)
+        
+@app.post(api_base+"/collections/{collection_id}/upsert")
+def upsert(collection_id:str, upsert: AddEmbedding):
+    print("Received upsert request")
+    return bkend.upsert(collection_id=_uuid(collection_id),embeddings=upsert.embeddings,metadatas=upsert.metadatas,documents=upsert.documents,ids=upsert.ids,increment_index=upsert.increment_index)
+        
+@app.post(api_base+"/collections/{collection_id}/get") 
+def get(self, collection_id: str, get: GetEmbedding) -> GetResult:
+    print("Received get request")
+    return bkend.get(collection_id=_uuid(collection_id), ids=get.ids, where=get.where,
+            where_document=get.where_document, sort=get.sort, limit=get.limit,
+            offset=get.offset, include=get.include)
+        
+@app.post(api_base+"/collections/{collection_id}/delete")
+def delete(collection_id: str, delete: DeleteEmbedding) -> List[UUID]:
+    print("Received delete request")
+    return bkend.delete(where=delete.where, ids=delete.ids,
+            collection_id=_uuid(collection_id), where_document=delete.where_document)
+        
+@app.get(api_base+"/collections/{collection_id}/count")
+def count(collection_id:str) ->int:
+    print("Received count request")
+    return bkend.count(_uuid(collection_id))
+        
+@app.post(api_base+"/collections/{collection_id}/query")
+def get_nearest_neighbors(collection_id: str, query: QueryEmbedding) -> QueryResult:
+    print("Received get_nearest_neighbors request")
+    return bkend.get_nearest_neighbors(collection_id=_uuid(collection_id), where=query.where,  where_document=query.where_document, 
+              query_embeddings=query.query_embeddings, n_results=query.n_results, include=query.include)
+        
+@app.post(api_base+"/collections/{collection_name}/create_index")
+def create_index(collection_name:str)-> bool:
+    print("Received create_index request")
+    return bkend.create_index(collection_name)
+        
+@app.get(api_base+"/collections/{collection_name}")
+def get_collection2( collection_name: str) -> Collection:
+    print("Received get_collection2 request")
+    return bkend.get_collection(collection_name)
+        
+@app.post(api_base+"/collections/{collection_id}")
+def modify(collection_id: str, collection: UpdateCollection) -> None:
+    print("Received modify(collection) request")
+    return bkend.modify(id=_uuid(collection_id), new_name=collection.new_name, new_metadata=collection.new_metadata)
+        
+@app.delete(api_base+"/collections/{collection_name}")
+def delete_collection(collection_name:str) -> None:
+    print("Received delete_collection request")
+    return bkend.delete_collection(collection_name)
 
 
 if __name__ == "__main__":

server.py

@@ -1,7 +1,8 @@
-from chromadb.config import Settings
 import chromadb
+from chromadb.config import Settings, System
+import time,json
 from pydantic import BaseModel
-from typing import List, Dict, Any, Generator, Optional, cast
+from typing import List, Dict, Any, Generator, Optional, cast, Callable
 from chromadb.api.types import (
     Documents,
     Embeddings,
@@ -44,8 +45,6 @@ from chromadb.telemetry import Telemetry
 from overrides import override
 
 
-import time
-
 class client():
     def __init__(self):
         self.db = chromadb.Client(Settings(
@@ -73,20 +72,83 @@ class client():
     def get_collection(
         self,
         name: str,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction()
     ) -> Collection:
         col=self.db.get_collection(name,embedding_function=embedding_function)
         print(col)
         return col
 
-    def reset():
+    def reset(self):
         return self.db.reset()
 
-    def version():
+    def version(self):
         return self.db.get_version()
 
-    def persist():
+    def persist(self):
         return self.db.persist()
 
-    def raw_sql(raw_sql: RawSql):
+    def raw_sql(self,raw_sql: RawSql):
         return self.db.raw_sql(raw_sql.raw_sql) 
+    
+    def add(self,ids: IDs,
+        collection_id: UUID,
+        embeddings: Embeddings,
+        metadatas: Optional[Metadatas] = None,
+        documents: Optional[Documents] = None,
+        increment_index: bool = True,
+    ) -> bool:
+        return self.db._add(collection_id=collection_id,embeddings=embeddings,
+            metadatas=metadatas,documents=documents,
+            ids=ids,increment_index=increment_index)
+    
+    def update( self, collection_id: UUID, ids: IDs,
+        embeddings: Optional[Embeddings] = None,
+        metadatas: Optional[Metadatas] = None,
+        documents: Optional[Documents] = None,
+    ) -> bool:
+        return self.db._update(ids=ids, collection_id=collection_id, embeddings=embeddings, documents=documents, metadatas=metadatas)
+    
+    def upsert( self, collection_id: UUID, ids: IDs,
+        embeddings: Embeddings,
+        metadatas: Optional[Metadatas] = None,
+        documents: Optional[Documents] = None,
+        increment_index: bool = True,
+    ) -> bool:
+        return self.db._upsert(collection_id=collection_id,embeddings=embeddings,metadatas=metadatas,documents=documents,ids=ids,increment_index=increment_index)
+    
+    def get( self, collection_id: UUID, ids: Optional[IDs] = None, where: Optional[Where] = {},
+        sort: Optional[str] = None, limit: Optional[int] = None, offset: Optional[int] = None,
+        page: Optional[int] = None, page_size: Optional[int] = None,
+        where_document: Optional[WhereDocument] = {},
+        include: Include = ["embeddings", "metadatas", "documents"],
+    ) -> GetResult:
+        return self.db._get(collection_id=collection_id, ids=ids, where=where,
+            where_document=where_document, sort=sort, limit=limit,
+            offset=offset, include=include)
+    def delete( self, collection_id: UUID, ids: Optional[IDs],
+        where: Optional[Where] = {}, where_document: Optional[WhereDocument] = {},
+    ) -> IDs:
+        return self.db._delete(where=where, ids=ids, collection_id=collection_id, where_document=where_document)
+
+    def count(self, collection_id: UUID) -> int:
+        return self.db._.count(collection_id)
+    
+    def get_nearest_neighbors( self, collection_id: UUID, query_embeddings: Embeddings,
+        n_results: int = 10, where: Where = {}, where_document: WhereDocument = {},
+        include: Include = ["embeddings", "metadatas", "documents", "distances"],
+    ) -> QueryResult:
+        return self.db._query(collection_id=collection_id, where=where,  where_document=where_document, 
+              query_embeddings=query_embeddings, n_results=n_results, include=include)
+    
+    def create_index(self, collection_name: str) -> bool:
+        return self.db.create_index(collection_name)
+    
+    def modify( self, id: UUID, new_name: Optional[str] = None,
+        new_metadata: Optional[CollectionMetadata] = None,
+    ) -> None:
+        """This is for updating the collection"""
+        return self.db._modify(id=id, new_name=new_name, new_metadata=new_metadata)
+    
+    def delete_collection( self, name: str,) -> None:
+        return self.db.delete_collection(name)
+