Class ToolServer

bmtools/tools/serve.py:63–101

This class host your own API backend.

        }

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: Dict = {}):
        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
        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)