Add models for search APIs (#52)

Search APIs didn't have models and they were directly returning the
generic response from the `_send_request` method.

compass_sdk/compass.py

@@ -35,8 +35,13 @@
     DEFAULT_MAX_RETRIES,
     DEFAULT_SLEEP_RETRY_SECONDS,
 )
-from compass_sdk.models import CreateDataSource, DataSource
-from compass_sdk.models.datasources import PaginatedList
+from compass_sdk.models import (
+    CreateDataSource,
+    DataSource,
+    PaginatedList,
+    SearchChunksResponse,
+    SearchDocumentsResponse,
+)
 
 
 @dataclass
@@ -567,31 +572,35 @@ def search_documents(
         query: str,
         top_k: int = 10,
         filters: Optional[List[SearchFilter]] = None,
-    ):
-        return self._search(
+    ) -> SearchDocumentsResponse:
+        result = self._search(
             api_name="search_documents",
             index_name=index_name,
             query=query,
             top_k=top_k,
             filters=filters,
         )
 
+        return SearchDocumentsResponse.model_validate(result.result)
+
     def search_chunks(
         self,
         *,
         index_name: str,
         query: str,
         top_k: int = 10,
         filters: Optional[List[SearchFilter]] = None,
-    ):
-        return self._search(
+    ) -> SearchChunksResponse:
+        result = self._search(
             api_name="search_chunks",
             index_name=index_name,
             query=query,
             top_k=top_k,
             filters=filters,
         )
 
+        return SearchChunksResponse.model_validate(result.result)
+
     def edit_group_authorization(self, *, index_name: str, group_auth_input: GroupAuthorizationInput):
         """
         Edit group authorization for an index

compass_sdk/models/__init__.py

@@ -1,2 +1,3 @@
 # import models into model package
-from compass_sdk.models.datasources import *
+from compass_sdk.models.datasources import *  # noqa: F403
+from compass_sdk.models.search import *  # noqa: F403

compass_sdk/models/search.py

@@ -0,0 +1,45 @@
+from typing import Any, Dict, List, Optional, TypeAlias
+
+from pydantic import BaseModel
+
+Content: TypeAlias = Dict[str, Any]
+
+
+class AssetInfo(BaseModel):
+    content_type: str
+    presigned_url: str
+
+
+class Chunk(BaseModel):
+    chunk_id: str
+    sort_id: int
+    parent_doc_id: str
+    content: Content
+    origin: Optional[Dict[str, Any]] = None
+    assets_info: Optional[list[AssetInfo]] = None
+    score: float
+
+
+class Document(BaseModel):
+    doc_id: str
+    path: str
+    parent_doc_id: str
+    content: Content
+    index_fields: Optional[List[str]] = None
+    authorized_groups: Optional[List[str]] = None
+    chunks: List[Chunk]
+    score: float
+
+
+class ChunkExtended(Chunk):
+    doc_id: str
+    path: str
+    index_fields: Optional[List[str]] = None
+
+
+class SearchDocumentsResponse(BaseModel):
+    hits: List[Document]
+
+
+class SearchChunksResponse(BaseModel):
+    hits: List[ChunkExtended]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "compass-sdk"
-version = "0.4.4"
+version = "0.5.0"
 authors = []
 description = "Compass SDK"