Chroma and chat updated

src/blackbox/chat.py

@@ -5,17 +5,22 @@ from fastapi.responses import JSONResponse
 
 from ..log.logging_time import logging_time
 from .blackbox import Blackbox
+from .chroma_query import ChromaQuery
 
 import requests
 import json
 from openai import OpenAI
 import re
 
-from injector import singleton
+from injector import singleton,inject
 
 @singleton
 class Chat(Blackbox):
     
+    @inject
+    def __init__(self, chroma_query: ChromaQuery):
+        self.chroma_query = chroma_query
+
     def __call__(self, *args, **kwargs):
         return self.processing(*args, **kwargs)
 
@@ -24,7 +29,7 @@ class Chat(Blackbox):
         return isinstance(data, list)
 
     # model_name有 Qwen1.5-14B-Chat , internlm2-chat-20b
-    #@logging_time()
+    # @logging_time()
     def processing(self, prompt: str, context: list, settings: dict)  -> str:
 
         if settings is None:
@@ -42,6 +47,9 @@ class Chat(Blackbox):
         user_presence_penalty = settings.get("presence_penalty")
         user_model_url = settings.get("model_url")
         user_model_key = settings.get("model_key")
+        
+        chroma_embedding_model = settings.get("chroma_embedding_model")
+        chroma_response = ''
 
         if user_context == None:
             user_context = []
@@ -82,6 +90,14 @@ class Chat(Blackbox):
         if user_model_key is None or user_model_key.isspace() or user_model_key == "":
             user_model_key = "YOUR_API_KEY"
 
+        if chroma_embedding_model != None:
+            chroma_response = self.chroma_query(user_question, settings)
+            print(chroma_response)
+
+        if chroma_response != None or chroma_response != '':
+            user_question = f"问题: {user_question}。- 根据知识库内的检索结果，以清晰简洁的表达方式回答问题。只从检索的内容中选取与问题相关信息。- 不要编造答案，如果答案不在经核实的资料中或无法从经核实的资料中得出，请回答“我无法回答您的问题。”检索内容：{chroma_response}"
+
+
         # 文心格式和openai的不一样，需要单独处理
         if re.search(r"ernie", user_model_name):
             # key = "24.22873ef3acf61fb343812681e4df251a.2592000.1719453781.282335-46723715" 没充钱，只有ernie-speed-128k能用
@@ -132,7 +148,6 @@ class Chat(Blackbox):
             header = {
                 'Content-Type': 'application/json',
             }
-              
 
         prompt_template = [
            {"role": "system", "content": user_template},