mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-11-03 22:32:27 +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,545 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from components.shared_state import (
|
||||
AmbiguousMutationError,
|
||||
ConfigStore,
|
||||
FenceLifecycleError,
|
||||
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_base_url(api_base_url: str) -> str:
|
||||
"""Reduce a configured base URL to ``scheme://host[:port]/path``.
|
||||
|
||||
A document mapping must name the upstream it was created on, so the stored
|
||||
form is canonical: scheme and host casing, a default port and a trailing
|
||||
slash never distinguish one deployment from another. An address without a
|
||||
host is stored stripped, so a later comparison still detects a change.
|
||||
"""
|
||||
url = httpx.URL(api_base_url)
|
||||
if not url.host:
|
||||
return api_base_url.rstrip("/")
|
||||
default_port = {'http': 80, 'https': 443}.get(url.scheme)
|
||||
port = f":{url.port}" if url.port is not None and url.port != default_port else ""
|
||||
return f"{url.scheme}://{url.host}{port}{url.path.rstrip('/')}"
|
||||
|
||||
|
||||
class RAGFlowConnector(ConfigStore, KnowledgeEngine):
|
||||
"""Knowledge Engine powered by RAGFlow.
|
||||
|
||||
Supports retrieval, document ingestion (upload + parse), and deletion
|
||||
via the RAGFlow HTTP 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 installation-bound configuration and validate dataset IDs."""
|
||||
logger.info(f"[RAGFlowKnowledgeEngine] Knowledge base created: {kb_id}")
|
||||
await self._save_config(kb_id, config)
|
||||
|
||||
# Validate dataset IDs
|
||||
api_base_url = config.get("api_base_url", "http://localhost:9380").rstrip("/")
|
||||
api_key = config.get("api_key")
|
||||
dataset_ids_str = config.get("dataset_ids", "")
|
||||
|
||||
if not api_key or not dataset_ids_str:
|
||||
return
|
||||
|
||||
dataset_ids = [did.strip() for did in dataset_ids_str.split(",") if did.strip()]
|
||||
if not dataset_ids:
|
||||
return
|
||||
|
||||
try:
|
||||
async with self.http_client() as client:
|
||||
resp = await client.get(
|
||||
f"{api_base_url}/api/v1/datasets",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
params={"page": 1, "page_size": 100},
|
||||
timeout=15.0,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
if data.get("code") != 0:
|
||||
logger.warning(
|
||||
f"[RAGFlowKnowledgeEngine] Failed to validate datasets: {data.get('message')}"
|
||||
)
|
||||
return
|
||||
|
||||
existing_ids = {ds["id"] for ds in data.get("data", [])}
|
||||
for did in dataset_ids:
|
||||
if did in existing_ids:
|
||||
logger.info(f"[RAGFlowKnowledgeEngine] Dataset {did} validated OK")
|
||||
else:
|
||||
logger.warning(
|
||||
f"[RAGFlowKnowledgeEngine] Dataset {did} NOT found in RAGFlow! "
|
||||
"Please check the dataset ID."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"[RAGFlowKnowledgeEngine] Dataset validation failed: {e}")
|
||||
|
||||
@serialized(allow_fenced=True)
|
||||
async def on_knowledge_base_delete(self, kb_id: str) -> None:
|
||||
"""Tombstone the stored configuration and clear any fence for this KB."""
|
||||
logger.info(f"[RAGFlowKnowledgeEngine] Knowledge base deleted: {kb_id}")
|
||||
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)
|
||||
|
||||
async def retrieve(self, context: RetrievalContext) -> RetrievalResponse:
|
||||
"""Execute retrieval against RAGFlow API."""
|
||||
config = context.creation_settings
|
||||
retrieval = context.retrieval_settings
|
||||
|
||||
api_base_url = config.get("api_base_url", "http://localhost:9380").rstrip("/")
|
||||
api_key = config.get("api_key")
|
||||
dataset_ids_str = config.get("dataset_ids", "")
|
||||
top_k = retrieval.get("top_k", 1024)
|
||||
similarity_threshold = retrieval.get("similarity_threshold", 0.2)
|
||||
vector_similarity_weight = retrieval.get("vector_similarity_weight", 0.3)
|
||||
page_size = retrieval.get("page_size", 30)
|
||||
keyword = retrieval.get("keyword", False)
|
||||
rerank_id = retrieval.get("rerank_id", "")
|
||||
use_kg = retrieval.get("use_kg", False)
|
||||
|
||||
if not api_key or not dataset_ids_str:
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] Missing required configuration. "
|
||||
f"Config keys: {list(config.keys())}"
|
||||
)
|
||||
return RetrievalResponse(results=[], total_found=0)
|
||||
|
||||
dataset_ids = [did.strip() for did in dataset_ids_str.split(",") if did.strip()]
|
||||
|
||||
if not dataset_ids:
|
||||
logger.error("[RAGFlowKnowledgeEngine] No valid dataset IDs provided")
|
||||
return RetrievalResponse(results=[], total_found=0)
|
||||
|
||||
url = f"{api_base_url}/api/v1/retrieval"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"question": context.query,
|
||||
"dataset_ids": dataset_ids,
|
||||
"page": 1,
|
||||
"page_size": int(page_size),
|
||||
"similarity_threshold": float(similarity_threshold),
|
||||
"vector_similarity_weight": float(vector_similarity_weight),
|
||||
"top_k": int(top_k),
|
||||
"keyword": keyword,
|
||||
"highlight": False,
|
||||
"use_kg": use_kg,
|
||||
}
|
||||
if rerank_id:
|
||||
payload["rerank_id"] = rerank_id
|
||||
|
||||
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()
|
||||
|
||||
if data.get("code") != 0:
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] API returned error code: {data.get('code')}"
|
||||
)
|
||||
return RetrievalResponse(results=[], total_found=0)
|
||||
|
||||
logger.info(
|
||||
f"[RAGFlowKnowledgeEngine] Request: vector_weight={vector_similarity_weight}, "
|
||||
f"top_k={top_k}, page_size={page_size}, threshold={similarity_threshold}, "
|
||||
f"keyword={keyword}, rerank_id={rerank_id or 'None'}, use_kg={use_kg}, "
|
||||
f"Response chunks: {len(data.get('data', {}).get('chunks', []))}"
|
||||
)
|
||||
|
||||
for chunk in data.get("data", {}).get("chunks", []):
|
||||
similarity = chunk.get("similarity")
|
||||
if similarity is None:
|
||||
similarity = 0.0
|
||||
|
||||
results.append(
|
||||
RetrievalResultEntry(
|
||||
id=chunk.get("id", ""),
|
||||
content=[ContentElement.from_text(chunk.get("content", ""))],
|
||||
metadata={
|
||||
"document_id": chunk.get("document_id", ""),
|
||||
"kb_id": chunk.get("kb_id", ""),
|
||||
"document_keyword": chunk.get("document_keyword", ""),
|
||||
"important_keywords": chunk.get("important_keywords", []),
|
||||
"term_similarity": chunk.get("term_similarity", 0.0),
|
||||
"vector_similarity": chunk.get("vector_similarity", 0.0),
|
||||
"image_id": chunk.get("image_id"),
|
||||
},
|
||||
distance=1.0 - float(similarity),
|
||||
score=float(similarity),
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(f"[RAGFlowKnowledgeEngine] Retrieved {len(results)} chunks from RAGFlow.")
|
||||
except Exception:
|
||||
logger.exception("[RAGFlowKnowledgeEngine] Error during retrieval")
|
||||
|
||||
return RetrievalResponse(results=results, total_found=len(results))
|
||||
|
||||
@serialized
|
||||
async def ingest(self, context: IngestionContext) -> IngestionResult:
|
||||
"""Upload a file to RAGFlow and trigger parsing."""
|
||||
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", "http://localhost:9380").rstrip("/")
|
||||
api_key = config.get("api_key")
|
||||
dataset_ids_str = config.get("dataset_ids", "")
|
||||
|
||||
if not api_key or not dataset_ids_str:
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] Missing required configuration for ingestion. "
|
||||
f"Config keys: {list(config.keys())}"
|
||||
)
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message="Missing api_key or dataset_ids in configuration.",
|
||||
)
|
||||
|
||||
dataset_ids = [did.strip() for did in dataset_ids_str.split(",") if did.strip()]
|
||||
if not dataset_ids:
|
||||
logger.error("[RAGFlowKnowledgeEngine] No valid dataset IDs provided for ingestion")
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message="No valid dataset IDs provided.",
|
||||
)
|
||||
|
||||
# Use the first dataset as the ingestion target
|
||||
target_dataset_id = dataset_ids[0]
|
||||
# The mapping names the canonical upstream it was created on, so a later
|
||||
# configuration change cannot be mistaken for the same deployment.
|
||||
upstream_target = _normalize_base_url(api_base_url)
|
||||
|
||||
# The Host knowledge-base id is the identity every entry point shares — the
|
||||
# create hook, deletion and the fence/lock scope — so the collection id,
|
||||
# which is a vector-store identity that need not equal it, is not used.
|
||||
kb_id = context.knowledge_base_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"[RAGFlowKnowledgeEngine] Failed to get file content: {e}")
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message=f"Could not read file: {e}",
|
||||
)
|
||||
|
||||
await self._save_document(kb_id, doc_id, {'upstream_id': '',
|
||||
'dataset_id': target_dataset_id, 'api_base_url': upstream_target,
|
||||
'status': 'pending'})
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
try:
|
||||
async with self.http_client() as client:
|
||||
# 2. Upload file to RAGFlow dataset
|
||||
upload_url = f"{api_base_url}/api/v1/datasets/{target_dataset_id}/documents"
|
||||
files = {"file": (filename, file_bytes)}
|
||||
|
||||
async def upload():
|
||||
response = await client.post(
|
||||
upload_url, headers=headers, files=files, timeout=60.0
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if data.get("code") != 0:
|
||||
# The provider answered: the upload was rejected, so no
|
||||
# Host record is written and the fence can be released.
|
||||
return None, data.get("message") or "Unknown upload error"
|
||||
# Extract the document ID returned by RAGFlow
|
||||
docs = data.get("data", [])
|
||||
ragflow_doc_id = docs[0].get("id") if docs and isinstance(docs[0], dict) else None
|
||||
if not isinstance(ragflow_doc_id, str) or not ragflow_doc_id:
|
||||
raise AmbiguousMutationError("RAGFlow upload omitted upstream document ID; outcome requires reconciliation")
|
||||
# The mapping is recorded before the fence is released: a
|
||||
# cancellation or a Host failure here must not leave an
|
||||
# uploaded document with no durable record of it.
|
||||
await self._save_document(kb_id, doc_id, {'upstream_id': ragflow_doc_id,
|
||||
'dataset_id': target_dataset_id, 'api_base_url': upstream_target,
|
||||
'status': 'created'})
|
||||
return ragflow_doc_id, None
|
||||
|
||||
# The intent is persisted before the upload leaves, and the fence
|
||||
# is released only once the response was validated and the mapping
|
||||
# recorded.
|
||||
ragflow_doc_id, upload_error = await self._state.dispatched(
|
||||
self.plugin, kb_id, upload, 'RAGFlow upload outcome unknown',
|
||||
)
|
||||
if upload_error is not None:
|
||||
logger.error(f"[RAGFlowKnowledgeEngine] Upload failed: {upload_error}")
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message=f"RAGFlow upload error: {upload_error}",
|
||||
)
|
||||
|
||||
# 3. Trigger parsing
|
||||
chunks_url = f"{api_base_url}/api/v1/datasets/{target_dataset_id}/chunks"
|
||||
|
||||
async def trigger_parse():
|
||||
response = await client.post(
|
||||
chunks_url,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
json={"document_ids": [ragflow_doc_id]},
|
||||
timeout=30.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if data.get("code") != 0:
|
||||
# The provider answered: parsing did not start, so the
|
||||
# fence can be released.
|
||||
return data.get("message") or "Unknown parsing error"
|
||||
return None
|
||||
|
||||
parse_error = await self._state.dispatched(
|
||||
self.plugin, kb_id,
|
||||
trigger_parse, 'RAGFlow parsing trigger outcome unknown',
|
||||
)
|
||||
if parse_error is not None:
|
||||
logger.warning(
|
||||
f"[RAGFlowKnowledgeEngine] Parsing trigger returned error: {parse_error}"
|
||||
)
|
||||
# Document was uploaded but parsing failed to start
|
||||
return IngestionResult(
|
||||
document_id=ragflow_doc_id,
|
||||
status=DocumentStatus.FAILED,
|
||||
error_message=f"RAGFlow parsing trigger error: {parse_error}",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[RAGFlowKnowledgeEngine] File '{filename}' uploaded and parsing triggered "
|
||||
f"(ragflow_doc_id={ragflow_doc_id})"
|
||||
)
|
||||
|
||||
# Auto-trigger GraphRAG construction if enabled
|
||||
auto_graphrag = config.get("auto_graphrag", False)
|
||||
if auto_graphrag:
|
||||
graphrag_url = f"{api_base_url}/api/v1/datasets/{target_dataset_id}/run_graphrag"
|
||||
|
||||
async def trigger_graphrag():
|
||||
response = await client.post(
|
||||
graphrag_url,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if data.get("code") != 0:
|
||||
# The provider answered: GraphRAG did not start. That
|
||||
# is a resolved answer, not an unresolved state.
|
||||
return None, data.get("message") or "unknown reason"
|
||||
task = (data.get("data") or {}).get("graphrag_task_id", "unknown")
|
||||
return task, None
|
||||
|
||||
# An unresolved GraphRAG outcome is not swallowed: it keeps
|
||||
# its own fence and the chain must not continue over it (the
|
||||
# RAPTOR dispatch would refuse it, and a later success must
|
||||
# never clear this fence).
|
||||
graphrag_task_id, graphrag_error = await self._state.dispatched(
|
||||
self.plugin, kb_id, trigger_graphrag,
|
||||
'RAGFlow GraphRAG trigger outcome unknown',
|
||||
)
|
||||
if graphrag_error is None:
|
||||
logger.info(
|
||||
f"[RAGFlowKnowledgeEngine] GraphRAG construction triggered "
|
||||
f"(task_id={graphrag_task_id})"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"[RAGFlowKnowledgeEngine] GraphRAG trigger returned: {graphrag_error}"
|
||||
)
|
||||
|
||||
# Auto-trigger RAPTOR construction if enabled
|
||||
auto_raptor = config.get("auto_raptor", False)
|
||||
if auto_raptor:
|
||||
raptor_url = f"{api_base_url}/api/v1/datasets/{target_dataset_id}/run_raptor"
|
||||
|
||||
async def trigger_raptor():
|
||||
response = await client.post(
|
||||
raptor_url,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if data.get("code") != 0:
|
||||
# The provider answered: RAPTOR did not start.
|
||||
return None, data.get("message") or "unknown reason"
|
||||
task = (data.get("data") or {}).get("raptor_task_id", "unknown")
|
||||
return task, None
|
||||
|
||||
raptor_task_id, raptor_error = await self._state.dispatched(
|
||||
self.plugin, kb_id, trigger_raptor,
|
||||
'RAGFlow RAPTOR trigger outcome unknown',
|
||||
)
|
||||
if raptor_error is None:
|
||||
logger.info(
|
||||
f"[RAGFlowKnowledgeEngine] RAPTOR construction triggered "
|
||||
f"(task_id={raptor_task_id})"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"[RAGFlowKnowledgeEngine] RAPTOR trigger returned: {raptor_error}"
|
||||
)
|
||||
|
||||
return IngestionResult(
|
||||
document_id=ragflow_doc_id,
|
||||
status=DocumentStatus.PROCESSING,
|
||||
)
|
||||
|
||||
except (AmbiguousMutationError, FenceLifecycleError):
|
||||
# An already-dispatched mutation with an unknown outcome, or a
|
||||
# mutation that was never dispatched: the caller must see it.
|
||||
raise
|
||||
except Exception as e:
|
||||
# An ambiguous failure already left its pre-dispatch fence in place.
|
||||
logger.error(f"[RAGFlowKnowledgeEngine] 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 RAGFlow."""
|
||||
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.error(
|
||||
f"[RAGFlowKnowledgeEngine] No stored config for kb_id={kb_id}, "
|
||||
"cannot delete document"
|
||||
)
|
||||
return False
|
||||
|
||||
api_base_url = config.get("api_base_url", "http://localhost:9380").rstrip("/")
|
||||
api_key = config.get("api_key")
|
||||
dataset_ids_str = config.get("dataset_ids", "")
|
||||
|
||||
if not api_key or not dataset_ids_str:
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] Missing api_key or dataset_ids for kb_id={kb_id}"
|
||||
)
|
||||
return False
|
||||
|
||||
dataset_ids = [did.strip() for did in dataset_ids_str.split(",") if did.strip()]
|
||||
if not dataset_ids:
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] No valid dataset IDs for kb_id={kb_id}"
|
||||
)
|
||||
return False
|
||||
|
||||
target_dataset_id = dataset_ids[0]
|
||||
|
||||
if mapping['dataset_id'] != target_dataset_id:
|
||||
# Local validation before any remote call: deterministic, no fence.
|
||||
raise RuntimeError('Dataset changed; reconcile before deletion')
|
||||
if mapping.get('api_base_url') != _normalize_base_url(api_base_url):
|
||||
# The recorded upstream is not the configured one, so the document id
|
||||
# does not necessarily exist here; mappings written before the target
|
||||
# was recorded are equally unresolvable.
|
||||
raise RuntimeError('Upstream target changed; reconcile before deletion')
|
||||
|
||||
url = f"{api_base_url}/api/v1/datasets/{target_dataset_id}/documents"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async def delete_upstream():
|
||||
async with self.http_client() as client:
|
||||
response = await client.request(
|
||||
"DELETE", url, headers=headers,
|
||||
json={"ids": [mapping['upstream_id']]},
|
||||
timeout=30.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
if data.get("code") != 0:
|
||||
# The provider answered: the delete was rejected, so there is
|
||||
# nothing to record and the fence can be released.
|
||||
return False, data.get("message") or "Unknown error"
|
||||
# The tombstone is recorded before the fence is released: a
|
||||
# cancellation or a Host failure here must not leave a deleted
|
||||
# document that the mapping still calls created.
|
||||
await self._save_document(kb_id, mapping['host_document_id'], {**mapping, 'status': 'deleted'})
|
||||
return True, None
|
||||
|
||||
try:
|
||||
deleted, error = await self._state.dispatched(
|
||||
self.plugin, kb_id, delete_upstream, 'RAGFlow delete outcome unknown',
|
||||
)
|
||||
|
||||
if not deleted:
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] Delete failed for doc={document_id}: {error}"
|
||||
)
|
||||
return False
|
||||
|
||||
logger.info(
|
||||
f"[RAGFlowKnowledgeEngine] Document {document_id} deleted from "
|
||||
f"dataset {target_dataset_id}"
|
||||
)
|
||||
return True
|
||||
|
||||
except (AmbiguousMutationError, FenceLifecycleError):
|
||||
# The delete was dispatched with an unknown outcome, or never
|
||||
# dispatched at all: the caller must not read this as a plain failure.
|
||||
raise
|
||||
except httpx.HTTPStatusError as e:
|
||||
# The provider answered, so this is not a transport ambiguity: a 4xx
|
||||
# rejection is a resolved failure whose fence was released, while a
|
||||
# 5xx keeps its fence (the provider or an intermediary may have
|
||||
# committed the delete and then failed), so the next mutation for
|
||||
# this knowledge base is still refused.
|
||||
logger.error(
|
||||
f"[RAGFlowKnowledgeEngine] Delete rejected for doc={document_id}: {e}"
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
# An ambiguous failure already left its pre-dispatch fence in place.
|
||||
logger.exception(
|
||||
f"[RAGFlowKnowledgeEngine] Error deleting document {document_id}"
|
||||
)
|
||||
return False
|
||||
@@ -0,0 +1,161 @@
|
||||
apiVersion: v1
|
||||
kind: KnowledgeEngine
|
||||
metadata:
|
||||
name: ragflow-connector
|
||||
label:
|
||||
en_US: RAGFlow Knowledge Base
|
||||
zh_Hans: RAGFlow 知识库
|
||||
ja_JP: RAGFlow 知識ベース
|
||||
description:
|
||||
en_US: Retrieve knowledge from or store files into RAGFlow knowledge bases
|
||||
zh_Hans: 从 RAGFlow 知识库中检索知识或向其中存入文件
|
||||
ja_JP: RAGFlow 知識ベースから知識を検索またはファイルを保存
|
||||
spec:
|
||||
creation_schema:
|
||||
- name: api_base_url
|
||||
label:
|
||||
en_US: API Base URL
|
||||
zh_Hans: API 基础地址
|
||||
ja_JP: API ベース URL
|
||||
description:
|
||||
en_US: 'Base URL for RAGFlow API (e.g., http://localhost:9380 for local, or your server URL)'
|
||||
zh_Hans: 'RAGFlow API 基础地址（本地版本使用 http://localhost:9380，或使用您的服务器地址）'
|
||||
ja_JP: 'RAGFlow API ベース URL（ローカル版は http://localhost:9380、またはサーバーの URL）'
|
||||
type: string
|
||||
default: http://localhost:9380
|
||||
required: true
|
||||
- name: api_key
|
||||
label:
|
||||
en_US: API Key
|
||||
zh_Hans: API 密钥
|
||||
ja_JP: API キー
|
||||
description:
|
||||
en_US: 'API key from your RAGFlow instance'
|
||||
zh_Hans: '从您的 RAGFlow 实例获取的 API Key'
|
||||
ja_JP: 'RAGFlow インスタンスから取得した API Key'
|
||||
type: password
|
||||
required: true
|
||||
- name: dataset_ids
|
||||
label:
|
||||
en_US: Dataset IDs
|
||||
zh_Hans: 数据集 IDs
|
||||
ja_JP: データセット ID
|
||||
description:
|
||||
en_US: 'Comma-separated dataset IDs to search (e.g., "id1,id2,id3")'
|
||||
zh_Hans: '要搜索的数据集 ID，用逗号分隔（例如："id1,id2,id3"）'
|
||||
ja_JP: '検索するデータセット ID（カンマ区切り、例："id1,id2,id3"）'
|
||||
type: string
|
||||
required: true
|
||||
- name: auto_graphrag
|
||||
label:
|
||||
en_US: Auto Build Knowledge Graph
|
||||
zh_Hans: 自动构建知识图谱
|
||||
ja_JP: ナレッジグラフ自動構築
|
||||
description:
|
||||
en_US: 'Automatically trigger GraphRAG construction after file ingestion'
|
||||
zh_Hans: '文件上传解析后自动触发知识图谱构建'
|
||||
ja_JP: 'ファイル取り込み後に自動的にナレッジグラフを構築'
|
||||
type: boolean
|
||||
default: false
|
||||
required: false
|
||||
- name: auto_raptor
|
||||
label:
|
||||
en_US: Auto Build RAPTOR
|
||||
zh_Hans: 自动构建 RAPTOR
|
||||
ja_JP: RAPTOR 自動構築
|
||||
description:
|
||||
en_US: 'Automatically trigger RAPTOR hierarchical summarization after file ingestion'
|
||||
zh_Hans: '文件上传解析后自动触发 RAPTOR 层次摘要构建'
|
||||
ja_JP: 'ファイル取り込み後に自動的に RAPTOR 階層要約を構築'
|
||||
type: boolean
|
||||
default: false
|
||||
required: false
|
||||
retrieval_schema:
|
||||
- name: vector_similarity_weight
|
||||
label:
|
||||
en_US: Vector Similarity Weight
|
||||
zh_Hans: 向量相似度权重
|
||||
ja_JP: ベクトル類似性の重み
|
||||
description:
|
||||
en_US: 'Balance between keyword and vector search (0=keyword only, 1=vector only, default 0.3)'
|
||||
zh_Hans: '关键词与向量检索的平衡（0=纯关键词，1=纯向量，默认 0.3）'
|
||||
ja_JP: 'キーワードとベクトル検索のバランス（0=キーワードのみ、1=ベクトルのみ、デフォルト 0.3）'
|
||||
type: float
|
||||
default: 0.3
|
||||
required: false
|
||||
- name: top_k
|
||||
label:
|
||||
en_US: Top K
|
||||
zh_Hans: 候选数量
|
||||
ja_JP: 候補数
|
||||
description:
|
||||
en_US: 'Maximum number of candidates to consider before pagination'
|
||||
zh_Hans: '分页前考虑的最大候选数量'
|
||||
ja_JP: 'ページネーション前に考慮する最大候補数'
|
||||
type: integer
|
||||
default: 1024
|
||||
required: false
|
||||
- name: page_size
|
||||
label:
|
||||
en_US: Page Size
|
||||
zh_Hans: 返回数量
|
||||
ja_JP: ページサイズ
|
||||
description:
|
||||
en_US: 'Number of results to return'
|
||||
zh_Hans: '返回的结果数量'
|
||||
ja_JP: '返す結果の数'
|
||||
type: integer
|
||||
default: 30
|
||||
required: false
|
||||
- name: similarity_threshold
|
||||
label:
|
||||
en_US: Similarity Threshold
|
||||
zh_Hans: 相似度阈值
|
||||
ja_JP: 類似度しきい値
|
||||
description:
|
||||
en_US: 'Minimum similarity score (0-1)'
|
||||
zh_Hans: '最小相似度分数 (0-1)'
|
||||
ja_JP: '最小類似性スコア (0-1)'
|
||||
type: float
|
||||
default: 0.2
|
||||
required: false
|
||||
- name: keyword
|
||||
label:
|
||||
en_US: LLM Keyword Extraction
|
||||
zh_Hans: LLM 关键词提取
|
||||
ja_JP: LLM キーワード抽出
|
||||
description:
|
||||
en_US: 'Use LLM to extract keywords from query to enhance retrieval (increases response time)'
|
||||
zh_Hans: '使用 LLM 从查询中提取关键词以增强检索（会增加响应时间）'
|
||||
ja_JP: 'LLM でクエリからキーワードを抽出して検索を強化（応答時間が増加）'
|
||||
type: boolean
|
||||
default: false
|
||||
required: false
|
||||
- name: rerank_id
|
||||
label:
|
||||
en_US: Rerank Model ID
|
||||
zh_Hans: 重排序模型 ID
|
||||
ja_JP: リランクモデル ID
|
||||
description:
|
||||
en_US: 'Rerank model ID configured in RAGFlow (e.g., BAAI/bge-reranker-v2-m3). Leave empty to disable'
|
||||
zh_Hans: '在 RAGFlow 中配置的重排序模型 ID（如 BAAI/bge-reranker-v2-m3），留空则不启用'
|
||||
ja_JP: 'RAGFlow で設定されたリランクモデル ID（例：BAAI/bge-reranker-v2-m3）、空欄で無効'
|
||||
type: string
|
||||
default: ''
|
||||
required: false
|
||||
- name: use_kg
|
||||
label:
|
||||
en_US: Knowledge Graph
|
||||
zh_Hans: 知识图谱
|
||||
ja_JP: ナレッジグラフ
|
||||
description:
|
||||
en_US: 'Enable knowledge graph retrieval (increases response time)'
|
||||
zh_Hans: '启用知识图谱检索（会增加响应时间）'
|
||||
ja_JP: 'ナレッジグラフ検索を有効化（応答時間が増加）'
|
||||
type: boolean
|
||||
default: false
|
||||
required: false
|
||||
execution:
|
||||
python:
|
||||
path: engine.py
|
||||
attr: RAGFlowConnector
|
||||
@@ -0,0 +1,502 @@
|
||||
"""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. An unreadable fence record refuses the operation instead of
|
||||
being read as "no fence", and such a read is retried rather than cached. 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. A fence that cannot be persisted aborts the mutation before
|
||||
it is dispatched, and a fence is released only after the provider answer was
|
||||
confirmed and its Host record written.
|
||||
"""
|
||||
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 operation is refused.
|
||||
|
||||
An unreadable fence record cannot be told apart from a knowledge base that
|
||||
is fenced after an ambiguous mutation, so an unknown fence state fails
|
||||
closed. The read is not cached: the next attempt re-reads the durable state,
|
||||
and nothing has been dispatched while the state stays unknown.
|
||||
"""
|
||||
|
||||
|
||||
class FenceLifecycleError(RuntimeError):
|
||||
"""A mutation was aborted before dispatch because its fence is not in place.
|
||||
|
||||
Both failures below mean "this call dispatched nothing": the pre-dispatch
|
||||
protection could not be persisted, or an earlier unresolved mutation is
|
||||
still fenced. Call sites must report them instead of folding them into an
|
||||
ordinary provider failure.
|
||||
"""
|
||||
|
||||
|
||||
class FencePersistError(FenceLifecycleError):
|
||||
"""The pre-dispatch fence record could not be persisted.
|
||||
|
||||
Raised by ``fence`` when the Host write of the fence record fails: the fence
|
||||
is not in place, so ``dispatched`` aborts before its request instead of
|
||||
dispatching a mutation whose protection does not exist.
|
||||
"""
|
||||
|
||||
|
||||
class FencedKnowledgeBaseError(FenceLifecycleError):
|
||||
"""A mutation was refused: an earlier mutation's outcome is unresolved.
|
||||
|
||||
Raised when a serialized operation, or a dispatch inside one, is attempted
|
||||
while the knowledge base (or the whole installation) is still fenced. A
|
||||
later successful operation therefore never clears an earlier unresolved
|
||||
operation's fence.
|
||||
"""
|
||||
|
||||
|
||||
def is_ambiguous_http_failure(exc):
|
||||
"""Whether a failed provider call had an unknown outcome.
|
||||
|
||||
Only a rejection the protocol makes unambiguous releases the caller's
|
||||
protection: a 4xx response means the provider received and refused the
|
||||
request, and a connection that was never established means nothing was
|
||||
dispatched. A 5xx or gateway response is NOT proof that the request was not
|
||||
applied — the provider or an intermediary may have committed the change and
|
||||
then failed — so it stays ambiguous together with timeouts, lost connections
|
||||
and unparsable responses.
|
||||
"""
|
||||
if isinstance(exc, httpx.HTTPStatusError):
|
||||
response = getattr(exc, "response", None)
|
||||
return response is None or response.status_code >= 500
|
||||
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 ``knowledge_base_id`` or ``get_collection_id()``
|
||||
(ingestion passes the whole context). The knowledge-base id is preferred: the
|
||||
collection id is a vector-store identity that need not equal it, and every
|
||||
other entry point of the connector is addressed by the knowledge-base id, so
|
||||
locking and fencing on the collection id would split one knowledge base into
|
||||
two independently protected scopes. 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():
|
||||
identity = getattr(value, 'knowledge_base_id', None)
|
||||
if isinstance(identity, str) and identity:
|
||||
return identity
|
||||
getter = getattr(value, 'get_collection_id', None)
|
||||
if callable(getter):
|
||||
try:
|
||||
identity = getter()
|
||||
except Exception:
|
||||
continue
|
||||
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()
|
||||
# binding -> single-flight lock and mutation counter for the fence read.
|
||||
self._fence_load_locks: dict = {}
|
||||
self._fence_generation: dict = {}
|
||||
|
||||
def _bump_fence_generation(self, binding):
|
||||
self._fence_generation[binding] = self._fence_generation.get(binding, 0) + 1
|
||||
|
||||
@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, or refuse the operation.
|
||||
|
||||
A failed read means the pending ambiguity of the installation is
|
||||
unknown, so it must not be mistaken for "no fences": the operation is
|
||||
refused with ``FenceStateUnavailableError`` and nothing is dispatched.
|
||||
The binding is not marked as loaded, so the next attempt re-reads and
|
||||
the knowledge base becomes usable again once storage answers. The read
|
||||
is single-flight per installation, and a snapshot read before a sibling
|
||||
operation changed the fence state is discarded instead of installed.
|
||||
"""
|
||||
if binding in self._fence_loaded:
|
||||
return
|
||||
lock = self._fence_load_locks.get(binding)
|
||||
if lock is None:
|
||||
lock = self._fence_load_locks[binding] = asyncio.Lock()
|
||||
async with lock:
|
||||
if binding in self._fence_loaded:
|
||||
return
|
||||
for _ in range(8):
|
||||
generation = self._fence_generation.get(binding, 0)
|
||||
snapshot = await self._read_fences(plugin, binding)
|
||||
if generation == self._fence_generation.get(binding, 0):
|
||||
self._fenced_keys(binding).update(snapshot)
|
||||
self._fence_loaded.add(binding)
|
||||
return
|
||||
raise FenceStateUnavailableError(
|
||||
f"Fence state for {binding!r} kept changing while it was being read"
|
||||
)
|
||||
|
||||
async def _read_fences(self, plugin, binding) -> set:
|
||||
"""Read the persisted fence keys for one installation.
|
||||
|
||||
An unreadable marker counts as fenced and a failed read refuses the
|
||||
operation with ``FenceStateUnavailableError``; the read is never cached.
|
||||
"""
|
||||
try:
|
||||
keys = await plugin.get_plugin_storage_keys()
|
||||
snapshot = set()
|
||||
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:
|
||||
snapshot.add(key)
|
||||
return snapshot
|
||||
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 read persisted fences for {binding!r}; pending ambiguous "
|
||||
f"mutations are unknown: {exc!r}"
|
||||
) from exc
|
||||
|
||||
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 FencedKnowledgeBaseError(
|
||||
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 one remote mutation under a pre-persisted fence.
|
||||
|
||||
``operation`` must cover the whole mutation: the provider request, the
|
||||
response validation and the durable result/mapping write. The fence is
|
||||
persisted before the request goes out — a failed persist aborts here,
|
||||
before any remote call — and it is released only once ``operation`` has
|
||||
returned, so an ambiguous outcome or a caller cancellation keeps the
|
||||
knowledge base fenced for reconciliation. A knowledge base that is
|
||||
already fenced is refused before anything is written or dispatched, so a
|
||||
later successful operation never clears an earlier unresolved
|
||||
operation's fence.
|
||||
"""
|
||||
binding = self.binding(plugin)
|
||||
if self.is_fenced(binding, kb_identity):
|
||||
# The enclosing serialized call loaded and validated the fence state,
|
||||
# so this is a real unresolved mutation rather than a stale view.
|
||||
raise FencedKnowledgeBaseError(
|
||||
f"Knowledge base {kb_identity!r} fenced by an unresolved mutation; "
|
||||
"reconcile with the provider before dispatching another mutation"
|
||||
)
|
||||
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):
|
||||
"""Persist one fence, then mark that knowledge base in this process.
|
||||
|
||||
The durable record is written first and a failed write is fatal
|
||||
(``FencePersistError``), so a caller that fences before dispatching a
|
||||
mutation aborts before its request instead of dispatching under a fence
|
||||
that does not exist. Only a persisted fence sets the in-process marker,
|
||||
so a failed write and a caller cancellation both leave no marker behind.
|
||||
The residual race — the Host write landed but its acknowledgement was
|
||||
lost — can make a later read load a fence for an operation that was
|
||||
never dispatched; that is the conservative direction and is accepted.
|
||||
"""
|
||||
binding = self.binding(plugin)
|
||||
fence_key = _kb_fence_key(kb_identity)
|
||||
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 cancelled:
|
||||
raise asyncio.CancelledError
|
||||
if failure is not None:
|
||||
raise FencePersistError(
|
||||
f"Could not persist fence for {kb_identity!r}: {failure!r}"
|
||||
) from failure
|
||||
self._fenced_keys(binding).add(fence_key)
|
||||
self._bump_fence_generation(binding)
|
||||
|
||||
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)
|
||||
self._bump_fence_generation(binding)
|
||||
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:
|
||||
try:
|
||||
await self.fence(plugin, kb_identity, f'Host write outcome unknown: {failure!r}')
|
||||
except FencePersistError:
|
||||
# The write failure is what the caller must see; a fence that
|
||||
# could not be persisted is reported alongside it, not instead.
|
||||
logger.error(
|
||||
"Could not persist fence for %r after a failed Host write",
|
||||
kb_identity, exc_info=True,
|
||||
)
|
||||
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)
|
||||
self._fence_load_locks.pop(binding, None)
|
||||
self._fence_generation.pop(binding, None)
|
||||
|
||||
|
||||
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.
|
||||
Configuration and document mappings are keyed by the Host knowledge-base id;
|
||||
a record written under another knowledge base's id is never adopted, and only
|
||||
an exact mapping for the requested knowledge base is honoured.
|
||||
"""
|
||||
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):
|
||||
"""Resolve this knowledge base's mapping, never another knowledge base's.
|
||||
|
||||
A record written under another knowledge base's id is never adopted:
|
||||
attributing historic collection-keyed data requires explicit
|
||||
reconciliation. An absent exact mapping therefore refuses the operation,
|
||||
and this path performs no writes.
|
||||
"""
|
||||
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 only for a record this
|
||||
# knowledge base itself wrote, 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