feat: Search Backend (#7426)

backend/alembic/versions/73e9983e5091_add_search_query_table.py

@@ -0,0 +1,47 @@
+"""add_search_query_table
+
+Revision ID: 73e9983e5091
+Revises: d1b637d7050a
+Create Date: 2026-01-14 14:16:52.837489
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "73e9983e5091"
+down_revision = "d1b637d7050a"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "search_query",
+        sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
+        sa.Column(
+            "user_id",
+            postgresql.UUID(as_uuid=True),
+            sa.ForeignKey("user.id"),
+            nullable=False,
+        ),
+        sa.Column("query", sa.String(), nullable=False),
+        sa.Column("query_expansions", postgresql.ARRAY(sa.String()), nullable=True),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            nullable=False,
+            server_default=sa.func.now(),
+        ),
+    )
+
+    op.create_index("ix_search_query_user_id", "search_query", ["user_id"])
+    op.create_index("ix_search_query_created_at", "search_query", ["created_at"])
+
+
+def downgrade() -> None:
+    op.drop_index("ix_search_query_created_at", table_name="search_query")
+    op.drop_index("ix_search_query_user_id", table_name="search_query")
+    op.drop_table("search_query")

backend/ee/onyx/db/search.py

@@ -0,0 +1,64 @@
+import uuid
+from datetime import timedelta
+from uuid import UUID
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from onyx.db.engine.time_utils import get_db_current_time
+from onyx.db.models import SearchQuery
+
+
+def create_search_query(
+    db_session: Session,
+    user_id: UUID,
+    query: str,
+    query_expansions: list[str] | None = None,
+) -> SearchQuery:
+    """Create and persist a `SearchQuery` row.
+
+    Notes:
+    - `SearchQuery.id` is a UUID PK without a server-side default, so we generate it.
+    - `created_at` is filled by the DB (server_default=now()).
+    """
+    search_query = SearchQuery(
+        id=uuid.uuid4(),
+        user_id=user_id,
+        query=query,
+        query_expansions=query_expansions,
+    )
+    db_session.add(search_query)
+    db_session.commit()
+    db_session.refresh(search_query)
+    return search_query
+
+
+def fetch_search_queries_for_user(
+    db_session: Session,
+    user_id: UUID,
+    filter_days: int | None = None,
+    limit: int | None = None,
+) -> list[SearchQuery]:
+    """Fetch `SearchQuery` rows for a user.
+
+    Args:
+        user_id: User UUID.
+        filter_days: Optional time filter. If provided, only rows created within
+            the last `filter_days` days are returned.
+        limit: Optional max number of rows to return.
+    """
+    if filter_days is not None and filter_days <= 0:
+        raise ValueError("filter_days must be > 0")
+
+    stmt = select(SearchQuery).where(SearchQuery.user_id == user_id)
+
+    if filter_days is not None and filter_days > 0:
+        cutoff = get_db_current_time(db_session) - timedelta(days=filter_days)
+        stmt = stmt.where(SearchQuery.created_at >= cutoff)
+
+    stmt = stmt.order_by(SearchQuery.created_at.desc())
+
+    if limit is not None:
+        stmt = stmt.limit(limit)
+
+    return list(db_session.scalars(stmt).all())

backend/ee/onyx/main.py

@@ -23,6 +23,7 @@ from ee.onyx.server.oauth.api import router as ee_oauth_router
 from ee.onyx.server.query_and_chat.query_backend import (
     basic_router as ee_query_router,
 )
+from ee.onyx.server.query_and_chat.search_backend import router as search_router
 from ee.onyx.server.query_history.api import router as query_history_router
 from ee.onyx.server.reporting.usage_export_api import router as usage_export_router
 from ee.onyx.server.seeding import seed_db
@@ -121,6 +122,7 @@ def get_application() -> FastAPI:
     # EE only backend APIs
     include_router_with_global_prefix_prepended(application, query_router)
     include_router_with_global_prefix_prepended(application, ee_query_router)
+    include_router_with_global_prefix_prepended(application, search_router)
     include_router_with_global_prefix_prepended(application, standard_answer_router)
     include_router_with_global_prefix_prepended(application, ee_oauth_router)
     include_router_with_global_prefix_prepended(application, ee_document_cc_pair_router)

backend/ee/onyx/prompts/search_flow_classification.py

@@ -0,0 +1,42 @@
+# ruff: noqa: E501, W605 start
+SEARCH_CLASS = "search"
+CHAT_CLASS = "chat"
+
+# Will note that with many larger LLMs the latency on running this prompt via third party APIs is as high as 2 seconds which is too slow for many
+# use cases.
+SEARCH_CHAT_PROMPT = f"""
+Determine if the following query is better suited for a search UI or a chat UI. Respond with "{SEARCH_CLASS}" or "{CHAT_CLASS}" literally and nothing else. \
+Do not provide any additional text or reasoning to your response. CRITICAL, IT MUST ONLY BE 1 SINGLE WORD - EITHER "{SEARCH_CLASS}" or "{CHAT_CLASS}".
+
+# Classification Guidelines:
+## {SEARCH_CLASS}
+- If the query consists entirely of keywords or query doesn't require any answer from the AI
+- If the query is a short statement that seems like a search query rather than a question
+- If the query feels nonsensical or is a short phrase that possibly describes a document or information that could be found in a internal document
+
+### Examples of {SEARCH_CLASS} queries:
+- Find me the document that goes over the onboarding process for a new hire
+- Pull requests since last week
+- Sales Runbook AMEA Region
+- Procurement process
+- Retrieve the PRD for project X
+
+## {CHAT_CLASS}
+- If the query is asking a question that requires an answer rather than a document
+- If the query is asking for a solution, suggestion, or general help
+- If the query is seeking information that is on the web and likely not in a company internal document
+- If the query should be answered without any context from additional documents or searches
+
+### Examples of {CHAT_CLASS} queries:
+- What led us to win the deal with company X? (seeking answer)
+- Google Drive not sync-ing files to my computer (seeking solution)
+- Review my email: <whatever the email is> (general help)
+- Write me a script to... (general help)
+- Cheap flights Europe to Tokyo (information likely found on the web, not internal)
+
+# User Query:
+{{user_query}}
+
+REMEMBER TO ONLY RESPOND WITH "{SEARCH_CLASS}" OR "{CHAT_CLASS}" AND NOTHING ELSE.
+""".strip()
+# ruff: noqa: E501, W605 end

backend/ee/onyx/secondary_llm_flows/search_flow_classification.py

@@ -0,0 +1,50 @@
+from ee.onyx.prompts.search_flow_classification import CHAT_CLASS
+from ee.onyx.prompts.search_flow_classification import SEARCH_CHAT_PROMPT
+from ee.onyx.prompts.search_flow_classification import SEARCH_CLASS
+from onyx.llm.interfaces import LLM
+from onyx.llm.models import LanguageModelInput
+from onyx.llm.models import ReasoningEffort
+from onyx.llm.models import UserMessage
+from onyx.llm.utils import llm_response_to_string
+from onyx.utils.logger import setup_logger
+from onyx.utils.timing import log_function_time
+
+logger = setup_logger()
+
+
+@log_function_time(print_only=True)
+def classify_is_search_flow(
+    query: str,
+    llm: LLM,
+) -> bool:
+    messages: LanguageModelInput = [
+        UserMessage(content=SEARCH_CHAT_PROMPT.format(user_query=query))
+    ]
+    response = llm.invoke(
+        prompt=messages,
+        reasoning_effort=ReasoningEffort.OFF,
+        # Nothing can happen in the UI until this call finishes so we need to be aggressive with the timeout
+        timeout_override=2,
+        # Well more than necessary but just to ensure completion and in case it succeeds with classifying but
+        # ends up rambling
+        max_tokens=20,
+    )
+
+    content = llm_response_to_string(response).strip().lower()
+    if not content:
+        logger.warning(
+            "Search flow classification returned empty response; defaulting to chat flow."
+        )
+        return False
+
+    # Prefer chat if both appear.
+    if CHAT_CLASS in content:
+        return False
+    if SEARCH_CLASS in content:
+        return True
+
+    logger.warning(
+        "Search flow classification returned unexpected response; defaulting to chat flow. Response=%r",
+        content,
+    )
+    return False

backend/ee/onyx/server/query_and_chat/models.py

@@ -22,3 +22,11 @@ class DocumentSearchRequest(BasicChunkRequest):
 
 class DocumentSearchResponse(BaseModel):
     top_documents: list[InferenceChunk]
+
+
+class SearchFlowClassificationRequest(BaseModel):
+    user_query: str
+
+
+class SearchFlowClassificationResponse(BaseModel):
+    is_search_flow: bool

backend/ee/onyx/server/query_and_chat/search_backend.py

@@ -0,0 +1,53 @@
+from fastapi import APIRouter
+from fastapi import Depends
+from sqlalchemy.orm import Session
+
+from ee.onyx.secondary_llm_flows.search_flow_classification import (
+    classify_is_search_flow,
+)
+from ee.onyx.server.query_and_chat.models import SearchFlowClassificationRequest
+from ee.onyx.server.query_and_chat.models import SearchFlowClassificationResponse
+from onyx.auth.users import current_user
+from onyx.db.engine.sql_engine import get_session
+from onyx.db.models import User
+from onyx.llm.factory import get_default_llm
+from onyx.server.usage_limits import check_llm_cost_limit_for_provider
+from onyx.utils.logger import setup_logger
+from shared_configs.contextvars import get_current_tenant_id
+
+logger = setup_logger()
+
+router = APIRouter(prefix="/search")
+
+
+@router.post("/search-flow-classification")
+def search_flow_classification(
+    request: SearchFlowClassificationRequest,
+    # This is added just to ensure this endpoint isn't spammed by non-authorized users since there's an LLM call underneath it
+    _: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> SearchFlowClassificationResponse:
+    query = request.user_query
+    # This is a heuristic that if the user is typing a lot of text, it's unlikely they're looking for some specific document
+    # Most likely something needs to be done with the text included so we'll just classify it as a chat flow
+    if len(query) > 200:
+        return SearchFlowClassificationResponse(is_search_flow=False)
+
+    llm = get_default_llm()
+
+    check_llm_cost_limit_for_provider(
+        db_session=db_session,
+        tenant_id=get_current_tenant_id(),
+        llm_provider_api_key=llm.config.api_key,
+    )
+
+    try:
+        is_search_flow = classify_is_search_flow(query=query, llm=llm)
+    except Exception as e:
+        logger.exception(
+            "Search flow classification failed; defaulting to chat flow",
+            exc_info=e,
+        )
+        is_search_flow = False
+
+    return SearchFlowClassificationResponse(is_search_flow=is_search_flow)

backend/onyx/db/models.py

@@ -2333,6 +2333,23 @@ class SearchDoc(Base):
     )
 
 
+class SearchQuery(Base):
+    # This table contains search queries for the Search UI. There are no followups and less is stored because the reply
+    # functionality is simply to rerun the search query again as things may have changed and this is more common for search.
+    __tablename__ = "search_query"
+    id: Mapped[UUID] = mapped_column(
+        PGUUID(as_uuid=True), primary_key=True, default=uuid4
+    )
+    user_id: Mapped[UUID] = mapped_column(PGUUID(as_uuid=True), ForeignKey("user.id"))
+    query: Mapped[str] = mapped_column(String)
+    query_expansions: Mapped[list[str] | None] = mapped_column(
+        postgresql.ARRAY(String), nullable=True
+    )
+    created_at: Mapped[datetime.datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now()
+    )
+
+
 """
 Feedback, Logging, Metrics Tables
 """