Add list/get/create collections

main.py

@@ -7,7 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
 from starlette.responses import StreamingResponse
 from pydantic import BaseModel
-from typing import List, Dict, Any, Generator, Optional
+from typing import List, Dict, Any, Generator, Optional, cast
 from server import client
 from chromadb.api.types import (
     Documents,
@@ -22,6 +22,18 @@ from chromadb.api.types import (
     QueryResult,
     CollectionMetadata,
 )
+from chromadb.api import API
+from chromadb.config import System
+import chromadb.utils.embedding_functions as ef
+import pandas as pd
+import requests
+import json
+from typing import Sequence
+from chromadb.api.models.Collection import Collection
+import chromadb.errors as errors
+from uuid import UUID
+from chromadb.telemetry import Telemetry
+from overrides import override
 
 
 app = fastapi.FastAPI(title="ChromaDB")
@@ -40,11 +52,31 @@ def heartbeat():
     print("Received heartbeat request")
     return bkend.heartbeat()
 
-@app.get(api_base+"/collection")
-def list_collection():
-    print("Received list_collection request")
+@app.get(api_base+"/collections")
+def list_collections():
+    print("Received list_collections request")
     return bkend.list_collections()
 
+@app.post(api_base+"/collections")
+def create_collection(
+        self,
+        name: str,
+        metadata: Optional[CollectionMetadata] = None,
+        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        get_or_create: bool = False,
+    ) -> Collection:
+    print("Received request to create_collection")
+    return bkend.create_collection(name,metadata=metadata,embedding_function=embedding_function,get_or_create=get_or_create)
+
+@app.get(api_base+"/collections")
+def get_collection(
+        self,
+        name: str,
+        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+    ) -> Collection:
+    return bkend.get_collection(name,embedding_function=embedding_function
+    )
+
 if __name__ == "__main__":
   uvicorn.run(app, host="0.0.0.0", port=8000)
 

server.py

@@ -1,5 +1,31 @@
 from chromadb.config import Settings
 import chromadb
+from typing import Optional, cast
+from chromadb.api import API
+from chromadb.config import System
+from chromadb.api.types import (
+    Documents,
+    Embeddings,
+    EmbeddingFunction,
+    IDs,
+    Include,
+    Metadatas,
+    Where,
+    WhereDocument,
+    GetResult,
+    QueryResult,
+    CollectionMetadata,
+)
+import chromadb.utils.embedding_functions as ef
+import pandas as pd
+import requests
+import json
+from typing import Sequence
+from chromadb.api.models.Collection import Collection
+import chromadb.errors as errors
+from uuid import UUID
+from chromadb.telemetry import Telemetry
+from overrides import override
 import time
 
 class client():
@@ -10,7 +36,24 @@ class client():
         ))
 
     def heartbeat(self):
-        return int(time.time_ns())
+        return {"nanosecond heartbeat":int(time.time_ns())}
 
     def list_collections(self):
-        return self.db.list_collections()
+        return self.db.list_collections()
+    
+    def create_collection(
+        self,
+        name: str,
+        metadata: Optional[CollectionMetadata] = None,
+        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        get_or_create: bool = False,
+    ) -> Collection:
+        return self.db.create_collection(name,metadata=metadata,embedding_function=embedding_function,get_or_create=get_or_create)
+    
+    def get_collection(
+        self,
+        name: str,
+        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+    ) -> Collection:
+        return self.db.get_collection(name,embedding_function=embedding_function)
+