1313 lines
56 KiB
Python
1313 lines
56 KiB
Python
"""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]
|