fix: chroma setting url and key

src/blackbox/chat.py

@@ -40,6 +40,9 @@ class Chat(Blackbox):
         user_stop = settings.get("stop")
         user_frequency_penalty = settings.get("frequency_penalty")
         user_presence_penalty = settings.get("presence_penalty")
+        user_model_url = settings.get("model_url")
+        user_model_key = settings.get("model_key")
+
         if user_context == None:
             user_context = []
 
@@ -73,14 +76,19 @@ class Chat(Blackbox):
         if user_presence_penalty is None or user_presence_penalty == "":
             user_presence_penalty = 0.8
  
+        if user_model_url is None or user_model_url.isspace() or user_model_url == "":
+            user_model_url = "http://120.196.116.194:48892/v1/chat/completions"
+
+        if user_model_key is None or user_model_key.isspace() or user_model_key == "":
+            user_model_key = "YOUR_API_KEY"
 
         # gpt-4， gpt-3.5-turbo
         if re.search(r"gpt", user_model_name):
             url = 'https://api.openai.com/v1/completions' 
             key = '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'
         else:
-            url = 'http://120.196.116.194:48892/v1/chat/completions'
-            key = 'YOUR_API_KEY'
+            url = user_model_url
+            key = user_model_key
               
         prompt_template = [
            {"role": "system", "content": user_template},

src/blackbox/chroma_chat.py

@@ -6,7 +6,10 @@ from .blackbox import Blackbox
 
 from .chat import Chat
 from .chroma_query import ChromaQuery
+from ..log.logging_time import logging_time
+import logging
 
+logger = logging.getLogger
 DEFAULT_COLLECTION_ID = "123"
 
 from injector import singleton,inject
@@ -25,18 +28,67 @@ class ChromaChat(Blackbox):
         data = args[0]
         return isinstance(data, list)
     
-    def processing(self, question, context: list) -> str:
+    @logging_time(logger=logger)
+    def processing(self, question: str, context: list, settings: dict) -> str:
+
+        if settings is None:
+            settings = {}
+
+        # # chat setting
+        user_model_name = settings.get("model_name")
+        user_context = context
+        user_question = question
+        user_template = settings.get("template")
+        user_temperature = settings.get("temperature")
+        user_top_p = settings.get("top_p")
+        user_n = settings.get("n")
+        user_max_tokens = settings.get("max_tokens")
+        user_stop = settings.get("stop")
+        user_frequency_penalty = settings.get("frequency_penalty")
+        user_presence_penalty = settings.get("presence_penalty")
+
+        # # 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 context == None:
             context = []
 
-        # load or create collection
-        collection_id = DEFAULT_COLLECTION_ID
+        if user_question is None:
+            return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
 
-        # query it
-        chroma_result = self.chroma_query(question, collection_id)
+        chroma_settings_json={
+            "chroma_embedding_model": chroma_embedding_model,
+            "chroma_host": chroma_host,
+            "chroma_port": chroma_port,
+            "chroma_collection_id": chroma_collection_id,
+            "chroma_n_results": chroma_n_results
+        }
 
-        fast_question = "问题: "+ question + "。根据问题，总结以下内容和来源：" + chroma_result
-        response = self.chat(model_name="Qwen1.5-14B-Chat", prompt=fast_question, template='', context=context, temperature=0.8, top_p=0.8, n=1, max_tokens=1024, stop=100,frequency_penalty=0.5,presence_penalty=0.8)
+        # chroma answer
+        chroma_result = self.chroma_query(user_question, chroma_settings_json)
+
+        # chat prompt 
+        fast_question = f"问题: {user_question}。根据问题，总结以下内容和来源：{chroma_result}"
+
+        chat_settings_json = {
+            "model_name": user_model_name,
+            "context": user_context,
+            "template": user_template,
+            "temperature": user_temperature,
+            "top_p": user_top_p,
+            "n": user_n,
+            "max_tokens": user_max_tokens,
+            "stop": user_stop,
+            "frequency_penalty": user_frequency_penalty,
+            "presence_penalty": user_presence_penalty
+        }
+
+        # chat answer
+        response = self.chat(fast_question, chat_settings_json)
 
         return response
 
@@ -49,10 +101,8 @@ class ChromaChat(Blackbox):
 
         user_question = data.get("question")
         user_context = data.get("context")
-
-        if user_question is None:
-            return JSONResponse(content={"error": "question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
+        setting: dict = data.get("settings")
 
         return JSONResponse(
-            content={"response": self.processing(user_question, user_context)},
+            content={"response": self.processing(user_question, user_context, setting)},
             status_code=status.HTTP_200_OK)

src/blackbox/chroma_query.py

@@ -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)