feat: init project

app/utils/add_async_route.py

@@ -0,0 +1,70 @@
+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)