mirror of
https://github.com/BoardWare-Genius/jarvis-models.git
synced 2025-12-13 16:53:24 +00:00
add chroma chat
This commit is contained in:
52
src/blackbox/chroma_chat.py
Executable file
52
src/blackbox/chroma_chat.py
Executable file
@ -0,0 +1,52 @@
|
|||||||
|
from typing import Any, Coroutine
|
||||||
|
|
||||||
|
from fastapi import Request, Response, status
|
||||||
|
from fastapi.responses import JSONResponse
|
||||||
|
from .blackbox import Blackbox
|
||||||
|
|
||||||
|
from injector import singleton
|
||||||
|
@singleton
|
||||||
|
class ChromaChat(Blackbox):
|
||||||
|
|
||||||
|
def __init__(self, fastchat, chroma_query):
|
||||||
|
self.fastchat = fastchat
|
||||||
|
self.chroma_query = chroma_query
|
||||||
|
|
||||||
|
def __call__(self, *args, **kwargs):
|
||||||
|
return self.processing(*args, **kwargs)
|
||||||
|
|
||||||
|
def valid(self, *args, **kwargs) -> bool:
|
||||||
|
data = args[0]
|
||||||
|
return isinstance(data, list)
|
||||||
|
|
||||||
|
def processing(self, question, context) -> str:
|
||||||
|
|
||||||
|
# load or create collection
|
||||||
|
if context is None:
|
||||||
|
collection_id = "123"
|
||||||
|
else:
|
||||||
|
collection_id = context["collections"][0]
|
||||||
|
# query it
|
||||||
|
chroma_result = self.chroma_query(question, collection_id)
|
||||||
|
|
||||||
|
fast_question = "问题: "+ question + "。根据问题,总结以下内容:" + chroma_result
|
||||||
|
response = self.fastchat(fast_question)
|
||||||
|
|
||||||
|
return response
|
||||||
|
|
||||||
|
|
||||||
|
async def fast_api_handler(self, request: Request) -> Response:
|
||||||
|
try:
|
||||||
|
data = await request.json()
|
||||||
|
except:
|
||||||
|
return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
return JSONResponse(
|
||||||
|
content={"response": self.processing(user_question, user_context)},
|
||||||
|
status_code=status.HTTP_200_OK)
|
||||||
Reference in New Issue
Block a user