add user metadata ask field (#2503)

nucliadb_models/src/nucliadb_models/search.py

@@ -17,6 +17,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+import json
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, TypeVar, Union
@@ -690,7 +691,33 @@ class MinScore(BaseModel):
     )
 
 
-class BaseSearchRequest(BaseModel):
+AUDIT_METADATA_MAX_BYTES = 1024 * 10  # 10KB
+
+
+class AuditMetadataBase(BaseModel):
+    audit_metadata: Optional[dict[str, str]] = Field(
+        default=None,
+        title="Audit metadata",
+        description=(
+            "A dictionary containing optional audit-specific metadata, such as user_id, environment, or other contextual information."
+            " This metadata can be leveraged for filtering and analyzing activity logs in future operations."
+            " Each key-value pair represents a piece of metadata relevant to the user's request."
+        ),
+        examples=[{"environment": "test", "user": "my-user-123"}],
+    )
+
+    @field_validator("audit_metadata", mode="after")
+    def check_audit_metadata_size(cls, value):
+        if value:
+            size = len(json.dumps(value).encode("utf-8"))
+            if size > AUDIT_METADATA_MAX_BYTES:
+                raise ValueError(
+                    f"Audit metadata size is too large: {size} bytes. Maximum size allowed: {AUDIT_METADATA_MAX_BYTES}"
+                )
+        return value
+
+
+class BaseSearchRequest(AuditMetadataBase):
     query: str = SearchParamDefaults.query.to_pydantic_field()
     fields: List[str] = SearchParamDefaults.fields.to_pydantic_field()
     filters: Union[List[str], List[Filter]] = Field(
@@ -1186,7 +1213,7 @@ class CustomPrompt(BaseModel):
     )
 
 
-class AskRequest(BaseModel):
+class AskRequest(AuditMetadataBase):
     query: str = SearchParamDefaults.chat_query.to_pydantic_field()
     top_k: int = Field(
         default=20,

nucliadb_models/tests/test_metadata.py

@@ -22,6 +22,7 @@
 from pydantic_core import ValidationError
 
 from nucliadb_models import metadata
+from nucliadb_models.search import AuditMetadataBase
 
 
 def test_relation_validator():
@@ -52,3 +53,9 @@ def test_relation_entity_model_validator():
         type=metadata.RelationNodeType.ENTITY,
         group="places",
     )
+
+
+def test_audit_metadata_base_max_size():
+    AuditMetadataBase(audit_metadata={"test": "test1"})
+    with pytest.raises(ValidationError):
+        AuditMetadataBase(audit_metadata={"test": "a" * 1024 * 10})

nucliadb_utils/src/nucliadb_utils/audit/stream.py

@@ -19,6 +19,7 @@
 #
 import asyncio
 import contextvars
+import json
 import time
 from datetime import datetime, timezone
 from typing import Callable, List, Optional
@@ -88,8 +89,12 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
         context.audit_request.trace_id = get_trace_id()
 
         if request.url.path.split("/")[-1] in ("ask", "search", "find"):
-            body = (await request.body()).decode()
-            context.audit_request.user_request = body
+            if request.method == "POST":
+                body = (await request.body()).decode()
+                context.audit_request.user_request = body
+            elif request.method == "GET":
+                query_params = json.dumps(dict(request.query_params))
+                context.audit_request.user_request = query_params
 
         response = await call_next(request)