[owl] Handle code execution upstream errors (#853)

haoshan98 · jiahuei · deafnv · haoshan98 · commit a3061e3df87d · 2025-12-17T11:42:56.000Z
* Replace upstream error with None for code execution

* Handle regen error propagation by saving error column to state

* Error state cleared when regen succeed

* Refactor error state handling (#865)

* [UI] fix build error; remove electron build (#861)

* [owl] Implement endpoints to update owner (#864)

* Don't log expected JamaiExceptions in code executor

* Refactor error state handling

* update

* Bug fix: `ModelConfigCreate.id` must not have a default value

* update

* [owl] Bug fixes and updates (#849)

Backend - owl (API server)

- Bug fixes:
  - Always update reasoning content in state
  - Only record reasoning time if there's reasoning content
  - Anthropic 4.1 and 4.5 models cannot specify both `temperature` and `top_p`
- Delete org secret upon org deletion
- Bump dependencies
- Deps: Pin Pydantic to `2.11.x`
- Test-LLM: Support reasoning content

* update

* Handle possible mid-stream provider error

---------

Co-authored-by: Okiniri &lt;okiniri@rokiniri.com&gt;

---------

Co-authored-by: Tan Jia Huei &lt;tanjiahuei@gmail.com&gt;
Co-authored-by: Okiniri &lt;okiniri@rokiniri.com&gt;
diff --git a/clients/python/src/jamaibase/types/db.py b/clients/python/src/jamaibase/types/db.py
@@ -720,6 +720,14 @@ def check_rerank_cost_per_ksearch(self) -> Self:
 
 class ModelConfigCreate(ModelConfigUpdate):
     # Overrides to make these field required in ModelConfigCreate.
+    id: SanitisedNonEmptyStr = Field(
+        description=(
+            "Unique identifier. "
+            "Users will specify this to select a model. "
+            "Must follow the following format: `{provider}/{model_id}`. "
+            "Examples=['openai/gpt-4o-mini', 'Qwen/Qwen2.5-0.5B']"
+        ),
+    )
     type: _ModelType = Field(
         description="Model type. Can be completion, chat, embed, or rerank.",
     )
diff --git a/services/api/src/owl/db/gen_executor.py b/services/api/src/owl/db/gen_executor.py
@@ -693,6 +693,8 @@ async def _execute_chat_task(self, task: Task, q: Queue[ResultT | None]) -> None
             pass
 
         # Perform completion
+        state_col = f"{task.output_column_name}_"
+        state: dict[str, Any] = self._column_dict.get(state_col, {})
         result = ""
         reasoning = ""
         references = None
@@ -765,7 +767,9 @@ async def _execute_chat_task(self, task: Task, q: Queue[ResultT | None]) -> None
                         )
                     )
                     if chunk.finish_reason == "error":
-                        self._error_columns.append(output_column)
+                        raise BadInputError(
+                            "LLM provider encountered unknown error during streaming."
+                        )
             else:
                 response = await self.lm.chat_completion(
                     messages=req.messages,
@@ -783,14 +787,15 @@ async def _execute_chat_task(self, task: Task, q: Queue[ResultT | None]) -> None
                 reasoning = response.reasoning_content
 
         except Exception as e:
+            result = f"[ERROR] {str(e)}"
             response_kwargs = dict(
                 id=self._request_id,
                 created=int(time()),
                 model="",
                 usage=ChatCompletionUsage(),
                 choices=[
                     ChatCompletionChoice(
-                        message=ChatCompletionMessage(content=f"[ERROR] {str(e)}"),
+                        message=ChatCompletionMessage(content=result),
                         index=0,
                         finish_reason="error",
                     )
@@ -811,17 +816,16 @@ async def _execute_chat_task(self, task: Task, q: Queue[ResultT | None]) -> None
                     row_id=self._row_id,
                 )
             )
-            result = response.content
-            reasoning = response.reasoning_content
+            state["error"] = {"message": result}
             self._error_columns.append(output_column)
             self.log_exception(
                 f'Table "{self._table_id}": Failed to generate completion for column "{output_column}": {repr(e)}',
                 e,
             )
+        else:
+            state.pop("error", None)
         finally:
             await q.put(None)
-            state_col = f"{task.output_column_name}_"
-            state = self._column_dict.get(state_col, {})
             # Always update state
             state["references"] = references.model_dump(mode="json") if references else None
             state["reasoning_content"] = reasoning if reasoning else None
@@ -911,21 +915,25 @@ async def _execute_code_task(self, task: Task, q: Queue[ResultT | None]) -> None
             pass
 
         # Perform code execution
+        state_col = f"{task.output_column_name}_"
+        state: dict[str, Any] = self._column_dict.get(state_col, {})
         result = ""
         try:
-            # Error circuit breaker
-            self._check_upstream_error([body.source_column])
+            error_cols = self._get_upstream_error([body.source_column])
             source_code = self._column_dict.get(body.source_column, "")
 
             # Extract bytes from ColumnDtype.AUDIO and ColumnDtype.IMAGE and put it into a dictionary
             row_data = self._column_dict.copy()
             self.table.postprocess_rows([row_data], include_state=False)
+            # Replace error columns with None value
+            for ec in error_cols:
+                row_data[ec] = None
             for k, v in row_data.items():
                 col = next((col for col in self.table.column_metadata if col.column_id == k), None)
                 if col and (col.dtype == ColumnDtype.AUDIO or col.dtype == ColumnDtype.IMAGE):
                     row_data[k] = await _load_uri_as_bytes(v)
 
-            if source_code and row_data:
+            if source_code:
                 result = await code_executor(
                     request=self.request,
                     organization_id=self.organization.id,
@@ -945,7 +953,7 @@ async def _execute_code_task(self, task: Task, q: Queue[ResultT | None]) -> None
                 usage=ChatCompletionUsage(),
                 choices=[
                     ChatCompletionChoice(
-                        message=ChatCompletionMessage(content=result),
+                        message=ChatCompletionMessage(content=str(result)),
                         index=0,
                     )
                 ],
@@ -966,10 +974,10 @@ async def _execute_code_task(self, task: Task, q: Queue[ResultT | None]) -> None
                     row_id=self._row_id,
                 )
             )
-
             self.log(f'Executed code for column "{output_column}": <{mask_string(result)}>.')
 
         except Exception as e:
+            result = None
             response_kwargs = dict(
                 id=self._request_id,
                 created=int(time()),
@@ -998,14 +1006,17 @@ async def _execute_code_task(self, task: Task, q: Queue[ResultT | None]) -> None
                     row_id=self._row_id,
                 )
             )
-            result = response.content
+            state["error"] = {"message": f"[ERROR] {str(e)}"}
             self._error_columns.append(output_column)
             self.log_exception(
                 f'Table "{self._table_id}": Failed to execute code for column "{output_column}": {repr(e)}',
                 e,
             )
+        else:
+            state.pop("error", None)
         finally:
             await q.put(None)
+            self._column_dict[state_col] = state
             await self._signal_task_completion(task, result)
 
     async def _execute_python_task(self, task: Task, q: Queue[ResultT | None]) -> None:
@@ -1052,21 +1063,27 @@ async def _execute_python_task(self, task: Task, q: Queue[ResultT | None]) -> No
             pass
 
         # Perform python fixed function execution
+        state_col = f"{task.output_column_name}_"
+        state: dict[str, Any] = self._column_dict.get(state_col, {})
         result = ""
         try:
-            # Error circuit breaker
-            # Extract all columns to the left and check for upstream errors
-            self._check_upstream_error(self._extract_all_upstream_columns(output_column))
+            # Extract all columns to the left and get upstream error columns
+            error_cols = self._get_upstream_error(
+                self._extract_all_upstream_columns(output_column)
+            )
 
             # Extract bytes from ColumnDtype.AUDIO and ColumnDtype.IMAGE and put it into a dictionary
             row_data = self._column_dict.copy()
             self.table.postprocess_rows([row_data], include_state=False)
+            # Replace error columns with None value
+            for ec in error_cols:
+                row_data[ec] = None
             for k, v in row_data.items():
                 col = next((col for col in self.table.column_metadata if col.column_id == k), None)
                 if col and (col.dtype == ColumnDtype.AUDIO or col.dtype == ColumnDtype.IMAGE):
                     row_data[k] = await _load_uri_as_bytes(v)
 
-            if body.python_code and row_data:
+            if body.python_code:
                 result = await code_executor(
                     request=self.request,
                     organization_id=self.organization.id,
@@ -1084,7 +1101,7 @@ async def _execute_python_task(self, task: Task, q: Queue[ResultT | None]) -> No
                 usage=ChatCompletionUsage(),
                 choices=[
                     ChatCompletionChoice(
-                        message=ChatCompletionMessage(content=result),
+                        message=ChatCompletionMessage(content=str(result)),
                         index=0,
                     )
                 ],
@@ -1104,12 +1121,12 @@ async def _execute_python_task(self, task: Task, q: Queue[ResultT | None]) -> No
                     row_id=self._row_id,
                 )
             )
-
             self.log(
                 f'Executed python code for column "{output_column}": <{mask_string(result)}>.'
             )
 
         except Exception as e:
+            result = None
             response_kwargs = dict(
                 id=self._request_id,
                 created=int(time()),
@@ -1138,14 +1155,17 @@ async def _execute_python_task(self, task: Task, q: Queue[ResultT | None]) -> No
                     row_id=self._row_id,
                 )
             )
-            result = response.content
+            state["error"] = {"message": f"[ERROR] {str(e)}"}
             self._error_columns.append(output_column)
             self.log_exception(
                 f'Table "{self._table_id}": Failed to execute python code for column "{output_column}": {repr(e)}',
                 e,
             )
+        else:
+            state.pop("error", None)
         finally:
             await q.put(None)
+            self._column_dict[state_col] = state
             await self._signal_task_completion(task, result)
 
     async def _signal_task_completion(self, task: Task, result: Any) -> None:
@@ -1197,12 +1217,25 @@ async def _load_files(self, message: ChatThreadEntry) -> ChatThreadEntry | ChatE
         # logger.warning(f"{message=}")
         return message
 
-    def _check_upstream_error(self, upstream_cols: list[str]) -> None:
+    def _get_upstream_error(self, upstream_cols: list[str]) -> list[str]:
         if not isinstance(upstream_cols, list):
             raise TypeError(f"`upstream_cols` must be a list, got: {type(upstream_cols)}")
-        error_cols = [f'"{col}"' for col in upstream_cols if col in self._error_columns]
+        error_cols = [
+            col
+            for col in upstream_cols
+            if col in self._error_columns
+            or (
+                f"{col}_" in self._column_dict
+                and self._column_dict[f"{col}_"].get("error", {}).get("message", None)
+            )
+        ]
+        return list(set(error_cols))
+
+    def _check_upstream_error(self, upstream_cols: list[str]) -> None:
+        error_cols = self._get_upstream_error(upstream_cols)
+        formatted_error_cols = [f'"{col}"' for col in error_cols]
         if len(error_cols) > 0:
-            raise UpStreamError(f"Upstream columns errored out: {', '.join(error_cols)}")
+            raise UpStreamError(f"Upstream columns errored out: {', '.join(formatted_error_cols)}")
 
     @classmethod
     async def setup_rag(
diff --git a/services/api/src/owl/utils/code.py b/services/api/src/owl/utils/code.py
@@ -18,7 +18,7 @@
 from owl.types import AUDIO_FILE_EXTENSIONS, IMAGE_FILE_EXTENSIONS, ColumnDtype
 from owl.utils.billing import OPENTELEMETRY_CLIENT
 from owl.utils.crypt import decrypt
-from owl.utils.exceptions import BadInputError
+from owl.utils.exceptions import BadInputError, JamaiException
 from owl.utils.io import s3_upload
 
 REQ_COUNTER = OPENTELEMETRY_CLIENT.get_counter("code_executor_requests_total")
@@ -126,7 +126,7 @@ async def code_executor(
     output_column: str,
     row_data: dict | None,
     dtype: str,
-) -> str:
+) -> str | None:
     async with observe_code_execution(
         organization_id=organization_id,
         project_id=project_id,
@@ -160,7 +160,7 @@ async def code_executor(
                     logger.info(
                         f"Code Executor: {request.state.id} - Python code execution completed for column {output_column}"
                     )
-                    return str(result)
+                    return None if result is None else str(result)
 
                 if not isinstance(result, bytes):
                     raise BadInputError(
@@ -195,6 +195,9 @@ async def code_executor(
                     )
                     return uri
 
+        except JamaiException:
+            # Don't log expected JamaiExceptions
+            raise
         except Exception as e:
             logger.error(
                 f"Code Executor: {request.state.id} - Python code execution encountered error for column {output_column} : {e}"
diff --git a/services/api/tests/gen_table/test_row_ops_v2.py b/services/api/tests/gen_table/test_row_ops_v2.py
