Repository navigation
Expand file tree
/
Copy pathfastapi4modelcache.py
More file actions
74 lines (64 loc) · 2.32 KB
/
Copy pathfastapi4modelcache.py
File metadata and controls
74 lines (64 loc) · 2.32 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
# -*- coding: utf-8 -*-
import asyncio
from contextlib import asynccontextmanager
import uvicorn
import json
from fastapi import FastAPI, Request, HTTPException
from modelcache.cache import Cache
from modelcache.embedding import EmbeddingModel
#创建一个FastAPI实例
@asynccontextmanager
async def startup_event(app: FastAPI):
global cache
cache, _ = await Cache.init(
sql_storage="mysql",
vector_storage="milvus",
embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
embedding_workers_num=1
)
yield
app = FastAPI(lifespan=startup_event)
cache: Cache = None
@app.get("/welcome")
async def first_fastapi():
return "hello, modelcache!"
@app.post("/modelcache")
async def user_backend(request: Request):
try:
raw_body = await request.body()
# 解析字符串为JSON对象
if isinstance(raw_body, bytes):
raw_body = raw_body.decode("utf-8")
if isinstance(raw_body, str):
try:
# 尝试将字符串解析为JSON对象
request_data = json.loads(raw_body)
except json.JSONDecodeError as e:
# 如果无法解析，返回格式错误
result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
"answer": ''}
cache.save_query_info(result, model='', query='', delta_time_log=0)
raise HTTPException(status_code=101, detail="Invalid JSON format")
else:
request_data = raw_body
# 确保request_data是字典对象
if isinstance(request_data, str):
try:
request_data = json.loads(request_data)
except json.JSONDecodeError:
raise HTTPException(status_code=101, detail="Invalid JSON format")
return await cache.handle_request(request_data)
except Exception as e:
request_data = raw_body if 'raw_body' in locals() else None
result = {
"errorCode": 103,
"errorDesc": str(e),
"cacheHit": False,
"delta_time": 0,
"hit_query": '',
"answer": '',
"para_dict": request_data
}
return result
if __name__ == '__main__':
uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio")