[owl] Validate row data before performing LLM generation (#836)

Backend - owl (API server)

- Validate row data before performing LLM generation
- Only log request start for mutating HTTP methods

Author: jiahuei

.github/workflows/ci.yml

@@ -127,7 +127,7 @@ jobs:
             # Replace the org with the key in the .env file
             sed -i "s/$org=.*/$org=$key/g" .env
           done
-          echo "OWL_DB_INIT=False" >> .env
+          echo "OWL_DB_INIT=0" >> .env
           echo "OWL_COMPUTE_STORAGE_PERIOD_SEC=15" >> .env
           echo "OWL_STRIPE_WEBHOOK_SECRET_TEST=${OWL_STRIPE_WEBHOOK_SECRET_TEST}" >> .env
           echo "OWL_STRIPE_PUBLISHABLE_KEY_TEST=${OWL_STRIPE_PUBLISHABLE_KEY_TEST}" >> .env

services/api/src/owl/entrypoints/api.py

@@ -247,7 +247,7 @@ async def log_request(request: Request, call_next):
 
     # Call request
     path = request.url.path
-    if "api/health" not in path:
+    if request.method in ("POST", "PATCH", "PUT", "DELETE"):
         logger.info(make_request_log_str(request))
     response = await call_next(request)
     response.headers["x-request-id"] = request_id

services/api/src/owl/routers/gen_table.py

@@ -553,18 +553,7 @@ async def add_rows(
     billing.has_db_storage_quota()
     billing.has_egress_quota()
     # Validate data
-    try:
-        [table.validate_row_data(d) for d in body.data]
-    except Exception as e:
-        logger.info(
-            (
-                "Row data validation failed. "
-                f'Table={table.schema_id}."{table.table_metadata.short_id}" '
-                f"Org={org.id} "
-                f"User={user.id} "
-                f"Error={repr(e)}"
-            )
-        )
+    [table.validate_row_data(d) for d in body.data]
     executor = MultiRowGenExecutor(
         request=request,
         table=table,
@@ -854,18 +843,7 @@ async def update_rows(
     billing.has_db_storage_quota()
     billing.has_egress_quota()
     # Validate data
-    try:
-        {row_id: table.validate_row_data(d) for row_id, d in body.data.items()}
-    except Exception as e:
-        logger.info(
-            (
-                "Row data validation failed. "
-                f'Table={table.schema_id}."{table.table_metadata.short_id}" '
-                f"Org={org.id} "
-                f"User={user.id} "
-                f"Error={repr(e)}"
-            )
-        )
+    {row_id: table.validate_row_data(d) for row_id, d in body.data.items()}
     await table.update_rows(body.data)
     return OkResponse()
 

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

@@ -384,7 +384,7 @@ async def _get_deployment(
             )
         try:
             logger.info(
-                f'{self.id} - Request started for model "{self.config.id}" ({provider=}, {routing_id=}).'
+                f'Request started for model "{self.config.id}" ({provider=}, {routing_id=}, {self.id}).'
             )
             t0 = perf_counter()
             self.request.state.model_start_time = t0
@@ -396,7 +396,9 @@ async def _get_deployment(
                 is_reasoning_model=is_reasoning_model,
             )
             self.request.state.timing["external_call"] = perf_counter() - t0
-            logger.info(f'{self.id} - Request completed for model "{self.config.id}".')
+            logger.info(
+                f'Request completed for model "{self.config.id}" ({provider=}, {routing_id=}, {self.id}).'
+            )
         except Exception as e:
             mapped_e = self._map_and_log_exception(e, deployment, api_key, **hyperparams)
             if isinstance(mapped_e, (ModelOverloadError, RateLimitExceedError)):

services/api/tests/gen_table/test_row_ops.py

@@ -987,22 +987,17 @@ def test_add_row_wrong_dtype(
         assert "concept" in row.columns
 
         # Test adding data with wrong dtype
-        response = add_table_rows(
-            client,
-            table_type,
-            table.id,
-            [dict(good="dummy1", words="dummy2", stars="dummy3", inputs=TEXT)],
-            stream=stream,
-        )
-        rows = list_table_rows(client, table_type, TABLE_ID_A)
-        assert len(rows.items) == 2
-        row = rows.items[-1]
-        assert row["good"]["value"] is None, row["good"]
-        assert row["good"]["original"] == "dummy1", row["good"]
-        assert row["words"]["value"] is None, row["words"]
-        assert row["words"]["original"] == "dummy2", row["words"]
-        assert row["stars"]["value"] is None, row["stars"]
-        assert row["stars"]["original"] == "dummy3", row["stars"]
+        with pytest.raises(BadInputError) as e:
+            add_table_rows(
+                client,
+                table_type,
+                table.id,
+                [dict(good="dummy1", words="dummy2", stars="dummy3", inputs=TEXT)],
+                stream=stream,
+            )
+        assert 'Column "good": Input should be a valid boolean' in str(e.value)
+        assert 'Column "words": Input should be a valid integer' in str(e.value)
+        assert 'Column "stars": Input should be a valid number' in str(e.value)
 
 
 @pytest.mark.parametrize("table_type", TABLE_TYPES)
@@ -1047,7 +1042,7 @@ def test_add_row_missing_columns(
             table_type,
             stream,
             TABLE_ID_A,
-            data=dict(good="dummy1", inputs=TEXT),
+            data=dict(good=True, inputs=TEXT),
         )
         if stream:
             responses = [r for r in response]
@@ -1057,8 +1052,7 @@ def test_add_row_missing_columns(
         rows = list_table_rows(client, table_type, TABLE_ID_A)
         assert len(rows.items) == 2
         row = rows.items[-1]
-        assert row["good"]["value"] is None, row["good"]
-        assert row["good"]["original"] == "dummy1", row["good"]
+        assert row["good"]["value"] is True, row["good"]
         assert row["words"]["value"] is None, row["words"]
         assert "original" not in row["words"], row["words"]
         assert row["stars"]["value"] is None, row["stars"]

services/api/tests/gen_table/test_row_ops_v2.py

@@ -617,26 +617,15 @@ def test_knowledge_table_embedding(
         row = rows.values[3]
         assert row["Title"] is None, row
         assert row["Text"] is None, row
-        # If embedding with invalid length is added, it will be coerced to None
-        # Original vector will be saved into state
-        response = add_table_rows(
-            client,
-            table_type,
-            table.id,
-            [{"Title": "test", "Title Embed": [1, 2, 3]}],
-            stream=stream,
-        )
-        # We currently dont return anything if LLM is not called
-        assert len(response.rows) == 0 if stream else 1
-        assert all(len(r.columns) == 0 for r in response.rows)
-        # Check the vectors
-        rows = list_table_rows(client, table_type, table.id)
-        assert rows.total == 5
-        row = rows.values[-1]
-        assert row["Title"] == "test", f"{row['Title']=}"
-        assert row["Title Embed"] is None, f"{row['Title Embed']=}"
-        assert row["Text"] is None, f"{row['Title']=}"
-        assert_is_vector_or_none(row["Text Embed"], allow_none=False)
+        # Embedding with invalid length will be rejected
+        with pytest.raises(BadInputError, match="Array input must have length 256"):
+            add_table_rows(
+                client,
+                table_type,
+                table.id,
+                [{"Title": "test", "Title Embed": [1, 2, 3]}],
+                stream=stream,
+            )
 
 
 @flaky(max_runs=3, min_passes=1)
@@ -1697,6 +1686,7 @@ def test_public_web_image(setup: ServingContext):
     - As input to model
     - Has valid raw and thumbnail URLs
     - Reject private URLs
+    - Reject malformed URL
     - Empty input is OK
     """
     table_type = TableType.ACTION
@@ -1729,18 +1719,21 @@ def test_public_web_image(setup: ServingContext):
         response = client.file.get_thumbnail_urls([image_uri])
         assert isinstance(response, GetURLResponse)
         assert response.urls[0] == image_uri
-        # Private URLs should be rejected
-        row = add_table_rows(
-            data=[dict(image="https://host.docker.internal:8080")], **kwargs
-        ).rows[0]
-        assert "cannot be opened" in row.columns["ocr"].content, row
-        row = add_table_rows(data=[dict(image="https://localhost")], **kwargs).rows[0]
-        assert "cannot be opened" in row.columns["ocr"].content, row
-        row = add_table_rows(data=[dict(image="https://192.168.0.1")], **kwargs).rows[0]
-        assert "cannot be opened" in row.columns["ocr"].content, row
         # Empty is OK
         row = add_table_rows(data=[dict()], **kwargs).rows[0]
         assert len(row.columns["ocr"].content) > 0, row
+        # Private URLs should be rejected
+        with pytest.raises(BadInputError, match="URL .+ invalid"):
+            add_table_rows(data=[dict(image="https://host.docker.internal:8080")], **kwargs)
+        with pytest.raises(BadInputError, match="URL .+ invalid"):
+            add_table_rows(data=[dict(image="https://localhost")], **kwargs)
+        with pytest.raises(BadInputError, match="URL .+ invalid"):
+            add_table_rows(data=[dict(image="https://192.168.0.1")], **kwargs)
+        # Malformed URL
+        with pytest.raises(BadInputError, match="URL .+ invalid"):
+            add_table_rows(
+                data=[dict(image='{"url": "https://host.docker.internal:8080"}')], **kwargs
+            )
 
 
 @pytest.mark.parametrize("table_type", TABLE_TYPES)
@@ -2130,20 +2123,15 @@ def test_update_row(
 
         # Test updating data with wrong dtype
         data = dict(ID="2", int="str", float="str", bool="str")
-        response = client.table.update_table_rows(
-            table_type,
-            MultiRowUpdateRequest(table_id=table.id, data={row["ID"]: data}),
-        )
-        assert isinstance(response, OkResponse)
-        _rows = list_table_rows(client, table_type, table.id)
-        assert len(_rows.items) == 1
-        _row = _rows.values[0]
-        t2 = datetime.fromisoformat(_row["Updated at"])
-        assert _row["int"] is None
-        assert _row["float"] is None
-        assert _row["bool"] is None
-        _assert_dict_equal(row, _row, exclude=["Updated at", "int", "float", "bool"])
-        assert t2 > t1
+        with pytest.raises(BadInputError) as e:
+            client.table.update_table_rows(
+                table_type,
+                MultiRowUpdateRequest(table_id=table.id, data={row["ID"]: data}),
+            )
+        assert 'Column "int": Input should be a valid integer' in str(e.value)
+        assert 'Column "float": Input should be a valid number' in str(e.value)
+        assert 'Column "bool": Input should be a valid boolean' in str(e.value)
+        _assert_dict_equal(_row, list_table_rows(client, table_type, table.id).values[0])
 
         if table_type == TableType.KNOWLEDGE:
             # Test updating embedding columns directly
@@ -2163,26 +2151,20 @@ def test_update_row(
             _rows = list_table_rows(client, table_type, table.id)
             assert len(_rows.items) == 1
             _row = _rows.values[0]
-            t3 = datetime.fromisoformat(_row["Updated at"])
+            t2 = datetime.fromisoformat(_row["Updated at"])
             assert sum(_row["Title Embed"]) == 0
             assert sum(_row["Text Embed"]) == len(row["Text Embed"])
-            assert t3 > t2
+            assert t2 > t1
             # Test updating embedding columns with wrong length
-            response = client.table.update_table_rows(
-                table_type,
-                MultiRowUpdateRequest(
-                    table_id=table.id,
-                    data={row["ID"]: {"Title Embed": [0], "Text Embed": [0]}},
-                ),
-            )
-            assert isinstance(response, OkResponse)
-            _rows = list_table_rows(client, table_type, table.id)
-            assert len(_rows.items) == 1
-            _row = _rows.values[0]
-            t4 = datetime.fromisoformat(_row["Updated at"])
-            assert _row["Title Embed"] is None
-            assert _row["Text Embed"] is None
-            assert t4 > t3
+            with pytest.raises(BadInputError, match="Array input must have length 256"):
+                client.table.update_table_rows(
+                    table_type,
+                    MultiRowUpdateRequest(
+                        table_id=table.id,
+                        data={row["ID"]: {"Title Embed": [0], "Text Embed": [0]}},
+                    ),
+                )
+            _assert_dict_equal(_row, list_table_rows(client, table_type, table.id).values[0])
 
 
 @pytest.mark.parametrize("stream", **STREAM_PARAMS)