55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
from fastapi import APIRouter, HTTPException, Request
|
||
from fastapi.responses import StreamingResponse
|
||
import json
|
||
from src.core import HistoryManager
|
||
from src.core.startup import startup
|
||
from src.utils.logging_config import setup_logger
|
||
|
||
chat = APIRouter(prefix="/chat")
|
||
logger = setup_logger("server-chat")
|
||
|
||
@chat.get("/")
|
||
async def chat_get():
|
||
return "Chat Get!"
|
||
|
||
@chat.post("/")
|
||
async def chat_post(request: Request):
|
||
request_data = await request.json()
|
||
query = request_data['query']
|
||
meta = request_data.get('meta')
|
||
history_manager = HistoryManager(request_data['history'])
|
||
|
||
new_query, refs = startup.retriever(query, history_manager.messages, meta)
|
||
|
||
messages = history_manager.get_history_with_msg(new_query, max_rounds=meta.get('history_round'))
|
||
history_manager.add_user(query)
|
||
logger.debug(f"Web history: {history_manager.messages}")
|
||
|
||
async def generate_response():
|
||
content = ""
|
||
for delta in startup.model.predict(messages, stream=True):
|
||
if not delta.content:
|
||
continue
|
||
|
||
if hasattr(delta, 'is_full') and delta.is_full:
|
||
content = delta.content
|
||
else:
|
||
content += delta.content
|
||
|
||
response_chunk = json.dumps({
|
||
"history": history_manager.update_ai(content),
|
||
"response": content,
|
||
"refs": refs # TODO: 优化 refs，不需要每次都返回
|
||
}, ensure_ascii=False).encode('utf8') + b'\n'
|
||
yield response_chunk
|
||
|
||
return StreamingResponse(generate_response(), media_type='application/json')
|
||
|
||
@chat.post("/call")
|
||
async def call(request: Request):
|
||
request_data = await request.json()
|
||
query = request_data['query']
|
||
response = startup.model.predict(query)
|
||
logger.debug({"query": query, "response": response.content})
|
||
|
||
return {"response": response.content} |