|
@@ -5,7 +5,7 @@ import inspect
|
|
import aiohttp
|
|
import aiohttp
|
|
import asyncio
|
|
import asyncio
|
|
|
|
|
|
-from typing import Any, Awaitable, Callable, get_type_hints, Dict, List, Union
|
|
|
|
|
|
+from typing import Any, Awaitable, Callable, get_type_hints, Dict, List, Union, Optional
|
|
from functools import update_wrapper, partial
|
|
from functools import update_wrapper, partial
|
|
|
|
|
|
|
|
|
|
@@ -348,40 +348,47 @@ async def get_tool_server_data(token: str, url: str) -> Dict[str, Any]:
|
|
return data
|
|
return data
|
|
|
|
|
|
|
|
|
|
-async def get_tool_servers_data(servers: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
|
|
- enabled_servers = [
|
|
|
|
- server for server in servers if server.get("config", {}).get("enable")
|
|
|
|
- ]
|
|
|
|
|
|
+async def get_tool_servers_data(
|
|
|
|
+ servers: List[Dict[str, Any]], session_token: Optional[str] = None
|
|
|
|
+) -> List[Dict[str, Any]]:
|
|
|
|
+ # Prepare list of enabled servers along with their original index
|
|
|
|
+ server_entries = []
|
|
|
|
+ for idx, server in enumerate(servers):
|
|
|
|
+ if server.get("config", {}).get("enable"):
|
|
|
|
+ url_path = server.get("path", "openapi.json")
|
|
|
|
+ full_url = f"{server.get('url')}/{url_path}"
|
|
|
|
|
|
- urls = [
|
|
|
|
- (
|
|
|
|
- server,
|
|
|
|
- f"{server.get('url')}/{server.get('path', 'openapi.json')}",
|
|
|
|
- server.get("key", ""),
|
|
|
|
- )
|
|
|
|
- for server in enabled_servers
|
|
|
|
- ]
|
|
|
|
|
|
+ auth_type = server.get("auth_type", "bearer")
|
|
|
|
+ token = None
|
|
|
|
|
|
- tasks = [get_tool_server_data(token, url) for _, url, token in urls]
|
|
|
|
|
|
+ if auth_type == "bearer":
|
|
|
|
+ token = server.get("key", "")
|
|
|
|
+ elif auth_type == "session":
|
|
|
|
+ token = session_token
|
|
|
|
+ server_entries.append((idx, server, full_url, token))
|
|
|
|
|
|
- results: List[Dict[str, Any]] = []
|
|
|
|
|
|
+ # Create async tasks to fetch data
|
|
|
|
+ tasks = [get_tool_server_data(token, url) for (_, _, url, token) in server_entries]
|
|
|
|
|
|
|
|
+ # Execute tasks concurrently
|
|
responses = await asyncio.gather(*tasks, return_exceptions=True)
|
|
responses = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
|
|
- for (server, _, _), response in zip(urls, responses):
|
|
|
|
|
|
+ # Build final results with index and server metadata
|
|
|
|
+ results = []
|
|
|
|
+ for (idx, server, url, _), response in zip(server_entries, responses):
|
|
if isinstance(response, Exception):
|
|
if isinstance(response, Exception):
|
|
- url_path = server.get("path", "openapi.json")
|
|
|
|
- full_url = f"{server.get('url')}/{url_path}"
|
|
|
|
- print(f"Failed to connect to {full_url} OpenAPI tool server")
|
|
|
|
- else:
|
|
|
|
- results.append(
|
|
|
|
- {
|
|
|
|
- "url": server.get("url"),
|
|
|
|
- "openapi": response["openapi"],
|
|
|
|
- "info": response["info"],
|
|
|
|
- "specs": response["specs"],
|
|
|
|
- }
|
|
|
|
- )
|
|
|
|
|
|
+ print(f"Failed to connect to {url} OpenAPI tool server")
|
|
|
|
+ continue
|
|
|
|
+
|
|
|
|
+ results.append(
|
|
|
|
+ {
|
|
|
|
+ "idx": idx,
|
|
|
|
+ "url": server.get("url"),
|
|
|
|
+ "openapi": response.get("openapi"),
|
|
|
|
+ "info": response.get("info"),
|
|
|
|
+ "specs": response.get("specs"),
|
|
|
|
+ }
|
|
|
|
+ )
|
|
|
|
|
|
return results
|
|
return results
|
|
|
|
|