"""Small, version-pinned adapter over WeKnora's official REST API.""" from __future__ import annotations import asyncio import hashlib import logging import os import time from collections import deque from collections.abc import Mapping, Sequence from typing import Any from urllib.parse import quote import httpx import regex logger = logging.getLogger(__name__) def _slow_diagnostic_threshold_ms() -> float: try: return max(0.0, min(float(os.getenv("DEERFLOW_SLOW_REQUEST_LOG_MS", "500")), 60_000.0)) except (TypeError, ValueError): return 500.0 _GREP_CONTRACT = "weknora-v0.7.2-grep-v1" _GREP_COMPATIBILITY_CONTRACT = "weknora-existing-api-compat-v1" _GREP_MAX_CANDIDATES = 500 _GREP_MAX_RESULTS = 30 _GREP_REGEX_TIMEOUT_SECONDS = 0.05 class WeKnoraError(RuntimeError): """Normalized upstream failure that is safe for DeerFlow to translate.""" def __init__(self, message: str, *, status_code: int = 502) -> None: super().__init__(message) self.status_code = status_code def _data(payload: Any) -> Any: if not isinstance(payload, dict): raise WeKnoraError("WeKnora returned an invalid response") if payload.get("success") is False: message = payload.get("message") or payload.get("error") or "WeKnora rejected the request" raise WeKnoraError(str(message), status_code=400) return payload.get("data", payload) def _as_list(value: Any) -> list[dict[str, Any]]: if isinstance(value, list): return [dict(item) for item in value if isinstance(item, Mapping)] if isinstance(value, dict): for key in ("items", "list", "knowledge_bases", "knowledge", "records"): nested = value.get(key) if isinstance(nested, list): return [dict(item) for item in nested if isinstance(item, Mapping)] return [] class WeKnoraClient: """Async client for a trusted-network WeKnora endpoint. The caller may inject an ``httpx.AsyncClient`` for connection pooling and tests. When omitted, each public method creates and closes a short-lived client. If admin credentials are configured, DeerFlow authenticates to WeKnora with a cached bearer token; otherwise it keeps the legacy trusted-network anonymous mode for deployments that allow DeerFlow through a reverse proxy. The knowledge-base directory is also cached briefly because several Gateway scopes request the same expensive upstream list at once. """ _KNOWLEDGE_BASES_CACHE_SECONDS = 15.0 def __init__( self, *, base_url: str, timeout_seconds: float = 30.0, http_client: httpx.AsyncClient | None = None, admin_email: str = "", admin_password: str = "", ) -> None: self.base_url = base_url.rstrip("/") self.timeout_seconds = timeout_seconds self._http_client = http_client self.admin_email = admin_email.strip() self.admin_password = admin_password self._auth_token = "" self._tenant_id = "" self._auth_expires_at = 0.0 self._auth_generation = 0 self._auth_lock = asyncio.Lock() self._knowledge_bases_cache: tuple[float, list[dict[str, Any]]] | None = None self._knowledge_bases_cache_generation = 0 self._knowledge_bases_task: tuple[int, asyncio.Task[list[dict[str, Any]]]] | None = None if not self.base_url: raise ValueError("WeKnora base URL is required") @property def _has_admin_credentials(self) -> bool: return bool(self.admin_email and self.admin_password) @staticmethod def _extract_login_token(payload: Any) -> tuple[str, str]: if not isinstance(payload, dict): return "", "" source = payload.get("data") if isinstance(payload.get("data"), dict) else payload token = str(source.get("token") or source.get("access_token") or payload.get("token") or payload.get("access_token") or "") tenant = source.get("tenant") if isinstance(source.get("tenant"), dict) else source.get("active_tenant") if isinstance(source.get("active_tenant"), dict) else payload.get("tenant") user = source.get("user") if isinstance(source.get("user"), dict) else payload.get("user") tenant_id = "" if isinstance(tenant, dict): tenant_id = str(tenant.get("id") or "") if not tenant_id and isinstance(user, dict): tenant_id = str(user.get("tenant_id") or "") return token, tenant_id async def _login(self, *, force_refresh: bool = False) -> None: if not self._has_admin_credentials: return now = time.time() if self._auth_token and not force_refresh and self._auth_expires_at > now + 60: return observed_generation = self._auth_generation async with self._auth_lock: now = time.time() token_is_fresh = bool(self._auth_token and self._auth_expires_at > now + 60) if token_is_fresh and (not force_refresh or self._auth_generation != observed_generation): return owns_client = self._http_client is None client = self._http_client or httpx.AsyncClient( timeout=httpx.Timeout(self.timeout_seconds), follow_redirects=False, trust_env=False, ) try: response = await client.post( f"{self.base_url}/api/v1/auth/login", json={"email": self.admin_email, "password": self.admin_password}, headers={"Accept": "application/json"}, ) except httpx.TimeoutException as exc: raise WeKnoraError("WeKnora admin login timed out", status_code=504) from exc except httpx.HTTPError as exc: raise WeKnoraError("WeKnora admin login is unavailable", status_code=502) from exc finally: if owns_client: await client.aclose() if response.status_code >= 400: raise WeKnoraError(f"WeKnora admin login failed ({response.status_code})", status_code=502) try: payload = response.json() except ValueError as exc: raise WeKnoraError("WeKnora admin login returned invalid JSON", status_code=502) from exc if isinstance(payload, dict) and payload.get("success") is False: raise WeKnoraError(str(payload.get("message") or "WeKnora admin login failed"), status_code=502) token, tenant_id = self._extract_login_token(payload) if not token: raise WeKnoraError("WeKnora admin login did not return a token", status_code=502) self._auth_token = token self._tenant_id = tenant_id self._auth_generation += 1 # WeKnora login responses do not consistently expose expiry. Refresh # periodically and retry once if upstream says the token is stale. self._auth_expires_at = now + 25 * 60 async def _headers(self, *, force_refresh: bool = False, accept: str = "application/json") -> dict[str, str]: headers = {"Accept": accept} if self._has_admin_credentials: await self._login(force_refresh=force_refresh) if self._auth_token: headers["Authorization"] = f"Bearer {self._auth_token}" if self._tenant_id: headers["X-Tenant-ID"] = self._tenant_id return headers async def _request( self, method: str, path: str, *, json: Any = None, params: Mapping[str, Any] | None = None, files: Any = None, data: Mapping[str, Any] | None = None, ) -> Any: diagnose_knowledge_base_list = method.upper() == "GET" and path == "/api/v1/knowledge-bases" diagnostic_started = time.perf_counter() if diagnose_knowledge_base_list else None auth_ms = 0.0 network_ms = 0.0 attempts = 0 owns_client = self._http_client is None client = self._http_client or httpx.AsyncClient( timeout=httpx.Timeout(self.timeout_seconds), follow_redirects=False, trust_env=False, ) try: for attempt in range(2): attempts += 1 auth_started = time.perf_counter() request_headers = await self._headers(force_refresh=attempt > 0) if diagnose_knowledge_base_list: auth_ms += (time.perf_counter() - auth_started) * 1000 request_started = time.perf_counter() response = await client.request( method, f"{self.base_url}{path}", headers=request_headers, json=json, params=params, files=files, data=data, ) if diagnose_knowledge_base_list: network_ms += (time.perf_counter() - request_started) * 1000 if response.status_code not in {401, 403} or not self._has_admin_credentials or attempt > 0: break except Exception as exc: if diagnose_knowledge_base_list: logger.warning( "PERF weknora_http_request operation=knowledge_base_directory method=GET total_ms=%.2f auth_ms=%.2f network_ms=%.2f attempts=%d outcome=error error_type=%s", (time.perf_counter() - diagnostic_started) * 1000, auth_ms, network_ms, attempts, type(exc).__name__, ) if isinstance(exc, httpx.TimeoutException): raise WeKnoraError("WeKnora request timed out", status_code=504) from exc if isinstance(exc, httpx.HTTPError): raise WeKnoraError("WeKnora is unavailable", status_code=502) from exc raise finally: if owns_client: await client.aclose() if diagnose_knowledge_base_list: total_ms = (time.perf_counter() - diagnostic_started) * 1000 if total_ms >= _slow_diagnostic_threshold_ms() or response.status_code >= 400: logger.warning( "PERF weknora_http_request operation=knowledge_base_directory method=GET total_ms=%.2f auth_ms=%.2f network_ms=%.2f attempts=%d status=%d outcome=%s", total_ms, auth_ms, network_ms, attempts, response.status_code, "http_error" if response.status_code >= 400 else "ok", ) if response.status_code == 404: raise WeKnoraError("WeKnora resource was not found", status_code=404) if response.status_code == 405: raise WeKnoraError("WeKnora endpoint does not support this method", status_code=405) if response.status_code in {401, 403}: if self._has_admin_credentials: raise WeKnoraError("WeKnora rejected DeerFlow's authenticated internal request", status_code=502) raise WeKnoraError( "WeKnora rejected DeerFlow's anonymous internal request; allow DeerFlow in WeKnora or its reverse proxy", status_code=502, ) if response.status_code >= 400: detail = "" try: body = response.json() if isinstance(body, dict): detail = str(body.get("message") or body.get("error") or body.get("detail") or "") except ValueError: detail = "" # Folder creation uses 409 as an optimistic-concurrency signal. # Preserve it so callers can re-read and adopt the concurrently # created folder even if a deployment localizes the error text. status_code = 409 if response.status_code == 409 else 502 raise WeKnoraError( detail or f"WeKnora returned HTTP {response.status_code}", status_code=status_code, ) if response.status_code == 204 or not response.content: return {"success": True} try: return response.json() except ValueError as exc: raise WeKnoraError("WeKnora returned invalid JSON") from exc async def health(self) -> dict[str, Any]: payload = await self._request("GET", "/health") return dict(payload) if isinstance(payload, dict) else {"status": "ok"} async def list_knowledge_bases(self, *, refresh: bool = False) -> list[dict[str, Any]]: now = time.monotonic() cached = self._knowledge_bases_cache if not refresh and cached is not None and cached[0] > now: return [dict(item) for item in cached[1]] # Share both success and failure with the current wave of callers. # A lock around an uncached failed read would make every waiter issue # another full timeout, serializing an unavailable service's fan-out. generation = self._knowledge_bases_cache_generation pending = self._knowledge_bases_task if pending is None or pending[0] != generation: task = asyncio.create_task(self._load_knowledge_bases(generation)) self._knowledge_bases_task = (generation, task) def clear_pending(completed: asyncio.Task[list[dict[str, Any]]]) -> None: if self._knowledge_bases_task is not None and self._knowledge_bases_task[1] is completed: self._knowledge_bases_task = None if not completed.cancelled(): completed.exception() task.add_done_callback(clear_pending) else: task = pending[1] rows = await asyncio.shield(task) return [dict(item) for item in rows] async def _load_knowledge_bases(self, generation: int) -> list[dict[str, Any]]: started = time.perf_counter() try: rows = _as_list(_data(await self._request("GET", "/api/v1/knowledge-bases"))) except Exception as exc: elapsed_ms = round((time.perf_counter() - started) * 1000, 2) logger.warning( "PERF weknora_knowledge_base_list upstream_ms=%.2f outcome=error error_type=%s", elapsed_ms, type(exc).__name__, ) raise elapsed_ms = round((time.perf_counter() - started) * 1000, 2) if elapsed_ms >= _slow_diagnostic_threshold_ms(): logger.warning( "PERF weknora_knowledge_base_list upstream_ms=%.2f item_count=%d outcome=ok", elapsed_ms, len(rows), ) # A mutation can finish while the read is in flight. Its old snapshot # must neither repopulate the cache nor be joined by post-write callers. if self._knowledge_bases_cache_generation == generation: self._knowledge_bases_cache = ( time.monotonic() + self._KNOWLEDGE_BASES_CACHE_SECONDS, [dict(item) for item in rows], ) return rows def _invalidate_knowledge_bases_cache(self) -> None: """Discard the directory after a successful local mutation.""" self._knowledge_bases_cache = None self._knowledge_bases_cache_generation += 1 async def get_knowledge_base(self, knowledge_base_id: str) -> dict[str, Any]: value = _data(await self._request("GET", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}")) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid knowledge base") return dict(value) async def list_models(self) -> list[dict[str, Any]]: """List models from WeKnora; model credentials are never returned.""" return _as_list(_data(await self._request("GET", "/api/v1/models"))) async def debug_model(self, model_id: str) -> dict[str, Any]: """Ask WeKnora to verify one saved model with a small real request.""" normalized_id = str(model_id or "").strip() if not normalized_id: raise WeKnoraError("WeKnora model ID is required", status_code=400) value = _data( await self._request( "POST", f"/api/v1/models/{quote(normalized_id, safe='')}/debug", files={"input": (None, "Reply with OK.")}, ) ) if not isinstance(value, dict) or not isinstance(value.get("ok"), bool): raise WeKnoraError("WeKnora returned an invalid model debug result") return dict(value) async def get_default_embedding_model_id(self) -> str: models = [item for item in await self.list_models() if str(item.get("type") or "").strip().lower() == "embedding"] preferred = next((item for item in models if item.get("is_default") is True), None) selected = preferred or next( (item for item in models if str(item.get("status") or "active").lower() == "active"), None, ) model_id = str((selected or {}).get("id") or "") if not model_id: raise WeKnoraError( "WeKnora has no available embedding model; configure its default model in WeKnora", status_code=400, ) return model_id async def get_default_knowledge_qa_model_id(self) -> str: """Resolve a live WeKnora synthesis model for Wiki indexing. WeKnora's ``active`` status means that the database configuration is enabled; it does not prove that the upstream provider still serves the model. Probe active candidates before binding a new Wiki so a disabled provider model cannot silently send every ingest task to dead letters. """ models = [ item for item in await self.list_models() if str(item.get("type") or "").strip().lower() == "knowledgeqa" and str(item.get("id") or "").strip() and str(item.get("status") or "active").strip().lower() == "active" ] if not models: raise WeKnoraError( "WeKnora has no available KnowledgeQA model; configure one in WeKnora before creating a Wiki", status_code=400, ) # Prefer an explicit database default. When no default exists, sort by # immutable persisted ID rather than trusting API row order or an # editable display name. models.sort( key=lambda item: ( 0 if item.get("is_default") is True else 1, str(item.get("id") or "").strip(), ) ) attempted: list[str] = [] for model in models: model_id = str(model.get("id") or "").strip() label = str(model.get("display_name") or model.get("name") or model_id).strip() attempted.append(label or model_id) # ``ok=false`` is a candidate-specific upstream-model failure. An # HTTP/auth/timeout failure from the debug endpoint is a WeKnora # service failure, so preserve it instead of multiplying probes # and misreporting the incident as bad model configuration. debug = await self.debug_model(model_id) if debug.get("ok") is True: return model_id attempted_models = ", ".join(attempted[:12]) if len(attempted) > 12: attempted_models = f"{attempted_models}, and {len(attempted) - 12} more" raise WeKnoraError( "WeKnora has no usable active KnowledgeQA model after connectivity checks" + (f": {attempted_models}" if attempted_models else "") + "; configure or repair a working default model before creating a Wiki", status_code=400, ) async def create_knowledge_base( self, *, name: str, description: str = "", kb_type: str = "document", embedding_model_id: str | None = None, knowledge_qa_model_id: str | None = None, wiki_enabled: bool = False, wiki_content_instructions: str = "", wiki_extraction_instructions: str = "", ) -> dict[str, Any]: model_id = embedding_model_id or await self.get_default_embedding_model_id() create_payload: dict[str, Any] = { "name": name, "description": description, "type": kb_type, "embedding_model_id": model_id, } if wiki_enabled: synthesis_model_id = knowledge_qa_model_id or await self.get_default_knowledge_qa_model_id() create_payload.update( { "summary_model_id": synthesis_model_id, "indexing_strategy": { "vector_enabled": True, "keyword_enabled": True, "wiki_enabled": True, "graph_enabled": False, }, "wiki_config": { "synthesis_model_id": synthesis_model_id, "extraction_granularity": "standard", "content_instructions": wiki_content_instructions.strip(), "extraction_instructions": wiki_extraction_instructions.strip(), }, } ) payload = await self._request( "POST", "/api/v1/knowledge-bases", json=create_payload, ) value = _data(payload) if not isinstance(value, dict) or not value.get("id"): raise WeKnoraError("WeKnora did not return the created knowledge base") self._invalidate_knowledge_bases_cache() return dict(value) async def update_knowledge_base(self, knowledge_base_id: str, changes: Mapping[str, Any]) -> dict[str, Any]: payload = dict(changes) wiki_config_changes = payload.pop("wiki_config", None) current: dict[str, Any] | None = None if not str(payload.get("name") or "").strip() or wiki_config_changes is not None: current = await self.get_knowledge_base(knowledge_base_id) if not str(payload.get("name") or "").strip(): assert current is not None current_name = str(current.get("name") or "").strip() if not current_name: raise WeKnoraError("WeKnora knowledge base name is required for updates") payload["name"] = current_name if wiki_config_changes is not None: if not isinstance(wiki_config_changes, Mapping): raise WeKnoraError("WeKnora Wiki instructions must be an object", status_code=422) assert current is not None current_wiki_config = current.get("wiki_config") merged_wiki_config = dict(current_wiki_config) if isinstance(current_wiki_config, Mapping) else {} for key in ("content_instructions", "extraction_instructions"): if key in wiki_config_changes: merged_wiki_config[key] = str(wiki_config_changes.get(key) or "").strip() payload["config"] = {"wiki_config": merged_wiki_config} value = _data( await self._request( "PUT", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}", json=payload, ) ) self._invalidate_knowledge_bases_cache() return dict(value) if isinstance(value, dict) else {} async def delete_knowledge_base(self, knowledge_base_id: str) -> None: await self._request("DELETE", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}") self._invalidate_knowledge_bases_cache() async def list_documents(self, knowledge_base_id: str, *, page: int = 1, page_size: int = 100) -> dict[str, Any]: payload = await self._request( "GET", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge", params={"page": page, "page_size": page_size}, ) rows = _as_list(_data(payload)) total = payload.get("total", len(rows)) if isinstance(payload, dict) else len(rows) return {"items": rows, "total": int(total or 0), "page": page, "page_size": page_size} async def upload_document( self, knowledge_base_id: str, *, filename: str, content: bytes, content_type: str = "application/octet-stream", ) -> dict[str, Any]: value = _data( await self._request( "POST", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge/file", files={"file": (filename, content, content_type)}, ) ) self._invalidate_knowledge_bases_cache() return dict(value) if isinstance(value, dict) else {} async def import_document_url(self, knowledge_base_id: str, *, url: str) -> dict[str, Any]: value = _data( await self._request( "POST", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge/url", json={"url": url}, ) ) self._invalidate_knowledge_bases_cache() return dict(value) if isinstance(value, dict) else {} async def create_manual_document( self, knowledge_base_id: str, *, title: str, content: str, ) -> dict[str, Any]: value = _data( await self._request( "POST", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge/manual", json={"title": title, "content": content, "status": "publish"}, ) ) self._invalidate_knowledge_bases_cache() return dict(value) if isinstance(value, dict) else {} async def delete_document(self, knowledge_id: str) -> None: await self._request("DELETE", f"/api/v1/knowledge/{quote(knowledge_id, safe='')}") self._invalidate_knowledge_bases_cache() async def get_document(self, knowledge_id: str) -> dict[str, Any]: value = _data(await self._request("GET", f"/api/v1/knowledge/{quote(knowledge_id, safe='')}")) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid document") return dict(value) async def get_document_preview(self, knowledge_id: str) -> tuple[bytes, str]: """Return the original document preview bytes and its media type.""" owns_client = self._http_client is None client = self._http_client or httpx.AsyncClient( timeout=httpx.Timeout(self.timeout_seconds), follow_redirects=True, trust_env=False, ) try: for attempt in range(2): response = await client.request( "GET", f"{self.base_url}/api/v1/knowledge/{quote(knowledge_id, safe='')}/preview", headers=await self._headers(force_refresh=attempt > 0, accept="*/*"), follow_redirects=True, ) if response.status_code not in {401, 403} or not self._has_admin_credentials or attempt > 0: break except httpx.TimeoutException as exc: raise WeKnoraError("WeKnora preview request timed out", status_code=504) from exc except httpx.HTTPError as exc: raise WeKnoraError("WeKnora preview is unavailable", status_code=502) from exc finally: if owns_client: await client.aclose() if response.status_code == 404: raise WeKnoraError("WeKnora document preview was not found", status_code=404) if response.status_code in {401, 403}: if self._has_admin_credentials: raise WeKnoraError("WeKnora rejected DeerFlow's authenticated preview request", status_code=502) raise WeKnoraError( "WeKnora rejected DeerFlow's anonymous internal request; allow DeerFlow in WeKnora or its reverse proxy", status_code=502, ) if response.status_code >= 400: raise WeKnoraError(f"WeKnora returned HTTP {response.status_code}", status_code=502) media_type = response.headers.get("content-type", "application/octet-stream").split(";", 1)[0].strip() return response.content, media_type or "application/octet-stream" async def get_knowledge_base_file( self, knowledge_base_id: str, file_path: str, ) -> tuple[bytes, str]: """Return one protected Wiki attachment using DeerFlow's WeKnora identity.""" owns_client = self._http_client is None client = self._http_client or httpx.AsyncClient( timeout=httpx.Timeout(self.timeout_seconds), follow_redirects=True, trust_env=False, ) try: for attempt in range(2): response = await client.request( "GET", f"{self.base_url}/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/files", params={"file_path": file_path}, headers=await self._headers(force_refresh=attempt > 0, accept="*/*"), follow_redirects=True, ) if response.status_code not in {401, 403} or not self._has_admin_credentials or attempt > 0: break except httpx.TimeoutException as exc: raise WeKnoraError("WeKnora Wiki image request timed out", status_code=504) from exc except httpx.HTTPError as exc: raise WeKnoraError("WeKnora Wiki image is unavailable", status_code=502) from exc finally: if owns_client: await client.aclose() if response.status_code == 404: raise WeKnoraError("WeKnora Wiki image was not found", status_code=404) if response.status_code in {401, 403}: if self._has_admin_credentials: raise WeKnoraError("WeKnora rejected DeerFlow's authenticated Wiki image request", status_code=502) raise WeKnoraError( "WeKnora rejected DeerFlow's anonymous internal request; allow DeerFlow in WeKnora or its reverse proxy", status_code=502, ) if response.status_code >= 400: raise WeKnoraError(f"WeKnora returned HTTP {response.status_code}", status_code=502) media_type = response.headers.get("content-type", "application/octet-stream").split(";", 1)[0].strip() return response.content, media_type or "application/octet-stream" async def reprocess_document(self, knowledge_id: str) -> dict[str, Any]: value = _data(await self._request("POST", f"/api/v1/knowledge/{quote(knowledge_id, safe='')}/reparse")) self._invalidate_knowledge_bases_cache() return dict(value) if isinstance(value, dict) else {} async def list_chunks(self, knowledge_id: str, *, page: int = 1, page_size: int = 20) -> dict[str, Any]: payload = await self._request( "GET", f"/api/v1/chunks/{quote(knowledge_id, safe='')}", params={"page": page, "page_size": page_size}, ) rows = _as_list(_data(payload)) total = payload.get("total", len(rows)) if isinstance(payload, dict) else len(rows) return {"items": rows, "total": int(total or 0), "page": page, "page_size": page_size} async def get_wiki_graph(self, knowledge_base_id: str, *, limit: int = 200) -> dict[str, Any]: value = _data( await self._request( "GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/graph", params={"mode": "overview", "limit": max(1, min(limit, 10000))}, ) ) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid wiki graph") return dict(value) async def get_wiki_stats(self, knowledge_base_id: str) -> dict[str, Any]: value = _data(await self._request("GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/stats")) if not isinstance(value, dict): raise WeKnoraError("Invalid Wiki generation status") return value async def create_wiki_folder(self, knowledge_base_id: str, *, name: str, parent_id: str = "") -> dict[str, Any]: value = _data(await self._request("POST", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/folders", json={"name": name, "parent_id": parent_id})) if not isinstance(value, dict) or not value.get("id"): raise WeKnoraError("Invalid Wiki folder response") return value async def list_wiki_folders( self, knowledge_base_id: str, *, parent_id: str = "", page_types: Sequence[str] | None = None, ) -> dict[str, Any]: params: dict[str, Any] = {"parent_id": parent_id or None} if page_types: params["page_types"] = ",".join(str(value) for value in page_types if str(value)) or None payload = await self._request( "GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/folders", params=params, ) value = _data(payload) source = value if isinstance(value, dict) else {} raw_folders = source.get("folders") if isinstance(value, dict) else value rows = _as_list(raw_folders) return { "parent_id": str(source.get("parent_id") or parent_id), "folders": rows, } async def update_wiki_folder( self, knowledge_base_id: str, folder_id: str, changes: Mapping[str, Any] | None = None, *, name: str | None = None, parent_id: str | None = None, move_parent: bool | None = None, ) -> dict[str, Any]: update = dict(changes or {}) if name is not None: update["name"] = name if parent_id is not None: update["parent_id"] = parent_id update.setdefault("move_parent", True) if move_parent is not None: update["move_parent"] = move_parent elif "parent_id" in update: update.setdefault("move_parent", True) value = _data( await self._request( "PUT", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/folders/{quote(folder_id, safe='')}", json=update, ) ) if not isinstance(value, dict): raise WeKnoraError("Invalid Wiki folder response") return dict(value) async def delete_wiki_folder(self, knowledge_base_id: str, folder_id: str) -> None: await self._request( "DELETE", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/folders/{quote(folder_id, safe='')}", ) async def move_wiki_page(self, knowledge_base_id: str, slug: str, folder_id: str) -> dict[str, Any]: value = _data( await self._request( "PUT", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/move-page", json={"slug": slug, "folder_id": folder_id}, ) ) if not isinstance(value, dict): raise WeKnoraError("Invalid Wiki page response") return dict(value) async def rebuild_wiki_links(self, knowledge_base_id: str) -> None: await self._request("POST", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/rebuild-links") async def list_wiki_pages( self, knowledge_base_id: str, *, page: int = 1, page_size: int = 50, query: str = "", folder_id: str | None = None, ) -> dict[str, Any]: params: dict[str, Any] = {"page": page, "page_size": page_size, "query": query or None} # ``None`` means no placement filter; an explicit empty string means # pages directly under the Wiki root. if folder_id is not None: params["folder_id"] = folder_id payload = await self._request( "GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages", params=params, ) value = _data(payload) source = value if isinstance(value, dict) else payload if isinstance(payload, dict) else {} rows = _as_list(source.get("pages") if isinstance(source, dict) else value) return { "pages": rows, "total": int(source.get("total") or len(rows)), "page": int(source.get("page") or page), "page_size": int(source.get("page_size") or page_size), "total_pages": int(source.get("total_pages") or max(1, (int(source.get("total") or len(rows)) + page_size - 1) // page_size)), } async def search_wiki_pages( self, knowledge_base_id: str, query: str, *, limit: int = 30, ) -> list[dict[str, Any]]: """Use WeKnora's native, ranked Wiki regex search. WeKnora v0.7.2 exposes this as a single-library ``GET`` endpoint with the regex in ``q`` (not the list endpoint's ``query`` parameter). It returns complete WikiPage rows, so callers do not need to enumerate a directory and issue one detail request per page. """ payload = await self._request( "GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/search", params={ "q": str(query or ""), "limit": max(1, min(int(limit), 50)), }, ) value = _data(payload) if isinstance(value, dict): return _as_list(value.get("pages")) return _as_list(value) async def get_wiki_page(self, knowledge_base_id: str, slug: str) -> dict[str, Any]: safe_slug = "/".join(quote(part, safe="") for part in slug.split("/")) value = _data( await self._request( "GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages/{safe_slug}", ) ) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid Wiki page") return dict(value) async def get_wiki_index( self, knowledge_base_id: str, *, types: Sequence[str], limit: int = 50, cursor: str = "", ) -> dict[str, Any]: params: dict[str, Any] = { "types": ",".join(value for value in types if value), "limit": max(1, min(limit, 500)), } if cursor: params["cursor"] = cursor value = _data( await self._request( "GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/index", params=params, ) ) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid Wiki index") return dict(value) async def create_wiki_page(self, knowledge_base_id: str, payload: Mapping[str, Any]) -> dict[str, Any]: value = _data( await self._request( "POST", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages", json=dict(payload), ) ) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid Wiki page") return dict(value) async def update_wiki_page(self, knowledge_base_id: str, slug: str, changes: Mapping[str, Any]) -> dict[str, Any]: safe_slug = "/".join(quote(part, safe="") for part in slug.split("/")) value = _data( await self._request( "PUT", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages/{safe_slug}", json=dict(changes), ) ) return dict(value) if isinstance(value, dict) else {} async def delete_wiki_page(self, knowledge_base_id: str, slug: str) -> None: safe_slug = "/".join(quote(part, safe="") for part in slug.split("/")) await self._request( "DELETE", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages/{safe_slug}", ) @staticmethod def _normalize_grep_chunk(row: Mapping[str, Any]) -> dict[str, Any]: """Normalize both the provider extension and compatibility rows.""" return { "chunk_id": str(row.get("id") or row.get("chunk_id") or ""), "content": str(row.get("content") or row.get("text") or row.get("match_snippet") or ""), "knowledge_id": str(row.get("knowledge_id") or ""), "knowledge_base_id": str(row.get("knowledge_base_id") or ""), "title": str( row.get("knowledge_title") or row.get("title") or row.get("knowledge_filename") or "" ), "chunk_index": int(row.get("chunk_index") or row.get("index") or 0), "match_count": int( row.get("matched_patterns") or row.get("match_count") or row.get("total_pattern_hits") or 0 ), "title_match": bool(row.get("title_match")), "score": float(row.get("match_score") or row.get("score") or 0.0), "metadata": row.get("metadata") if isinstance(row.get("metadata"), dict) else {}, } @staticmethod def _grep_matches(expression: regex.Pattern, value: str) -> list[regex.Match]: try: return list( expression.finditer( value, timeout=_GREP_REGEX_TIMEOUT_SECONDS, ) ) except TimeoutError as exc: raise WeKnoraError( "WeKnora compatibility grep regular expression timed out", status_code=422, ) from exc async def _grep_chunks_compatibility( self, query: str, knowledge_base_ids: Sequence[str], ) -> dict[str, Any]: """Bounded, fair fallback for older providers without ``chunks/grep``. The native endpoint remains authoritative. This path round-robins both knowledge bases and documents so one large first scope cannot consume the global 500-chunk compatibility budget. """ ids = list(dict.fromkeys(str(value).strip() for value in knowledge_base_ids if str(value).strip())) try: expression = regex.compile(query, regex.IGNORECASE | regex.VERSION1) except regex.error as exc: raise WeKnoraError(f"Invalid grep regular expression: {exc}", status_code=422) from exc per_scope_document_cap = max(1, (_GREP_MAX_CANDIDATES + len(ids) - 1) // max(1, len(ids))) async def load_documents(knowledge_base_id: str) -> tuple[str, list[dict[str, Any]] | Exception]: rows: list[dict[str, Any]] = [] page = 1 try: while len(rows) < per_scope_document_cap: page_size = min(100, per_scope_document_cap - len(rows)) payload = await self.list_documents( knowledge_base_id, page=page, page_size=page_size, ) batch = [ dict(item) for item in payload.get("items") or [] if isinstance(item, Mapping) ] rows.extend(batch) if not batch or page * page_size >= int(payload.get("total") or len(rows)): break page += 1 return knowledge_base_id, rows except Exception as exc: # noqa: BLE001 - returned as structured scope failure return knowledge_base_id, exc loaded = await asyncio.gather(*(load_documents(knowledge_base_id) for knowledge_base_id in ids)) documents_by_scope: dict[str, list[dict[str, Any]]] = {} warnings: list[dict[str, Any]] = [ { "code": "WEKNORA_GREP_COMPATIBILITY_FALLBACK", "message": "Provider chunks/grep is unavailable; used bounded existing-API compatibility scan.", } ] failed_scope_ids: set[str] = set() for knowledge_base_id, outcome in loaded: if isinstance(outcome, Exception): failed_scope_ids.add(knowledge_base_id) warnings.append( { "code": "PROVIDER_SCOPE_FAILED", "knowledge_base_id": knowledge_base_id, "stage": "list_documents", } ) else: documents_by_scope[knowledge_base_id] = outcome if not documents_by_scope: raise WeKnoraError("WeKnora compatibility grep failed for every knowledge base") # Seed the work queue in KB/document round-robin order. Subsequent # chunk pages are re-appended to the tail, preserving document fairness. work: deque[tuple[str, dict[str, Any], int]] = deque() max_documents = max((len(rows) for rows in documents_by_scope.values()), default=0) for index in range(max_documents): for knowledge_base_id in ids: rows = documents_by_scope.get(knowledge_base_id) or [] if index < len(rows): work.append((knowledge_base_id, rows[index], 1)) candidates: list[dict[str, Any]] = [] seen_content: set[str] = set() scanned_chunks = 0 chunk_successes = {knowledge_base_id: 0 for knowledge_base_id in documents_by_scope} chunk_failures = {knowledge_base_id: 0 for knowledge_base_id in documents_by_scope} while work and scanned_chunks < _GREP_MAX_CANDIDATES: knowledge_base_id, document, page = work.popleft() knowledge_id = str(document.get("id") or document.get("knowledge_id") or "") if not knowledge_id: continue page_size = min(20, _GREP_MAX_CANDIDATES - scanned_chunks) try: payload = await self.list_chunks( knowledge_id, page=page, page_size=page_size, ) except Exception: # noqa: BLE001 - one document must not fail the whole fallback chunk_failures[knowledge_base_id] += 1 continue chunk_successes[knowledge_base_id] += 1 chunks = [ dict(item) for item in payload.get("items") or [] if isinstance(item, Mapping) ] title = str( document.get("title") or document.get("name") or document.get("filename") or "" ) title_matches = self._grep_matches(expression, title[:5000]) for chunk in chunks: if scanned_chunks >= _GREP_MAX_CANDIDATES: break scanned_chunks += 1 content = str(chunk.get("content") or chunk.get("text") or "") content_matches = self._grep_matches(expression, content[:50_000]) if not title_matches and not content_matches: continue signature = hashlib.sha256(" ".join(content.split()).encode()).hexdigest() if signature in seen_content: continue seen_content.add(signature) match_count = min(99, len(title_matches) * 4 + len(content_matches)) candidates.append( self._normalize_grep_chunk( { **chunk, "knowledge_id": knowledge_id, "knowledge_base_id": knowledge_base_id, "knowledge_title": title, "match_count": match_count, "title_match": bool(title_matches), "match_score": float(match_count), } ) ) if chunks and page * page_size < int(payload.get("total") or len(chunks)): work.append((knowledge_base_id, document, page + 1)) successful_scope_ids: list[str] = [] for knowledge_base_id, documents in documents_by_scope.items(): # An empty KB was successfully scanned. A non-empty KB for which # every chunk read failed is a failed provider scope. if not documents or chunk_successes[knowledge_base_id] > 0: successful_scope_ids.append(knowledge_base_id) else: failed_scope_ids.add(knowledge_base_id) if chunk_failures[knowledge_base_id]: warnings.append( { "code": "PROVIDER_DOCUMENT_READ_FAILED", "knowledge_base_id": knowledge_base_id, "failed_documents": chunk_failures[knowledge_base_id], } ) if not successful_scope_ids: raise WeKnoraError("WeKnora compatibility grep failed for every knowledge base") candidates.sort( key=lambda item: ( bool(item.get("title_match")), int(item.get("match_count") or 0), float(item.get("score") or 0), -int(item.get("chunk_index") or 0), ), reverse=True, ) return { "items": candidates[:_GREP_MAX_RESULTS], "contract": _GREP_COMPATIBILITY_CONTRACT, "compatibility_fallback": True, "partial": bool(failed_scope_ids or any(chunk_failures.values())), "warnings": warnings, "successful_knowledge_base_ids": successful_scope_ids, "failed_knowledge_base_ids": [value for value in ids if value in failed_scope_ids], "scanned_chunks": scanned_chunks, "scan_truncated": scanned_chunks >= _GREP_MAX_CANDIDATES, } async def grep_chunks( self, query: str, knowledge_base_ids: Sequence[str], ) -> dict[str, Any]: """Run the v0.7.2 cross-KB grep extension, with a marked fallback.""" ids = list(dict.fromkeys(str(value).strip() for value in knowledge_base_ids if str(value).strip())) if not ids: return { "items": [], "contract": _GREP_CONTRACT, "compatibility_fallback": False, "partial": False, "warnings": [], "successful_knowledge_base_ids": [], "failed_knowledge_base_ids": [], "scanned_chunks": 0, "scan_truncated": False, } try: payload = await self._request( "POST", "/api/v1/chunks/grep", json={"query": query, "knowledge_base_ids": ids}, ) except WeKnoraError as exc: if exc.status_code not in {404, 405}: raise return await self._grep_chunks_compatibility(query, ids) value = _data(payload) source = value if isinstance(value, dict) else {} contract = str( source.get("contract") or (payload.get("contract") if isinstance(payload, dict) else "") or "" ) if contract != _GREP_CONTRACT: raise WeKnoraError( f"WeKnora chunks/grep contract mismatch (expected {_GREP_CONTRACT})" ) raw_rows = ( source.get("chunk_results") or source.get("results") or source.get("items") or [] ) if not isinstance(raw_rows, list): raise WeKnoraError("WeKnora chunks/grep returned invalid results") rows = [ self._normalize_grep_chunk(row) for row in raw_rows if isinstance(row, Mapping) ] warnings = [ dict(item) for item in source.get("warnings") or [] if isinstance(item, Mapping) ] failed_ids = [ str(value) for value in source.get("failed_knowledge_base_ids") or [] if str(value) in ids ] explicit_successes = source.get("successful_knowledge_base_ids") successful_ids = ( [str(value) for value in explicit_successes if str(value) in ids] if isinstance(explicit_successes, list) else [value for value in ids if value not in failed_ids] ) warned_scope_ids = { str(item.get("knowledge_base_id") or "") for item in warnings } warnings.extend( { "code": "PROVIDER_SCOPE_FAILED", "knowledge_base_id": knowledge_base_id, } for knowledge_base_id in failed_ids if knowledge_base_id not in warned_scope_ids ) if source.get("partial") and not warnings: warnings.append({"code": "PROVIDER_PARTIAL_RESPONSE"}) return { "items": rows[:_GREP_MAX_RESULTS], "contract": contract, "compatibility_fallback": False, "partial": bool(source.get("partial")) or bool(warnings) or bool(failed_ids), "warnings": warnings, "successful_knowledge_base_ids": successful_ids, "failed_knowledge_base_ids": failed_ids, "scanned_chunks": int(source.get("scanned_chunks") or len(rows)), "scan_truncated": bool(source.get("scan_truncated")), } async def search(self, query: str, knowledge_base_ids: Sequence[str]) -> list[dict[str, Any]]: ids = [str(value) for value in knowledge_base_ids if str(value)] if not ids: return [] rows = _as_list( _data( await self._request( "POST", "/api/v1/knowledge-search", json={"query": query, "knowledge_base_ids": ids}, ) ) ) normalized: list[dict[str, Any]] = [] for row in rows: normalized.append( { "chunk_id": str(row.get("id") or row.get("chunk_id") or ""), "content": str(row.get("content") or ""), "knowledge_id": str(row.get("knowledge_id") or ""), "knowledge_base_id": str(row.get("knowledge_base_id") or ""), "title": str(row.get("knowledge_title") or row.get("title") or row.get("knowledge_filename") or ""), "filename": str(row.get("knowledge_filename") or ""), "score": float(row.get("score") or 0), "chunk_index": int(row.get("chunk_index") or 0), "metadata": row.get("metadata") if isinstance(row.get("metadata"), dict) else {}, } ) return normalized async def get_chunk(self, chunk_id: str) -> dict[str, Any]: value = _data(await self._request("GET", f"/api/v1/chunks/by-id/{quote(chunk_id, safe='')}")) if not isinstance(value, dict): raise WeKnoraError("WeKnora returned an invalid chunk") return dict(value) async def get_chunk_context(self, chunk_id: str, *, expected_knowledge_base_id: str, radius: int = 2) -> list[dict[str, Any]]: center = await self.get_chunk(chunk_id) if str(center.get("knowledge_base_id") or "") != expected_knowledge_base_id: raise WeKnoraError("Chunk does not belong to the authorized knowledge base", status_code=404) before: list[dict[str, Any]] = [] cursor = center for _ in range(max(0, min(radius, 4))): previous_id = str(cursor.get("pre_chunk_id") or "") if not previous_id: break previous = await self.get_chunk(previous_id) if str(previous.get("knowledge_base_id") or "") != expected_knowledge_base_id: break before.append(previous) cursor = previous after: list[dict[str, Any]] = [] cursor = center for _ in range(max(0, min(radius, 4))): next_id = str(cursor.get("next_chunk_id") or "") if not next_id: break following = await self.get_chunk(next_id) if str(following.get("knowledge_base_id") or "") != expected_knowledge_base_id: break after.append(following) cursor = following return [*reversed(before), center, *after]