mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-11-01 05:15:22 +08:00
feat(plugins): consolidate official plugins with independent releases (LANG-180) (#2797)
Co-authored-by: multica-agent <[email protected]>
This commit is contained in:
RockChinQ
co-authored by
multica-agent
multica-agent
parent
d099033d79
commit
3a77306c3c
@@ -0,0 +1,136 @@
|
||||
apiVersion: v1
|
||||
kind: KnowledgeEngine
|
||||
metadata:
|
||||
name: dify-connector
|
||||
label:
|
||||
en_US: Dify Knowledge Base
|
||||
zh_Hans: Dify 知识库
|
||||
ja_JP: Dify 知識ベース
|
||||
description:
|
||||
en_US: Retrieve knowledge from or store files into Dify knowledge bases
|
||||
zh_Hans: 从 Dify 知识库中检索知识或向其中存入文件
|
||||
ja_JP: Dify 知識ベースから知識を検索または保存
|
||||
spec:
|
||||
creation_schema:
|
||||
- name: api_base_url
|
||||
label:
|
||||
en_US: API Base URL
|
||||
zh_Hans: API 基础地址
|
||||
description:
|
||||
en_US: 'Base URL for Dify API (e.g., https://api.dify.ai/v1 for cloud, or your self-hosted URL)'
|
||||
zh_Hans: 'Dify API 基础地址（云端版本使用 https://api.dify.ai/v1，自托管版本使用您的服务器地址）'
|
||||
ja_JP: 'Dify API ベース URL（クラウド版は https://api.dify.ai/v1、セルフホストの場合はサーバーの URL）'
|
||||
type: string
|
||||
default: https://api.dify.ai/v1
|
||||
required: true
|
||||
- name: dify_apikey
|
||||
label:
|
||||
en_US: API Key
|
||||
zh_Hans: API 密钥
|
||||
description:
|
||||
en_US: 'API key from your Dify instance'
|
||||
zh_Hans: '从您的 Dify 实例获取的 API Key'
|
||||
ja_JP: 'Dify インスタンスから取得した API Key'
|
||||
type: password
|
||||
required: true
|
||||
- name: dataset_id
|
||||
label:
|
||||
en_US: Dataset ID
|
||||
zh_Hans: 数据集 ID
|
||||
ja_JP: データセット ID
|
||||
description:
|
||||
en_US: 'The UUID of the Dify Dataset to retrieve from'
|
||||
zh_Hans: '要检索的 Dify 数据集 UUID'
|
||||
ja_JP: 'Dify ナレッジベース/データセットの ID'
|
||||
type: string
|
||||
required: true
|
||||
retrieval_schema:
|
||||
- name: search_method
|
||||
label:
|
||||
en_US: Search Method
|
||||
zh_Hans: 检索模式
|
||||
ja_JP: 検索方法
|
||||
description:
|
||||
en_US: 'The search method to use'
|
||||
zh_Hans: '使用的搜索方法'
|
||||
ja_JP: '使用する検索方法'
|
||||
type: select
|
||||
default: semantic_search
|
||||
options:
|
||||
- name: keyword_search
|
||||
label:
|
||||
en_US: Keyword Search
|
||||
zh_Hans: 关键词检索
|
||||
ja_JP: キーワード検索
|
||||
- name: semantic_search
|
||||
label:
|
||||
en_US: Semantic Search
|
||||
zh_Hans: 语义检索
|
||||
ja_JP: セマンティック検索
|
||||
- name: full_text_search
|
||||
label:
|
||||
en_US: Full Text Search
|
||||
zh_Hans: 全文检索
|
||||
ja_JP: 全文検索
|
||||
- name: hybrid_search
|
||||
label:
|
||||
en_US: Hybrid Search
|
||||
zh_Hans: 混合检索
|
||||
ja_JP: ハイブリッド検索
|
||||
required: false
|
||||
- name: top_k
|
||||
label:
|
||||
en_US: Top K
|
||||
zh_Hans: 返回数量
|
||||
ja_JP: 取得件数
|
||||
description:
|
||||
en_US: 'Maximum number of results to return'
|
||||
zh_Hans: '返回的最大结果数量'
|
||||
ja_JP: '返す結果の最大数'
|
||||
type: integer
|
||||
default: 5
|
||||
required: false
|
||||
- name: score_threshold_enabled
|
||||
label:
|
||||
en_US: Enable Score Threshold
|
||||
zh_Hans: 启用分数阈值
|
||||
ja_JP: スコア閾値を有効化
|
||||
description:
|
||||
en_US: 'Filter results below the minimum score'
|
||||
zh_Hans: '过滤低于最低分数的结果'
|
||||
ja_JP: '最小スコア以下の結果をフィルタリング'
|
||||
type: boolean
|
||||
default: false
|
||||
required: false
|
||||
- name: score_threshold
|
||||
label:
|
||||
en_US: Score Threshold
|
||||
zh_Hans: 分数阈值
|
||||
ja_JP: スコア閾値
|
||||
description:
|
||||
en_US: 'Minimum relevance score (0-1)'
|
||||
zh_Hans: '最小相关性分数 (0-1)'
|
||||
ja_JP: '最小関連性スコア (0-1)'
|
||||
type: float
|
||||
default: 0.5
|
||||
required: false
|
||||
show_if:
|
||||
field: score_threshold_enabled
|
||||
operator: eq
|
||||
value: true
|
||||
- name: reranking_enable
|
||||
label:
|
||||
en_US: Enable Reranking
|
||||
zh_Hans: 启用重排序
|
||||
ja_JP: リランキングを有効化
|
||||
description:
|
||||
en_US: 'Enable reranking using the model configured in your Dify dataset settings (configure it in Dify console first)'
|
||||
zh_Hans: '启用重排序，将自动使用 Dify 数据集中配置的重排序模型（请先在 Dify 控制台配置）'
|
||||
ja_JP: 'Dify データセット設定で構成されたリランキングモデルを使用（先に Dify コンソールで設定してください）'
|
||||
type: boolean
|
||||
default: false
|
||||
required: false
|
||||
execution:
|
||||
python:
|
||||
path: engine.py
|
||||
attr: DifyDatasetsConnector
|
||||
@@ -0,0 +1,373 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
from components.shared_state import (
|
||||
AmbiguousMutationError,
|
||||
ConfigStore,
|
||||
serialized,
|
||||
)
|
||||
|
||||
import httpx
|
||||
|
||||
from langbot_plugin.api.definition.components.knowledge_engine import KnowledgeEngine, KnowledgeEngineCapability
|
||||
from langbot_plugin.api.entities.builtin.rag import (
|
||||
IngestionContext,
|
||||
IngestionResult,
|
||||
DocumentStatus,
|
||||
RetrievalContext,
|
||||
RetrievalResultEntry,
|
||||
RetrievalResponse,
|
||||
)
|
||||
from langbot_plugin.api.entities.builtin.provider.message import ContentElement
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _normalize_api_base_url(url: str) -> str:
|
||||
"""Canonical provider target: ``scheme://host[:port]/path``, no trailing slash.
|
||||
|
||||
Used to bind a stored document mapping to the upstream instance that created
|
||||
it, so a later configuration change is detected before any remote call.
|
||||
"""
|
||||
parts = urlsplit(url or '')
|
||||
return urlunsplit((parts.scheme.lower(), parts.netloc.lower(), parts.path.rstrip('/'), '', ''))
|
||||
|
||||
|
||||
class DifyDatasetsConnector(ConfigStore, KnowledgeEngine):
|
||||
"""Knowledge Engine powered by Dify Datasets.
|
||||
|
||||
Supports retrieval, document ingestion (file upload), and deletion
|
||||
via the Dify Dataset API.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_capabilities(cls) -> list[str]:
|
||||
return [KnowledgeEngineCapability.DOC_INGESTION, KnowledgeEngineCapability.DOC_PARSING]
|
||||
|
||||
# ========== Lifecycle Hooks ==========
|
||||
|
||||
@serialized
|
||||
async def on_knowledge_base_create(self, kb_id: str, config: dict) -> None:
|
||||
"""Persist the knowledge-base config for deletion after worker restart."""
|
||||
await self._save_config(kb_id, config)
|
||||
logger.info(f"[DifyDatasetsConnector] Knowledge base created: {kb_id}")
|
||||
|
||||
@serialized(allow_fenced=True)
|
||||
async def on_knowledge_base_delete(self, kb_id: str) -> None:
|
||||
"""Tombstone the stored config and clear any fence for this knowledge base."""
|
||||
await self._save_config(kb_id, None)
|
||||
# Deleting the knowledge base is the operator path out of a fence: the
|
||||
# pending ambiguity is moot once it no longer exists.
|
||||
await self._state.clear_fence(self.plugin, kb_id)
|
||||
logger.info(f"[DifyDatasetsConnector] Knowledge base deleted: {kb_id}")
|
||||
|
||||
# ========== Helper Methods ==========
|
||||
|
||||
async def _fetch_dataset_reranking_config(
|
||||
self, api_base_url: str, api_key: str, dataset_id: str
|
||||
) -> dict:
|
||||
"""Fetch the reranking model config from the Dify dataset settings."""
|
||||
url = f"{api_base_url}/datasets/{dataset_id}"
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
try:
|
||||
async with self.http_client() as client:
|
||||
response = await client.get(url, headers=headers, timeout=15.0)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
retrieval_model = data.get("retrieval_model_dict") or data.get("retrieval_model") or {}
|
||||
reranking = retrieval_model.get("reranking_model", {})
|
||||
mode = retrieval_model.get("reranking_mode")
|
||||
logger.info(
|
||||
f"[DifyDatasetsConnector] Fetched dataset reranking config: "
|
||||
f"mode={mode}, provider={reranking.get('reranking_provider_name')}, "
|
||||
f"model={reranking.get('reranking_model_name')}"
|
||||
)
|
||||
return {
|
||||
"reranking_mode": mode or "reranking_model",
|
||||
"reranking_model": reranking,
|
||||
}
|
||||
except Exception:
|
||||
logger.exception("[DifyDatasetsConnector] Failed to fetch dataset reranking config")
|
||||
return {
|
||||
"reranking_mode": "reranking_model",
|
||||
"reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
|
||||
}
|
||||
|
||||
# ========== Core Methods ==========
|
||||
|
||||
async def retrieve(self, context: RetrievalContext) -> RetrievalResponse:
|
||||
"""Execute retrieval against Dify Dataset API."""
|
||||
config = context.creation_settings
|
||||
retrieval = context.retrieval_settings
|
||||
|
||||
api_base_url = config.get("api_base_url", "https://api.dify.ai/v1").rstrip("/")
|
||||
api_key = config.get("dify_apikey")
|
||||
dataset_id = config.get("dataset_id")
|
||||
top_k = retrieval.get("top_k", 5)
|
||||
score_threshold_enabled = retrieval.get("score_threshold_enabled", False)
|
||||
score_threshold = retrieval.get("score_threshold", 0.5)
|
||||
search_method = retrieval.get("search_method", "semantic_search")
|
||||
reranking_enable = retrieval.get("reranking_enable", False)
|
||||
|
||||
if not api_key or not dataset_id:
|
||||
logger.error(
|
||||
f"[DifyDatasetsConnector] Missing required configuration. "
|
||||
f"Config keys: {list(config.keys())}"
|
||||
)
|
||||
return RetrievalResponse(results=[], total_found=0)
|
||||
|
||||
# Fetch reranking model config from Dify dataset when reranking is enabled
|
||||
reranking_mode = None
|
||||
reranking_model = {"reranking_provider_name": "", "reranking_model_name": ""}
|
||||
if reranking_enable:
|
||||
reranking_config = await self._fetch_dataset_reranking_config(
|
||||
api_base_url, api_key, dataset_id
|
||||
)
|
||||
reranking_mode = reranking_config["reranking_mode"]
|
||||
reranking_model = reranking_config["reranking_model"]
|
||||
|
||||
url = f"{api_base_url}/datasets/{dataset_id}/retrieve"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"query": context.query,
|
||||
"retrieval_model": {
|
||||
"search_method": search_method,
|
||||
"reranking_enable": reranking_enable,
|
||||
"reranking_mode": reranking_mode,
|
||||
"reranking_model": reranking_model,
|
||||
"weights": None,
|
||||
"top_k": int(top_k),
|
||||
"score_threshold_enabled": score_threshold_enabled,
|
||||
"score_threshold": float(score_threshold),
|
||||
},
|
||||
}
|
||||
|
||||
results: list[RetrievalResultEntry] = []
|
||||
try:
|
||||
async with self.http_client() as client:
|
||||
response = await client.post(url, json=payload, headers=headers, timeout=30.0)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
logger.info(
|
||||
f"[DifyDatasetsConnector] Request: search_method={search_method}, top_k={top_k}, "
|
||||
f"score_threshold_enabled={score_threshold_enabled}, score_threshold={score_threshold}, "
|
||||
f"reranking_enable={reranking_enable}, reranking_mode={reranking_mode}, "
|
||||
f"reranking_model={reranking_model}, "
|
||||
f"Response records count: {len(data.get('records', []))}"
|
||||
)
|
||||
|
||||
for record in data.get("records", []):
|
||||
segment = record.get("segment", {})
|
||||
document = segment.get("document", {})
|
||||
|
||||
score = record.get("score")
|
||||
if score is None:
|
||||
score = 0.0
|
||||
|
||||
results.append(
|
||||
RetrievalResultEntry(
|
||||
id=segment.get("id", ""),
|
||||
content=[ContentElement.from_text(segment.get("content", ""))],
|
||||
metadata={
|
||||
"document_id": segment.get("document_id", ""),
|
||||
"document_name": document.get("name", ""),
|
||||
"segment_id": segment.get("id", ""),
|
||||
"keywords": segment.get("keywords", []),
|
||||
"answer": segment.get("answer"),
|
||||
},
|
||||
distance=1.0 - float(score),
|
||||
score=float(score),
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(f"[DifyDatasetsConnector] Retrieved {len(results)} chunks from Dify.")
|
||||
except Exception:
|
||||
logger.exception("[DifyDatasetsConnector] Error during retrieval")
|
||||
|
||||
return RetrievalResponse(results=results, total_found=len(results))
|
||||
|
||||
@serialized
|
||||
async def ingest(self, context: IngestionContext) -> IngestionResult:
|
||||
"""Upload a file to a Dify dataset for indexing."""
|
||||
doc_id = context.file_object.metadata.document_id
|
||||
filename = context.file_object.metadata.filename
|
||||
config = context.creation_settings
|
||||
|
||||
api_base_url = config.get("api_base_url", "https://api.dify.ai/v1").rstrip("/")
|
||||
upstream_target = _normalize_api_base_url(api_base_url)
|
||||
api_key = config.get("dify_apikey")
|
||||
dataset_id = config.get("dataset_id")
|
||||
|
||||
if not api_key or not dataset_id:
|
||||
logger.error(
|
||||
f"[DifyDatasetsConnector] Missing required configuration for ingestion. "
|
||||
f"Config keys: {list(config.keys())}"
|
||||
)
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message="Missing required config: dify_apikey or dataset_id.",
|
||||
)
|
||||
|
||||
# Persist config for deletion after worker restart
|
||||
kb_id = context.get_collection_id()
|
||||
await self._save_config(kb_id, config)
|
||||
if await self._load_document(kb_id, doc_id) is not None:
|
||||
raise RuntimeError('Existing upload intent; reconcile before retry')
|
||||
|
||||
# 1. Read file content from Host
|
||||
try:
|
||||
file_bytes = await self.plugin.get_knowledge_file_stream(context.file_object.storage_path)
|
||||
except Exception as e:
|
||||
logger.error(f"[DifyDatasetsConnector] Failed to read file content: {e}")
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message=f"Could not read file: {e}",
|
||||
)
|
||||
|
||||
# Intent is durable before remote I/O; an ambiguous timeout cannot create
|
||||
# a second document on retry without explicit upstream reconciliation.
|
||||
await self._save_document(kb_id, doc_id, {'upstream_id': '', 'dataset_id': dataset_id,
|
||||
'status': 'pending'})
|
||||
# 2. Upload to Dify dataset
|
||||
url = f"{api_base_url}/datasets/{dataset_id}/document/create-by-file"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
try:
|
||||
data_payload = json.dumps({
|
||||
"indexing_technique": "high_quality",
|
||||
"process_rule": {"mode": "automatic"},
|
||||
})
|
||||
|
||||
async def upload():
|
||||
async with self.http_client() as client:
|
||||
response = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
files={"file": (filename, file_bytes)},
|
||||
data={"data": data_payload},
|
||||
timeout=120.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
resp_data = await self._state.dispatched(
|
||||
self.plugin, kb_id, upload, 'Dify upload outcome unknown')
|
||||
|
||||
dify_document = resp_data.get("document", {})
|
||||
dify_doc_id = dify_document.get("id")
|
||||
if not isinstance(dify_doc_id, str) or not dify_doc_id:
|
||||
raise AmbiguousMutationError("Dify upload omitted upstream document ID; outcome requires reconciliation")
|
||||
await self._save_document(kb_id, doc_id, {'upstream_id': dify_doc_id,
|
||||
'dataset_id': dataset_id, 'status': 'created',
|
||||
'api_base_url': upstream_target})
|
||||
|
||||
logger.info(
|
||||
f"[DifyDatasetsConnector] File uploaded: {filename} -> "
|
||||
f"dify_doc_id={dify_doc_id}, status={dify_document.get('indexing_status')}"
|
||||
)
|
||||
|
||||
return IngestionResult(
|
||||
document_id=dify_doc_id,
|
||||
status=DocumentStatus.PROCESSING,
|
||||
)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
# The provider answered, so this rejection is deterministic: no fence.
|
||||
error_body = e.response.text
|
||||
logger.error(
|
||||
f"[DifyDatasetsConnector] Dify API error during ingestion: "
|
||||
f"status={e.response.status_code}, body={error_body}"
|
||||
)
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message=f"Dify API error {e.response.status_code}: {error_body}",
|
||||
)
|
||||
except AmbiguousMutationError:
|
||||
# Already-dispatched upload with an unknown outcome: fence this KB.
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"[DifyDatasetsConnector] Ingestion failed for {filename}: {e}")
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message=str(e),
|
||||
)
|
||||
|
||||
@serialized
|
||||
async def delete_document(self, kb_id: str, document_id: str) -> bool:
|
||||
"""Delete a document from a Dify dataset."""
|
||||
mapping = await self._load_document(kb_id, document_id)
|
||||
if not mapping or not mapping['upstream_id'] or mapping['status'] != 'created':
|
||||
return False
|
||||
config = await self._load_config(kb_id)
|
||||
if not config:
|
||||
logger.warning(
|
||||
f"[DifyDatasetsConnector] No stored config for kb_id={kb_id}. "
|
||||
"Cannot delete document without API credentials."
|
||||
)
|
||||
return False
|
||||
|
||||
api_base_url = config.get("api_base_url", "https://api.dify.ai/v1").rstrip("/")
|
||||
api_key = config.get("dify_apikey")
|
||||
dataset_id = config.get("dataset_id")
|
||||
|
||||
if not api_key or not dataset_id:
|
||||
logger.error(
|
||||
f"[DifyDatasetsConnector] Missing required configuration for deletion. "
|
||||
f"Config keys: {list(config.keys())}"
|
||||
)
|
||||
return False
|
||||
|
||||
if mapping['dataset_id'] != dataset_id:
|
||||
raise RuntimeError('Dataset changed; reconcile before deletion')
|
||||
# The mapping is bound to the upstream instance that created it: a changed
|
||||
# base URL (or a legacy mapping without one) must not be sent to a
|
||||
# different provider. Reject before dispatching any remote request.
|
||||
if mapping.get('api_base_url') != _normalize_api_base_url(api_base_url):
|
||||
raise RuntimeError('Upstream target changed; reconcile before deletion')
|
||||
url = f"{api_base_url}/datasets/{dataset_id}/documents/{mapping['upstream_id']}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
async def remove():
|
||||
async with self.http_client() as client:
|
||||
response = await client.delete(url, headers=headers, timeout=30.0)
|
||||
if response.status_code != 204:
|
||||
# Some Dify versions may return 200 instead of 204.
|
||||
response.raise_for_status()
|
||||
return response.status_code
|
||||
|
||||
try:
|
||||
status_code = await self._state.dispatched(
|
||||
self.plugin, kb_id, remove, 'Dify delete outcome unknown')
|
||||
await self._save_document(kb_id, mapping['host_document_id'], {**mapping, 'status': 'deleted'})
|
||||
logger.info(
|
||||
f"[DifyDatasetsConnector] Document deleted: {document_id} "
|
||||
f"from dataset {dataset_id} (status={status_code})"
|
||||
)
|
||||
return True
|
||||
except httpx.HTTPStatusError as e:
|
||||
# The provider answered, so the delete was rejected deterministically.
|
||||
logger.error(
|
||||
f"[DifyDatasetsConnector] Failed to delete document {document_id}: {e}"
|
||||
)
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[DifyDatasetsConnector] Failed to delete document {document_id}: {e}"
|
||||
)
|
||||
return False
|
||||
|
||||
@@ -0,0 +1,381 @@
|
||||
"""Installation- and knowledge-base-scoped state using the SDK-bound Host proxy.
|
||||
|
||||
Bundled in each independent archive. No filesystem/global tenant registry.
|
||||
|
||||
Locks and fences are keyed by ``(installation binding, knowledge-base identity)``
|
||||
so an ambiguous failure for one knowledge base never blocks the other knowledge
|
||||
bases that share the same installation; the platform storage key is the same for
|
||||
a whole installation, hence the knowledge-base id is always folded into our keys.
|
||||
Fences are persisted in installation-bound Host storage (key includes the
|
||||
knowledge-base id) with an in-process cache in front, so a restarted worker still
|
||||
recognises them. Only dispatched mutations whose outcome is unknown fence; a
|
||||
deterministic validation error raised before any remote call and ordinary
|
||||
cancellation before a call is dispatched do not.
|
||||
"""
|
||||
import asyncio
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
from functools import wraps
|
||||
from contextlib import asynccontextmanager
|
||||
import httpx
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
HTTP_TOTAL_TIMEOUT = 150.0
|
||||
FENCE_KEY_PREFIX = "ke.fence.v1."
|
||||
# Fallback fence scope when the wrapped method carries no knowledge-base identity.
|
||||
_INSTALLATION_FENCE_KEY = FENCE_KEY_PREFIX + "installation"
|
||||
_IDENTITY_ARGUMENTS = ('kb_id', 'collection_id', 'knowledge_base_id')
|
||||
|
||||
|
||||
class AmbiguousMutationError(RuntimeError):
|
||||
"""A dispatched mutation's outcome is unknown and requires reconciliation.
|
||||
|
||||
Raise this from a call site when a remote mutation may have taken effect but
|
||||
its result (or its durable record) is unknown. Only these failures fence a
|
||||
knowledge base; deterministic errors raised before any remote call, and
|
||||
ordinary cancellation before a call is dispatched, do not fence.
|
||||
"""
|
||||
|
||||
|
||||
class FenceStateUnavailableError(RuntimeError):
|
||||
"""Persisted fence state could not be read, so the fence cannot be confirmed.
|
||||
|
||||
A failed read is not evidence that no fence exists, so the operation fails
|
||||
closed rather than proceeding as if the knowledge base were clear. The read
|
||||
is retried on the next call instead of being cached as loaded.
|
||||
"""
|
||||
|
||||
|
||||
def is_ambiguous_http_failure(exc):
|
||||
"""Whether a failed provider call had an unknown outcome.
|
||||
|
||||
A response status means the provider answered, so the failure is
|
||||
deterministic. A connection that was never established means nothing was
|
||||
dispatched. Everything else (timeout, lost connection, unparsable response)
|
||||
was dispatched with an unknown outcome.
|
||||
"""
|
||||
if isinstance(exc, httpx.HTTPStatusError):
|
||||
return False
|
||||
if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout)):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _settle_commit(task):
|
||||
"""Wait for a dispatched commit and report its outcome.
|
||||
|
||||
Repeated caller cancellation never abandons the wait, so a Host commit that
|
||||
was already dispatched settles before its lock is released. Returns
|
||||
``(caller_cancelled, failure)``; ``failure`` is the commit's exception, or
|
||||
``asyncio.CancelledError`` when the commit itself was cancelled.
|
||||
"""
|
||||
cancelled = False
|
||||
while not task.done():
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
except Exception:
|
||||
# The commit raised; its outcome is read from the state below.
|
||||
pass
|
||||
if task.cancelled():
|
||||
return cancelled, asyncio.CancelledError()
|
||||
return cancelled, task.exception()
|
||||
|
||||
|
||||
def _kb_identity(method, instance, args, kwargs):
|
||||
"""Extract the knowledge-base identity from a serialized method's arguments.
|
||||
|
||||
Reads a ``kb_id`` / ``collection_id`` / ``knowledge_base_id`` argument first,
|
||||
then an argument exposing ``get_collection_id()`` or ``knowledge_base_id``
|
||||
(ingestion passes the whole context). Returns ``None`` when the wrapped
|
||||
method carries no knowledge-base identity, in which case the caller falls
|
||||
back to an installation-wide lock and fence.
|
||||
"""
|
||||
try:
|
||||
bound = inspect.signature(method).bind_partial(instance, *args, **kwargs)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
for name in _IDENTITY_ARGUMENTS:
|
||||
value = bound.arguments.get(name)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
for value in bound.arguments.values():
|
||||
getter = getattr(value, 'get_collection_id', None)
|
||||
if callable(getter):
|
||||
try:
|
||||
identity = getter()
|
||||
except Exception:
|
||||
continue
|
||||
if isinstance(identity, str) and identity:
|
||||
return identity
|
||||
identity = getattr(value, 'knowledge_base_id', None)
|
||||
if isinstance(identity, str) and identity:
|
||||
return identity
|
||||
return None
|
||||
|
||||
|
||||
def _kb_fence_key(kb_identity):
|
||||
if kb_identity is None:
|
||||
return _INSTALLATION_FENCE_KEY
|
||||
return FENCE_KEY_PREFIX + hashlib.sha256(kb_identity.encode()).hexdigest()
|
||||
|
||||
|
||||
class SerialState:
|
||||
def __init__(self):
|
||||
self._locks = WeakValueDictionary()
|
||||
# binding -> set of fence keys; cache in front of persisted fences.
|
||||
self._fenced = {}
|
||||
self._fence_loaded = set()
|
||||
|
||||
@staticmethod
|
||||
def binding(plugin):
|
||||
try:
|
||||
binding = plugin.get_installation_binding()
|
||||
except Exception:
|
||||
# Revocation has no invocation context; fall back to dedicated scope.
|
||||
binding = None
|
||||
return binding if binding is not None else ('dedicated',)
|
||||
|
||||
@staticmethod
|
||||
def _scope(binding, kb_identity):
|
||||
return (binding, kb_identity)
|
||||
|
||||
def _fenced_keys(self, binding):
|
||||
return self._fenced.setdefault(binding, set())
|
||||
|
||||
def is_fenced(self, binding, kb_identity):
|
||||
fenced = self._fenced_keys(binding)
|
||||
return (
|
||||
_kb_fence_key(kb_identity) in fenced
|
||||
or _INSTALLATION_FENCE_KEY in fenced
|
||||
)
|
||||
|
||||
async def _load_fences(self, plugin, binding):
|
||||
"""Read persisted fences once per binding.
|
||||
|
||||
A failed read cannot be treated as "no fence": it means the fence state
|
||||
is unknown, so the caller fails closed and retries on the next call
|
||||
rather than caching the failure as loaded.
|
||||
"""
|
||||
if binding in self._fence_loaded:
|
||||
return
|
||||
try:
|
||||
keys = await plugin.get_plugin_storage_keys()
|
||||
fenced = self._fenced_keys(binding)
|
||||
for key in keys:
|
||||
if not key.startswith(FENCE_KEY_PREFIX):
|
||||
continue
|
||||
try:
|
||||
marker = json.loads(await plugin.get_plugin_storage(key))
|
||||
except Exception:
|
||||
# Unreadable marker: fail closed on this knowledge base.
|
||||
marker = {'reason': 'unreadable fence record'}
|
||||
if marker is not None:
|
||||
fenced.add(key)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("Could not read persisted fences for %r", binding, exc_info=True)
|
||||
raise FenceStateUnavailableError(
|
||||
f"Could not confirm persisted fence state for {binding!r}: {exc!r}"
|
||||
) from exc
|
||||
self._fence_loaded.add(binding)
|
||||
|
||||
async def run(self, plugin, kb_identity, operation, *, allow_fenced=False):
|
||||
binding = self.binding(plugin)
|
||||
key = self._scope(binding, kb_identity)
|
||||
lock = self._locks.get(key)
|
||||
if lock is None:
|
||||
lock = self._locks[key] = asyncio.Lock()
|
||||
async with lock:
|
||||
await self._load_fences(plugin, binding)
|
||||
if not allow_fenced and self.is_fenced(binding, kb_identity):
|
||||
raise RuntimeError(
|
||||
f"Knowledge base {kb_identity!r} fenced after an ambiguous mutation; "
|
||||
"reconcile with the provider before retrying"
|
||||
)
|
||||
try:
|
||||
return await operation()
|
||||
except AmbiguousMutationError:
|
||||
await self.fence(plugin, kb_identity, 'ambiguous remote mutation')
|
||||
raise
|
||||
|
||||
async def dispatched(self, plugin, kb_identity, operation, reason):
|
||||
"""Dispatch a remote mutation under a pre-persisted intent.
|
||||
|
||||
The fence is written before the request goes out, so a cancellation after
|
||||
dispatch (or an ambiguous failure) cannot leave a mutation with no marker
|
||||
for reconciliation; a deterministic answer clears it again.
|
||||
"""
|
||||
await self.fence(plugin, kb_identity, reason)
|
||||
try:
|
||||
result = await operation()
|
||||
except asyncio.CancelledError:
|
||||
# Dispatched, outcome unknown: keep the fence and let the caller see it.
|
||||
raise
|
||||
except Exception as exc:
|
||||
if not is_ambiguous_http_failure(exc):
|
||||
await self.clear_fence(plugin, kb_identity)
|
||||
raise
|
||||
await self.clear_fence(plugin, kb_identity)
|
||||
return result
|
||||
|
||||
async def fence(self, plugin, kb_identity, reason):
|
||||
"""Mark one knowledge base as requiring reconciliation.
|
||||
|
||||
The in-process marker is set synchronously; persisting it to Host storage
|
||||
is best effort so that a caller cancellation cannot lose the fence.
|
||||
"""
|
||||
binding = self.binding(plugin)
|
||||
fence_key = _kb_fence_key(kb_identity)
|
||||
self._fenced_keys(binding).add(fence_key)
|
||||
payload = json.dumps({'kb_id': kb_identity, 'reason': reason}).encode()
|
||||
cancelled, failure = await _settle_commit(
|
||||
asyncio.create_task(plugin.set_plugin_storage(fence_key, payload))
|
||||
)
|
||||
if failure is not None:
|
||||
logger.warning("Could not persist fence for %r: %r", kb_identity, failure)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def clear_fence(self, plugin, kb_identity):
|
||||
"""Drop the fence for one knowledge base, in process and in storage.
|
||||
|
||||
Used when a knowledge base is deleted: its pending ambiguity is moot, so
|
||||
the operator path out of a fence is deletion. A lingering persisted
|
||||
marker is inert because the configuration tombstone is already written.
|
||||
"""
|
||||
binding = self.binding(plugin)
|
||||
fence_key = _kb_fence_key(kb_identity)
|
||||
self._fenced_keys(binding).discard(fence_key)
|
||||
cancelled, failure = await _settle_commit(
|
||||
asyncio.create_task(plugin.set_plugin_storage(fence_key, b'null'))
|
||||
)
|
||||
if failure is not None:
|
||||
logger.warning("Could not clear fence for %r: %r", kb_identity, failure)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def write(self, plugin, kb_identity, key, value):
|
||||
# A caller cancellation must not abandon a dispatched Host commit while
|
||||
# its lock is released; settle it and fence only a real ambiguity.
|
||||
task = asyncio.create_task(plugin.set_plugin_storage(key, value))
|
||||
cancelled, failure = await _settle_commit(task)
|
||||
if failure is not None:
|
||||
await self.fence(plugin, kb_identity, f'Host write outcome unknown: {failure!r}')
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
if failure is not None:
|
||||
raise failure
|
||||
return task.result()
|
||||
|
||||
def release_binding(self, binding):
|
||||
"""Drop process-local state keyed by one installation binding."""
|
||||
for key in [k for k in self._locks if k[0] == binding]:
|
||||
self._locks.pop(key, None)
|
||||
self._fenced.pop(binding, None)
|
||||
self._fence_loaded.discard(binding)
|
||||
|
||||
|
||||
def serialized(method=None, *, allow_fenced=False):
|
||||
def decorate(func):
|
||||
@wraps(func)
|
||||
async def call(self, *args, **kwargs):
|
||||
return await self._state.run(
|
||||
self.plugin,
|
||||
_kb_identity(func, self, args, kwargs),
|
||||
lambda: func(self, *args, **kwargs),
|
||||
allow_fenced=allow_fenced,
|
||||
)
|
||||
return call
|
||||
return decorate if method is None else decorate(method)
|
||||
|
||||
|
||||
class ConfigStore:
|
||||
"""No credential cache: every delete resolves durable installation storage.
|
||||
|
||||
A JSON null tombstone makes missing KBs explicit without ambiguous delete RPCs.
|
||||
The one mutation lock also orders ingest/delete/create inside this worker.
|
||||
"""
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._state = SerialState()
|
||||
self._http_slots = asyncio.Semaphore(4)
|
||||
|
||||
async def initialize(self) -> None:
|
||||
await super().initialize()
|
||||
# The runtime dispatches installation revocation to the plugin instance,
|
||||
# which holds no back-reference to its components; hand it this
|
||||
# component's state so the revocation hook can drop per-binding locks and
|
||||
# fences. Without the hook (langbot-plugin 0.7.4) this attribute is simply
|
||||
# never read.
|
||||
plugin = getattr(self, 'plugin', None)
|
||||
if plugin is not None:
|
||||
plugin.knowledge_engine_serial_state = self._state
|
||||
|
||||
@asynccontextmanager
|
||||
async def http_client(self):
|
||||
# Include queueing and the whole response body in the deadline, not just
|
||||
# individual socket reads. Clients/credentials never survive a call.
|
||||
async with asyncio.timeout(HTTP_TOTAL_TIMEOUT):
|
||||
async with self._http_slots:
|
||||
async with httpx.AsyncClient() as client:
|
||||
yield client
|
||||
|
||||
@staticmethod
|
||||
def _key(kb_id):
|
||||
return "ke.config.v1." + hashlib.sha256(kb_id.encode()).hexdigest()
|
||||
|
||||
async def _save_config(self, kb_id, config):
|
||||
value = json.dumps(config, ensure_ascii=False).encode()
|
||||
if len(value) > 65536:
|
||||
raise ValueError("Knowledge-base configuration exceeds 64 KiB")
|
||||
await self._state.write(self.plugin, kb_id, self._key(kb_id), value)
|
||||
|
||||
async def _load_config(self, kb_id):
|
||||
key = self._key(kb_id)
|
||||
# Only an authoritative keys response establishes absence; RPC errors propagate.
|
||||
if key not in await self.plugin.get_plugin_storage_keys():
|
||||
return None
|
||||
value = json.loads(await self.plugin.get_plugin_storage(key))
|
||||
if value is not None and not isinstance(value, dict):
|
||||
raise ValueError("Invalid persisted knowledge-base configuration")
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _document_key(kb_id, host_document_id):
|
||||
return "ke.document.v1." + hashlib.sha256(
|
||||
json.dumps([kb_id, host_document_id]).encode()
|
||||
).hexdigest()
|
||||
|
||||
async def _load_document(self, kb_id, host_document_id):
|
||||
key = self._document_key(kb_id, host_document_id)
|
||||
keys = await self.plugin.get_plugin_storage_keys()
|
||||
if key in keys:
|
||||
value = json.loads(await self.plugin.get_plugin_storage(key))
|
||||
else:
|
||||
# Current Host passes the durable upstream ID on delete; older Host
|
||||
# revisions pass its file UUID. Resolve either without trusting an
|
||||
# arbitrary unrecorded identifier.
|
||||
matches = []
|
||||
for candidate in keys:
|
||||
if candidate.startswith('ke.document.v1.'):
|
||||
item = json.loads(await self.plugin.get_plugin_storage(candidate))
|
||||
if item.get('kb_id') == kb_id and item.get('upstream_id') == host_document_id:
|
||||
matches.append(item)
|
||||
if len(matches) > 1:
|
||||
raise ValueError('Ambiguous upstream document ID')
|
||||
return matches[0] if matches else None
|
||||
if not isinstance(value, dict) or not isinstance(value.get('upstream_id'), str):
|
||||
raise ValueError('Invalid upstream document mapping; reconcile installation')
|
||||
return value
|
||||
|
||||
async def _save_document(self, kb_id, host_document_id, value):
|
||||
value = {**value, 'kb_id': kb_id, 'host_document_id': host_document_id}
|
||||
await self._state.write(self.plugin, kb_id, self._document_key(kb_id, host_document_id),
|
||||
json.dumps(value).encode())
|
||||
Reference in New Issue
Block a user