From 31dbe29b8bd0e053d7ee0546c11e10d0ab850832 Mon Sep 17 00:00:00 2001 From: ACBBZ Date: Tue, 7 May 2024 06:24:35 +0000 Subject: [PATCH] update chroma chat --- src/blackbox/chroma_chat.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/blackbox/chroma_chat.py b/src/blackbox/chroma_chat.py index c1e47dd..560712c 100755 --- a/src/blackbox/chroma_chat.py +++ b/src/blackbox/chroma_chat.py @@ -4,6 +4,9 @@ from fastapi import Request, Response, status from fastapi.responses import JSONResponse from .blackbox import Blackbox +from ..utils import chroma_setting +DEFAULT_COLLECTION_ID = "123" + from injector import singleton @singleton class ChromaChat(Blackbox): @@ -19,13 +22,13 @@ class ChromaChat(Blackbox): data = args[0] return isinstance(data, list) - def processing(self, question, context) -> str: + def processing(self, question, setting: chroma_setting) -> str: # load or create collection - if context is None: - collection_id = "123" + if setting is None: + collection_id = DEFAULT_COLLECTION_ID else: - collection_id = context["collections"][0] + collection_id = setting.ChromaSetting.collection_ids[0] # query it chroma_result = self.chroma_query(question, collection_id)