george
73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
"""Async docs-api client for the WebUI."""
|
|
import os
|
|
from typing import Any, Dict, Optional
|
|
|
|
from httpx import AsyncClient, Timeout
|
|
|
|
|
|
class DocsAPIClient:
|
|
"""Small async HTTP client for the docs-api backend."""
|
|
|
|
def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None):
|
|
self.base_url = (base_url or os.environ.get("DOCS_API_URL", "http://docs-api:8787")).rstrip("/")
|
|
self.api_key = api_key if api_key is not None else os.environ.get("WEBUI_API_KEY")
|
|
self.headers = {"X-API-Key": self.api_key} if self.api_key else {}
|
|
self._client: Optional[AsyncClient] = None
|
|
|
|
async def _get_client(self) -> AsyncClient:
|
|
if self._client is None or self._client.is_closed:
|
|
self._client = AsyncClient(
|
|
base_url=self.base_url,
|
|
headers=self.headers,
|
|
timeout=Timeout(120.0),
|
|
)
|
|
return self._client
|
|
|
|
async def request(self, method: str, path: str, **kwargs: Any) -> Dict[str, Any]:
|
|
client = await self._get_client()
|
|
resp = await client.request(method, path, **kwargs)
|
|
if resp.status_code >= 400:
|
|
raise RuntimeError(f"{method} {path} failed: {resp.status_code} {resp.text}")
|
|
if resp.headers.get("content-type", "").startswith("application/json"):
|
|
data = resp.json()
|
|
return data if isinstance(data, dict) else {"data": data}
|
|
return {"data": resp.text}
|
|
|
|
async def get(self, path: str, **kwargs: Any) -> Dict[str, Any]:
|
|
return await self.request("GET", path, **kwargs)
|
|
|
|
async def post(self, path: str, **kwargs: Any) -> Dict[str, Any]:
|
|
return await self.request("POST", path, **kwargs)
|
|
|
|
async def delete(self, path: str, **kwargs: Any) -> Dict[str, Any]:
|
|
return await self.request("DELETE", path, **kwargs)
|
|
|
|
async def health(self) -> Dict[str, Any]:
|
|
try:
|
|
return await self.get("/health")
|
|
except Exception as e:
|
|
return {"status": "error", "message": str(e)}
|
|
|
|
async def upload_file(self, library_id: str, filename: str, content: bytes) -> Dict[str, Any]:
|
|
files = {"file": (filename, content)}
|
|
return await self.post(f"/api/v1/upload/{library_id}", files=files)
|
|
|
|
async def close(self) -> None:
|
|
if self._client is not None and not self._client.is_closed:
|
|
await self._client.aclose()
|
|
|
|
|
|
_client_instance: Optional[DocsAPIClient] = None
|
|
|
|
|
|
async def get_client() -> DocsAPIClient:
|
|
global _client_instance
|
|
if _client_instance is None:
|
|
_client_instance = DocsAPIClient()
|
|
return _client_instance
|
|
|
|
|
|
async def close_client() -> None:
|
|
if _client_instance is not None:
|
|
await _client_instance.close()
|