sharedkernel 1.0.0__tar.gz → 1.1.0__tar.gz

{sharedkernel-1.0.0 → sharedkernel-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sharedkernel
-Version: 1.0.0
+Version: 1.1.0
 Summary: sharekernel is an shared package between all python projects
 Author: Smilinno
 Requires-Dist: numpy
@@ -8,3 +8,5 @@ Requires-Dist: requests
 Requires-Dist: pymongo
 Requires-Dist: fastapi==0.89.1
 Requires-Dist: PyJWT
+Requires-Dist: pymilvus
+Requires-Dist: chromadb

{sharedkernel-1.0.0 → sharedkernel-1.1.0}/setup.py

@@ -4,10 +4,10 @@ setup(
     # Needed to silence warnings (and to be a worthwhile package)
     name='sharedkernel',
     author='Smilinno',
-    packages=['sharedkernel','sharedkernel.database','sharedkernel.enum','sharedkernel.exception','sharedkernel.objects'],
+    packages=['sharedkernel','sharedkernel.database','sharedkernel.database.vector_database_repository','sharedkernel.enum','sharedkernel.exception','sharedkernel.objects'],
     # Needed for dependencies
-    install_requires=['numpy','requests','pymongo','fastapi==0.89.1','PyJWT'],
+    install_requires=['numpy','requests','pymongo','fastapi==0.89.1','PyJWT','pymilvus','chromadb'],
     # *strongly* suggested for sharing
-    version='1.0.0',
+    version='1.1.0',
     description='sharekernel is an shared package between all python projects',
 )

sharedkernel-1.1.0/sharedkernel/database/__init__.py

@@ -0,0 +1,2 @@
+from .vector_database_repository import VectorRepository
+# from .mongo_repository_base import MongoRepositoryBase

sharedkernel-1.0.0/sharedkernel/database/repository_base.py → sharedkernel-1.1.0/sharedkernel/database/mongo_repository_base.py

@@ -19,7 +19,7 @@ ResultT = TypeVar("ResultT")
 
 T = TypeVar("T", bound=BaseModel)
 
-class RepositoryBase(Generic[T]):
+class MongoRepositoryBase(Generic[T]):
     
     def __init__(self,model):
       self.model = model

sharedkernel-1.1.0/sharedkernel/database/vector_database_repository/__init__.py

@@ -0,0 +1,3 @@
+from .chroma_startegy import ChromaStrategy
+from .milvus_strategy import MilvusStrategy
+from .vector_database_repository import VectorRepository

sharedkernel-1.1.0/sharedkernel/database/vector_database_repository/chroma_startegy.py

@@ -0,0 +1,39 @@
+import chromadb
+import numpy as np
+from chromadb.config import Settings
+from .vector_database_strategy import VectorDatabaseStrategy
+import uuid
+
+
+class ChromaStrategy(VectorDatabaseStrategy):
+    def __init__(self, collection_name: str):
+        self.collection_name = collection_name
+        self.collection = None
+
+    def connect(self, host: str = "localhost", port: int = 8000):
+        client = chromadb.Client(
+            Settings(
+                chroma_api_impl="rest",
+                chroma_server_host=host,
+                chroma_server_http_port=port,
+            )
+        )
+        self.collection = client.get_or_create_collection(self.collection_name)
+
+    def insert_vector(self, vector: np.ndarray, metadata: dict) -> str:
+        id = str(uuid.uuid4())
+        self.collection.upsert(ids=id, embeddings=vector.tolist(), metadatas=[metadata])
+
+        return id
+
+    def search_vector(self, vector: np.ndarray, top_k: int):
+        results = self.collection.query(vectors=[vector.tolist()], n_results=top_k)
+        return results
+
+    def get_vector_by_id(self, id: str):
+        result = self.collection.get(ids=id)
+
+        return result
+
+    def delete_vector(self, id: str):
+        self.collection.delete(ids=id)

sharedkernel-1.1.0/sharedkernel/database/vector_database_repository/milvus_strategy.py

@@ -0,0 +1,50 @@
+import numpy as np
+from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
+
+from .vector_database_strategy import VectorDatabaseStrategy
+import uuid
+
+
+class MilvusStrategy(VectorDatabaseStrategy):
+    def __init__(self, collection_name: str):
+        self.collection_name = collection_name
+        self.collection = None
+
+    def connect(self, host: str = "localhost", port: str = "19530"):
+        connections.connect(alias="default", host=host, port=port)
+        # Define fields
+        fields = [
+            FieldSchema(
+                name="id", dtype=DataType.VARCHAR, max_length=36, is_primary=True
+            ),
+            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=128),
+        ]
+        schema = CollectionSchema(
+            fields, description="Vector collection", enable_dynamic_field=True
+        )
+        self.collection = Collection(name=self.collection_name, schema=schema)
+
+        if not self.collection.has_index():
+             self.collection.create_index(field_name="vector", index_params={"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 128}})
+
+        self.collection.load()
+
+    def insert_vector(self, vector: np.ndarray, metadata: dict) -> str:
+        id = str(uuid.uuid4())
+        self.collection.insert(data={"id": id, "vector": vector.tolist()})
+ 
+        return id
+
+    def search_vector(self, vector: np.ndarray, top_k: int):
+        search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
+        results = self.collection.search(
+            [vector.tolist()], "vector", search_params, top_k
+        )
+        return results
+
+    def get_vector_by_id(self, id: str):
+        result = self.collection.query(expr=f"id=='{id}'",output_fields=["vector"])
+        return result
+
+    def delete_vector(self, id: str):
+        self.collection.delete(expr=f"id=='{id}'")

sharedkernel-1.1.0/sharedkernel/database/vector_database_repository/vector_database_repository.py

@@ -0,0 +1,28 @@
+import numpy as np
+from .vector_database_strategy import VectorDatabaseStrategy
+from .milvus_strategy import MilvusStrategy
+from .chroma_startegy import ChromaStrategy
+from sharedkernel.enum.vector_database_type import VectorDatabaseType
+
+class VectorRepository:
+    def __init__(self, database_type: VectorDatabaseType, collection_name: str, **connection_params):
+        self.strategy = self._get_strategy(database_type, collection_name)
+        self.strategy.connect(**connection_params)
+    
+    def _get_strategy(self, database_type: VectorDatabaseType, collection_name: str) -> VectorDatabaseStrategy:
+        if database_type == VectorDatabaseType.MILVUS:
+            return MilvusStrategy(collection_name)
+        else:
+            return ChromaStrategy(collection_name)
+    
+    def add_vector(self, vector: np.ndarray, metadata: dict) -> str:
+        return self.strategy.insert_vector(vector, metadata)
+    
+    def find_similar_vectors(self, vector: np.ndarray, top_k: int):
+        return self.strategy.search_vector(vector, top_k)
+    
+    def remove_vector(self, id: str):
+        self.strategy.delete_vector(id)
+
+    def get_vector_by_id(self, id: str):
+       return self.strategy.get_vector_by_id(id)

sharedkernel-1.1.0/sharedkernel/database/vector_database_repository/vector_database_strategy.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+import numpy as np
+
+class VectorDatabaseStrategy(ABC):
+    @abstractmethod
+    def connect(self, **kwargs):
+        pass
+    @abstractmethod
+    def insert_vector(self, vector: np.ndarray, metadata: dict) -> str:
+        pass
+    
+    @abstractmethod
+    def search_vector(self, vector: np.ndarray, top_k: int):
+        pass
+    
+    @abstractmethod
+    def delete_vector(self, id: str):
+        pass
+    
+    @abstractmethod
+    def get_vector_by_id(self, id: str):
+        pass

sharedkernel-1.1.0/sharedkernel/enum/__init__.py

@@ -0,0 +1,2 @@
+from .error_code import ErrorCode
+from .vector_database_type import VectorDatabaseType

sharedkernel-1.1.0/sharedkernel/enum/vector_database_type.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class VectorDatabaseType(Enum):
+    MILVUS = "milvus"
+    CHROMA = "chroma"

{sharedkernel-1.0.0 → sharedkernel-1.1.0}/sharedkernel.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sharedkernel
-Version: 1.0.0
+Version: 1.1.0
 Summary: sharekernel is an shared package between all python projects
 Author: Smilinno
 Requires-Dist: numpy
@@ -8,3 +8,5 @@ Requires-Dist: requests
 Requires-Dist: pymongo
 Requires-Dist: fastapi==0.89.1
 Requires-Dist: PyJWT
+Requires-Dist: pymilvus
+Requires-Dist: chromadb

{sharedkernel-1.0.0 → sharedkernel-1.1.0}/sharedkernel.egg-info/SOURCES.txt

@@ -9,14 +9,21 @@ sharedkernel.egg-info/SOURCES.txt
 sharedkernel.egg-info/dependency_links.txt
 sharedkernel.egg-info/requires.txt
 sharedkernel.egg-info/top_level.txt
+sharedkernel/database/__init__.py
 sharedkernel/database/adapter.py
-sharedkernel/database/repository_base.py
+sharedkernel/database/mongo_repository_base.py
+sharedkernel/database/vector_database_repository/__init__.py
+sharedkernel/database/vector_database_repository/chroma_startegy.py
+sharedkernel/database/vector_database_repository/milvus_strategy.py
+sharedkernel/database/vector_database_repository/vector_database_repository.py
+sharedkernel/database/vector_database_repository/vector_database_strategy.py
+sharedkernel/enum/__init__.py
 sharedkernel/enum/error_code.py
+sharedkernel/enum/vector_database_type.py
 sharedkernel/exception/__init__.py
 sharedkernel/exception/exception.py
 sharedkernel/exception/exception_handlers.py
 sharedkernel/objects/__init__.py
 sharedkernel/objects/base_document.py
 sharedkernel/objects/jwt_model.py
-sharedkernel/objects/result.py
-test/test.py
+sharedkernel/objects/result.py

{sharedkernel-1.0.0 → sharedkernel-1.1.0}/sharedkernel.egg-info/requires.txt

@@ -3,3 +3,5 @@ requests
 pymongo
 fastapi==0.89.1
 PyJWT
+pymilvus
+chromadb

sharedkernel-1.0.0/test/test.py

@@ -1,12 +0,0 @@
-from fastapi import FastAPI,Depends
-
-from sharedkernel import jwt_service
-from sharedkernel.objects import JwtModel
-from sharedkernel import config
-app = FastAPI(title="Domain Apis",dependencies=[Depends(jwt_service.JWTBearer(JwtModel(secret_key=config.JWT_SECRETKEY,
-                                                                                       algorithms=config.JWT_ALGORITHM,
-                                                                                       audience=config.JWT_AUDIENCE,
-                                                                                       issuer=config.JWT_ISSURE)))])
-
-# history = api_client.bots()
-# print("Chat history:", history)