diff --git a/backend/alembic/versions/73e9983e5091_add_search_query_table.py b/backend/alembic/versions/73e9983e5091_add_search_query_table.py new file mode 100644 index 0000000000..505021ea22 --- /dev/null +++ b/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") diff --git a/backend/ee/onyx/db/search.py b/backend/ee/onyx/db/search.py new file mode 100644 index 0000000000..9ce4055155 --- /dev/null +++ b/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()) diff --git a/backend/ee/onyx/main.py b/backend/ee/onyx/main.py index 3a7e912c87..17282cd511 100644 --- a/backend/ee/onyx/main.py +++ b/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) diff --git a/backend/ee/onyx/prompts/__init__.py b/backend/ee/onyx/prompts/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/ee/onyx/prompts/search_flow_classification.py b/backend/ee/onyx/prompts/search_flow_classification.py new file mode 100644 index 0000000000..58d4d78372 --- /dev/null +++ b/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: (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 diff --git a/backend/ee/onyx/secondary_llm_flows/__init__.py b/backend/ee/onyx/secondary_llm_flows/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/ee/onyx/secondary_llm_flows/search_flow_classification.py b/backend/ee/onyx/secondary_llm_flows/search_flow_classification.py new file mode 100644 index 0000000000..9e4db38a88 --- /dev/null +++ b/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 diff --git a/backend/ee/onyx/server/query_and_chat/models.py b/backend/ee/onyx/server/query_and_chat/models.py index 60da4d7160..aed7a59115 100644 --- a/backend/ee/onyx/server/query_and_chat/models.py +++ b/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 diff --git a/backend/ee/onyx/server/query_and_chat/search_backend.py b/backend/ee/onyx/server/query_and_chat/search_backend.py new file mode 100644 index 0000000000..afb124ea74 --- /dev/null +++ b/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) diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index f625b7db0d..89c3367083 100644 --- a/backend/onyx/db/models.py +++ b/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 """