.

Author: edwin-onyx

backend/onyx/agents/agent_search/orchestration/states.py

@@ -1,6 +1,6 @@
 from pydantic import BaseModel
 
-from onyx.chat.prompt_builder.answer_prompt_builder import PromptSnapshot
+from onyx.chat.prompt_builder.schemas import PromptSnapshot
 from onyx.tools.message import ToolCallSummary
 from onyx.tools.models import SearchToolOverrideKwargs
 from onyx.tools.models import ToolCallFinalResult

backend/onyx/background/celery/apps/app_base.py

@@ -24,15 +24,15 @@
 from onyx.background.celery.apps.task_formatters import CeleryTaskPlainFormatter
 from onyx.background.celery.celery_utils import celery_is_worker_primary
 from onyx.background.celery.celery_utils import make_probe_path
-# from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_PREFIX
-# from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_TASKSET_KEY
+from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_PREFIX
+from onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_TASKSET_KEY
 from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX
 from onyx.configs.constants import OnyxRedisLocks
 from onyx.db.engine.sql_engine import get_sqlalchemy_engine
 from onyx.document_index.vespa.shared_utils.utils import wait_for_vespa_with_timeout
 from onyx.httpx.httpx_pool import HttpxPool
-# from onyx.redis.redis_connector import RedisConnector
-# from onyx.redis.redis_connector_delete import RedisConnectorDelete
+from onyx.redis.redis_connector import RedisConnector
+from onyx.redis.redis_connector_delete import RedisConnectorDelete
 from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync
 from onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync
 from onyx.redis.redis_connector_prune import RedisConnectorPrune
@@ -157,9 +157,9 @@ def on_task_postrun(
     # NOTE: we want to remove the `Redis*` classes, prefer to just have functions to
     # do these things going forward. In short, things should generally be like the doc
     # sync task rather than the others below
-    # if task_id.startswith(DOCUMENT_SYNC_PREFIX):
-    #     r.srem(DOCUMENT_SYNC_TASKSET_KEY, task_id)
-    #     return
+    if task_id.startswith(DOCUMENT_SYNC_PREFIX):
+        r.srem(DOCUMENT_SYNC_TASKSET_KEY, task_id)
+        return
 
     if task_id.startswith(RedisDocumentSet.PREFIX):
         document_set_id = RedisDocumentSet.get_id_from_task_id(task_id)
@@ -175,33 +175,33 @@ def on_task_postrun(
             r.srem(rug.taskset_key, task_id)
         return
 
-    # if task_id.startswith(RedisConnectorDelete.PREFIX):
-    #     cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
-    #     if cc_pair_id is not None:
-    #         RedisConnectorDelete.remove_from_taskset(int(cc_pair_id), task_id, r)
-    #     return
-
-    # if task_id.startswith(RedisConnectorPrune.SUBTASK_PREFIX):
-    #     cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
-    #     if cc_pair_id is not None:
-    #         RedisConnectorPrune.remove_from_taskset(int(cc_pair_id), task_id, r)
-    #     return
-
-    # if task_id.startswith(RedisConnectorPermissionSync.SUBTASK_PREFIX):
-    #     cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
-    #     if cc_pair_id is not None:
-    #         RedisConnectorPermissionSync.remove_from_taskset(
-    #             int(cc_pair_id), task_id, r
-    #         )
-    #     return
-
-    # if task_id.startswith(RedisConnectorExternalGroupSync.SUBTASK_PREFIX):
-    #     cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
-    #     if cc_pair_id is not None:
-    #         RedisConnectorExternalGroupSync.remove_from_taskset(
-    #             int(cc_pair_id), task_id, r
-    #         )
-    #     return
+    if task_id.startswith(RedisConnectorDelete.PREFIX):
+        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            RedisConnectorDelete.remove_from_taskset(int(cc_pair_id), task_id, r)
+        return
+
+    if task_id.startswith(RedisConnectorPrune.SUBTASK_PREFIX):
+        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            RedisConnectorPrune.remove_from_taskset(int(cc_pair_id), task_id, r)
+        return
+
+    if task_id.startswith(RedisConnectorPermissionSync.SUBTASK_PREFIX):
+        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            RedisConnectorPermissionSync.remove_from_taskset(
+                int(cc_pair_id), task_id, r
+            )
+        return
+
+    if task_id.startswith(RedisConnectorExternalGroupSync.SUBTASK_PREFIX):
+        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            RedisConnectorExternalGroupSync.remove_from_taskset(
+                int(cc_pair_id), task_id, r
+            )
+        return
 
 
 def on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:

backend/onyx/background/celery/tasks/docfetching/tasks.py

@@ -3,6 +3,7 @@
 import time
 import traceback
 from time import sleep
+from typing import Any
 
 import sentry_sdk
 from celery import Celery
@@ -21,7 +22,7 @@
 from onyx.background.indexing.job_client import SimpleJob
 from onyx.background.indexing.job_client import SimpleJobClient
 from onyx.background.indexing.job_client import SimpleJobException
-# from onyx.background.indexing.run_docfetching import run_docfetching_entrypoint
+from onyx.background.indexing.run_docfetching import run_docfetching_entrypoint
 from onyx.configs.constants import CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.connectors.exceptions import ConnectorValidationError
@@ -32,7 +33,7 @@
 from onyx.db.index_attempt import mark_attempt_canceled
 from onyx.db.index_attempt import mark_attempt_failed
 from onyx.db.indexing_coordination import IndexingCoordination
-# from onyx.redis.redis_connector import RedisConnector
+from onyx.redis.redis_connector import RedisConnector
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import global_version
 from shared_configs.configs import SENTRY_DSN
@@ -211,14 +212,14 @@ def _docfetching_task(
         )
 
         # This is where the heavy/real work happens
-        # run_docfetching_entrypoint(
-        #     app,
-        #     index_attempt_id,
-        #     tenant_id,
-        #     cc_pair_id,
-        #     is_ee,
-        #     callback=callback,
-        # )
+        run_docfetching_entrypoint(
+            app,
+            index_attempt_id,
+            tenant_id,
+            cc_pair_id,
+            is_ee,
+            callback=callback,
+        )
 
     except ConnectorValidationError:
         raise SimpleJobException(
@@ -259,7 +260,7 @@ def process_job_result(
     job: SimpleJob,
     connector_source: str | None,
     index_attempt_id: int,
-    log_builder: ConnectorIndexingLogBuilder,
+    log_builder: Any,
 ) -> SimpleJobResult:
     result = SimpleJobResult()
     result.connector_source = connector_source

backend/onyx/background/celery/tasks/docprocessing/tasks.py

@@ -86,7 +86,6 @@
 from onyx.file_store.document_batch_storage import get_document_batch_storage
 from onyx.httpx.httpx_pool import HttpxPool
 from onyx.indexing.embedder import DefaultIndexingEmbedder
-from onyx.indexing.indexing_pipeline import run_indexing_pipeline
 from onyx.natural_language_processing.search_nlp_models import EmbeddingModel
 from onyx.natural_language_processing.search_nlp_models import (
     InformationContentClassificationModel,
@@ -1268,6 +1267,8 @@ def _docprocessing_task(
     tenant_id: str,
     batch_num: int,
 ) -> None:
+    from onyx.indexing.indexing_pipeline import run_indexing_pipeline
+
     start_time = time.monotonic()
 
     if tenant_id:

backend/onyx/background/indexing/run_docfetching.py

@@ -28,7 +28,6 @@
 from onyx.connectors.connector_runner import ConnectorRunner
 from onyx.connectors.exceptions import ConnectorValidationError
 from onyx.connectors.exceptions import UnexpectedValidationError
-from onyx.connectors.factory import instantiate_connector
 from onyx.connectors.interfaces import CheckpointedConnector
 from onyx.connectors.models import ConnectorFailure
 from onyx.connectors.models import ConnectorStopSignal
@@ -66,7 +65,6 @@
 from onyx.httpx.httpx_pool import HttpxPool
 from onyx.indexing.embedder import DefaultIndexingEmbedder
 from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
-from onyx.indexing.indexing_pipeline import run_indexing_pipeline
 from onyx.natural_language_processing.search_nlp_models import (
     InformationContentClassificationModel,
 )
@@ -100,6 +98,8 @@ def _get_connector_runner(
     are the complete list of existing documents of the connector. If the task
     of type LOAD_STATE, the list will be considered complete and otherwise incomplete.
     """
+    from onyx.connectors.factory import instantiate_connector
+
     task = attempt.connector_credential_pair.connector.input_type
 
     try:
@@ -283,6 +283,8 @@ def _run_indexing(
     2. Embed and index these documents into the chosen datastore (vespa)
     3. Updates Postgres to record the indexed documents + the outcome of this run
     """
+    from onyx.indexing.indexing_pipeline import run_indexing_pipeline
+
     start_time = time.monotonic()  # jsut used for logging
 
     with get_session_with_current_tenant() as db_session_temp:

backend/onyx/chat/prompt_builder/answer_prompt_builder.py

@@ -4,7 +4,6 @@
 from langchain_core.messages import BaseMessage
 from langchain_core.messages import HumanMessage
 from langchain_core.messages import SystemMessage
-from pydantic import BaseModel
 from pydantic.v1 import BaseModel as BaseModel__v1
 
 from onyx.chat.models import PromptConfig
@@ -196,10 +195,6 @@ def build(self) -> list[BaseMessage]:
 
 
 # Stores some parts of a prompt builder as needed for tool calls
-class PromptSnapshot(BaseModel):
-    raw_message_history: list[PreviousMessage]
-    raw_user_query: str
-    built_prompt: list[BaseMessage]
 
 
 # TODO: rename this? AnswerConfig maybe?

backend/onyx/chat/prompt_builder/schemas.py

@@ -0,0 +1,10 @@
+from langchain_core.messages import BaseMessage
+from pydantic import BaseModel
+
+from onyx.llm.models import PreviousMessage
+
+
+class PromptSnapshot(BaseModel):
+    raw_message_history: list[PreviousMessage]
+    raw_user_query: str
+    built_prompt: list[BaseMessage]

backend/onyx/chat/tool_handling/tool_response_handler.py

@@ -7,7 +7,7 @@
 from onyx.chat.models import ResponsePart
 from onyx.chat.prompt_builder.answer_prompt_builder import AnswerPromptBuilder
 from onyx.chat.prompt_builder.answer_prompt_builder import LLMCall
-from onyx.chat.prompt_builder.answer_prompt_builder import PromptSnapshot
+from onyx.chat.prompt_builder.schemas import PromptSnapshot
 from onyx.llm.interfaces import LLM
 from onyx.tools.force import ForceUseTool
 from onyx.tools.message import build_tool_message

backend/onyx/db/chat.py

@@ -56,7 +56,7 @@
 from onyx.server.query_and_chat.models import ChatMessageDetail
 from onyx.server.query_and_chat.models import SubQueryDetail
 from onyx.server.query_and_chat.models import SubQuestionDetail
-from onyx.tools.tool_runner import ToolCallFinalResult
+from onyx.tools.models import ToolCallFinalResult
 from onyx.utils.logger import setup_logger
 from onyx.utils.special_types import JSON_ro
 

backend/onyx/db/document.py

@@ -22,7 +22,7 @@
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import null
 
-# from onyx.agents.agent_search.kb_search.models import KGEntityDocInfo
+from onyx.agents.agent_search.kb_search.models import KGEntityDocInfo
 from onyx.configs.constants import DEFAULT_BOOST
 from onyx.configs.constants import DocumentSource
 from onyx.configs.kg_configs import KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES
@@ -33,7 +33,7 @@
 from onyx.db.entities import delete_from_kg_entities_extraction_staging__no_commit
 from onyx.db.enums import AccessType
 from onyx.db.enums import ConnectorCredentialPairStatus
-# from onyx.db.feedback import delete_document_feedback_for_documents__no_commit
+from onyx.db.feedback import delete_document_feedback_for_documents__no_commit
 from onyx.db.models import Connector
 from onyx.db.models import ConnectorCredentialPair
 from onyx.db.models import Credential
@@ -747,9 +747,9 @@ def delete_documents_complete__no_commit(
     )
 
     delete_documents_by_connector_credential_pair__no_commit(db_session, document_ids)
-    # delete_document_feedback_for_documents__no_commit(
-        # document_ids=document_ids, db_session=db_session
-    # )
+    delete_document_feedback_for_documents__no_commit(
+        document_ids=document_ids, db_session=db_session
+    )
     delete_document_tags_for_documents__no_commit(
         document_ids=document_ids, db_session=db_session
     )
@@ -1199,35 +1199,35 @@ def get_skipped_kg_documents(db_session: Session) -> list[str]:
     return list(db_session.scalars(stmt).all())
 
 
-# def get_kg_doc_info_for_entity_name(
-#     db_session: Session, document_id: str, entity_type: str
-# ) -> KGEntityDocInfo:
-#     """
-#     Get the semantic ID and the link for an entity name.
-#     """
-
-#     result = (
-#         db_session.query(Document.semantic_id, Document.link)
-#         .filter(Document.id == document_id)
-#         .first()
-#     )
-
-#     if result is None:
-#         return KGEntityDocInfo(
-#             doc_id=None,
-#             doc_semantic_id=None,
-#             doc_link=None,
-#             semantic_entity_name=f"{entity_type}:{document_id}",
-#             semantic_linked_entity_name=f"{entity_type}:{document_id}",
-#         )
-
-#     return KGEntityDocInfo(
-#         doc_id=document_id,
-#         doc_semantic_id=result[0],
-#         doc_link=result[1],
-#         semantic_entity_name=f"{entity_type.upper()}:{result[0]}",
-#         semantic_linked_entity_name=f"[{entity_type.upper()}:{result[0]}]({result[1]})",
-#     )
+def get_kg_doc_info_for_entity_name(
+    db_session: Session, document_id: str, entity_type: str
+) -> KGEntityDocInfo:
+    """
+    Get the semantic ID and the link for an entity name.
+    """
+
+    result = (
+        db_session.query(Document.semantic_id, Document.link)
+        .filter(Document.id == document_id)
+        .first()
+    )
+
+    if result is None:
+        return KGEntityDocInfo(
+            doc_id=None,
+            doc_semantic_id=None,
+            doc_link=None,
+            semantic_entity_name=f"{entity_type}:{document_id}",
+            semantic_linked_entity_name=f"{entity_type}:{document_id}",
+        )
+
+    return KGEntityDocInfo(
+        doc_id=document_id,
+        doc_semantic_id=result[0],
+        doc_link=result[1],
+        semantic_entity_name=f"{entity_type.upper()}:{result[0]}",
+        semantic_linked_entity_name=f"[{entity_type.upper()}:{result[0]}]({result[1]})",
+    )
 
 
 def check_for_documents_needing_kg_processing(