File size: 3,015 Bytes
6284167 | 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 | """Virtual API server for StableToolBench.
Cache-only server. For cached requests returns cached response.
For uncached requests returns simulated response.
For full simulation, deploy stabletoolbench/MirrorAPI.
"""
import json, os, re
from typing import Union, Optional
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
app = FastAPI()
def standardize(name):
res = re.compile(r"[^\u4e00-\u9fa5a-zA-Z0-9_]")
name = res.sub("_", name)
name = re.sub(r"(_)\1+", "_", name).lower().strip("_")
if name and name[0].isdigit(): name = "get_" + name
return name
def change_name(name):
if name in ["from", "class", "return", "false", "true", "id", "and"]: return "is_" + name
return name
CACHE_FOLDER = os.environ.get("CACHE_FOLDER", os.path.join(os.path.dirname(__file__), "api_cache"))
TOOLS_FOLDER = os.environ.get("TOOLS_FOLDER", os.path.join(os.path.dirname(__file__), "..", "toolenv2404_filtered"))
class Info(BaseModel):
category: str
tool_name: str
api_name: str
tool_input: Union[str, dict]
strip: str
toolbench_key: str = ""
def prepare_tool_name_and_url(info):
category = info.category.replace(" ", "_").replace(",", "_").replace("/", "_").replace("__", "_")
tool_name = info.tool_name
api_name = change_name(standardize(info.api_name)).split(f"_for_{tool_name}")[0]
if not tool_name.endswith(f"_for_{category}"):
tool_name = standardize(info.tool_name) + f"_for_{category}"
return tool_name, category, api_name
@app.post('/virtual')
def get_virtual_response(info: Info):
tool_name, standard_category, api_name = prepare_tool_name_and_url(info)
tool_input = info.tool_input
try:
if isinstance(tool_input, str): tool_input = json.loads(tool_input) if tool_input else {}
except: tool_input = {}
cache_path = os.path.join(CACHE_FOLDER, standard_category, tool_name, api_name + ".json")
if os.path.exists(cache_path):
try:
with open(cache_path) as f: cache = json.load(f)
if str(tool_input) in cache: return cache[str(tool_input)]
except: pass
api_doc = _load_api_doc(standard_category, info.tool_name, api_name)
if api_doc:
return {"error": "", "response": json.dumps({"message": f"Simulated response for {api_name}", "input": tool_input, "status": "success"})}
return {"error": "API not working error...", "response": ""}
def _load_api_doc(category, tool_name, api_name):
json_path = os.path.join(TOOLS_FOLDER, category, standardize(tool_name) + ".json")
if not os.path.exists(json_path): return None
try:
with open(json_path) as f: return json.load(f)
except: return None
@app.get('/health')
def health(): return {"status": "ok"}
def main():
port = int(os.environ.get("API_SERVER_PORT", 8080))
print(f"Starting virtual API server on port {port}\nCache: {CACHE_FOLDER}\nTools: {TOOLS_FOLDER}")
uvicorn.run(app, host="0.0.0.0", port=port)
if __name__ == "__main__": main()
|