ivanwu/jarvis-models
1
0
Fork 0
mirror of https://github.com/BoardWare-Genius/jarvis-models.git synced 2025-12-13 16:53:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: chroma setting url and key

Browse Source
This commit is contained in:
ACBBZ
2024-05-27 08:42:58 +00:00
parent 8fee786a41
commit 76971c87f5
3 changed files with 124 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
src/blackbox/chroma_query.py
View File

@ -6,8 +6,11 @@ from .blackbox import Blackbox
import chromadb
from chromadb.utils import embedding_functions
from ..utils import chroma_setting
import logging
from ..log.logging_time import logging_time
import re
logger = logging.getLogger
DEFAULT_COLLECTION_ID = "123"
from injector import singleton
@ -16,10 +19,11 @@ class ChromaQuery(Blackbox):
def __init__(self, *args, **kwargs) -> None:
# config = read_yaml(args[0])
# load embedding model
self.embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/model/Weight/BAAI/bge-small-en-v1.5", device = "cuda")
# load chromadb
self.client = chromadb.HttpClient(host='10.6.82.192', port=8000)
# load chromadb and embedding model
self.embedding_model_1 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/model/Weight/BAAI/bge-small-en-v1.5", device = "cuda")
# self.embedding_model_2 = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="/model/Weight/BAAI/bge-small-en-v1.5", device = "cuda")
self.client_1 = chromadb.HttpClient(host='10.6.82.192', port=8000)
# self.client_2 = chromadb.HttpClient(host='10.6.82.192', port=8000)
def __call__(self, *args, **kwargs):
return self.processing(*args, **kwargs)
@ -28,14 +32,52 @@ class ChromaQuery(Blackbox):
data = args[0]
return isinstance(data, list)
def processing(self, question: str, collection_id) -> str:
@logging_time(logger=logger)
def processing(self, question: str, settings: dict) -> str:
# load or create collection
collection = self.client.get_collection(collection_id, embedding_function=self.embedding_model)
if settings is None:
settings = {}
usr_question = question
# # chroma_query settings
chroma_embedding_model = settings.get("chroma_embedding_model")
chroma_host = settings.get("chroma_host")
chroma_port = settings.get("chroma_port")
chroma_collection_id = settings.get("chroma_collection_id")
chroma_n_results = settings.get("chroma_n_results")
if usr_question is None:
return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
if chroma_embedding_model is None or chroma_embedding_model.isspace() or chroma_embedding_model == "":
chroma_embedding_model = "bge-small-en-v1.5"
if chroma_host is None or chroma_host.isspace() or chroma_host == "":
chroma_host = "10.6.82.192"
if chroma_port is None or chroma_port.isspace() or chroma_port == "":
chroma_port = "8000"
if chroma_collection_id is None or chroma_collection_id.isspace() or chroma_collection_id == "":
chroma_collection_id = DEFAULT_COLLECTION_ID
if chroma_n_results is None or chroma_n_results == "":
chroma_n_results = 3
# load client
if re.search(r"10.6.82.192", chroma_host) and re.search(r"8000", chroma_port):
client = self.client_1
if re.search(r"bge-small-en-v1.5", chroma_embedding_model):
embedding_model = self.embedding_model_1
# load collection
collection = client.get_collection(chroma_collection_id, embedding_function=embedding_model)
# query it
results = collection.query(
query_texts=[question],
query_texts=[usr_question],
n_results=3,
)
@ -50,14 +92,8 @@ class ChromaQuery(Blackbox):
return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)
user_question = data.get("question")
user_collection_id = data.get("collection_id")
if user_question is None:
return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
if user_collection_id is None:
user_collection_id = DEFAULT_COLLECTION_ID
setting = data.get("settings")
return JSONResponse(
content={"response": self.processing(user_question, user_collection_id)},
content={"response": self.processing(user_question, setting)},
status_code=status.HTTP_200_OK)