From 494ee31c5a1586800b18c7e926b5582b821ca8d8 Mon Sep 17 00:00:00 2001
From: ACBBZ <tpcacbbz@gmail.com>
Date: Thu, 23 May 2024 04:02:50 +0000
Subject: [PATCH] update chroma_chat

---
 src/blackbox/chroma_chat.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/blackbox/chroma_chat.py b/src/blackbox/chroma_chat.py
index 33b9840..01685f3 100755
--- a/src/blackbox/chroma_chat.py
+++ b/src/blackbox/chroma_chat.py
@@ -4,8 +4,7 @@ from fastapi import Request, Response, status
 from fastapi.responses import JSONResponse
 from .blackbox import Blackbox
 
-from ..utils import chroma_setting
-from .fastchat import Fastchat
+from .chat import Chat
 from .chroma_query import ChromaQuery
 
 DEFAULT_COLLECTION_ID = "123"
@@ -15,8 +14,8 @@ from injector import singleton,inject
 class ChromaChat(Blackbox):
 
     @inject
-    def __init__(self, fastchat: Fastchat, chroma_query: ChromaQuery):
-        self.fastchat = fastchat
+    def __init__(self, chat: Chat, chroma_query: ChromaQuery):
+        self.chat = chat
         self.chroma_query = chroma_query
 
     def __call__(self, *args, **kwargs):
@@ -26,18 +25,18 @@ class ChromaChat(Blackbox):
         data = args[0]
         return isinstance(data, list)
 
-    def processing(self, question, setting: chroma_setting) -> str:
+    def processing(self, question, context: list) -> str:
+        if context == None:
+            context = []
 
         # load or create collection
-        if setting is None:
-            collection_id = DEFAULT_COLLECTION_ID
-        else:
-            collection_id = setting.ChromaSetting.collection_ids[0]
+        collection_id = DEFAULT_COLLECTION_ID
+
         # query it
         chroma_result = self.chroma_query(question, collection_id)
 
         fast_question = "问题: "+ question + "。根据问题，总结以下内容和来源：" + chroma_result
-        response = self.fastchat(model_name="Qwen1.5-14B-Chat", prompt=fast_question, template='回答限制50字.', context=None, temperature=0.8, top_p=0.8, top_k=-1, n=1, max_tokens=1024)
+        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)
 
         return response