Files
wehub-resource-sync fed8b2eed7
Backend release / release (push) Waiting to run
Bandit Security Scan / bandit_scan (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / manifest (push) Blocked by required conditions
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Waiting to run
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Waiting to run
Build and push DocsGPT FE Docker image for development / manifest (push) Blocked by required conditions
Python linting / ruff (push) Waiting to run
Run python tests with pytest / Run tests and count coverage (3.12) (push) Waiting to run
React Widget Build / build (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:28:29 +08:00

2591 lines
99 KiB
Python
Executable File

import datetime
import json
import logging
import mimetypes
import os
import shutil
import string
import tempfile
import threading
from typing import Any, Dict
import zipfile
import uuid
from collections import Counter
from urllib.parse import urljoin, urlsplit
import requests
from application.core.settings import settings
from application.events.publisher import publish_user_event
from application.parser.chunking_creator import ChunkerCreator
from application.parser.connectors.connector_creator import ConnectorCreator
from application.parser.embedding_pipeline import (
assert_index_complete,
embed_and_store_documents,
)
from application.parser.file.bulk import SimpleDirectoryReader, get_default_file_extractor
from application.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS
from application.parser.remote.remote_creator import (
RemoteCreator,
normalize_remote_data,
)
from application.parser.schema.base import Document
from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.attachments import AttachmentsRepository
from application.storage.db.repositories.ingest_chunk_progress import (
IngestChunkProgressRepository,
)
from application.storage.db.repositories.sources import SourcesRepository
from application.storage.db.repositories.wiki_pages import (
WikiPagesRepository,
_content_hash,
rebuild_wiki_directory_structure,
)
from application.storage.db.session import db_readonly, db_session
from application.storage.db.source_config import SourceConfig
from application.storage.storage_creator import StorageCreator
from application.utils import count_tokens_docs, num_tokens_from_string, safe_filename
# Constants
MIN_TOKENS = 150
MAX_TOKENS = 1250
RECURSION_DEPTH = 2
INGEST_HEARTBEAT_INTERVAL_SECONDS = 30
def graph_extraction_key(source_id, updated_at) -> str:
"""Build the extract_graph idempotency key for a source's current state.
The key embeds ``updated_at`` (the trigger-maintained column) so it
changes whenever re-extraction is warranted — a re-ingest or re-enable
bumps ``updated_at`` and yields a fresh key, bypassing the 24h completed
cache so the worker re-runs. Two enqueues for the same state share a key
and dedupe. ``updated_at`` is stringified deterministically.
"""
return f"extract-graph:{source_id}:{updated_at}"
def _source_updated_at(source) -> str:
"""Return the source's ``updated_at`` (falling back to ``date``) as a string."""
if not source:
return ""
stamp = source.get("updated_at") or source.get("date")
return str(stamp) if stamp is not None else ""
def _reset_graph_for_source(source_id) -> None:
"""Drop a source's existing graph so a re-enable/re-ingest rebuilds from scratch.
Clears nodes, edges, node→chunk links and the ingest checkpoint. Run at the
enqueue site (not inside the worker) so a broker redelivery of an interrupted
build still resumes from its checkpoint rather than restarting from zero.
"""
from application.graphrag.store import GraphStore
GraphStore().delete_by_source(str(source_id))
def _publish_graph_event(user, source_id, event_type, payload) -> None:
"""Publish a graph-extraction SSE event, scoped to the source. Never raises."""
if not user:
return
try:
publish_user_event(
user,
event_type,
payload,
scope={"kind": "source", "id": str(source_id)},
)
except Exception as e:
logging.debug(f"Failed to publish graph event {event_type}: {e}")
def _maybe_enqueue_graph_extraction(cfg, source_id, user):
"""Reset and re-enqueue graph extraction after embed for a graphrag source.
The graph lights up asynchronously once chunks are embedded, so ClassicRAG
works immediately. A no-op for non-graphrag sources or when GraphRAG is
unavailable. The prior graph is cleared first so a re-ingest rebuilds rather
than accumulating stale nodes. The work is isolated so a broker hiccup can
never fail an otherwise-successful ingest.
"""
if cfg.kind != "graphrag":
return
from application.graphrag import graphrag_available
if not graphrag_available():
return
source_id = str(source_id)
try:
from application.api.user.tasks import extract_graph
with db_readonly() as conn:
source = SourcesRepository(conn).get_any(source_id, user)
_reset_graph_for_source(source_id)
key = graph_extraction_key(source_id, _source_updated_at(source))
extract_graph.delay(source_id, user, idempotency_key=key)
except Exception as e:
logging.warning(
f"Failed to enqueue graph extraction for {source_id}: {e}",
exc_info=True,
)
# Re-exported here for backward-compatible imports
# (``from application.worker import _derive_source_id`` /
# ``DOCSGPT_INGEST_NAMESPACE``) from tests and any other in-tree callers.
# New code should import from ``application.storage.db.source_ids``
# directly to avoid pulling this Celery worker module into the API
# process at import time.
from application.storage.db.source_ids import ( # noqa: E402, F401
DOCSGPT_INGEST_NAMESPACE,
derive_source_id as _derive_source_id,
)
def _ingest_heartbeat_loop(source_id, stop_event, interval=INGEST_HEARTBEAT_INTERVAL_SECONDS):
"""Bump ``ingest_chunk_progress.last_updated`` until ``stop_event`` is set."""
while not stop_event.wait(interval):
try:
with db_session() as conn:
IngestChunkProgressRepository(conn).bump_heartbeat(source_id)
except Exception as e:
logging.warning(
f"Heartbeat failed for {source_id}: {e}", exc_info=True
)
def _start_ingest_heartbeat(source_id):
"""Spawn the heartbeat daemon and return ``(thread, stop_event)``."""
stop_event = threading.Event()
thread = threading.Thread(
target=_ingest_heartbeat_loop,
args=(str(source_id), stop_event),
daemon=True,
name=f"ingest-heartbeat-{source_id}",
)
thread.start()
return thread, stop_event
def _stop_ingest_heartbeat(thread, stop_event):
"""Signal the heartbeat daemon to exit and wait briefly for it."""
if stop_event is not None:
stop_event.set()
if thread is not None:
thread.join(timeout=5)
def _make_parse_progress_callback(task, user, source_id, start_pct, end_pct):
"""Build a ``load_data`` callback mapping parse progress to
``[start_pct, end_pct]`` via ``update_state`` + a throttled
``stage='parsing'`` SSE event.
"""
span = end_pct - start_pct
source_id_str = str(source_id)
state = {"last_pct": -1}
def _callback(files_done, total_files):
if not total_files:
return
pct = start_pct + int((files_done / total_files) * span)
task.update_state(
state="PROGRESS",
meta={"current": pct, "status": "Parsing files"},
)
if user and pct > state["last_pct"]:
publish_user_event(
user,
"source.ingest.progress",
{
"current": pct,
"total": total_files,
"files_done": files_done,
"stage": "parsing",
},
scope={"kind": "source", "id": source_id_str},
)
state["last_pct"] = pct
return _callback
# Define a function to extract metadata from a given filename.
def metadata_from_filename(title):
return {"title": title}
def _normalize_file_name_map(file_name_map):
if not file_name_map:
return {}
if isinstance(file_name_map, str):
try:
file_name_map = json.loads(file_name_map)
except Exception:
return {}
return file_name_map if isinstance(file_name_map, dict) else {}
def _get_display_name(file_name_map, rel_path):
if not file_name_map or not rel_path:
return None
if rel_path in file_name_map:
return file_name_map[rel_path]
base_name = os.path.basename(rel_path)
return file_name_map.get(base_name)
def _apply_display_names_to_structure(structure, file_name_map, prefix=""):
if not isinstance(structure, dict) or not file_name_map:
return structure
for name, node in structure.items():
if isinstance(node, dict) and "type" in node and "size_bytes" in node:
rel_path = f"{prefix}/{name}" if prefix else name
display_name = _get_display_name(file_name_map, rel_path)
if display_name:
node["display_name"] = display_name
elif isinstance(node, dict):
next_prefix = f"{prefix}/{name}" if prefix else name
_apply_display_names_to_structure(node, file_name_map, next_prefix)
return structure
def _download_source_files_to_dir(storage, source_file_path, temp_dir):
"""Mirror a source's stored files into ``temp_dir``, preserving structure."""
if not storage.is_directory(source_file_path):
return
for storage_file_path in storage.list_files(source_file_path):
if storage.is_directory(storage_file_path):
continue
rel_path = os.path.relpath(storage_file_path, source_file_path)
local_file_path = os.path.join(temp_dir, rel_path)
os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
try:
file_data = storage.get_file(storage_file_path)
with open(local_file_path, "wb") as f:
f.write(file_data.read())
except Exception as e:
logging.error(f"Error downloading file {storage_file_path}: {e}")
continue
# Define a function to generate a random string of a given length.
def generate_random_string(length):
return "".join([string.ascii_letters[i % 52] for i in range(length)])
current_dir = os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
# Zip extraction security limits
MAX_UNCOMPRESSED_SIZE = 500 * 1024 * 1024 # 500 MB max uncompressed size
MAX_FILE_COUNT = 10000 # Maximum number of files to extract
MAX_COMPRESSION_RATIO = 100 # Maximum compression ratio (to detect zip bombs)
class ZipExtractionError(Exception):
"""Raised when zip extraction fails due to security constraints."""
pass
def _is_path_safe(base_path: str, target_path: str) -> bool:
"""
Check if target_path is safely within base_path (prevents zip slip attacks).
Args:
base_path: The base directory where extraction should occur.
target_path: The full path where a file would be extracted.
Returns:
True if the path is safe, False otherwise.
"""
# Resolve to absolute paths and check containment
base_resolved = os.path.realpath(base_path)
target_resolved = os.path.realpath(target_path)
return target_resolved.startswith(base_resolved + os.sep) or target_resolved == base_resolved
def _validate_zip_safety(zip_path: str, extract_to: str) -> None:
"""
Validate a zip file for security issues before extraction.
Checks for:
- Zip bombs (excessive compression ratio or uncompressed size)
- Too many files
- Path traversal attacks (zip slip)
Args:
zip_path: Path to the zip file.
extract_to: Destination directory.
Raises:
ZipExtractionError: If the zip file fails security validation.
"""
try:
with zipfile.ZipFile(zip_path, "r") as zip_ref:
# Get compressed size
compressed_size = os.path.getsize(zip_path)
# Calculate total uncompressed size and file count
total_uncompressed = 0
file_count = 0
for info in zip_ref.infolist():
file_count += 1
# Check file count limit
if file_count > MAX_FILE_COUNT:
raise ZipExtractionError(
f"Zip file contains too many files (>{MAX_FILE_COUNT}). "
"This may be a zip bomb attack."
)
# Accumulate uncompressed size
total_uncompressed += info.file_size
# Check total uncompressed size
if total_uncompressed > MAX_UNCOMPRESSED_SIZE:
raise ZipExtractionError(
f"Zip file uncompressed size exceeds limit "
f"({total_uncompressed / (1024*1024):.1f} MB > "
f"{MAX_UNCOMPRESSED_SIZE / (1024*1024):.1f} MB). "
"This may be a zip bomb attack."
)
# Check for path traversal (zip slip)
target_path = os.path.join(extract_to, info.filename)
if not _is_path_safe(extract_to, target_path):
raise ZipExtractionError(
f"Zip file contains path traversal attempt: {info.filename}"
)
# Check compression ratio (only if compressed size is meaningful)
if compressed_size > 0 and total_uncompressed > 0:
compression_ratio = total_uncompressed / compressed_size
if compression_ratio > MAX_COMPRESSION_RATIO:
raise ZipExtractionError(
f"Zip file has suspicious compression ratio ({compression_ratio:.1f}:1 > "
f"{MAX_COMPRESSION_RATIO}:1). This may be a zip bomb attack."
)
except zipfile.BadZipFile as e:
raise ZipExtractionError(f"Invalid or corrupted zip file: {e}")
def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5):
"""
Recursively extract zip files with security protections.
Security measures:
- Limits recursion depth to prevent infinite loops
- Validates uncompressed size to prevent zip bombs
- Limits number of files to prevent resource exhaustion
- Checks compression ratio to detect zip bombs
- Validates paths to prevent zip slip attacks
Args:
zip_path (str): Path to the zip file to be extracted.
extract_to (str): Destination path for extracted files.
current_depth (int): Current depth of recursion.
max_depth (int): Maximum allowed depth of recursion to prevent infinite loops.
"""
if current_depth > max_depth:
logging.warning(f"Reached maximum recursion depth of {max_depth}")
return
try:
# Validate zip file safety before extraction
_validate_zip_safety(zip_path, extract_to)
# Safe to extract
with zipfile.ZipFile(zip_path, "r") as zip_ref:
zip_ref.extractall(extract_to)
os.remove(zip_path) # Remove the zip file after extracting
except ZipExtractionError as e:
logging.error(f"Zip security validation failed for {zip_path}: {e}")
# Remove the potentially malicious zip file
try:
os.remove(zip_path)
except OSError:
pass
return
except Exception as e:
logging.error(f"Error extracting zip file {zip_path}: {e}", exc_info=True)
return
# Check for nested zip files and extract them
for root, dirs, files in os.walk(extract_to):
for file in files:
if file.endswith(".zip"):
# If a nested zip file is found, extract it recursively
file_path = os.path.join(root, file)
extract_zip_recursive(file_path, root, current_depth + 1, max_depth)
def download_file(url, params, dest_path):
try:
response = requests.get(url, params=params, timeout=100)
response.raise_for_status()
with open(dest_path, "wb") as f:
f.write(response.content)
except requests.RequestException as e:
logging.error(f"Error downloading file: {e}")
raise
def upload_index(full_path, file_data):
files = None
try:
headers = {}
if settings.INTERNAL_KEY:
headers["X-Internal-Key"] = settings.INTERNAL_KEY
if settings.VECTOR_STORE == "faiss":
faiss_path = full_path + "/index.faiss"
pkl_path = full_path + "/index.pkl"
if not os.path.exists(faiss_path):
logging.error(f"FAISS index file not found: {faiss_path}")
raise FileNotFoundError(f"FAISS index file not found: {faiss_path}")
if not os.path.exists(pkl_path):
logging.error(f"FAISS pickle file not found: {pkl_path}")
raise FileNotFoundError(f"FAISS pickle file not found: {pkl_path}")
files = {
"file_faiss": open(faiss_path, "rb"),
"file_pkl": open(pkl_path, "rb"),
}
response = requests.post(
urljoin(settings.API_URL, "/api/upload_index"),
files=files,
data=file_data,
headers=headers,
timeout=100,
)
else:
response = requests.post(
urljoin(settings.API_URL, "/api/upload_index"),
data=file_data,
headers=headers,
timeout=100,
)
response.raise_for_status()
except (requests.RequestException, FileNotFoundError) as e:
logging.error(f"Error uploading index: {e}")
raise
finally:
if settings.VECTOR_STORE == "faiss" and files is not None:
for file in files.values():
file.close()
# Define the main function for ingesting and processing documents.
def ingest_worker(
self,
directory,
formats,
job_name,
file_path,
filename,
user,
retriever="classic",
file_name_map=None,
config=None,
idempotency_key=None,
source_id=None,
):
"""
Ingest and process documents.
Args:
self: Reference to the instance of the task.
directory (str): Specifies the directory for ingesting ('inputs' or 'temp').
formats (list of str): List of file extensions to consider for ingestion (e.g., [".rst", ".md"]).
job_name (str): Name of the job for this ingestion task (original, unsanitized).
file_path (str): Complete file path to use consistently throughout the pipeline.
filename (str): Original unsanitized filename provided by the user.
user (str): Identifier for the user initiating the ingestion (original, unsanitized).
retriever (str): Type of retriever to use for processing the documents.
file_name_map (dict|str|None): Optional mapping of safe relative paths to original filenames.
config (dict|None): Per-source ``SourceConfig`` dict. ``None``/``{}`` →
classic defaults (byte-identical to prior behavior).
idempotency_key (str|None): When provided, the ``source_id`` is derived
deterministically from the key so a retried task reuses the same
source row instead of duplicating it.
source_id (str|None): UUID minted by the HTTP route and returned in
its response. When supplied, the worker uses it verbatim so SSE
envelopes carry the same id the frontend already has — required
for non-idempotent uploads where the route can't predict
``_derive_source_id(idempotency_key)``.
Returns:
dict: Information about the completed ingestion task, including input parameters and a "limited" flag.
"""
input_files = None
recursive = True
limit = None
exclude = True
sample = False
storage = StorageCreator.get_storage()
logging.info(f"Ingest path: {file_path}", extra={"user": user, "job": job_name})
# Source id resolution order:
# 1. Caller-supplied ``source_id`` (HTTP route minted + returned to
# the frontend) — keeps the route response and the SSE event
# payloads in lockstep on the non-idempotent path.
# 2. Deterministic uuid5 from ``idempotency_key`` — retried tasks
# reuse the original source row instead of duplicating it.
# 3. Fresh uuid4 (caller has neither) — opaque, single-shot only.
if source_id:
source_uuid = uuid.UUID(source_id)
else:
source_uuid = _derive_source_id(idempotency_key)
source_id_for_events = str(source_uuid)
# Only emit ``queued`` on the original attempt. Celery retries re-run
# the body, and re-publishing here would oscillate the toast through
# ``queued`` again between ``failed`` and ``completed``.
if self.request.retries == 0:
publish_user_event(
user,
"source.ingest.queued",
{
"job_name": job_name,
"filename": filename,
"source_id": source_id_for_events,
"operation": "upload",
},
scope={"kind": "source", "id": source_id_for_events},
)
# Wrap the entire body in try/except so a failure between the
# ``queued`` publish above and the inner work (e.g. tempdir
# creation, OS-level resource exhaustion) still emits a terminal
# ``failed`` event rather than leaving the toast wedged on
# 'training' until the polling fallback rescues it 30s later.
try:
with tempfile.TemporaryDirectory() as temp_dir:
os.makedirs(temp_dir, exist_ok=True)
if storage.is_directory(file_path):
# Handle directory case
logging.info(f"Processing directory: {file_path}")
files_list = storage.list_files(file_path)
for storage_file_path in files_list:
if storage.is_directory(storage_file_path):
continue
# Create relative path structure in temp directory
rel_path = os.path.relpath(storage_file_path, file_path)
local_file_path = os.path.join(temp_dir, rel_path)
os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
# Download file
try:
file_data = storage.get_file(storage_file_path)
with open(local_file_path, "wb") as f:
f.write(file_data.read())
except Exception as e:
logging.error(
f"Error downloading file {storage_file_path}: {e}"
)
continue
else:
# Handle single file case
temp_filename = os.path.basename(file_path)
temp_file_path = os.path.join(temp_dir, temp_filename)
file_data = storage.get_file(file_path)
with open(temp_file_path, "wb") as f:
f.write(file_data.read())
# Handle zip files
if temp_filename.endswith(".zip"):
logging.info(f"Extracting zip file: {temp_filename}")
extract_zip_recursive(
temp_file_path,
temp_dir,
current_depth=0,
max_depth=RECURSION_DEPTH,
)
self.update_state(state="PROGRESS", meta={"current": 1})
if sample:
logging.info(f"Sample mode enabled. Using {limit} documents.")
reader = SimpleDirectoryReader(
input_dir=temp_dir,
input_files=input_files,
recursive=recursive,
required_exts=formats,
exclude_hidden=exclude,
file_metadata=metadata_from_filename,
)
# Parsing/OCR owns 1-50% of the bar; embedding takes 50-100%.
raw_docs = reader.load_data(
progress_callback=_make_parse_progress_callback(
self, user, source_uuid, start_pct=1, end_pct=50,
)
)
directory_structure = getattr(reader, "directory_structure", {})
logging.info(f"Directory structure from reader: {directory_structure}")
file_name_map = _normalize_file_name_map(file_name_map)
if file_name_map:
for doc in raw_docs:
extra_info = getattr(doc, "extra_info", None)
if not isinstance(extra_info, dict):
continue
rel_path = extra_info.get("source") or extra_info.get("file_path")
display_name = _get_display_name(file_name_map, rel_path)
if display_name:
display_name = str(display_name)
extra_info["filename"] = display_name
extra_info["file_name"] = display_name
extra_info["title"] = display_name
directory_structure = _apply_display_names_to_structure(
directory_structure, file_name_map
)
cfg = SourceConfig.parse(config)
chunker = ChunkerCreator.create_chunker(
cfg.chunking.strategy,
chunking_strategy=cfg.chunking.strategy,
max_tokens=cfg.chunking.max_tokens,
min_tokens=cfg.chunking.min_tokens,
duplicate_headers=cfg.chunking.duplicate_headers,
)
raw_docs = chunker.chunk(documents=raw_docs)
docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]
vector_store_path = os.path.join(temp_dir, "vector_store")
os.makedirs(vector_store_path, exist_ok=True)
heartbeat_thread, heartbeat_stop = _start_ingest_heartbeat(source_uuid)
try:
embed_and_store_documents(
docs, vector_store_path, source_uuid, self,
attempt_id=getattr(self.request, "id", None),
user_id=user,
progress_start=50, progress_end=100,
)
finally:
_stop_ingest_heartbeat(heartbeat_thread, heartbeat_stop)
# Defense-in-depth: chunk-progress is the authoritative
# record of how many chunks landed; mismatch raises so the
# task fails loud rather than caching a partial index.
assert_index_complete(source_uuid)
tokens = count_tokens_docs(docs)
self.update_state(state="PROGRESS", meta={"current": 100})
if sample:
for i in range(min(5, len(raw_docs))):
logging.info(f"Sample document {i}: {raw_docs[i]}")
file_data = {
"name": job_name,
"file": filename,
"user": user,
"tokens": tokens,
"retriever": retriever,
"id": source_id_for_events,
"type": "local",
"file_path": file_path,
"directory_structure": json.dumps(directory_structure),
}
if file_name_map:
file_data["file_name_map"] = json.dumps(file_name_map)
if config:
file_data["config"] = json.dumps(config)
upload_index(vector_store_path, file_data)
publish_user_event(
user,
"source.ingest.completed",
{
"source_id": source_id_for_events,
"filename": filename,
"tokens": tokens,
"operation": "upload",
# Forward-looking contract: ``limited`` is always
# ``False`` today but is carried on the wire so a
# future token-cap detection path can flip it and
# the frontend slice / UploadToast already react.
"limited": False,
},
scope={"kind": "source", "id": source_id_for_events},
)
_maybe_enqueue_graph_extraction(cfg, source_id_for_events, user)
except Exception as e:
logging.error(f"Error in ingest_worker: {e}", exc_info=True)
publish_user_event(
user,
"source.ingest.failed",
{
"source_id": source_id_for_events,
"filename": filename,
"operation": "upload",
"error": str(e)[:1024],
},
scope={"kind": "source", "id": source_id_for_events},
)
raise
return {
"directory": directory,
"formats": formats,
"name_job": job_name, # Use original job_name
"filename": filename,
"user": user, # Use original user
"limited": False,
}
def reingest_source_worker(self, source_id, user):
"""
Re-ingestion worker that handles incremental updates by:
1. Adding chunks from newly added files
2. Removing chunks from deleted files
Args:
self: Task instance
source_id: ID of the source to re-ingest
user: User identifier
Returns:
dict: Information about the re-ingestion task
Note:
Reingest does its own ``vector_store.add_chunk`` work rather
than going through ``embed_and_store_documents`` so it does
*not* emit per-percent SSE progress events — only ``queued``,
``completed`` (carrying ``chunks_added`` / ``chunks_deleted``),
or ``failed``. v1 limitation; revisit if reingest gains a
progress-driven UI.
"""
# Declared at the function scope so the outer except can include
# ``name`` in the failed event payload when the failure happens
# after the source lookup. Empty string until the lookup succeeds.
source_name = ""
# Tracks inner-block failures so a ``completed`` event reflects
# partial-success accurately rather than masking it.
inner_warnings: list[str] = []
try:
from application.vectorstore.vector_creator import VectorCreator
self.update_state(
state="PROGRESS",
meta={"current": 10, "status": "Initializing re-ingestion scan"},
)
with db_readonly() as conn:
source = SourcesRepository(conn).get_any(source_id, user)
if not source:
raise ValueError(f"Source {source_id} not found or access denied")
source_id = str(source["id"])
source_name = source.get("name") or ""
# Re-chunk with the source's persisted config (classic defaults
# when absent/empty), keeping reingest consistent with ingest.
cfg = SourceConfig.parse(source.get("config"))
# Publish ``queued`` *after* canonicalising ``source_id`` so the
# event references the same id as the source row. Trade-off
# documented: a Celery-backend or PG-lookup hiccup before this
# publish means the toast may see only a ``failed`` event with
# no preceding ``queued`` — acceptable for v1 since both
# conditions also imply broader system trouble. Gate on first
# attempt only so Celery retries don't re-emit ``queued`` after
# a prior attempt already published ``failed``.
if self.request.retries == 0:
publish_user_event(
user,
"source.ingest.queued",
{
"source_id": source_id,
"name": source_name,
# ``filename`` labels the upload toast on auto-create.
"filename": source_name,
"operation": "reingest",
},
scope={"kind": "source", "id": source_id},
)
storage = StorageCreator.get_storage()
source_file_path = source.get("file_path", "")
file_name_map = _normalize_file_name_map(source.get("file_name_map"))
self.update_state(
state="PROGRESS", meta={"current": 20, "status": "Scanning current files"}
)
with tempfile.TemporaryDirectory() as temp_dir:
_download_source_files_to_dir(storage, source_file_path, temp_dir)
reader = SimpleDirectoryReader(
input_dir=temp_dir,
recursive=True,
required_exts=list(SUPPORTED_SOURCE_EXTENSIONS),
exclude_hidden=True,
file_metadata=metadata_from_filename,
)
reader.load_data()
directory_structure = reader.directory_structure
logging.info(
f"Directory structure built with token counts: {directory_structure}"
)
try:
old_directory_structure = source.get("directory_structure") or {}
if isinstance(old_directory_structure, str):
try:
old_directory_structure = json.loads(old_directory_structure)
except Exception:
old_directory_structure = {}
def _flatten_directory_structure(struct, prefix=""):
files = set()
if isinstance(struct, dict):
for name, meta in struct.items():
current_path = (
os.path.join(prefix, name) if prefix else name
)
if isinstance(meta, dict) and (
"type" in meta and "size_bytes" in meta
):
files.add(current_path)
elif isinstance(meta, dict):
files |= _flatten_directory_structure(
meta, current_path
)
return files
old_files = _flatten_directory_structure(old_directory_structure)
new_files = _flatten_directory_structure(directory_structure)
added_files = sorted(new_files - old_files)
removed_files = sorted(old_files - new_files)
if added_files:
logging.info(f"Files added since last ingest: {added_files}")
else:
logging.info("No files added since last ingest.")
if removed_files:
logging.info(f"Files removed since last ingest: {removed_files}")
else:
logging.info("No files removed since last ingest.")
except Exception as e:
logging.error(
f"Error comparing directory structures: {e}", exc_info=True
)
added_files = []
removed_files = []
try:
if not added_files and not removed_files:
logging.info("No changes detected.")
publish_user_event(
user,
"source.ingest.completed",
{
"source_id": source_id,
"name": source_name,
"filename": source_name,
"operation": "reingest",
"no_changes": True,
"chunks_added": 0,
"chunks_deleted": 0,
},
scope={"kind": "source", "id": source_id},
)
return {
"source_id": source_id,
"user": user,
"status": "no_changes",
"added_files": [],
"removed_files": [],
}
vector_store = VectorCreator.create_vectorstore(
settings.VECTOR_STORE,
source_id,
settings.EMBEDDINGS_KEY,
)
self.update_state(
state="PROGRESS",
meta={"current": 40, "status": "Processing file changes"},
)
# 1) Delete chunks from removed files
deleted = 0
if removed_files:
try:
for ch in vector_store.get_chunks() or []:
metadata = (
ch.get("metadata", {})
if isinstance(ch, dict)
else getattr(ch, "metadata", {})
)
raw_source = metadata.get("source")
source_file = str(raw_source) if raw_source else ""
if source_file in removed_files:
cid = ch.get("doc_id")
if cid:
try:
vector_store.delete_chunk(cid)
deleted += 1
except Exception as de:
logging.error(
f"Failed deleting chunk {cid}: {de}"
)
logging.info(
f"Deleted {deleted} chunks from {len(removed_files)} removed files"
)
except Exception as e:
logging.error(
f"Error during deletion of removed file chunks: {e}",
exc_info=True,
)
inner_warnings.append(
f"deletion failed: {str(e)[:200]}"
)
# 2) Add chunks from new files
added = 0
if added_files:
try:
# Build list of local files for added files only
added_local_files = []
for rel_path in added_files:
local_path = os.path.join(temp_dir, rel_path)
if os.path.isfile(local_path):
added_local_files.append(local_path)
if added_local_files:
reader_new = SimpleDirectoryReader(
input_files=added_local_files,
exclude_hidden=True,
errors="ignore",
file_metadata=metadata_from_filename,
)
raw_docs_new = reader_new.load_data()
chunker_new = ChunkerCreator.create_chunker(
cfg.chunking.strategy,
chunking_strategy=cfg.chunking.strategy,
max_tokens=cfg.chunking.max_tokens,
min_tokens=cfg.chunking.min_tokens,
duplicate_headers=cfg.chunking.duplicate_headers,
)
chunked_new = chunker_new.chunk(documents=raw_docs_new)
for (
file_path,
token_count,
) in reader_new.file_token_counts.items():
try:
rel_path = os.path.relpath(
file_path, start=temp_dir
)
path_parts = rel_path.split(os.sep)
current_dir = directory_structure
for part in path_parts[:-1]:
if part in current_dir and isinstance(
current_dir[part], dict
):
current_dir = current_dir[part]
else:
break
filename = path_parts[-1]
if filename in current_dir and isinstance(
current_dir[filename], dict
):
current_dir[filename][
"token_count"
] = token_count
logging.info(
f"Updated token count for {rel_path}: {token_count}"
)
except Exception as e:
logging.warning(
f"Could not update token count for {file_path}: {e}"
)
for d in chunked_new:
meta = dict(d.extra_info or {})
try:
raw_src = meta.get("source")
if isinstance(raw_src, str) and os.path.isabs(
raw_src
):
meta["source"] = os.path.relpath(
raw_src, start=temp_dir
)
except Exception:
pass
display_name = _get_display_name(
file_name_map, meta.get("source")
)
if display_name:
display_name = str(display_name)
meta["filename"] = display_name
meta["file_name"] = display_name
meta["title"] = display_name
vector_store.add_chunk(d.text, metadata=meta)
added += 1
logging.info(
f"Added {added} chunks from {len(added_files)} new files"
)
except Exception as e:
logging.error(
f"Error during ingestion of new files: {e}", exc_info=True
)
inner_warnings.append(
f"add failed: {str(e)[:200]}"
)
# 3) Update source directory structure timestamp
try:
total_tokens = sum(reader.file_token_counts.values())
directory_structure = _apply_display_names_to_structure(
directory_structure, file_name_map
)
now = datetime.datetime.now(datetime.timezone.utc)
with db_session() as conn:
SourcesRepository(conn).update(
source_id, user,
{
"directory_structure": directory_structure,
"date": now,
"tokens": total_tokens,
},
)
except Exception as e:
logging.error(
f"Error updating directory_structure in DB: {e}", exc_info=True
)
self.update_state(
state="PROGRESS",
meta={"current": 100, "status": "Re-ingestion completed"},
)
completed_payload: dict = {
"source_id": source_id,
"name": source_name,
"filename": source_name,
"operation": "reingest",
"chunks_added": added,
"chunks_deleted": deleted,
"tokens": int(total_tokens) if "total_tokens" in locals() else 0,
}
if inner_warnings:
# Surface the per-block failures so the toast can warn
# rather than claim a clean success.
completed_payload["warnings"] = inner_warnings
publish_user_event(
user,
"source.ingest.completed",
completed_payload,
scope={"kind": "source", "id": source_id},
)
_maybe_enqueue_graph_extraction(cfg, source_id, user)
return {
"source_id": source_id,
"user": user,
"status": "completed",
"added_files": added_files,
"removed_files": removed_files,
"chunks_added": added,
"chunks_deleted": deleted,
}
except Exception as e:
logging.error(
f"Error while processing file changes: {e}", exc_info=True
)
raise
except Exception as e:
logging.error(f"Error in reingest_source_worker: {e}", exc_info=True)
publish_user_event(
user,
"source.ingest.failed",
{
"source_id": str(source_id),
"name": source_name,
"filename": source_name,
"operation": "reingest",
"error": str(e)[:1024],
},
scope={"kind": "source", "id": str(source_id)},
)
raise
def remote_worker(
self,
source_data,
name_job,
user,
loader,
directory="temp",
retriever="classic",
sync_frequency="never",
operation_mode="upload",
doc_id=None,
config=None,
idempotency_key=None,
source_id=None,
):
safe_user = safe_filename(user)
full_path = os.path.join(directory, safe_user, uuid.uuid4().hex)
os.makedirs(full_path, exist_ok=True)
# Source id resolution order matches ``ingest_worker``:
# 1. ``operation_mode == "sync"`` reuses the existing source's ``doc_id``.
# 2. Caller-supplied ``source_id`` (the HTTP route minted it and
# already returned it to the frontend) — keeps the route
# response and the SSE event payloads in lockstep on the
# no-idempotency-key path.
# 3. Deterministic uuid5 from ``idempotency_key`` — retried tasks
# reuse the original source row instead of duplicating it.
# 4. Fresh uuid4 — opaque, single-shot only.
if operation_mode == "sync" and doc_id:
source_uuid = str(doc_id)
elif source_id:
source_uuid = uuid.UUID(source_id)
else:
source_uuid = _derive_source_id(idempotency_key)
source_id_for_events = str(source_uuid)
# Emit the queued event before any work that could fail (including
# ``update_state``) so the toast UI always sees a queued envelope
# before any subsequent failed event. Gated on first attempt so
# Celery retries don't re-emit ``queued`` after a prior ``failed``.
if self.request.retries == 0:
publish_user_event(
user,
"source.ingest.queued",
{
"source_id": source_id_for_events,
"job_name": name_job,
"loader": loader,
"operation": operation_mode,
},
scope={"kind": "source", "id": source_id_for_events},
)
# Wrap ``update_state`` plus the entire body so any pre-loader
# failure (Celery backend down, OS resource issue) still emits a
# terminal ``failed`` event rather than wedging the toast.
try:
self.update_state(state="PROGRESS", meta={"current": 1})
logging.info("Initializing remote loader with type: %s", loader)
remote_loader = RemoteCreator.create_loader(loader)
raw_docs = remote_loader.load_data(source_data)
cfg = SourceConfig.parse(config)
chunker = ChunkerCreator.create_chunker(
cfg.chunking.strategy,
chunking_strategy=cfg.chunking.strategy,
max_tokens=cfg.chunking.max_tokens,
min_tokens=cfg.chunking.min_tokens,
duplicate_headers=cfg.chunking.duplicate_headers,
)
raw_docs = chunker.chunk(documents=raw_docs)
docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]
tokens = count_tokens_docs(docs)
logging.info("Total tokens calculated: %d", tokens)
# Build directory structure from loaded documents
# Format matches local file uploads: nested structure with type, size_bytes, token_count
directory_structure = {}
for doc in raw_docs:
# Get the file path from extra_info
# For crawlers: file_path is a virtual path like "guides/setup.md"
# For other remotes: use key or title as fallback
file_path = ""
if doc.extra_info:
file_path = (
doc.extra_info.get("file_path", "")
or doc.extra_info.get("key", "")
or doc.extra_info.get("title", "")
)
if not file_path:
file_path = doc.doc_id or ""
if file_path:
# Calculate token count
token_count = num_tokens_from_string(doc.text) if doc.text else 0
# Estimate size in bytes from text content
size_bytes = len(doc.text.encode("utf-8")) if doc.text else 0
# Guess mime type from extension
file_name = (
file_path.split("/")[-1] if "/" in file_path else file_path
)
ext = os.path.splitext(file_name)[1].lower()
mime_types = {
".txt": "text/plain",
".md": "text/markdown",
".pdf": "application/pdf",
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
".doc": "application/msword",
".html": "text/html",
".json": "application/json",
".csv": "text/csv",
".xml": "application/xml",
".py": "text/x-python",
".js": "text/javascript",
".ts": "text/typescript",
".jsx": "text/jsx",
".tsx": "text/tsx",
}
file_type = mime_types.get(ext, "application/octet-stream")
# Build nested directory structure from path
# e.g., "guides/setup.md" -> {"guides": {"setup.md": {...}}}
path_parts = file_path.split("/")
current_level = directory_structure
for i, part in enumerate(path_parts):
if i == len(path_parts) - 1:
# Last part is the file
current_level[part] = {
"type": file_type,
"size_bytes": size_bytes,
"token_count": token_count,
}
else:
# Intermediate parts are directories
if part not in current_level:
current_level[part] = {}
current_level = current_level[part]
logging.info(
f"Built directory structure with {len(directory_structure)} files: "
f"{list(directory_structure.keys())}"
)
if operation_mode == "upload":
embed_and_store_documents(
docs, full_path, source_uuid, self,
attempt_id=getattr(self.request, "id", None),
user_id=user,
)
assert_index_complete(source_uuid)
elif operation_mode == "sync":
if not doc_id:
logging.error("Invalid doc_id provided for sync operation: %s", doc_id)
raise ValueError("doc_id must be provided for sync operation.")
embed_and_store_documents(
docs, full_path, source_uuid, self,
attempt_id=getattr(self.request, "id", None),
user_id=user,
)
assert_index_complete(source_uuid)
self.update_state(state="PROGRESS", meta={"current": 100})
# Serialize remote_data as JSON if it's a dict (for S3, Reddit, etc.)
remote_data_serialized = (
json.dumps(source_data) if isinstance(source_data, dict) else source_data
)
file_data = {
"name": name_job,
"user": user,
"tokens": tokens,
"retriever": retriever,
"id": source_id_for_events,
"type": loader,
"remote_data": remote_data_serialized,
"sync_frequency": sync_frequency,
"directory_structure": json.dumps(directory_structure),
}
if config:
file_data["config"] = json.dumps(config)
if operation_mode == "sync":
last_sync_now = datetime.datetime.now(datetime.timezone.utc)
file_data["last_sync"] = last_sync_now
try:
with db_session() as conn:
repo = SourcesRepository(conn)
src = repo.get_any(source_id_for_events, user)
if src is not None:
repo.update(str(src["id"]), user, {"date": last_sync_now})
except Exception as upd_err:
logging.warning(
f"Failed to update last_sync for source {source_id_for_events}: {upd_err}"
)
upload_index(full_path, file_data)
publish_user_event(
user,
"source.ingest.completed",
{
"source_id": source_id_for_events,
"job_name": name_job,
"loader": loader,
"operation": operation_mode,
"tokens": tokens,
# Forward-looking contract: see ingest_worker.
"limited": False,
},
scope={"kind": "source", "id": source_id_for_events},
)
_maybe_enqueue_graph_extraction(cfg, source_id_for_events, user)
except Exception as e:
logging.error("Error in remote_worker task: %s", str(e), exc_info=True)
publish_user_event(
user,
"source.ingest.failed",
{
"source_id": source_id_for_events,
"job_name": name_job,
"loader": loader,
"operation": operation_mode,
"error": str(e)[:1024],
},
scope={"kind": "source", "id": source_id_for_events},
)
raise
finally:
if os.path.exists(full_path):
shutil.rmtree(full_path)
logging.info("remote_worker task completed successfully")
return {
"id": source_id_for_events,
"urls": source_data,
"name_job": name_job,
"user": user,
"limited": False,
}
def sync(
self,
source_data,
name_job,
user,
loader,
sync_frequency,
retriever,
doc_id=None,
directory="temp",
):
try:
remote_worker(
self,
source_data,
name_job,
user,
loader,
directory,
retriever,
sync_frequency,
"sync",
doc_id,
)
except Exception as e:
logging.error(f"Error during sync: {e}", exc_info=True)
return {"status": "error", "error": str(e)}
return {"status": "success"}
def sync_worker(self, frequency):
from sqlalchemy import text as sql_text
sync_counts = Counter()
with db_readonly() as conn:
result = conn.execute(
sql_text(
"SELECT id, name, user_id, type, remote_data, retriever "
"FROM sources WHERE sync_frequency = :freq"
),
{"freq": frequency},
)
rows = result.fetchall()
for row in rows:
doc = dict(row._mapping)
name = doc.get("name")
user = doc.get("user_id")
source_type = doc.get("type")
retriever = doc.get("retriever")
doc_id = str(doc.get("id"))
sync_counts["total_sync_count"] += 1
# Connector sources have no RemoteCreator loader and need an OAuth
# token to sync, which a scheduled task lacks — skip them.
if source_type and source_type.startswith("connector"):
sync_counts["sync_skipped"] += 1
continue
source_data = normalize_remote_data(source_type, doc.get("remote_data"))
if not source_data:
# No syncable URL/config — skip instead of dispatching a sync
# that can only fail (and emit a spurious failed event).
sync_counts["sync_skipped"] += 1
continue
resp = sync(
self, source_data, name, user, source_type, frequency, retriever, doc_id
)
sync_counts[
"sync_success" if resp["status"] == "success" else "sync_failure"
] += 1
return {
key: sync_counts[key]
for key in [
"total_sync_count", "sync_success", "sync_failure", "sync_skipped",
]
}
def attachment_worker(self, file_info, user):
"""
Process and store a single attachment without vectorization.
"""
filename = file_info["filename"]
attachment_id = file_info["attachment_id"]
relative_path = file_info["path"]
metadata = file_info.get("metadata", {})
publish_user_event(
user,
"attachment.queued",
{"attachment_id": str(attachment_id), "filename": filename},
scope={"kind": "attachment", "id": str(attachment_id)},
)
try:
self.update_state(state="PROGRESS", meta={"current": 10})
storage = StorageCreator.get_storage()
self.update_state(
state="PROGRESS", meta={"current": 30, "status": "Processing content"}
)
publish_user_event(
user,
"attachment.progress",
{
"attachment_id": str(attachment_id),
"filename": filename,
"current": 30,
"stage": "processing",
},
scope={"kind": "attachment", "id": str(attachment_id)},
)
file_extractor = get_default_file_extractor(
ocr_enabled=settings.DOCLING_OCR_ATTACHMENTS_ENABLED
)
attachment_document = storage.process_file(
relative_path,
lambda local_path, **kwargs: SimpleDirectoryReader(
input_files=[local_path],
exclude_hidden=True,
errors="ignore",
file_extractor=file_extractor,
file_metadata=metadata_from_filename,
)
.load_data()[0],
)
content = attachment_document.text
parser_metadata = {
key: value
for key, value in (attachment_document.extra_info or {}).items()
if key.startswith("transcript_")
}
if parser_metadata:
metadata = {**metadata, **parser_metadata}
token_count = num_tokens_from_string(content)
if token_count > 100000:
content = content[:250000]
token_count = num_tokens_from_string(content)
self.update_state(
state="PROGRESS", meta={"current": 80, "status": "Storing in database"}
)
publish_user_event(
user,
"attachment.progress",
{
"attachment_id": str(attachment_id),
"filename": filename,
"current": 80,
"stage": "storing",
},
scope={"kind": "attachment", "id": str(attachment_id)},
)
mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
# The upload route produces a UUID-shaped ``attachment_id`` (stored
# in the storage path) but the PG ``attachments.id`` is generated
# by the DB. Keep ``attachment_id`` as the caller-visible handle
# used for the storage path, and stash it in ``legacy_mongo_id``
# so the attachment row is resolvable via that handle too.
with db_session() as conn:
AttachmentsRepository(conn).create(
user, filename, relative_path,
mime_type=mime_type,
content=content,
token_count=token_count,
metadata=metadata,
legacy_mongo_id=str(attachment_id),
)
logging.info(
f"Stored attachment with ID: {attachment_id}", extra={"user": user}
)
self.update_state(state="PROGRESS", meta={"current": 100, "status": "Complete"})
publish_user_event(
user,
"attachment.completed",
{
"attachment_id": str(attachment_id),
"filename": filename,
"token_count": token_count,
"mime_type": mime_type,
},
scope={"kind": "attachment", "id": str(attachment_id)},
)
return {
"filename": filename,
"path": relative_path,
"token_count": token_count,
"attachment_id": attachment_id,
"mime_type": mime_type,
"metadata": metadata,
}
except Exception as e:
logging.error(
f"Error processing file {filename}: {e}",
extra={"user": user},
exc_info=True,
)
publish_user_event(
user,
"attachment.failed",
{
"attachment_id": str(attachment_id),
"filename": filename,
"error": str(e)[:1024],
},
scope={"kind": "attachment", "id": str(attachment_id)},
)
raise
def parse_document_worker(self, artifact_id, parent, user_id, options):
"""Thin Celery-task wrapper; delegates to the process-agnostic ``run_parse_document``."""
return run_parse_document(artifact_id, parent, user_id, options)
def run_parse_document(artifact_id, parent, user_id, options):
"""Parse an input artifact's bytes to a shaped result; runnable inline OR on the parsing queue.
Security: the artifact is re-resolved through the run-scoped gate here (never trusting a
raw storage path) so authz is enforced independently, in addition to the pre-enqueue check
in the tool. ``read_document`` calls this directly (in-process) when it already runs inside a
Celery worker; the web process dispatches ``parse_document`` to the parsing queue, which lands
here via ``parse_document_worker``.
"""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.parser.document_reader import bound_parse_payload, parse_document_bytes
options = options or {}
parent = parent or {}
conversation_id = parent.get("conversation_id")
workflow_run_id = parent.get("workflow_run_id")
if conversation_id is None and workflow_run_id is None:
return {"status": "error", "error": "parse_document requires a conversation_id or workflow_run_id."}
# Re-resolve through the parent-scoped gate so a forged/cross-run id is rejected
# here too; resolve a short ref to an id within this parent only.
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
resolved_id = resolve_artifact_id(
repo, artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
artifact = (
repo.get_artifact_in_parent(
resolved_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
if resolved_id is not None
else None
)
if artifact is None:
return {"status": "error", "error": f"input artifact {artifact_id} not found in this conversation/run."}
version = repo.get_version(resolved_id, artifact["current_version"])
except Exception:
logging.error("run_parse_document: failed to resolve input artifact", exc_info=True)
return {"status": "error", "error": f"failed to load input artifact {artifact_id}."}
if not version or not version.get("storage_path"):
return {"status": "error", "error": f"input artifact {artifact_id} has no stored content."}
display_name = version.get("filename") or artifact.get("title") or str(resolved_id)
filename = safe_filename(display_name)
try:
data = StorageCreator.get_storage().get_file(version["storage_path"]).read()
except Exception:
logging.error("run_parse_document: failed to read input artifact bytes", exc_info=True)
return {"status": "error", "error": f"failed to read input artifact {artifact_id}."}
# Parse returns the FULL content (no max_chars/window here) so the persisted artifact is the
# complete parse; all view-bounding is applied by ``bound_parse_payload`` below.
result = parse_document_bytes(
data,
filename,
output=options.get("output", "markdown"),
ocr=options.get("ocr", "auto"),
pages=options.get("pages"),
engine=options.get("engine", "auto"),
include_tables=bool(options.get("include_tables", True)),
)
if result.get("error"):
return {"status": "error", "error": result["error"]}
payload = {"status": "ok", **result}
if options.get("persist"):
# Persist the FULL shaped result by reference (bytes live in the artifact); only the
# bounded view computed below rides back through the Redis result backend.
artifact_ref = _persist_parse_result(result, display_name, user_id, parent, options)
if isinstance(artifact_ref, dict) and artifact_ref.get("error"):
payload["artifact_error"] = artifact_ref["error"]
elif artifact_ref is not None:
payload["artifact"] = artifact_ref
# Bound the Redis-backed VIEW across all shapes: content is capped by max_chars (else the
# default head+tail window), chunks are count/length-capped, and structured (needed for
# json_schema validation) rides back as-is. The FULL result already lives in the artifact.
payload = bound_parse_payload(payload, max_chars=options.get("max_chars"))
return payload
def _persist_parse_result(result, title, user_id, parent, options):
"""Persist the full shaped parse result as an owner/parent-scoped ``data`` artifact; return its ref."""
from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact
try:
data = json.dumps(result).encode("utf-8")
except (TypeError, ValueError):
logging.error("parse_document_worker: parse result is not JSON-serializable", exc_info=True)
return {"error": "parse result is not JSON-serializable."}
base = safe_filename(title) or "document"
filename = f"{base}.parsed.json"
try:
return persist_new_artifact(
user_id=user_id,
kind="data",
data=data,
filename=filename,
mime_type="application/json",
title=f"{title} (parsed)",
conversation_id=parent.get("conversation_id"),
workflow_run_id=parent.get("workflow_run_id"),
message_id=parent.get("message_id"),
produced_by={"tool": "read_document", "action": "read_document", "tool_id": options.get("tool_id")},
)
except QuotaExceeded as exc:
return {"error": str(exc)}
def agent_webhook_worker(self, agent_id, payload):
"""Process the webhook payload for an agent.
Raises on failure: Celery treats a returned dict as success and
would skip retries, leaving the caller with a stale 200.
"""
self.update_state(state="PROGRESS", meta={"current": 1})
try:
with db_readonly() as conn:
repo = AgentsRepository(conn)
agent_config = None
if looks_like_uuid(str(agent_id)):
# Access without user scoping — webhooks authenticate via
# the incoming token, not a user context.
from sqlalchemy import text as sql_text
from application.storage.db.base_repository import row_to_dict
result = conn.execute(
sql_text("SELECT * FROM agents WHERE id = CAST(:id AS uuid)"),
{"id": str(agent_id)},
)
row = result.fetchone()
if row is not None:
agent_config = row_to_dict(row)
if agent_config is None:
agent_config = repo.get_by_legacy_id(str(agent_id))
if not agent_config:
raise ValueError(f"Agent with ID {agent_id} not found.")
input_data = json.dumps(payload)
except Exception as e:
logging.error(f"Error processing agent webhook: {e}", exc_info=True)
raise
self.update_state(state="PROGRESS", meta={"current": 50})
try:
# Shared headless path with the scheduler; approval-gated tools auto-deny.
from application.agents.headless_runner import run_agent_headless
outcome = run_agent_headless(
agent_config,
input_data,
tool_allowlist=_webhook_tool_allowlist(agent_config),
endpoint="webhook",
)
result = {
"answer": outcome.get("answer", ""),
"sources": outcome.get("sources", []),
"tool_calls": outcome.get("tool_calls", []),
"thought": outcome.get("thought", ""),
}
except Exception as e:
logging.error(f"Error running agent logic: {e}", exc_info=True)
raise
else:
logging.info(
f"Webhook processed for agent {agent_id}", extra={"agent_id": agent_id}
)
return {"status": "success", "result": result}
finally:
self.update_state(state="PROGRESS", meta={"current": 100})
def _webhook_tool_allowlist(agent_config):
"""Deny-all on approval-gated tools for webhooks (per-agent opt-in is TBD)."""
return []
def ingest_connector(
self,
job_name: str,
user: str,
source_type: str,
session_token=None,
file_ids=None,
folder_ids=None,
recursive=True,
retriever: str = "classic",
operation_mode: str = "upload",
doc_id=None,
sync_frequency: str = "never",
config=None,
idempotency_key=None,
source_id=None,
) -> Dict[str, Any]:
"""
Ingestion for internal knowledge bases (GoogleDrive, etc.).
Args:
job_name: Name of the ingestion job
user: User identifier
source_type: Type of remote source ("google_drive", "dropbox", etc.)
session_token: Authentication token for the service
file_ids: List of file IDs to download
folder_ids: List of folder IDs to download
recursive: Whether to recursively download folders
retriever: Type of retriever to use
operation_mode: "upload" for initial ingestion, "sync" for incremental sync
doc_id: Document ID for sync operations (required when operation_mode="sync")
sync_frequency: How often to sync ("never", "daily", "weekly", "monthly")
config: Per-source ``SourceConfig`` dict. ``None``/``{}`` → classic
defaults (byte-identical to prior behavior).
idempotency_key: When provided, the ``source_id`` is derived
deterministically so a retried upload reuses the same source row.
source_id: When supplied, the worker uses it verbatim so SSE envelopes
carry the same id the HTTP route already returned to the frontend
— required for non-idempotent uploads where the route can't
predict ``_derive_source_id(idempotency_key)``.
"""
logging.info(
f"Starting remote ingestion from {source_type} for user: {user}, job: {job_name}"
)
# Source id resolution mirrors ``ingest_worker`` / ``remote_worker``:
# sync mode reuses ``doc_id``; otherwise the caller-supplied
# ``source_id`` (minted by the HTTP route and already echoed to the
# client) wins; fall back to ``_derive_source_id`` only when neither
# is supplied. Without rule (2) the no-idempotency-key path would
# mint a fresh uuid4 here that the frontend has no way to correlate
# SSE envelopes to.
if operation_mode == "sync" and doc_id:
source_uuid = str(doc_id)
elif source_id:
source_uuid = uuid.UUID(source_id)
else:
source_uuid = _derive_source_id(idempotency_key)
source_id_for_events = str(source_uuid)
# First-attempt gate: Celery retries re-run the body, and a
# repeated ``queued`` here would oscillate the toast through
# ``queued`` again between ``failed`` and ``completed``.
if self.request.retries == 0:
publish_user_event(
user,
"source.ingest.queued",
{
"source_id": source_id_for_events,
"job_name": job_name,
"loader": source_type,
"operation": operation_mode,
},
scope={"kind": "source", "id": source_id_for_events},
)
self.update_state(state="PROGRESS", meta={"current": 1})
try:
with tempfile.TemporaryDirectory() as temp_dir:
# Step 1: Initialize the appropriate loader
self.update_state(
state="PROGRESS",
meta={"current": 10, "status": "Initializing connector"},
)
if not session_token:
raise ValueError(f"{source_type} connector requires session_token")
if not ConnectorCreator.is_supported(source_type):
raise ValueError(
f"Unsupported connector type: {source_type}. Supported types: {ConnectorCreator.get_supported_connectors()}"
)
remote_loader = ConnectorCreator.create_connector(
source_type, session_token
)
# Create a clean config for storage
api_source_config = {
"file_ids": file_ids or [],
"folder_ids": folder_ids or [],
"recursive": recursive,
}
# Step 2: Download files to temp directory
self.update_state(
state="PROGRESS", meta={"current": 20, "status": "Downloading files"}
)
download_info = remote_loader.download_to_directory(
temp_dir, api_source_config
)
if download_info.get("empty_result", False) or not download_info.get(
"files_downloaded", 0
):
logging.warning(f"No files were downloaded from {source_type}")
# Connector returned no files — surface as a benign
# ``completed`` event with zero tokens so the toast
# closes out cleanly instead of waiting on polling.
publish_user_event(
user,
"source.ingest.completed",
{
"source_id": source_id_for_events,
"job_name": job_name,
"loader": source_type,
"operation": operation_mode,
"tokens": 0,
"no_changes": True,
},
scope={"kind": "source", "id": source_id_for_events},
)
# Create empty result directly instead of calling a separate method
return {
"name": job_name,
"user": user,
"tokens": 0,
"type": source_type,
"source_config": api_source_config,
"directory_structure": "{}",
}
# Step 3: Use SimpleDirectoryReader to process downloaded files
self.update_state(
state="PROGRESS", meta={"current": 40, "status": "Processing files"}
)
reader = SimpleDirectoryReader(
input_dir=temp_dir,
recursive=True,
required_exts=list(SUPPORTED_SOURCE_EXTENSIONS),
exclude_hidden=True,
file_metadata=metadata_from_filename,
)
# Parsing/OCR fills 40-60% of the bar; embedding takes 60-100%.
raw_docs = reader.load_data(
progress_callback=_make_parse_progress_callback(
self, user, source_uuid, start_pct=40, end_pct=60,
)
)
directory_structure = getattr(reader, "directory_structure", {})
# Step 4: Process documents (chunking, embedding, etc.)
cfg = SourceConfig.parse(config)
chunker = ChunkerCreator.create_chunker(
cfg.chunking.strategy,
chunking_strategy=cfg.chunking.strategy,
max_tokens=cfg.chunking.max_tokens,
min_tokens=cfg.chunking.min_tokens,
duplicate_headers=cfg.chunking.duplicate_headers,
)
raw_docs = chunker.chunk(documents=raw_docs)
# Preserve source information in document metadata
for doc in raw_docs:
if hasattr(doc, "extra_info") and doc.extra_info:
source = doc.extra_info.get("source")
if source and os.path.isabs(source):
# Convert absolute path to relative path
doc.extra_info["source"] = os.path.relpath(
source, start=temp_dir
)
docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]
# Validate operation_mode here too (the source_uuid path
# at the top of the function only branches on the
# sync+doc_id combination; surfacing the wrong-mode error
# this far in matches the legacy behaviour).
if operation_mode == "sync" and not doc_id:
logging.error(
"Invalid doc_id provided for sync operation: %s", doc_id
)
raise ValueError("doc_id must be provided for sync operation.")
if operation_mode not in ("upload", "sync"):
raise ValueError(f"Invalid operation_mode: {operation_mode}")
vector_store_path = os.path.join(temp_dir, "vector_store")
os.makedirs(vector_store_path, exist_ok=True)
self.update_state(
state="PROGRESS", meta={"current": 60, "status": "Storing documents"}
)
embed_and_store_documents(
docs, vector_store_path, source_uuid, self,
attempt_id=getattr(self.request, "id", None),
user_id=user,
progress_start=60, progress_end=100,
)
assert_index_complete(source_uuid)
tokens = count_tokens_docs(docs)
# Step 6: Upload index files
file_data = {
"user": user,
"name": job_name,
"tokens": tokens,
"retriever": retriever,
"id": source_id_for_events,
"type": "connector:file",
"remote_data": json.dumps(
{"provider": source_type, **api_source_config}
),
"directory_structure": json.dumps(directory_structure),
"sync_frequency": sync_frequency,
}
if config:
file_data["config"] = json.dumps(config)
file_data["last_sync"] = datetime.datetime.now(datetime.timezone.utc)
if operation_mode == "sync":
try:
with db_session() as conn:
repo = SourcesRepository(conn)
src = repo.get_any(source_id_for_events, user)
if src is not None:
repo.update(
str(src["id"]), user,
{"date": file_data["last_sync"]},
)
except Exception as upd_err:
logging.warning(
"Failed to update last_sync for source %s: %s",
source_id_for_events,
upd_err,
)
upload_index(vector_store_path, file_data)
# Ensure we mark the task as complete
self.update_state(
state="PROGRESS", meta={"current": 100, "status": "Complete"}
)
logging.info(f"Remote ingestion completed: {job_name}")
publish_user_event(
user,
"source.ingest.completed",
{
"source_id": source_id_for_events,
"job_name": job_name,
"loader": source_type,
"operation": operation_mode,
"tokens": tokens,
},
scope={"kind": "source", "id": source_id_for_events},
)
_maybe_enqueue_graph_extraction(cfg, source_id_for_events, user)
return {
"user": user,
"name": job_name,
"tokens": tokens,
"type": source_type,
"id": source_id_for_events,
"status": "complete",
}
except Exception as e:
logging.error(f"Error during remote ingestion: {e}", exc_info=True)
publish_user_event(
user,
"source.ingest.failed",
{
"source_id": source_id_for_events,
"job_name": job_name,
"loader": source_type,
"operation": operation_mode,
"error": str(e)[:1024],
},
scope={"kind": "source", "id": source_id_for_events},
)
raise
def mcp_oauth(self, config: Dict[str, Any], user_id: str = None) -> Dict[str, Any]:
"""Worker to handle MCP OAuth flow asynchronously.
Publishes SSE events at each phase boundary so the frontend can
drive the OAuth popup directly from the push channel. The
``mcp.oauth.awaiting_redirect`` envelope carries the
``authorization_url`` once the upstream OAuth client surfaces it,
eliminating the prior polling-only path for that URL.
"""
# Bind ``task_id`` and the publish helpers OUTSIDE the outer try so
# the ``except`` handler at the bottom can reach them even when an
# early statement raises. Without this, ``publish_oauth`` would
# UnboundLocalError on top of the original failure.
task_id = self.request.id if getattr(self, "request", None) else None
def publish_oauth(event_type: str, payload: Dict[str, Any]) -> None:
# MCP OAuth can be invoked without a route-bound user_id by
# legacy paths. Skip the SSE publish in that case \u2014 the caller
# has no per-user channel to subscribe to, and the status is
# surfaced via the task's return value.
if not user_id or task_id is None:
return
publish_user_event(
user_id,
event_type,
{"task_id": task_id, **payload},
scope={"kind": "mcp_oauth", "id": task_id},
)
def publish_awaiting_redirect(authorization_url: str) -> None:
"""Callback invoked by ``DocsGPTOAuth.redirect_handler`` once
the OAuth client has minted the authorization URL.
Carrying the URL on the SSE envelope lets the frontend open the
popup directly from the event \u2014 the prior polling-only path
for the URL is gone.
"""
publish_oauth(
"mcp.oauth.awaiting_redirect",
{
"message": "Awaiting OAuth redirect...",
"authorization_url": authorization_url,
},
)
try:
import asyncio
from application.agents.tools.mcp_tool import MCPTool
publish_oauth("mcp.oauth.in_progress", {"message": "Starting OAuth..."})
tool_config = config.copy()
tool_config["oauth_task_id"] = task_id
# Inject the awaiting-redirect publish callback. ``MCPTool`` pops
# it out of the config and threads it into ``DocsGPTOAuth`` so
# the publish fires synchronously from inside
# ``redirect_handler`` \u2014 the only point where the URL is known.
tool_config["oauth_redirect_publish"] = publish_awaiting_redirect
mcp_tool = MCPTool(tool_config, user_id)
async def run_oauth_discovery():
if not mcp_tool._client:
mcp_tool._setup_client()
return await mcp_tool._execute_with_client("list_tools")
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(run_oauth_discovery())
tools = mcp_tool.get_actions_metadata()
publish_oauth(
"mcp.oauth.completed",
{"tools": tools, "tools_count": len(tools)},
)
return {"success": True, "tools": tools, "tools_count": len(tools)}
except Exception as e:
error_msg = f"OAuth failed: {str(e)}"
logging.error("MCP OAuth discovery failed: %s", error_msg, exc_info=True)
publish_oauth("mcp.oauth.failed", {"error": error_msg[:1024]})
return {"success": False, "error": error_msg}
finally:
loop.close()
except Exception as e:
error_msg = f"OAuth init failed: {str(e)}"
logging.error("MCP OAuth init failed: %s", error_msg, exc_info=True)
publish_oauth("mcp.oauth.failed", {"error": error_msg[:1024]})
return {"success": False, "error": error_msg}
def reembed_wiki_page_worker(self, source_id, path, content_hash, user):
"""Re-embed one wiki page after an edit, or purge its chunks on delete.
Targeted delete of the page's existing chunks runs first, so a deleted
page becomes a pure purge and an edited page is re-embedded cleanly. The
chunk metadata shape matches the reingest path (``source``/``title``/
``filename``) so retrieval and future targeted deletes line up.
Args:
self: Celery task instance.
source_id: Source the page belongs to.
path: Page path; doubles as the chunk ``source`` metadata key.
content_hash: The edit's content hash (idempotency anchor).
user: Owner identifier used to load the source row.
Returns:
A status dict: ``{"status", "added", "deleted"}`` on re-embed, or
``{"status": "deleted", "deleted": n}`` when the page is gone.
"""
from application.vectorstore.vector_creator import VectorCreator
source_id = str(source_id)
with db_readonly() as conn:
source = SourcesRepository(conn).get_any(source_id, user)
if not source:
raise ValueError(f"Source {source_id} not found or access denied")
source_id = str(source["id"])
cfg = SourceConfig.parse(source.get("config"))
store = VectorCreator.create_vectorstore(
settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
)
deleted = store.delete_chunks_by_source_path(path)
with db_readonly() as conn:
page = WikiPagesRepository(conn).get_by_path(source_id, path)
if page is None:
return {"status": "deleted", "deleted": deleted}
try:
title = page.get("title") or path
chunker = ChunkerCreator.create_chunker(
cfg.chunking.strategy,
chunking_strategy=cfg.chunking.strategy,
max_tokens=cfg.chunking.max_tokens,
min_tokens=cfg.chunking.min_tokens,
duplicate_headers=cfg.chunking.duplicate_headers,
)
chunks = chunker.chunk(
documents=[
Document(
text=page["content"],
extra_info={
"source": path,
"title": title,
"filename": path,
},
)
]
)
added = 0
for chunk in chunks:
store.add_chunk(
chunk.text,
metadata={"source": path, "title": title, "filename": path},
)
added += 1
with db_session() as conn:
WikiPagesRepository(conn).set_embed_status(source_id, path, "embedded")
except Exception:
with db_session() as conn:
WikiPagesRepository(conn).set_embed_status(source_id, path, "failed")
raise
return {"status": "embedded", "added": added, "deleted": deleted}
def _wiki_page_path_from_rel(rel_path):
"""Build a validated leading-slash wiki page path, or None if invalid.
Runs the derived path through ``validate_tool_path`` so a stored filename
carrying traversal (``..``) never lands as a ``wiki_pages.path``.
"""
from application.agents.tools.path_utils import validate_tool_path
raw = "/" + rel_path.replace(os.sep, "/").lstrip("/")
return validate_tool_path(raw)
def _chunk_text(chunk) -> str:
if isinstance(chunk, dict):
return chunk.get("text") or chunk.get("page_content") or ""
return getattr(chunk, "text", None) or getattr(chunk, "page_content", "") or ""
def _chunk_metadata(chunk) -> dict:
if isinstance(chunk, dict):
meta = chunk.get("metadata")
else:
meta = getattr(chunk, "metadata", None)
return meta if isinstance(meta, dict) else {}
def _chunk_doc_id(chunk):
if isinstance(chunk, dict):
return chunk.get("doc_id")
return getattr(chunk, "doc_id", None)
def _url_to_virtual_path(value: str) -> str:
parts = urlsplit(value)
path = parts.path.strip("/")
if path:
return path
if parts.netloc:
return parts.netloc
return value.replace("://", "/")
def _chunk_page_path(metadata) -> str:
for key in ("file_path", "file_name", "filename", "title", "source"):
value = metadata.get(key)
if value:
value = str(value)
if "://" in value:
value = _url_to_virtual_path(value)
return value
return ""
def _chunk_order_hint(metadata):
for key in ("chunk", "chunk_index", "index", "start", "start_index", "offset"):
value = metadata.get(key)
if isinstance(value, (int, float)):
return value
if isinstance(value, str):
try:
return float(value)
except ValueError:
continue
return None
MIN_CHUNK_OVERLAP_TRIM = 32
def _join_chunk_texts(texts: list[str]) -> str:
"""Concatenate chunk texts, trimming an exact overlapping boundary.
Only an exact suffix/prefix overlap of at least ``MIN_CHUNK_OVERLAP_TRIM``
characters between consecutive chunks is removed (the common
``chunk_overlap`` case); shorter, possibly legitimate repeats are kept.
"""
parts: list[str] = []
for text in texts:
if parts:
prev = parts[-1]
limit = min(len(prev), len(text))
overlap = 0
for size in range(limit, MIN_CHUNK_OVERLAP_TRIM - 1, -1):
if prev.endswith(text[:size]):
overlap = size
break
if overlap:
text = text[overlap:]
if text:
parts.append(text)
return "\n\n".join(parts)
def convert_source_to_wiki_worker(self, source_id, user):
"""Convert an ingested source into a wiki by reassembling pages from chunks.
Groups the source's existing vector-store chunks by their page path
(``metadata.source``), concatenating each group into one wiki page. This
works uniformly for every source type, including crawler/remote/connector
sources that hold no stored originals. Chunks with a missing or invalid
path are skipped and reported. Embedding is enqueued per page (async); the
source's ``config.kind`` flips to ``wiki`` only after at least one page is
materialized — a source with no usable chunks keeps its original kind and
directory structure.
Args:
self: Celery task instance.
source_id: ID of the source to convert.
user: Owner identifier used to load the source row and author pages.
Returns:
A summary dict ``{"status": "converted", "pages_created", "skipped"}``,
``{"status": "no_pages", ...}`` when nothing was reassembled, or
``{"status": "already_wiki"}`` when the source is already a wiki.
"""
from application.api.user.tasks import reembed_wiki_page
from application.vectorstore.vector_creator import VectorCreator
source_id = str(source_id)
with db_readonly() as conn:
source = SourcesRepository(conn).get_any(source_id, user)
if not source:
raise ValueError(f"Source {source_id} not found or access denied")
source_id = str(source["id"])
cfg = SourceConfig.parse(source.get("config"))
if cfg.kind == "wiki":
return {"status": "already_wiki", "pages_created": 0, "skipped": []}
file_name_map = _normalize_file_name_map(source.get("file_name_map"))
created_pages: list[tuple[str, str]] = []
skipped: list[dict] = []
store = VectorCreator.create_vectorstore(
settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
)
grouped: dict[str, list[tuple[object, str]]] = {}
original_doc_ids: list[object] = []
for chunk in store.get_chunks() or []:
# Track every original chunk so a skipped one (empty/no-path/invalid)
# is purged on convert too, not left orphaned in the vector store after
# the source flips to wiki. Only deleted once pages are created below.
doc_id = _chunk_doc_id(chunk)
if doc_id is not None:
original_doc_ids.append(doc_id)
metadata = _chunk_metadata(chunk)
rel_path = _chunk_page_path(metadata)
text = _chunk_text(chunk)
if not text.strip():
continue
if not rel_path:
skipped.append({"file": "", "reason": "missing path"})
continue
if _wiki_page_path_from_rel(rel_path) is None:
skipped.append({"file": rel_path, "reason": "invalid path"})
continue
grouped.setdefault(rel_path, []).append((_chunk_order_hint(metadata), text))
with db_session() as conn:
repo = WikiPagesRepository(conn)
for rel_path in sorted(grouped):
page_path = _wiki_page_path_from_rel(rel_path)
entries = grouped[rel_path]
if any(hint is not None for hint, _ in entries):
entries = sorted(
enumerate(entries),
key=lambda item: (
item[1][0] if item[1][0] is not None else float("inf"),
item[0],
),
)
entries = [entry for _, entry in entries]
content = _join_chunk_texts([text for _, text in entries])
title = _get_display_name(file_name_map, rel_path) or os.path.basename(
rel_path
)
repo.upsert(
source_id,
page_path,
content,
title=title,
updated_by=user,
updated_via="agent",
)
created_pages.append((page_path, _content_hash(content)))
if not created_pages:
return {
"status": "no_pages",
"pages_created": 0,
"skipped": skipped,
}
with db_session() as conn:
rebuild_wiki_directory_structure(conn, source_id, user)
for doc_id in original_doc_ids:
try:
store.delete_chunk(doc_id)
except Exception as exc:
logging.error(f"Failed deleting original chunk {doc_id}: {exc}")
for page_path, content_hash in created_pages:
reembed_wiki_page.delay(
source_id,
page_path,
content_hash,
user=user,
idempotency_key=f"reembed-wiki:{source_id}:{page_path}:{content_hash}",
)
with db_session() as conn:
SourcesRepository(conn).update(
source_id, user, {"config": cfg.wiki_enabled()}
)
return {
"status": "converted",
"pages_created": len(created_pages),
"skipped": skipped,
}
def extract_graph_worker(self, source_id, user):
"""Build a graphrag source's knowledge graph from its embedded chunks.
Loads the source, fetches its chunks from the vector store, and runs the
per-chunk LLM extraction pipeline. The chunks carry ``doc_id`` + ``text``,
matching the retrievable ids the extraction pipeline links against. No-ops
cleanly when GraphRAG is unavailable or the source has no chunks yet.
Streams ``graph.extract.progress`` SSE events as chunks are processed and a
terminal ``graph.extract.completed``/``graph.extract.failed`` on exit.
Args:
self: Celery task instance.
source_id: Source whose graph is being built.
user: Owner identifier used to load the source and attribute token usage.
Returns:
``{"status": "unavailable"}`` when GraphRAG is off, otherwise the
extraction summary ``{nodes, edges, chunks_processed, ...}``.
"""
from application.graphrag import graphrag_available
from application.graphrag.extraction import extract_graph_for_source
from application.vectorstore.vector_creator import VectorCreator
source_id = str(source_id)
if not graphrag_available():
return {"status": "unavailable"}
with db_readonly() as conn:
source = SourcesRepository(conn).get_any(source_id, user)
if not source:
raise ValueError(f"Source {source_id} not found or access denied")
source_id = str(source["id"])
cfg = SourceConfig.parse(source.get("config"))
store = VectorCreator.create_vectorstore(
settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
)
chunks = store.get_chunks() or []
total = len(chunks)
# Throttle: at most ~20 progress events regardless of chunk count.
step = max(1, total // 20)
def _progress(info):
current = int(info.get("current", 0))
if current and current % step != 0 and current != info.get("total"):
return
_publish_graph_event(
user,
source_id,
"graph.extract.progress",
{
"source_id": source_id,
"current": current,
"total": int(info.get("total", total)),
"nodes": int(info.get("nodes", 0)),
"edges": int(info.get("edges", 0)),
},
)
_publish_graph_event(
user,
source_id,
"graph.extract.progress",
{
"source_id": source_id,
"current": 0,
"total": total,
"nodes": 0,
"edges": 0,
},
)
try:
summary = extract_graph_for_source(
source_id,
user,
chunks,
config=cfg,
request_id=getattr(self.request, "id", None),
progress_cb=_progress,
)
except Exception as e:
_publish_graph_event(
user,
source_id,
"graph.extract.failed",
{"source_id": source_id, "error": str(e)[:1024]},
)
raise
_publish_graph_event(
user,
source_id,
"graph.extract.completed",
{"source_id": source_id, **summary},
)
return summary