starred/onyx
1
0 Fork 0
mirror of https://github.com/onyx-dot-app/onyx.git synced 2026-10-19 01:50:51 -05:00
Code Issues Packages Projects Releases 10 Wiki Activity

chore: kg stuff in celery (#7908)

Browse Source
This commit is contained in:
Yuhong Sun authored and GitHub committed 2026-01-28 03:36:31 +00:00
1 parent bbba32b989
commit b75b1019f3
19 files changed
+20 -1016
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/onyx/background/celery/apps/background.py
-1
View File
@@ -121,7 +121,6 @@ celery_app.autodiscover_tasks(
[
# Original background worker tasks
"onyx.background.celery.tasks.pruning",
"onyx.background.celery.tasks.kg_processing",
"onyx.background.celery.tasks.monitoring",
"onyx.background.celery.tasks.user_file_processing",
"onyx.background.celery.tasks.llm_model_update",
backend/onyx/background/celery/apps/kg_processing.py
-109
View File
@@ -1,109 +0,0 @@
from typing import Any
from typing import cast
from celery import Celery
from celery import signals
from celery import Task
from celery.apps.worker import Worker
from celery.signals import celeryd_init
from celery.signals import worker_init
from celery.signals import worker_process_init
from celery.signals import worker_ready
from celery.signals import worker_shutdown
import onyx.background.celery.apps.app_base as app_base
from onyx.configs.constants import POSTGRES_CELERY_WORKER_KG_PROCESSING_APP_NAME
from onyx.db.engine.sql_engine import SqlEngine
from onyx.utils.logger import setup_logger
from shared_configs.configs import MULTI_TENANT
logger = setup_logger()
celery_app = Celery(__name__)
celery_app.config_from_object("onyx.background.celery.configs.kg_processing")
celery_app.Task = app_base.TenantAwareTask # type: ignore [misc]
@signals.task_prerun.connect
def on_task_prerun(
sender: Any | None = None,
task_id: str | None = None,
task: Task | None = None,
args: tuple | None = None,
kwargs: dict | None = None,
**kwds: Any,
) -> None:
app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)
@signals.task_postrun.connect
def on_task_postrun(
sender: Any | None = None,
task_id: str | None = None,
task: Task | None = None,
args: tuple | None = None,
kwargs: dict | None = None,
retval: Any | None = None,
state: str | None = None,
**kwds: Any,
) -> None:
app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)
@celeryd_init.connect
def on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:
app_base.on_celeryd_init(sender, conf, **kwargs)
@worker_init.connect
def on_worker_init(sender: Worker, **kwargs: Any) -> None:
logger.info("worker_init signal received.")
SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_KG_PROCESSING_APP_NAME)
pool_size = cast(int, sender.concurrency) # type: ignore
SqlEngine.init_engine(pool_size=pool_size, max_overflow=8)
app_base.wait_for_redis(sender, **kwargs)
app_base.wait_for_db(sender, **kwargs)
app_base.wait_for_vespa_or_shutdown(sender, **kwargs)
# Less startup checks in multi-tenant case
if MULTI_TENANT:
return
app_base.on_secondary_worker_init(sender, **kwargs)
@worker_ready.connect
def on_worker_ready(sender: Any, **kwargs: Any) -> None:
app_base.on_worker_ready(sender, **kwargs)
@worker_shutdown.connect
def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
app_base.on_worker_shutdown(sender, **kwargs)
@worker_process_init.connect
def init_worker(**kwargs: Any) -> None:
SqlEngine.reset_engine()
@signals.setup_logging.connect
def on_setup_logging(
loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
) -> None:
app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
base_bootsteps = app_base.get_bootsteps()
for bootstep in base_bootsteps:
celery_app.steps["worker"].add(bootstep)
celery_app.autodiscover_tasks(
[
"onyx.background.celery.tasks.kg_processing",
]
)
backend/onyx/background/celery/apps/primary.py
-1
View File
@@ -323,7 +323,6 @@ celery_app.autodiscover_tasks(
"onyx.background.celery.tasks.shared",
"onyx.background.celery.tasks.vespa",
"onyx.background.celery.tasks.llm_model_update",
"onyx.background.celery.tasks.kg_processing",
"onyx.background.celery.tasks.user_file_processing",
]
)
backend/onyx/background/celery/configs/kg_processing.py
-21
View File
@@ -1,21 +0,0 @@
import onyx.background.celery.configs.base as shared_config
from onyx.configs.app_configs import CELERY_WORKER_KG_PROCESSING_CONCURRENCY
broker_url = shared_config.broker_url
broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup
broker_pool_limit = shared_config.broker_pool_limit
broker_transport_options = shared_config.broker_transport_options
redis_socket_keepalive = shared_config.redis_socket_keepalive
redis_retry_on_timeout = shared_config.redis_retry_on_timeout
redis_backend_health_check_interval = shared_config.redis_backend_health_check_interval
result_backend = shared_config.result_backend
result_expires = shared_config.result_expires # 86400 seconds is the default
task_default_priority = shared_config.task_default_priority
task_acks_late = shared_config.task_acks_late
worker_concurrency = CELERY_WORKER_KG_PROCESSING_CONCURRENCY
worker_pool = "threads"
worker_prefetch_multiplier = 1
backend/onyx/background/celery/tasks/beat_schedule.py
-18
View File
@@ -57,24 +57,6 @@ beat_task_templates: list[dict] = [
"expires": BEAT_EXPIRES_DEFAULT,
},
},
{
"name": "check-for-kg-processing",
"task": OnyxCeleryTask.CHECK_KG_PROCESSING,
"schedule": timedelta(seconds=60),
"options": {
"priority": OnyxCeleryPriority.MEDIUM,
"expires": BEAT_EXPIRES_DEFAULT,
},
},
{
"name": "check-for-kg-processing-clustering-only",
"task": OnyxCeleryTask.CHECK_KG_PROCESSING_CLUSTERING_ONLY,
"schedule": timedelta(seconds=600),
"options": {
"priority": OnyxCeleryPriority.LOW,
"expires": BEAT_EXPIRES_DEFAULT,
},
},
{
"name": "check-for-indexing",
"task": OnyxCeleryTask.CHECK_FOR_INDEXING,
backend/onyx/background/celery/tasks/kg_processing/kg_indexing.py
-77
View File
@@ -1,77 +0,0 @@
from onyx.background.celery.apps.app_base import task_logger
from onyx.background.celery.apps.client import celery_app
from onyx.background.celery.tasks.kg_processing.utils import is_kg_processing_blocked
from onyx.background.celery.tasks.kg_processing.utils import (
is_kg_processing_requirements_met,
)
from onyx.configs.constants import OnyxCeleryPriority
from onyx.configs.constants import OnyxCeleryQueues
from onyx.configs.constants import OnyxCeleryTask
def try_creating_kg_processing_task(
tenant_id: str,
) -> bool:
"""Schedules the KG Processing for a tenant immediately. Will not schedule if
the tenant is not ready for KG processing.
"""
try:
if not is_kg_processing_requirements_met():
return False
# Send the KG processing task
result = celery_app.send_task(
OnyxCeleryTask.KG_PROCESSING,
kwargs=dict(
tenant_id=tenant_id,
),
queue=OnyxCeleryQueues.KG_PROCESSING,
priority=OnyxCeleryPriority.MEDIUM,
)
if not result:
task_logger.error("send_task for kg processing failed.")
return bool(result)
except Exception:
task_logger.exception(
f"try_creating_kg_processing_task - Unexpected exception for tenant={tenant_id}"
)
return False
def try_creating_kg_source_reset_task(
tenant_id: str,
source_name: str | None,
index_name: str,
) -> bool:
"""Schedules the KG Source Reset for a tenant immediately. Will not do anything if
the tenant is currently being processed.
"""
try:
if is_kg_processing_blocked():
return False
# Send the KG source reset task
result = celery_app.send_task(
OnyxCeleryTask.KG_RESET_SOURCE_INDEX,
kwargs=dict(
tenant_id=tenant_id,
source_name=source_name,
index_name=index_name,
),
queue=OnyxCeleryQueues.KG_PROCESSING,
priority=OnyxCeleryPriority.MEDIUM,
)
if not result:
task_logger.error("send_task for kg source reset failed.")
return bool(result)
except Exception:
task_logger.exception(
f"try_creating_kg_source_reset_task - Unexpected exception for tenant={tenant_id}"
)
return False
backend/onyx/background/celery/tasks/kg_processing/tasks.py
-253
View File
@@ -1,253 +0,0 @@
import time
from celery import shared_task
from celery import Task
from celery.exceptions import SoftTimeLimitExceeded
from redis.lock import Lock as RedisLock
from onyx.background.celery.apps.app_base import task_logger
from onyx.background.celery.tasks.kg_processing.utils import (
is_kg_clustering_only_requirements_met,
)
from onyx.background.celery.tasks.kg_processing.utils import (
is_kg_processing_requirements_met,
)
from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT
from onyx.configs.constants import OnyxCeleryPriority
from onyx.configs.constants import OnyxCeleryQueues
from onyx.configs.constants import OnyxCeleryTask
from onyx.configs.constants import OnyxRedisLocks
from onyx.db.engine.sql_engine import get_session_with_current_tenant
from onyx.db.search_settings import get_current_search_settings
from onyx.kg.clustering.clustering import kg_clustering
from onyx.kg.extractions.extraction_processing import kg_extraction
from onyx.kg.resets.reset_source import reset_source_kg_index
from onyx.redis.redis_pool import get_redis_client
from onyx.redis.redis_pool import redis_lock_dump
from onyx.utils.logger import setup_logger
logger = setup_logger()
@shared_task(
name=OnyxCeleryTask.CHECK_KG_PROCESSING,
soft_time_limit=300,
bind=True,
)
def check_for_kg_processing(self: Task, *, tenant_id: str) -> None:
"""a lightweight task used to kick off kg processing tasks."""
time_start = time.monotonic()
task_logger.warning("check_for_kg_processing - Starting")
try:
if not is_kg_processing_requirements_met():
return
task_logger.info(
f"Found documents needing KG processing for tenant {tenant_id}"
)
self.app.send_task(
OnyxCeleryTask.KG_PROCESSING,
kwargs={
"tenant_id": tenant_id,
},
queue=OnyxCeleryQueues.KG_PROCESSING,
priority=OnyxCeleryPriority.MEDIUM,
)
except SoftTimeLimitExceeded:
task_logger.info(
"Soft time limit exceeded, task is being terminated gracefully."
)
except Exception:
task_logger.exception("Unexpected exception during kg processing check")
time_elapsed = time.monotonic() - time_start
task_logger.info(f"check_for_kg_processing finished: elapsed={time_elapsed:.2f}")
@shared_task(
name=OnyxCeleryTask.CHECK_KG_PROCESSING_CLUSTERING_ONLY,
soft_time_limit=300,
bind=True,
)
def check_for_kg_processing_clustering_only(self: Task, *, tenant_id: str) -> None:
"""a lightweight task used to kick off kg clustering tasks."""
time_start = time.monotonic()
task_logger.warning("check_for_kg_processing_clustering_only - Starting")
try:
if not is_kg_clustering_only_requirements_met():
return
task_logger.info(
f"Found documents needing KG clustering for tenant {tenant_id}"
)
self.app.send_task(
OnyxCeleryTask.KG_CLUSTERING_ONLY,
kwargs={
"tenant_id": tenant_id,
},
queue=OnyxCeleryQueues.KG_PROCESSING,
priority=OnyxCeleryPriority.MEDIUM,
)
except SoftTimeLimitExceeded:
task_logger.info(
"Soft time limit exceeded, task is being terminated gracefully."
)
except Exception:
task_logger.exception("Unexpected exception during kg clustering-only check")
time_elapsed = time.monotonic() - time_start
task_logger.info(
f"check_for_kg_processing_clustering_only finished: elapsed={time_elapsed:.2f}"
)
@shared_task(
name=OnyxCeleryTask.KG_PROCESSING,
bind=True,
)
def kg_processing(self: Task, *, tenant_id: str) -> None:
"""a task for doing kg extraction and clustering."""
task_logger.warning(f"kg_processing - Starting for tenant {tenant_id}")
redis_client = get_redis_client()
lock_beat: RedisLock = redis_client.lock(
OnyxRedisLocks.KG_PROCESSING_LOCK,
timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
)
# these tasks should never overlap
if not lock_beat.acquire(blocking=False):
return
try:
with get_session_with_current_tenant() as db_session:
search_settings = get_current_search_settings(db_session)
index_str = search_settings.index_name
task_logger.info(f"KG processing set to in progress for tenant {tenant_id}")
kg_extraction(
tenant_id=tenant_id,
index_name=index_str,
lock=lock_beat,
processing_chunk_batch_size=8,
)
kg_clustering(
tenant_id=tenant_id,
index_name=index_str,
lock=lock_beat,
processing_chunk_batch_size=8,
)
except Exception:
task_logger.exception("Error during kg processing")
finally:
if lock_beat.owned():
lock_beat.release()
else:
task_logger.error(
"kg_processing - Lock not owned on completion: " f"tenant={tenant_id}"
)
redis_lock_dump(lock_beat, redis_client)
task_logger.debug("Completed kg processing task!")
@shared_task(
name=OnyxCeleryTask.KG_CLUSTERING_ONLY,
bind=True,
)
def kg_clustering_only(self: Task, *, tenant_id: str) -> None:
"""a task for doing kg clustering only."""
task_logger.warning(f"kg_clustering_only - Starting for tenant {tenant_id}")
redis_client = get_redis_client()
lock_beat: RedisLock = redis_client.lock(
OnyxRedisLocks.KG_PROCESSING_LOCK,
timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
)
# these tasks should never overlap
if not lock_beat.acquire(blocking=False):
return
try:
with get_session_with_current_tenant() as db_session:
search_settings = get_current_search_settings(db_session)
index_str = search_settings.index_name
task_logger.info(
f"KG clustering-only set to in progress for tenant {tenant_id}"
)
kg_clustering(
tenant_id=tenant_id,
index_name=index_str,
lock=lock_beat,
processing_chunk_batch_size=8,
)
except Exception:
task_logger.exception("Error during kg clustering-only")
finally:
if lock_beat.owned():
lock_beat.release()
else:
task_logger.error(
"kg_clustering_only - Lock not owned on completion: "
f"tenant={tenant_id}"
)
redis_lock_dump(lock_beat, redis_client)
task_logger.debug("Completed kg clustering-only task!")
@shared_task(
name=OnyxCeleryTask.KG_RESET_SOURCE_INDEX,
bind=True,
)
def kg_reset_source_index(
self: Task, *, tenant_id: str, source_name: str, index_name: str
) -> None:
"""a task for KG reset of a source."""
task_logger.warning(f"kg_reset_source_index - Starting for tenant {tenant_id}")
redis_client = get_redis_client()
lock_beat: RedisLock = redis_client.lock(
OnyxRedisLocks.KG_PROCESSING_LOCK,
timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
)
# these tasks should never overlap
if not lock_beat.acquire(blocking=False):
return
try:
reset_source_kg_index(
source_name=source_name,
tenant_id=tenant_id,
index_name=index_name,
lock=lock_beat,
)
except Exception:
task_logger.exception("Error during kg reset")
finally:
if lock_beat.owned():
lock_beat.release()
else:
task_logger.error(
"kg_reset_source_index - Lock not owned on completion: "
f"tenant={tenant_id}"
)
redis_lock_dump(lock_beat, redis_client)
task_logger.debug("Completed kg reset task!")
backend/onyx/background/celery/tasks/kg_processing/utils.py
-78
View File
@@ -1,78 +0,0 @@
import time
from redis.lock import Lock as RedisLock
from onyx.configs.constants import OnyxRedisLocks
from onyx.db.document import check_for_documents_needing_kg_processing
from onyx.db.engine.sql_engine import get_session_with_current_tenant
from onyx.db.kg_config import get_kg_config_settings
from onyx.db.kg_config import is_kg_config_settings_enabled_valid
from onyx.db.models import KGEntityExtractionStaging
from onyx.db.models import KGRelationshipExtractionStaging
from onyx.redis.redis_pool import get_redis_client
def is_kg_processing_blocked() -> bool:
"""Checks if there are any KG tasks in progress."""
redis_client = get_redis_client()
lock_beat: RedisLock = redis_client.lock(OnyxRedisLocks.KG_PROCESSING_LOCK)
return lock_beat.locked()
def is_kg_processing_requirements_met() -> bool:
"""Checks that there are no other KG tasks in progress, KG is enabled, valid,
and there are documents that need KG processing."""
if is_kg_processing_blocked():
return False
kg_config = get_kg_config_settings()
if not is_kg_config_settings_enabled_valid(kg_config):
return False
with get_session_with_current_tenant() as db_session:
has_staging_entities = (
db_session.query(KGEntityExtractionStaging).first() is not None
)
has_staging_relationships = (
db_session.query(KGRelationshipExtractionStaging).first() is not None
)
return (
check_for_documents_needing_kg_processing(
db_session,
kg_config.KG_COVERAGE_START_DATE,
kg_config.KG_MAX_COVERAGE_DAYS,
)
or has_staging_entities
or has_staging_relationships
)
def is_kg_clustering_only_requirements_met() -> bool:
"""Checks that there are no other KG tasks in progress, KG is enabled, valid,
and there are documents that need KG clustering."""
if is_kg_processing_blocked():
return False
kg_config = get_kg_config_settings()
if not is_kg_config_settings_enabled_valid(kg_config):
return False
# Check if there are any entries in the staging tables
with get_session_with_current_tenant() as db_session:
has_staging_entities = (
db_session.query(KGEntityExtractionStaging).first() is not None
)
has_staging_relationships = (
db_session.query(KGRelationshipExtractionStaging).first() is not None
)
return has_staging_entities or has_staging_relationships
def extend_lock(lock: RedisLock, timeout: int, last_lock_time: float) -> float:
current_time = time.monotonic()
if current_time - last_lock_time >= (timeout / 4):
lock.reacquire()
last_lock_time = current_time
return last_lock_time
backend/onyx/background/celery/versioned_apps/kg_processing.py
-18
View File
@@ -1,18 +0,0 @@
"""Factory stub for running celery worker for knowledge graph processing.
This code is different from the primary/beat stubs because there is no EE version to
fetch. Port over the code in those files if we add an EE version of this worker."""
from celery import Celery
from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable
set_is_ee_based_on_env_variable()
def get_app() -> Celery:
from onyx.background.celery.apps.kg_processing import celery_app
return celery_app
app = get_app()
backend/onyx/chat/chat_utils.py
+1 -40
View File
@@ -9,12 +9,6 @@ from fastapi.datastructures import Headers
from sqlalchemy.orm import Session
from onyx.auth.users import is_user_admin
from onyx.background.celery.tasks.kg_processing.kg_indexing import (
try_creating_kg_processing_task,
)
from onyx.background.celery.tasks.kg_processing.kg_indexing import (
try_creating_kg_source_reset_task,
)
from onyx.chat.models import ChatLoadedFile
from onyx.chat.models import ChatMessageSimple
from onyx.chat.models import PersonaOverrideConfig
@@ -37,7 +31,6 @@ from onyx.db.models import Tool
from onyx.db.models import User
from onyx.db.models import UserFile
from onyx.db.projects import check_project_ownership
from onyx.db.search_settings import get_current_search_settings
from onyx.file_processing.extract_file_text import extract_file_text
from onyx.file_store.file_store import get_default_file_store
from onyx.file_store.models import ChatFileType
@@ -359,39 +352,7 @@ def process_kg_commands(
if not is_kg_config_settings_enabled_valid(kg_config_settings):
return
# get Vespa index
search_settings = get_current_search_settings(db_session)
index_str = search_settings.index_name
if message == "kg_p":
success = try_creating_kg_processing_task(tenant_id)
if success:
raise KGException("KG processing scheduled")
else:
raise KGException(
"Cannot schedule another KG processing if one is already running "
"or there are no documents to process"
)
elif message.startswith("kg_rs_source"):
msg_split = [x for x in message.split(":")]
if len(msg_split) > 2:
raise KGException("Invalid format for a source reset command")
elif len(msg_split) == 2:
source_name = msg_split[1].strip()
elif len(msg_split) == 1:
source_name = None
else:
raise KGException("Invalid format for a source reset command")
success = try_creating_kg_source_reset_task(tenant_id, source_name, index_str)
if success:
source_name = source_name or "all"
raise KGException(f"KG index reset for source '{source_name}' scheduled")
else:
raise KGException("Cannot reset index while KG processing is running")
elif message == "kg_setup":
if message == "kg_setup":
populate_missing_default_entity_types__commit(db_session=db_session)
raise KGException("KG setup done")
backend/onyx/configs/app_configs.py
+1 -5
View File
@@ -413,7 +413,7 @@ CELERY_WORKER_PRIMARY_POOL_OVERFLOW = int(
os.environ.get("CELERY_WORKER_PRIMARY_POOL_OVERFLOW") or 4
)
# Consolidated background worker (light, docprocessing, docfetching, heavy, kg_processing, monitoring, user_file_processing)
# Consolidated background worker (light, docprocessing, docfetching, heavy, monitoring, user_file_processing)
# separate workers' defaults: light=24, docprocessing=6, docfetching=1, heavy=4, kg=2, monitoring=1, user_file=2
# Total would be 40, but we use a more conservative default of 20 for the consolidated worker
CELERY_WORKER_BACKGROUND_CONCURRENCY = int(
@@ -425,10 +425,6 @@ CELERY_WORKER_HEAVY_CONCURRENCY = int(
os.environ.get("CELERY_WORKER_HEAVY_CONCURRENCY") or 4
)
CELERY_WORKER_KG_PROCESSING_CONCURRENCY = int(
os.environ.get("CELERY_WORKER_KG_PROCESSING_CONCURRENCY") or 2
)
CELERY_WORKER_MONITORING_CONCURRENCY = int(
os.environ.get("CELERY_WORKER_MONITORING_CONCURRENCY") or 1
)
backend/onyx/configs/constants.py
-17
View File
@@ -80,7 +80,6 @@ POSTGRES_CELERY_WORKER_DOCFETCHING_APP_NAME = "celery_worker_docfetching"
POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME = "celery_worker_indexing_child"
POSTGRES_CELERY_WORKER_BACKGROUND_APP_NAME = "celery_worker_background"
POSTGRES_CELERY_WORKER_HEAVY_APP_NAME = "celery_worker_heavy"
POSTGRES_CELERY_WORKER_KG_PROCESSING_APP_NAME = "celery_worker_kg_processing"
POSTGRES_CELERY_WORKER_MONITORING_APP_NAME = "celery_worker_monitoring"
POSTGRES_CELERY_WORKER_USER_FILE_PROCESSING_APP_NAME = (
"celery_worker_user_file_processing"
@@ -383,9 +382,6 @@ class OnyxCeleryQueues:
# Monitoring queue
MONITORING = "monitoring"
# KG processing queue
KG_PROCESSING = "kg_processing"
# Sandbox processing queue
SANDBOX = "sandbox"
@@ -423,9 +419,6 @@ class OnyxRedisLocks:
CLOUD_BEAT_TASK_GENERATOR_LOCK = "da_lock:cloud_beat_task_generator"
CLOUD_CHECK_ALEMBIC_BEAT_LOCK = "da_lock:cloud_check_alembic"
# KG processing
KG_PROCESSING_LOCK = "da_lock:kg_processing"
# User file processing
USER_FILE_PROCESSING_BEAT_LOCK = "da_lock:check_user_file_processing_beat"
USER_FILE_PROCESSING_LOCK_PREFIX = "da_lock:user_file_processing"
@@ -457,9 +450,6 @@ class OnyxRedisSignals:
"signal:block_validate_connector_deletion_fences"
)
# KG processing
CHECK_KG_PROCESSING_BEAT_LOCK = "da_lock:check_kg_processing_beat"
class OnyxRedisConstants:
ACTIVE_FENCES = "active_fences"
@@ -559,13 +549,6 @@ class OnyxCeleryTask:
EXPORT_QUERY_HISTORY_TASK = "export_query_history_task"
EXPORT_QUERY_HISTORY_CLEANUP_TASK = "export_query_history_cleanup_task"
# KG processing
CHECK_KG_PROCESSING = "check_kg_processing"
KG_PROCESSING = "kg_processing"
KG_CLUSTERING_ONLY = "kg_clustering_only"
CHECK_KG_PROCESSING_CLUSTERING_ONLY = "check_kg_processing_clustering_only"
KG_RESET_SOURCE_INDEX = "kg_reset_source_index"
# Sandbox cleanup
CLEANUP_IDLE_SANDBOXES = "cleanup_idle_sandboxes"
CLEANUP_OLD_SNAPSHOTS = "cleanup_old_snapshots"
backend/onyx/kg/clustering/clustering.py
+1 -1
View File
@@ -7,7 +7,6 @@ from redis.lock import Lock as RedisLock
from sqlalchemy import func
from sqlalchemy import text
from onyx.background.celery.tasks.kg_processing.utils import extend_lock
from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT
from onyx.configs.kg_configs import KG_CLUSTERING_RETRIEVE_THRESHOLD
from onyx.configs.kg_configs import KG_CLUSTERING_THRESHOLD
@@ -32,6 +31,7 @@ from onyx.document_index.vespa.kg_interactions import (
from onyx.document_index.vespa.kg_interactions import update_kg_chunks_vespa_info
from onyx.kg.models import KGGroundingType
from onyx.kg.utils.formatting_utils import make_relationship_id
from onyx.kg.utils.lock_utils import extend_lock
from onyx.utils.logger import setup_logger
from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
backend/onyx/kg/extractions/extraction_processing.py
+1 -1
View File
@@ -3,7 +3,6 @@ from typing import Any
from redis.lock import Lock as RedisLock
from onyx.background.celery.tasks.kg_processing.utils import extend_lock
from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT
from onyx.db.connector import get_kg_enabled_connectors
from onyx.db.document import get_document_updated_at
@@ -40,6 +39,7 @@ from onyx.kg.utils.formatting_utils import extract_relationship_type_id
from onyx.kg.utils.formatting_utils import get_entity_type
from onyx.kg.utils.formatting_utils import split_entity_id
from onyx.kg.utils.formatting_utils import split_relationship_id
from onyx.kg.utils.lock_utils import extend_lock
from onyx.utils.logger import setup_logger
from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel
backend/onyx/kg/resets/reset_vespa.py
+1 -1
View File
@@ -3,7 +3,6 @@ from typing import Any
from redis.lock import Lock as RedisLock
from onyx.background.celery.tasks.kg_processing.utils import extend_lock
from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT
from onyx.configs.constants import DocumentSource
from onyx.db.document import get_num_chunks_for_document
@@ -15,6 +14,7 @@ from onyx.document_index.document_index_utils import get_uuid_from_chunk_info
from onyx.document_index.vespa.index import KGVespaChunkUpdateRequest
from onyx.document_index.vespa.index import VespaIndex
from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT
from onyx.kg.utils.lock_utils import extend_lock
from onyx.utils.logger import setup_logger
from shared_configs.configs import MULTI_TENANT
backend/onyx/kg/utils/lock_utils.py
+12
View File
@@ -0,0 +1,12 @@
import time
from redis.lock import Lock as RedisLock
def extend_lock(lock: RedisLock, timeout: int, last_lock_time: float) -> float:
current_time = time.monotonic()
if current_time - last_lock_time >= (timeout / 4):
lock.reacquire()
last_lock_time = current_time
return last_lock_time
backend/scripts/dev_run_background_jobs.py
+1 -15
View File
@@ -96,7 +96,7 @@ def run_jobs() -> None:
"--loglevel=INFO",
"--hostname=background@%n",
"-Q",
"connector_pruning,connector_doc_permissions_sync,connector_external_group_sync,csv_generation,kg_processing,monitoring,user_file_processing,user_file_project_sync,user_file_delete",
"connector_pruning,connector_doc_permissions_sync,connector_external_group_sync,csv_generation,monitoring,user_file_processing,user_file_project_sync,user_file_delete",
]
background_workers = [("BACKGROUND", cmd_worker_background)]
else:
@@ -114,19 +114,6 @@ def run_jobs() -> None:
"-Q",
"connector_pruning,sandbox",
]
cmd_worker_kg_processing = [
"celery",
"-A",
"onyx.background.celery.versioned_apps.kg_processing",
"worker",
"--pool=threads",
"--concurrency=2",
"--prefetch-multiplier=1",
"--loglevel=INFO",
"--hostname=kg_processing@%n",
"-Q",
"kg_processing",
]
cmd_worker_monitoring = [
"celery",
"-A",
@@ -155,7 +142,6 @@ def run_jobs() -> None:
]
background_workers = [
("HEAVY", cmd_worker_heavy),
("KG_PROCESSING", cmd_worker_kg_processing),
("MONITORING", cmd_worker_monitoring),
("USER_FILE_PROCESSING", cmd_worker_user_file_processing),
]
backend/supervisord.conf
+2 -17
View File
@@ -57,12 +57,12 @@ autostart=%(ENV_USE_SEPARATE_BACKGROUND_WORKERS)s
# Lightweight mode: single consolidated background worker
# Used when USE_LIGHTWEIGHT_BACKGROUND_WORKER=true (default)
# Consolidates: light, docprocessing, docfetching, heavy, kg_processing, monitoring, user_file_processing
# Consolidates: light, docprocessing, docfetching, heavy, monitoring, user_file_processing
[program:celery_worker_background]
command=celery -A onyx.background.celery.versioned_apps.background worker
--loglevel=INFO
--hostname=background@%%n
-Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,sandbox,docprocessing,connector_doc_fetching,connector_pruning,connector_doc_permissions_sync,connector_external_group_sync,csv_generation,kg_processing,monitoring,user_file_processing,user_file_project_sync
-Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,sandbox,docprocessing,connector_doc_fetching,connector_pruning,connector_doc_permissions_sync,connector_external_group_sync,csv_generation,monitoring,user_file_processing,user_file_project_sync
stdout_logfile=/var/log/celery_worker_background.log
stdout_logfile_maxbytes=16MB
redirect_stderr=true
@@ -140,20 +140,6 @@ startsecs=10
stopasgroup=true
autostart=%(ENV_USE_SEPARATE_BACKGROUND_WORKERS)s
[program:celery_worker_kg_processing]
command=celery -A onyx.background.celery.versioned_apps.kg_processing worker
--loglevel=INFO
--hostname=kg_processing@%%n
-Q kg_processing
stdout_logfile=/var/log/celery_worker_kg_processing.log
stdout_logfile_maxbytes=16MB
redirect_stderr=true
autorestart=true
startsecs=10
stopasgroup=true
autostart=%(ENV_USE_SEPARATE_BACKGROUND_WORKERS)s
# Job scheduler for periodic tasks
[program:celery_beat]
@@ -213,7 +199,6 @@ command=tail -qF
/var/log/celery_worker_background.log
/var/log/celery_worker_heavy.log
/var/log/celery_worker_docprocessing.log
/var/log/celery_worker_kg_processing.log
/var/log/celery_worker_monitoring.log
/var/log/celery_worker_user_file_processing.log
/var/log/celery_worker_docfetching.log
backend/tests/integration/tests/kg/test_kg_processing.py
-343
View File
@@ -1,343 +0,0 @@
import time
import pytest
from onyx.background.celery.tasks.kg_processing.kg_indexing import (
try_creating_kg_processing_task,
)
from onyx.background.celery.tasks.kg_processing.utils import (
is_kg_processing_blocked,
)
from onyx.configs.constants import DocumentSource
from onyx.connectors.models import InputType
from onyx.db.engine.sql_engine import get_session_with_current_tenant
from onyx.db.kg_config import get_kg_config_settings
from onyx.db.kg_config import set_kg_config_settings
from onyx.db.models import Connector
from onyx.db.models import Document
from onyx.db.models import KGEntity
from onyx.db.models import KGEntityExtractionStaging
from onyx.db.models import KGEntityType
from onyx.db.models import KGRelationship
from onyx.db.models import KGRelationshipExtractionStaging
from onyx.db.models import KGStage
from onyx.kg.models import KGAttributeEntityOption
from onyx.kg.models import KGAttributeImplicationProperty
from onyx.kg.models import KGAttributeProperty
from onyx.kg.models import KGAttributeTrackInfo
from onyx.kg.models import KGAttributeTrackType
from onyx.kg.models import KGEntityTypeAttributes
from onyx.kg.models import KGEntityTypeDefinition
from onyx.kg.models import KGGroundingType
from shared_configs.contextvars import get_current_tenant_id
from tests.integration.common_utils.managers.api_key import APIKeyManager
from tests.integration.common_utils.managers.cc_pair import CCPairManager
from tests.integration.common_utils.managers.document import DocumentManager
from tests.integration.common_utils.managers.llm_provider import LLMProviderManager
from tests.integration.common_utils.managers.user import UserManager
from tests.integration.common_utils.reset import reset_all
from tests.integration.common_utils.test_models import DATestUser
@pytest.fixture(autouse=True)
def reset_for_test() -> None:
"""Reset all data before each test."""
reset_all()
kg_config_settings = get_kg_config_settings()
kg_config_settings.KG_EXPOSED = True
kg_config_settings.KG_ENABLED = True
kg_config_settings.KG_VENDOR = "Test"
kg_config_settings.KG_VENDOR_DOMAINS = ["example.com", "tester.ai"]
kg_config_settings.KG_IGNORE_EMAIL_DOMAINS = ["gmail.com"]
kg_config_settings.KG_COVERAGE_START = "2020-01-01"
set_kg_config_settings(kg_config_settings)
@pytest.fixture()
def kg_test_docs() -> tuple[list[str], int, list[KGEntityType]]:
# create admin user
admin_user: DATestUser = UserManager.create(email="admin@example.com")
# create a minimal file connector
cc_pair = CCPairManager.create_from_scratch(
name="KG-Test-FileConnector",
source=DocumentSource.FILE,
input_type=InputType.LOAD_STATE,
connector_specific_config={
"file_locations": [],
"file_names": [],
"zip_metadata": {},
},
user_performing_action=admin_user,
)
api_key = APIKeyManager.create(user_performing_action=admin_user)
api_key.headers.update(admin_user.headers)
LLMProviderManager.create(user_performing_action=admin_user)
# create test document
# semantic id = Test Document {document_id}
# source type = FILE
doc1 = DocumentManager.seed_doc_with_content(
cc_pair=cc_pair,
content="Dummy content for KG doc A",
document_id="docA",
metadata={
"teamname": "Team1",
"assignees": ["john@gmail.com", "dane@tester.ai"],
"parent": "Test Document docB",
"key": "Test Document docA",
"stuff": ["a", "b", "c", "d", "e", "f", "g"],
},
api_key=api_key,
)
doc2 = DocumentManager.seed_doc_with_content(
cc_pair=cc_pair,
content="Dummy content for KG doc BS",
document_id="docB",
metadata={
"teamname": "Team2",
"assignees": ["david@outsider.com"],
"key": "Test Document docB",
"stuff": ["h", "i", "j", "k", "l", "m", "n"],
},
api_key=api_key,
)
# create entity type
entity_type_definitions = {
"TEST": KGEntityTypeDefinition(
description="A test entity",
attributes=KGEntityTypeAttributes(
metadata_attribute_conversion={
"teamname": KGAttributeProperty(name="team", keep=True),
"stuff": KGAttributeProperty(name="stuff", keep=True),
"assignees": KGAttributeProperty(
name="assignees",
keep=False,
implication_property=KGAttributeImplicationProperty(
implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,
implied_relationship_name="is_assignee_of",
),
),
"parent": KGAttributeProperty(
name="parent", # will also create has_subcomponent
keep=True,
implication_property=KGAttributeImplicationProperty(
implied_entity_type="TEST",
implied_relationship_name="is_parent_of",
),
),
"key": KGAttributeProperty(name="key", keep=True),
},
),
grounding=KGGroundingType.GROUNDED,
grounded_source_name=DocumentSource.FILE,
active=True,
),
"ACCOUNT": KGEntityTypeDefinition(
description=("A company"),
attributes=KGEntityTypeAttributes(
entity_filter_attributes={"object_type": "Account"},
),
grounding=KGGroundingType.GROUNDED,
grounded_source_name=DocumentSource.SALESFORCE,
active=True,
),
"EMPLOYEE": KGEntityTypeDefinition(
description="An employee",
grounding=KGGroundingType.GROUNDED,
active=True,
grounded_source_name=None,
),
}
kg_entity_types: list[KGEntityType] = []
with get_session_with_current_tenant() as db_session:
for (
entity_type_id_name,
entity_type_definition,
) in entity_type_definitions.items():
grounded_source_name = (
entity_type_definition.grounded_source_name.value
if entity_type_definition.grounded_source_name
else None
)
kg_entity_type = KGEntityType(
id_name=entity_type_id_name,
description=entity_type_definition.description,
attributes=entity_type_definition.attributes.model_dump(),
grounding=entity_type_definition.grounding,
grounded_source_name=grounded_source_name,
active=entity_type_definition.active,
)
kg_entity_types.append(kg_entity_type)
db_session.add(kg_entity_type)
db_session.commit()
# enable connector kg processing
with get_session_with_current_tenant() as db_session:
db_session.query(Connector).filter(Connector.id == cc_pair.connector_id).update(
{"kg_processing_enabled": True}
)
db_session.commit()
return ([doc1.id, doc2.id], cc_pair.id, kg_entity_types)
def wait_until_kg_processing_done(timeout: float = 60) -> bool:
"""
Wait until KG processing is done. Returns True if the KG processing
finished before the timeout, False otherwise.
"""
start = time.monotonic()
while time.monotonic() - start < timeout:
time.sleep(1)
if not is_kg_processing_blocked():
return True
return False
def test_kg_processing(
kg_test_docs: tuple[list[str], int, KGEntityType],
) -> None:
"""Test KG processing."""
# run extraction and clustering
tenant_id = get_current_tenant_id()
result = try_creating_kg_processing_task(tenant_id)
assert result
assert wait_until_kg_processing_done()
# check entities
with get_session_with_current_tenant() as db_session:
entities_list = db_session.query(KGEntity).all()
entities = {entity.name: entity for entity in entities_list}
# should have the two TEST entities, the ACCOUNT and EMPLOYEE entities
# should not have gmail.com as it is in the IGNORE_EMAIL_DOMAINS
# the two docb entities should be clustered
assert len(entities_list) == 4
assert len(entities) == 4
assert "test document doca" in entities
assert "test document docb" in entities
assert "outsider.com" in entities
assert "dane" in entities
# check properties of each entity
doca_entity = entities["test document doca"]
doca_id = doca_entity.id_name
assert doca_entity.id_name.startswith("TEST::")
assert doca_entity.entity_type_id_name == "TEST"
assert doca_entity.document_id == "docA"
assert doca_entity.entity_key == "Test Document docA"
assert doca_entity.parent_key == "Test Document docB"
assert doca_entity.attributes == {
"team": "Team1",
"stuff": ["a", "b", "c", "d", "e", "f", "g"],
}
docb_entity = entities["test document docb"]
docb_id = docb_entity.id_name
assert docb_entity.id_name.startswith("TEST::")
assert docb_entity.entity_type_id_name == "TEST"
assert docb_entity.document_id == "docB"
assert docb_entity.entity_key == "Test Document docB"
assert docb_entity.parent_key is None
assert docb_entity.attributes == {
"team": "Team2",
"stuff": ["h", "i", "j", "k", "l", "m", "n"],
}
account_entity = entities["outsider.com"]
account_id = account_entity.id_name
assert account_entity.id_name.startswith("ACCOUNT::")
assert account_entity.entity_type_id_name == "ACCOUNT"
assert account_entity.document_id is None
assert account_entity.entity_key is None
assert account_entity.parent_key is None
assert account_entity.attributes == {}
employee_entity = entities["dane"]
employee_id = employee_entity.id_name
assert employee_entity.id_name.startswith("EMPLOYEE::")
assert employee_entity.entity_type_id_name == "EMPLOYEE"
assert employee_entity.document_id is None
assert employee_entity.entity_key is None
assert employee_entity.parent_key is None
assert employee_entity.attributes == {}
# check relationships
with get_session_with_current_tenant() as db_session:
relationships_list = db_session.query(KGRelationship).all()
relationships = {
relationship.relationship_type_id_name: relationship
for relationship in relationships_list
}
assert len(relationships_list) == 4
assert len(relationships) == 4
assert "ACCOUNT__is_assignee_of__TEST" in relationships
assert "EMPLOYEE__is_assignee_of__TEST" in relationships
assert "TEST__is_parent_of__TEST" in relationships
assert "TEST__has_subcomponent__TEST" in relationships
# check properties of each relationship
act_assignee_rel = relationships["ACCOUNT__is_assignee_of__TEST"]
assert act_assignee_rel.source_node == account_id
assert act_assignee_rel.target_node == docb_id
assert act_assignee_rel.source_node_type == "ACCOUNT"
assert act_assignee_rel.target_node_type == "TEST"
assert act_assignee_rel.type == "is_assignee_of"
emp_assignee_rel = relationships["EMPLOYEE__is_assignee_of__TEST"]
assert emp_assignee_rel.source_node == employee_id
assert emp_assignee_rel.target_node == doca_id
assert emp_assignee_rel.source_node_type == "EMPLOYEE"
assert emp_assignee_rel.target_node_type == "TEST"
assert emp_assignee_rel.type == "is_assignee_of"
parent_rel = relationships["TEST__is_parent_of__TEST"]
assert parent_rel.source_node == docb_id
assert parent_rel.target_node == doca_id
assert parent_rel.source_node_type == "TEST"
assert parent_rel.target_node_type == "TEST"
assert parent_rel.type == "is_parent_of"
subcomponent_rel = relationships["TEST__has_subcomponent__TEST"]
assert subcomponent_rel.source_node == docb_id
assert subcomponent_rel.target_node == doca_id
assert subcomponent_rel.source_node_type == "TEST"
assert subcomponent_rel.target_node_type == "TEST"
assert subcomponent_rel.type == "has_subcomponent"
# check staging tables are empty
with get_session_with_current_tenant() as db_session:
assert db_session.query(KGEntityExtractionStaging).count() == 0
assert db_session.query(KGRelationshipExtractionStaging).count() == 0
# check document kg_stage
with get_session_with_current_tenant() as db_session:
docs = db_session.query(Document).all()
assert len(docs) == 2
assert docs[0].kg_stage == KGStage.NORMALIZED
assert docs[1].kg_stage == KGStage.NORMALIZED
# check entity type attribute extraction
with get_session_with_current_tenant() as db_session:
entity_types_list = db_session.query(KGEntityType).all()
entity_types = {
entity_type.id_name: entity_type for entity_type in entity_types_list
}
assert len(entity_types_list) == 3
assert len(entity_types) == 3
assert "TEST" in entity_types
assert "ACCOUNT" in entity_types
assert "EMPLOYEE" in entity_types
test_type_attr = entity_types["TEST"].parsed_attributes.attribute_values
assert test_type_attr == {
"team": KGAttributeTrackInfo(
type=KGAttributeTrackType.VALUE,
values={"Team1", "Team2"},
),
"stuff": KGAttributeTrackInfo(
type=KGAttributeTrackType.LIST,
values=None,
),
}
account_type_attr = entity_types["ACCOUNT"].parsed_attributes.attribute_values
assert account_type_attr == {}
employee_type_attr = entity_types["EMPLOYEE"].parsed_attributes.attribute_values
assert employee_type_attr == {}