github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-13 22:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: integrate Weaviate vector database in DB-GPT.

Browse Source 
1.Weaviate default schema update
2.Weaviate database config
3.requirement
This commit is contained in:
aries-ckt 2023-06-19 09:56:54 +08:00
parent 398aa57d46
commit b95084b89f
5 changed files with 53 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pilot/configs/config.py
View File

@ -150,6 +150,9 @@ class Config(metaclass=Singleton):
self.MILVUS_USERNAME = os.getenv("MILVUS_USERNAME", None)
self.MILVUS_PASSWORD = os.getenv("MILVUS_PASSWORD", None)
self.WEAVIATE_URL = os.getenv("WEAVIATE_URL", "http://127.0.0.1:8080")
# QLoRA
self.QLoRA = os.getenv("QUANTIZE_QLORA", "True")

2
pilot/scene/chat_knowledge/custom/chat.py
View File

@ -53,7 +53,7 @@ class ChatNewKnowledge(BaseChat):
docs = self.knowledge_embedding_client.similar_search(
self.current_user_input, CFG.KNOWLEDGE_SEARCH_TOP_SIZE
)
context = [d.page_content for d in docs]
context = [d["page_content"] for d in docs]
context = context[:2000]
input_values = {"context": context, "question": self.current_user_input}
return input_values

5
pilot/vector_store/connector.py
View File

@ -1,8 +1,9 @@
from pilot.vector_store.chroma_store import ChromaStore
# from pilot.vector_store.milvus_store import MilvusStore
from pilot.vector_store.milvus_store import MilvusStore
from pilot.vector_store.weaviate_store import WeaviateStore
connector = {"Chroma": ChromaStore, "Milvus": None}
connector = {"Chroma": ChromaStore, "Milvus": MilvusStore, "Weaviate": WeaviateStore}
class VectorStoreConnector:

81
pilot/vector_store/weaviate_store.py
View File

@ -2,15 +2,19 @@ import os
import json
import weaviate
from langchain.vectorstores import Weaviate
from pilot.configs.config import Config
from pilot.configs.model_config import KNOWLEDGE_UPLOAD_ROOT_PATH
from pilot.logs import logger
from pilot.vector_store.vector_store_base import VectorStoreBase
CFG = Config()
class WeaviateStore(VectorStoreBase):
"""Weaviate database"""
def __init__(self, ctx: dict, weaviate_url: str) -> None:
def __init__(self, ctx: dict) -> None:
"""Initialize with Weaviate client."""
try:
import weaviate
@ -21,9 +25,11 @@ class WeaviateStore(VectorStoreBase):
)
self.ctx = ctx
self.weaviate_url = weaviate_url
self.weaviate_url = CFG.WEAVIATE_URL
self.embedding = ctx.get("embeddings", None)
self.vector_name = ctx["vector_store_name"]
self.persist_dir = os.path.join(
KNOWLEDGE_UPLOAD_ROOT_PATH, ctx["vector_store_name"] + ".vectordb"
KNOWLEDGE_UPLOAD_ROOT_PATH, self.vector_name + ".vectordb"
)
self.vector_store_client = weaviate.Client(self.weaviate_url)
@ -31,26 +37,26 @@ class WeaviateStore(VectorStoreBase):
def similar_search(self, text: str, topk: int) -> None:
"""Perform similar search in Weaviate"""
logger.info("Weaviate similar search")
nearText = {
"concepts": [text],
"distance": 0.75, # prior to v1.14 use "certainty" instead of "distance"
}
# nearText = {
# "concepts": [text],
# "distance": 0.75, # prior to v1.14 use "certainty" instead of "distance"
# }
# vector = self.embedding.embed_query(text)
response = (
self.vector_store_client.query.get("Document", ["metadata", "text"])
.with_near_vector({"vector": nearText})
self.vector_store_client.query.get(self.vector_name, ["metadata", "page_content"])
# .with_near_vector({"vector": vector})
.with_limit(topk)
.with_additional(["distance"])
.do()
)
return json.dumps(response, indent=2)
docs = response['data']['Get'][list(response['data']['Get'].keys())[0]]
return docs
def vector_name_exists(self) -> bool:
"""Check if a vector name exists for a given class in Weaviate.
Returns:
bool: True if the vector name exists, False otherwise.
"""
if self.vector_store_client.schema.get("Document"):
if self.vector_store_client.schema.get(self.vector_name):
return True
return False
@ -62,39 +68,39 @@ class WeaviateStore(VectorStoreBase):
schema = {
"classes": [
{
"class": "Document",
"class": self.vector_name,
"description": "A document with metadata and text",
"moduleConfig": {
"text2vec-transformers": {
"poolingStrategy": "masked_mean",
"vectorizeClassName": False,
}
},
# "moduleConfig": {
# "text2vec-transformers": {
# "poolingStrategy": "masked_mean",
# "vectorizeClassName": False,
# }
# },
"properties": [
{
"dataType": ["text"],
"moduleConfig": {
"text2vec-transformers": {
"skip": False,
"vectorizePropertyName": False,
}
},
# "moduleConfig": {
# "text2vec-transformers": {
# "skip": False,
# "vectorizePropertyName": False,
# }
# },
"description": "Metadata of the document",
"name": "metadata",
},
{
"dataType": ["text"],
"moduleConfig": {
"text2vec-transformers": {
"skip": False,
"vectorizePropertyName": False,
}
},
# "moduleConfig": {
# "text2vec-transformers": {
# "skip": False,
# "vectorizePropertyName": False,
# }
# },
"description": "Text content of the document",
"name": "text",
},
"name": "page_content",
}
],
"vectorizer": "text2vec-transformers",
# "vectorizer": "text2vec-transformers",
}
]
}
@ -114,6 +120,7 @@ class WeaviateStore(VectorStoreBase):
# Batch import all documents
for i in range(len(texts)):
properties = {"metadata": metadatas[i], "text": texts[i]}
properties = {"metadata": metadatas[i]['source'], "page_content": texts[i]}
self.vector_store_client.batch.add_data_object(properties, "Document")
self.vector_store_client.batch.add_data_object(data_object=properties, class_name=self.vector_name)
self.vector_store_client.batch.flush()

3
requirements.txt
View File

@ -59,12 +59,13 @@ nltk
python-dotenv==1.0.0
# pymilvus==2.2.1
vcrpy
chromadb
chromadb=0.3.22
markdown2
colorama
playsound
distro
pypdf
weaviate-client
# Testing dependencies
pytest