Add virtual_api_server.py"

pipeline/virtual_api_server.py

@@ -0,0 +1,78 @@
+"""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()