[owl] Make RAG more robust: Upon error use fallback search query and skip reranking (#828)

Backend - owl (API server)

- Make RAG more robust: Upon error use fallback search query and skip reranking
- If `search_columns` is not provided, default to all string columns

Author: jiahuei

services/api/src/owl/db/gen_executor.py

@@ -1224,16 +1224,28 @@ async def setup_rag(
             project_id=project.id, table_id=kt_id, request_id=request_id
         )
         kt_cols = {c.column_id for c in kt.column_metadata if not c.is_state_column}
-        t0 = perf_counter()
-        fts_query, vs_query = await lm.generate_search_query(
-            messages=body.messages,
-            rag_params=body.rag_params,
-            **body.hyperparams,
-        )
-        cls._log(
-            f'Query rewrite using "{body.model}" took t={(perf_counter() - t0) * 1e3:,.2f} ms.',
-            request_id=request_id,
-        )
+        try:
+            t0 = perf_counter()
+            fts_query, vs_query = await lm.generate_search_query(
+                messages=body.messages,
+                rag_params=body.rag_params,
+                **body.hyperparams,
+            )
+            cls._log(
+                f'Query rewrite using "{body.model}" took t={(perf_counter() - t0) * 1e3:,.2f} ms.',
+                request_id=request_id,
+            )
+        except Exception as e:
+            cls._log(
+                f"Query rewrite failed with error: {repr(e)}. Using last user message as query.",
+                request_id=request_id,
+            )
+            # Fallback: use last user message
+            for msg in reversed(body.messages):
+                if msg.role == ChatRole.USER:
+                    fts_query = msg.text_content
+                    vs_query = msg.text_content
+                    break
         rows = await kt.hybrid_search(
             fts_query=fts_query,
             vs_query=vs_query,
@@ -1270,14 +1282,20 @@ async def setup_rag(
             chunk.metadata["project_id"] = project.id
             chunk.metadata["table_id"] = body.rag_params.table_id
         if len(rows) > 0 and body.rag_params.reranking_model is not None:
-            order = (
-                await lm.rerank_documents(
-                    model=body.rag_params.reranking_model,
-                    query=vs_query,
-                    documents=kt.rows_to_documents(rows),
+            try:
+                order = (
+                    await lm.rerank_documents(
+                        model=body.rag_params.reranking_model,
+                        query=vs_query,
+                        documents=kt.rows_to_documents(rows),
+                    )
+                ).results
+                chunks = [chunks[i.index] for i in order]
+            except Exception as e:
+                cls._log(
+                    f"Reranking failed with error: {repr(e)}. Proceeding with original order.",
+                    request_id=request_id,
                 )
-            ).results
-            chunks = [chunks[i.index] for i in order]
         chunks = chunks[: body.rag_params.k]
         references = References(chunks=chunks, search_query=vs_query)
         if body.messages[-1].role == ChatRole.USER:
@@ -1286,7 +1304,7 @@ async def setup_rag(
             replacement_idx = -2
         else:
             raise BadInputError("The message list should end with user or assistant message.")
-        rag_prompt = await lm.generate_rag_prompt(
+        rag_prompt = await lm.make_rag_prompt(
             messages=body.messages,
             references=references,
             inline_citations=body.rag_params.inline_citations,

services/api/src/owl/db/models/oss.py

@@ -199,7 +199,9 @@ def _search_query_filter(
         search_columns: list[str] | None,
     ) -> SelectBase:
         # Apply search filters
-        if search_query and search_columns:
+        if search_query:
+            if not search_columns:
+                search_columns = cls.str_cols()
             search_conditions = []
             for column_name in search_columns:
                 if (column := getattr(cls, column_name, None)) is not None:

services/api/src/owl/utils/lm.py

@@ -1585,7 +1585,7 @@ async def generate_search_query(
                 queries[q_type] = generated_query
         return queries["fts"], queries["vs"]
 
-    async def generate_rag_prompt(
+    async def make_rag_prompt(
         self,
         *,
         messages: list[ChatEntry],