71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
import json
|
|
import time
|
|
from typing import Any, List
|
|
|
|
from fastapi import APIRouter, FastAPI
|
|
from langchain_core.messages import AIMessage
|
|
from langchain_core.runnables import Runnable
|
|
from starlette.responses import StreamingResponse
|
|
|
|
|
|
def format_ai_message(ai_message: AIMessage):
|
|
return {
|
|
"choices": [{
|
|
"finish_reason": "stop",
|
|
"index": 0,
|
|
"message": {
|
|
"content": ai_message.content,
|
|
"role": "assistant"
|
|
}
|
|
}],
|
|
"created": int(time.time()),
|
|
"id": ai_message.id,
|
|
"usage": ai_message.response_metadata.get('token_usage')
|
|
}
|
|
|
|
|
|
def add_async_route(
|
|
app: FastAPI,
|
|
runnable: Runnable,
|
|
path: str,
|
|
input_type: Any = None, # 设置接口的传入参数类型，input_type实际上是一个pydantic类
|
|
):
|
|
router = APIRouter(prefix=path, tags=[path])
|
|
|
|
_input_type = input_type or runnable.input_schema
|
|
|
|
@router.post('/ainvoke')
|
|
async def ainvoke(input: _input_type):
|
|
ai_message = await runnable.ainvoke(input.model_dump())
|
|
return format_ai_message(ai_message)
|
|
|
|
@router.post('/abatch')
|
|
async def abatch(inputs: List[_input_type]):
|
|
ai_message_list = await runnable.abatch([input.model_dump() for input in inputs])
|
|
return [format_ai_message(ai_message) for ai_message in ai_message_list]
|
|
|
|
@router.post('/astream')
|
|
async def astream(input: _input_type):
|
|
async def generator_function():
|
|
result_template = {
|
|
"choices": [{"delta": {}, "index": 0}],
|
|
"created": time.time(),
|
|
"id": "",
|
|
"usage": None
|
|
}
|
|
async for chunk in runnable.astream(input.model_dump()):
|
|
result_template["choices"][0]["delta"]['content'] = chunk.content
|
|
result_template["choices"][0]["delta"]['role'] = 'assistant'
|
|
|
|
if chunk.response_metadata.get('finish_reason') is not None:
|
|
result_template["choices"][0]["delta"]['finish_reason'] = chunk.response_metadata.get('finish_reason')
|
|
|
|
result_template['id'] = chunk.id
|
|
result_template['created'] = int(time.time())
|
|
yield f"data: {json.dumps(result_template, ensure_ascii=False)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
return StreamingResponse(generator_function(), media_type="text/event-stream")
|
|
|
|
app.include_router(router)
|