-
Notifications
You must be signed in to change notification settings - Fork 105
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add milvus vector database for rag recipe
Signed-off-by: Shreyanand <[email protected]> Co-authored-by: Michael Clifford <[email protected]> Co-authored-by: greg pereira <[email protected]>
- Loading branch information
1 parent
e2ead10
commit 705b5d1
Showing
10 changed files
with
182 additions
and
101 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
81 changes: 81 additions & 0 deletions
81
recipes/natural_language_processing/rag/app/manage_vectordb.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,81 @@ | ||
from langchain_community.vectorstores import Chroma | ||
from chromadb import HttpClient | ||
from chromadb.config import Settings | ||
import chromadb.utils.embedding_functions as embedding_functions | ||
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings | ||
from langchain_community.vectorstores import Milvus | ||
from pymilvus import MilvusClient | ||
from pymilvus import connections, utility | ||
|
||
class VectorDB: | ||
def __init__(self, vector_vendor, host, port, collection_name, embedding_model): | ||
self.vector_vendor = vector_vendor | ||
self.host = host | ||
self.port = port | ||
self.collection_name = collection_name | ||
self.embedding_model = embedding_model | ||
|
||
def connect(self): | ||
# Connection logic | ||
print(f"Connecting to {self.host}:{self.port}...") | ||
if self.vector_vendor == "chromadb": | ||
self.client = HttpClient(host=self.host, | ||
port=self.port, | ||
settings=Settings(allow_reset=True,)) | ||
elif self.vector_vendor == "milvus": | ||
self.client = MilvusClient(uri=f"http://{self.host}:{self.port}") | ||
return self.client | ||
|
||
def populate_db(self, documents): | ||
# Logic to populate the VectorDB with vectors | ||
e = SentenceTransformerEmbeddings(model_name=self.embedding_model) | ||
print(f"Populating VectorDB with vectors...") | ||
if self.vector_vendor == "chromadb": | ||
embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=self.embedding_model) | ||
collection = self.client.get_or_create_collection(self.collection_name, | ||
embedding_function=embedding_func) | ||
if collection.count() < 1: | ||
db = Chroma.from_documents( | ||
documents=documents, | ||
embedding=e, | ||
collection_name=self.collection_name, | ||
client=self.client | ||
) | ||
print("DB populated") | ||
else: | ||
db = Chroma(client=self.client, | ||
collection_name=self.collection_name, | ||
embedding_function=e, | ||
) | ||
print("DB already populated") | ||
|
||
elif self.vector_vendor == "milvus": | ||
connections.connect(host=self.host, port=self.port) | ||
if not utility.has_collection(self.collection_name): | ||
print("Populating VectorDB with vectors...") | ||
db = Milvus.from_documents( | ||
documents, | ||
e, | ||
collection_name=self.collection_name, | ||
connection_args={"host": self.host, "port": self.port}, | ||
) | ||
print("DB populated") | ||
else: | ||
print("DB already populated") | ||
db = Milvus( | ||
e, | ||
collection_name=self.collection_name, | ||
connection_args={"host": self.host, "port": self.port}, | ||
) | ||
return db | ||
|
||
def clear_db(self): | ||
print(f"Clearing VectorDB...") | ||
try: | ||
if self.vector_vendor == "chromadb": | ||
self.client.delete_collection(self.collection_name) | ||
elif self.vector_vendor == "milvus": | ||
self.client.drop_collection(self.collection_name) | ||
print("Cleared DB") | ||
except: | ||
print("Couldn't clear the collection possibly because it doesn't exist") |
36 changes: 0 additions & 36 deletions
36
recipes/natural_language_processing/rag/app/populate_vectordb.py
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
FROM docker.io/milvusdb/milvus:master-20240426-bed6363f | ||
ADD embedEtcd.yaml /milvus/configs/embedEtcd.yaml |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
REGISTRY ?= quay.io | ||
REGISTRY_ORG ?= ai-lab | ||
COMPONENT = vector_dbs | ||
|
||
IMAGE ?= $(REGISTRY)/$(REGISTRY_ORG)/$(COMPONENT)/milvus:latest | ||
|
||
ARCH ?= $(shell uname -m) | ||
PLATFORM ?= linux/$(ARCH) | ||
|
||
gRCP_PORT := 19530 | ||
REST_PORT := 9091 | ||
CLIENT_PORT := 2379 | ||
|
||
LIB_MILVUS_DIR_MOUNTPATH := $(shell pwd)/volumes/milvus | ||
|
||
.PHONY: build | ||
build: | ||
podman build --platform $(PLATFORM) -f Containerfile -t ${IMAGE} . | ||
|
||
.PHONY: run | ||
run: | ||
podman run -d \ | ||
--name milvus-standalone \ | ||
--security-opt seccomp:unconfined \ | ||
-e ETCD_USE_EMBED=true \ | ||
-e ETCD_CONFIG_PATH=/milvus/configs/embedEtcd.yaml \ | ||
-e COMMON_STORAGETYPE=local \ | ||
-v $(LIB_MILVUS_DIR_MOUNTPATH):/var/lib/milvus \ | ||
-p $(gRCP_PORT):$(gRCP_PORT) \ | ||
-p $(REST_PORT):$(REST_PORT) \ | ||
-p $(CLIENT_PORT):$(CLIENT_PORT) \ | ||
--health-cmd="curl -f http://localhost:$(REST_PORT)/healthz" \ | ||
--health-interval=30s \ | ||
--health-start-period=90s \ | ||
--health-timeout=20s \ | ||
--health-retries=3 \ | ||
$(IMAGE) \ | ||
milvus run standalone 1> /dev/null | ||
|
||
.PHONY: stop | ||
stop: | ||
-podman stop milvus-standalone | ||
|
||
.PHONY: delete | ||
delete: | ||
-podman rm milvus-standalone -f | ||
|
||
.PHONY: podman-clean | ||
podman-clean: | ||
@container_ids=$$(podman ps --format "{{.ID}} {{.Image}}" | awk '$$2 == "$(IMAGE)" {print $$1}'); \ | ||
echo "removing all containers with IMAGE=$(IMAGE)"; \ | ||
for id in $$container_ids; do \ | ||
echo "Removing container: $$id,"; \ | ||
podman rm -f $$id; \ | ||
done |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
listen-client-urls: http://0.0.0.0:2379 | ||
advertise-client-urls: http://0.0.0.0:2379 | ||
quota-backend-bytes: 4294967296 | ||
auto-compaction-mode: revision | ||
auto-compaction-retention: '1000' |