Teknium
d4b4e8b0de
- session_search_tool_common/_discover folded back into session_search_tool (single module, no re-export shim) - MemoryStore: one locked _mutate pipeline for add/replace/remove/batch; _reload_target/_read_entries_checked/_commit/_batch_error/_previews inlined - memory_tool: _STORE_ACTIONS table now carries store call + gate text; _target_label/_missing_old_text_error inlined - session_search: _discovery_entry (title + FTS hits), _same_lineage, _loud (error-level DB failure -> tool_error), _session_end_reason/_normalize_title_query inlined - graph client: iterate_pages folded into collect_paginated (its only caller), _decode_json_or, try/else retry loop - docstrings compacted by hand keeping every WHY/invariant; tool schemas byte-identical; SQL untouched
207 lines
10 KiB
Python
207 lines
10 KiB
Python
"""Reusable Microsoft Graph REST client helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
import httpx
|
|
|
|
from agent.retry_utils import parse_retry_after_seconds
|
|
from tools.microsoft_graph_auth import GraphCredentials, MicrosoftGraphTokenProvider, format_graph_error
|
|
|
|
|
|
DEFAULT_GRAPH_BASE_URL = "https://graph.microsoft.com/v1.0"
|
|
|
|
Headers = dict[str, str] | None
|
|
Params = dict[str, Any] | None
|
|
|
|
|
|
class MicrosoftGraphClientError(RuntimeError):
|
|
"""Base class for Graph client failures."""
|
|
|
|
|
|
class MicrosoftGraphAPIError(MicrosoftGraphClientError):
|
|
"""Raised when a Graph API request fails."""
|
|
|
|
def __init__(
|
|
self, status_code: int, method: str, url: str, message: str, *,
|
|
retry_after_seconds: float | None = None, payload: Any = None) -> None:
|
|
self.status_code, self.method, self.url = status_code, method, url
|
|
self.retry_after_seconds, self.payload = retry_after_seconds, payload
|
|
super().__init__(f"Microsoft Graph API error {status_code} for {method} {url}: {message}")
|
|
|
|
|
|
class MicrosoftGraphClient:
|
|
"""Minimal async Graph client. Retry policy (JSON requests and streaming downloads
|
|
alike): transport errors back off exponentially; 401 clears the token cache and
|
|
refetches; 429/5xx honor ``Retry-After``. Each attempt uses a fresh ``AsyncClient``."""
|
|
|
|
def __init__(
|
|
self, token_provider: MicrosoftGraphTokenProvider, *,
|
|
base_url: str = DEFAULT_GRAPH_BASE_URL, timeout: float = 60.0, max_retries: int = 3,
|
|
transport: httpx.AsyncBaseTransport | None = None,
|
|
sleep: Callable[[float], Awaitable[None]] | None = None,
|
|
user_agent: str = "Hermes-Agent/graph-client") -> None:
|
|
self.token_provider, self.base_url, self.timeout = token_provider, base_url.rstrip("/"), timeout
|
|
self.max_retries, self.user_agent = max(0, int(max_retries)), user_agent
|
|
self._transport, self._sleep = transport, sleep or asyncio.sleep
|
|
|
|
@classmethod
|
|
def from_env(cls, **kwargs: Any) -> "MicrosoftGraphClient":
|
|
return cls(MicrosoftGraphTokenProvider(GraphCredentials.from_env()), **kwargs)
|
|
|
|
async def get_json(self, path: str, *, params: Params = None, headers: Headers = None) -> Any:
|
|
return self._decode_json(await self._request("GET", path, params=params, headers=headers))
|
|
|
|
async def post_json(self, path: str, *, json_body: Any | None = None, headers: Headers = None) -> Any:
|
|
return self._decode_json(await self._request("POST", path, json_body=json_body, headers=headers))
|
|
|
|
async def patch_json(self, path: str, *, json_body: Any | None = None, headers: Headers = None) -> Any:
|
|
response = await self._request("PATCH", path, json_body=json_body, headers=headers)
|
|
return self._decode_json_or(response, {})
|
|
|
|
async def delete(self, path: str, *, headers: Headers = None) -> dict[str, Any]:
|
|
response = await self._request("DELETE", path, headers=headers)
|
|
return self._decode_json_or(response, {"deleted": True, "status_code": response.status_code})
|
|
|
|
def _decode_json_or(self, response: httpx.Response, empty: Any) -> Any:
|
|
"""*empty* for a 204 / bodiless response, else the decoded JSON body."""
|
|
return empty if response.status_code == 204 or not response.content else self._decode_json(response)
|
|
|
|
async def collect_paginated(
|
|
self, path: str, *, params: Params = None, headers: Headers = None) -> list[Any]:
|
|
"""Follow ``@odata.nextLink`` and concatenate every page's ``value`` list."""
|
|
items: list[Any] = []
|
|
# Query params go on the first request only; @odata.nextLink already embeds them.
|
|
next_url: str | None = self._resolve_url(path)
|
|
next_params = dict(params or {})
|
|
while next_url:
|
|
response = await self._request("GET", next_url, params=next_params or None, headers=headers)
|
|
payload = self._decode_json(response)
|
|
if not isinstance(payload, dict):
|
|
raise MicrosoftGraphClientError(
|
|
f"Expected paginated Graph response dict, got {type(payload).__name__}.")
|
|
if isinstance(payload.get("value"), list):
|
|
items.extend(payload["value"])
|
|
next_url, next_params = payload.get("@odata.nextLink"), {}
|
|
return items
|
|
|
|
async def download_to_file(
|
|
self, path: str, destination: str | Path, *, headers: Headers = None, chunk_size: int = 65536
|
|
) -> dict[str, Any]:
|
|
"""Stream a Graph resource to disk chunk-by-chunk (large recordings never
|
|
fit in memory); written to ``.part`` and renamed into place only on success."""
|
|
url = self._resolve_url(path)
|
|
target = Path(destination)
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp_target = target.with_suffix(target.suffix + ".part")
|
|
|
|
async def perform(client: httpx.AsyncClient, request_headers: dict[str, str]):
|
|
try:
|
|
async with client.stream("GET", url, headers=request_headers) as response:
|
|
if response.status_code >= 400:
|
|
# Materialize the (small) error body so the message is meaningful.
|
|
await response.aread()
|
|
return response, None
|
|
with tmp_target.open("wb") as handle:
|
|
async for chunk in response.aiter_bytes(chunk_size=chunk_size):
|
|
if chunk:
|
|
handle.write(chunk)
|
|
return response, response.headers.get("content-type")
|
|
except httpx.HTTPError:
|
|
tmp_target.unlink(missing_ok=True)
|
|
raise
|
|
|
|
content_type = await self._with_retries("GET", url, "*/*", None, headers, perform, "download")
|
|
os.replace(tmp_target, target)
|
|
return {"path": str(target), "size_bytes": target.stat().st_size, "content_type": content_type}
|
|
|
|
async def _request(
|
|
self, method: str, path_or_url: str, *,
|
|
params: Params = None, json_body: Any | None = None, headers: Headers = None) -> httpx.Response:
|
|
url = self._resolve_url(path_or_url)
|
|
|
|
async def perform(client: httpx.AsyncClient, request_headers: dict[str, str]):
|
|
response = await client.request(method, url, params=params, json=json_body, headers=request_headers)
|
|
return response, response
|
|
|
|
return await self._with_retries(method, url, "application/json", json_body, headers, perform, "request")
|
|
|
|
async def _with_retries(
|
|
self, method: str, url: str, accept: str, json_body: Any | None, headers: Headers,
|
|
perform: Callable[[httpx.AsyncClient, dict[str, str]], Awaitable[tuple[httpx.Response, Any]]],
|
|
kind: str) -> Any:
|
|
"""Run ``perform`` (-> ``(response, result)``) under the retry policy. ``kind``
|
|
only labels transport-failure messages. Raises ``MicrosoftGraphAPIError`` once
|
|
retries are exhausted or the status is not retryable; only 401 forces a token refresh."""
|
|
attempt = 0
|
|
last_error: Exception | None = None
|
|
|
|
while attempt <= self.max_retries:
|
|
token = await self.token_provider.get_access_token(
|
|
force_refresh=attempt > 0
|
|
and isinstance(last_error, MicrosoftGraphAPIError)
|
|
and last_error.status_code == 401)
|
|
request_headers = {"Authorization": f"Bearer {token}", "Accept": accept, "User-Agent": self.user_agent}
|
|
if json_body is not None:
|
|
request_headers["Content-Type"] = "application/json"
|
|
if headers:
|
|
request_headers.update(headers)
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(self.timeout), transport=self._transport) as client:
|
|
response, result = await perform(client, request_headers)
|
|
except httpx.HTTPError as exc:
|
|
last_error, response = exc, None
|
|
if attempt >= self.max_retries:
|
|
raise MicrosoftGraphClientError(
|
|
f"Microsoft Graph {kind} failed for {method} {url}: {exc}") from exc
|
|
else:
|
|
if response.status_code < 400:
|
|
return result
|
|
last_error = self._build_api_error(method, url, response)
|
|
status = response.status_code
|
|
if attempt >= self.max_retries or not (status in (401, 429) or 500 <= status < 600):
|
|
raise last_error
|
|
if status == 401:
|
|
self.token_provider.clear_cache()
|
|
await self._sleep(self._retry_delay(response, attempt))
|
|
attempt += 1
|
|
|
|
raise MicrosoftGraphClientError(f"Microsoft Graph {kind} exhausted retries for {method} {url}.")
|
|
|
|
def _resolve_url(self, path_or_url: str) -> str:
|
|
if path_or_url.startswith(("http://", "https://")):
|
|
return path_or_url
|
|
path = path_or_url if path_or_url.startswith("/") else f"/{path_or_url}"
|
|
return f"{self.base_url}{path}"
|
|
|
|
@staticmethod
|
|
def _decode_json(response: httpx.Response) -> Any:
|
|
try:
|
|
return response.json()
|
|
except ValueError as exc:
|
|
raise MicrosoftGraphClientError(
|
|
"Microsoft Graph response was not valid JSON for "
|
|
f"{response.request.method} {response.request.url}") from exc
|
|
|
|
@staticmethod
|
|
def _retry_delay(response: httpx.Response | None, attempt: int) -> float:
|
|
retry_after = parse_retry_after_seconds(response.headers) if response is not None else None
|
|
return min(8.0, 0.5 * (2 ** attempt)) if retry_after is None else retry_after
|
|
|
|
@staticmethod
|
|
def _build_api_error(method: str, url: str, response: httpx.Response) -> MicrosoftGraphAPIError:
|
|
try:
|
|
payload: Any = response.json()
|
|
except ValueError:
|
|
payload = None
|
|
detail = format_graph_error(payload.get("error")) if isinstance(payload, dict) else None
|
|
return MicrosoftGraphAPIError(
|
|
response.status_code, method, url,
|
|
detail if detail is not None else (response.text.strip() or "unknown error"),
|
|
retry_after_seconds=parse_retry_after_seconds(response.headers), payload=payload)
|