feat: index records with suggestions for search engine

src/argilla/server/contexts/datasets.py

@@ -619,7 +619,7 @@ async def create_records(
     async with db.begin_nested():
         db.add_all(records)
         await db.flush(records)
-        await _load_users_from_record_responses(records)
+        await _preload_records_relationships_before_index(db, records)
         await search_engine.index_records(dataset, records)
 
     await db.commit()
@@ -788,6 +788,23 @@ async def _update_record(
     return params, suggestions, vectors, needs_search_engine_update, caches
 
 
+async def _preload_records_relationships_before_index(db: "AsyncSession", records: List[Record]) -> None:
+    for record in records:
+        await _preload_record_relationships_before_index(db, record)
+
+
+async def _preload_record_relationships_before_index(db: "AsyncSession", record: Record) -> None:
+    await db.execute(
+        select(Record)
+        .filter_by(id=record.id)
+        .options(
+            selectinload(Record.responses).selectinload(Response.user),
+            selectinload(Record.suggestions).selectinload(Suggestion.question),
+            selectinload(Record.vectors),
+        )
+    )
+
+
 async def update_records(
     db: "AsyncSession", search_engine: "SearchEngine", dataset: Dataset, records_update: "RecordsUpdate"
 ) -> None:
@@ -866,6 +883,7 @@ async def update_records(
                 include=RecordIncludeParam(keys=[RecordInclude.vectors], vectors=None),
             )
             await dataset.awaitable_attrs.vectors_settings
+            await _preload_records_relationships_before_index(db, records)
             await search_engine.index_records(dataset, records)
 
     await db.commit()
@@ -905,6 +923,7 @@ async def update_record(
 
         if needs_search_engine_update:
             await record.dataset.awaitable_attrs.vectors_settings
+            await _preload_record_relationships_before_index(db, record)
             await search_engine.index_records(record.dataset, [record])
 
     await db.commit()

src/argilla/server/models/database.py

@@ -24,7 +24,14 @@
 from sqlalchemy.ext.mutable import MutableDict, MutableList
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
-from argilla.server.enums import DatasetStatus, MetadataPropertyType, ResponseStatus, SuggestionType, UserRole
+from argilla.server.enums import (
+    DatasetStatus,
+    MetadataPropertyType,
+    QuestionType,
+    ResponseStatus,
+    SuggestionType,
+    UserRole,
+)
 from argilla.server.models.base import DatabaseModel
 from argilla.server.models.metadata_properties import MetadataPropertySettings
 from argilla.server.models.mixins import inserted_at_current_value
@@ -237,6 +244,10 @@ class Question(DatabaseModel):
     def parsed_settings(self) -> QuestionSettings:
         return parse_obj_as(QuestionSettings, self.settings)
 
+    @property
+    def type(self) -> QuestionType:
+        return QuestionType(self.settings["type"])
+
     def __repr__(self):
         return (
             f"Question(id={str(self.id)!r}, name={self.name!r}, required={self.required!r}, "

src/argilla/server/search_engine/base.py

@@ -15,7 +15,6 @@
 from abc import ABCMeta, abstractmethod
 from contextlib import asynccontextmanager
 from typing import (
-    Annotated,
     Any,
     AsyncGenerator,
     ClassVar,
@@ -46,7 +45,6 @@
 
 __all__ = [
     "SearchEngine",
-    "UserResponse",
     "TextQuery",
     "MetadataFilter",
     "TermsMetadataFilter",
@@ -125,11 +123,6 @@ def __init__(self, scope: FilterScope, order: SortOrder):
         self.order = order
 
 
-class UserResponse(BaseModel):
-    values: Optional[Dict[str, Any]]
-    status: ResponseStatus
-
-
 class TextQuery(BaseModel):
     q: str
     field: Optional[str] = None