Spaces:
Runtime error
Runtime error
File size: 2,980 Bytes
b044427 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import fastapi
import uvicorn
from .registry import build_tool, list_tools
from .retriever import Retriever
from typing import List
from pydantic import BaseModel
class RetrieveRequest(BaseModel):
query: str
topk: int = 3
def _bind_tool_server(t : "ToolServer"):
""" Add property API to ToolServer.
t.api is a FastAPI object
"""
@t.api.get("/")
def health():
return {
"status": "ok",
}
@t.api.get("/list")
def get_tools_list():
return {
"tools": t.list_tools(),
}
@t.api.get("/loaded")
def get_loaded_tools():
return {
"tools": list(t.loaded_tools),
}
@t.api.get("/.well-known/ai-plugin.json", include_in_schema=False)
def get_api_info():
return {
"schema_version": "v1",
"name_for_human": "BMTools",
"name_for_model": "BMTools",
"description_for_human": "tools to big models",
"description_for_model": "tools to big models",
"auth": {
"type": "none",
},
"api": {
"type": "openapi",
"url": "/openapi.json",
"is_user_authenticated": False,
},
"logo_url": None,
"contact_email": "",
"legal_info_url": "",
}
@t.api.post("/retrieve")
def retrieve(request: RetrieveRequest):
tool_list = t.retrieve(request.query, request.topk)
return {
"tools": tool_list
}
class ToolServer:
""" This class host your own API backend.
"""
def __init__(self) -> None:
# define the root API server
self.api = fastapi.FastAPI(
title="BMTools",
description="Tools for bigmodels",
)
self.loaded_tools = dict()
self.retriever = Retriever()
_bind_tool_server(self)
def load_tool(self, name : str, config = {}):
if self.is_loaded(name):
raise ValueError(f"Tool {name} is already loaded")
try:
tool = build_tool(name, config)
except BaseException as e:
print(f"Cannot load tool {name}: {repr(e)}")
return
# tool = build_tool(name, config)
self.loaded_tools[name] = tool.api_info
self.retriever.add_tool(name, tool.api_info)
# mount sub API server to the root API server, thus can mount all urls of sub API server to /tools/{name} route
self.api.mount(f"/tools/{name}", tool, name)
return
def is_loaded(self, name : str):
return name in self.loaded_tools
def serve(self, host : str = "0.0.0.0", port : int = 8079):
uvicorn.run(self.api, host=host, port=port)
def list_tools(self) -> List[str]:
return list_tools()
def retrieve(self, query: str, topk: int = 3) -> List[str]:
return self.retriever.query(query, topk) |