feat: direct search (#99)

Added methods `direct_search` and `direct_search_scroll`.
Added example for scrolling through all chunks in
`scroll_all_chunks.py`.

Author: jeffh-cohere

cohere/compass/clients/compass.py

@@ -49,6 +49,9 @@
     CompassSdkStage,
     CreateDataSource,
     DataSource,
+    DirectSearchInput,
+    DirectSearchResponse,
+    DirectSearchScrollInput,
     Document,
     DocumentStatus,
     ParseableDocument,
@@ -115,7 +118,6 @@ def __init__(
             )
         self.default_max_retries = default_max_retries
         self.default_sleep_retry_seconds = default_sleep_retry_seconds
-
         self.api_method = {
             "create_index": self.session.put,
             "list_indexes": self.session.get,
@@ -130,6 +132,8 @@ def __init__(
             "refresh": self.session.post,
             "upload_documents": self.session.post,
             "update_group_authorization": self.session.post,
+            "direct_search": self.session.post,
+            "direct_search_scroll": self.session.post,
             # Data Sources APIs
             "create_datasource": self.session.post,
             "list_datasources": self.session.get,
@@ -152,6 +156,8 @@ def __init__(
             "refresh": "/api/v1/indexes/{index_name}/_refresh",
             "upload_documents": "/api/v1/indexes/{index_name}/documents/_upload",
             "update_group_authorization": "/api/v1/indexes/{index_name}/group_authorization",  # noqa: E501
+            "direct_search": "/api/v1/indexes/{index_name}/_direct_search",
+            "direct_search_scroll": "/api/v1/indexes/_direct_search/scroll",
             # Data Sources APIs
             "create_datasource": "/api/v1/datasources",
             "list_datasources": "/api/v1/datasources",
@@ -881,6 +887,77 @@ def update_group_authorization(
             raise CompassError(result.error)
         return PutDocumentsResponse.model_validate(result.result)
 
+    def direct_search(
+        self,
+        *,
+        index_name: str,
+        query: dict[str, Any],
+        size: int = 100,
+        scroll: Optional[str] = None,
+        max_retries: Optional[int] = None,
+        sleep_retry_seconds: Optional[int] = None,
+    ) -> DirectSearchResponse:
+        """
+        Perform a direct search query against the Compass API.
+
+        :param index_name: the name of the index
+        :param query: the direct search query (e.g. {"match_all": {}})
+        :param size: the number of results to return
+        :param scroll: the scroll duration (e.g. "1m" for 1 minute)
+        :param max_retries: the maximum number of times to retry the request
+        :param sleep_retry_seconds: the number of seconds to sleep between retries
+
+        :returns: the direct search results
+        :raises CompassError: if the search fails
+        """
+        data = DirectSearchInput(query=query, size=size, scroll=scroll)
+
+        result = self._send_request(
+            api_name="direct_search",
+            index_name=index_name,
+            data=data,
+            max_retries=max_retries,
+            sleep_retry_seconds=sleep_retry_seconds,
+        )
+
+        if result.error:
+            raise CompassError(result.error)
+
+        return DirectSearchResponse.model_validate(result.result)
+
+    def direct_search_scroll(
+        self,
+        *,
+        scroll_id: str,
+        scroll: str = "1m",
+        max_retries: Optional[int] = None,
+        sleep_retry_seconds: Optional[int] = None,
+    ) -> DirectSearchResponse:
+        """
+        Continue a search using a scroll ID from a previous direct_search call.
+
+        :param scroll_id: the scroll ID from a previous direct_search call
+        :param scroll: the scroll duration (e.g. "1m" for 1 minute)
+        :param max_retries: the maximum number of times to retry the request
+        :param sleep_retry_seconds: the number of seconds to sleep between retries
+
+        :returns: the next batch of search results
+        :raises CompassError: if the scroll search fails
+        """
+        data = DirectSearchScrollInput(scroll_id=scroll_id, scroll=scroll)
+
+        result = self._send_request(
+            api_name="direct_search_scroll",
+            data=data,
+            max_retries=max_retries,
+            sleep_retry_seconds=sleep_retry_seconds,
+        )
+
+        if result.error:
+            raise CompassError(result.error)
+
+        return DirectSearchResponse.model_validate(result.result)
+
     # todo Simplify this method so we don't have to ignore the C901 complexity warning.
     def _send_request(  # noqa: C901
         self,

cohere/compass/models/search.py

@@ -84,3 +84,25 @@ class SearchInput(BaseModel):
     query: str
     top_k: int
     filters: Optional[list[SearchFilter]] = None
+
+
+class DirectSearchInput(BaseModel):
+    """Input to direct search APIs."""
+
+    query: dict[str, Any]
+    size: int
+    scroll: Optional[str] = None
+
+
+class DirectSearchScrollInput(BaseModel):
+    """Input to direct search scroll API."""
+
+    scroll_id: str
+    scroll: str
+
+
+class DirectSearchResponse(BaseModel):
+    """Response object for direct search APIs."""
+
+    hits: list[RetrievedChunkExtended]
+    scroll_id: Optional[str] = None

examples/compass_sdk_examples/scroll_all_chunks.py

@@ -0,0 +1,88 @@
+import argparse
+import json
+
+from compass_sdk_examples.utils import get_compass_api
+
+
+def parse_args():
+    """
+    Parse the user arguments using argparse.
+    """
+    parser = argparse.ArgumentParser(
+        description="""
+This script retrieves all chunks from an existing index in Compass using pagination.
+""".strip(),
+        add_help=True,
+    )
+
+    parser.add_argument(
+        "--index-name",
+        type=str,
+        help="Specify the name of the index to retrieve chunks from.",
+        required=True,
+    )
+    parser.add_argument(
+        "--query",
+        type=str,
+        help='JSON string of the query to use (default: {"match_all": {}})',
+        default='{"match_all": {}}',
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        help="Number of documents to retrieve per batch (default: 100)",
+        default=100,
+    )
+    parser.add_argument(
+        "--scroll",
+        type=str,
+        help="Scroll duration (default: '1m')",
+        default="1m",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    index_name = args.index_name
+
+    try:
+        query = json.loads(args.query)
+    except json.JSONDecodeError:
+        print("Error: Invalid JSON in query argument")
+        return
+
+    client = get_compass_api()
+
+    # Inline scroll_all_chunks functionality
+    if query is None:
+        query = {"match_all": {}}  # type: ignore
+
+    response = client.direct_search(
+        index_name=index_name,
+        query=query,  # type: ignore
+        size=args.batch_size,
+        scroll=args.scroll,
+    )
+
+    all_chunks = response.hits
+
+    while response.hits and response.scroll_id:
+        response = client.direct_search_scroll(
+            scroll_id=response.scroll_id,
+            scroll=args.scroll,
+        )
+        all_chunks.extend(response.hits)
+
+    results = all_chunks
+
+    print(f"Retrieved {len(results)} total documents")
+
+    if results:
+        print("\nPreview of first document:")
+        print(results[0])
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "compass-sdk"
-version = "0.17.0"
+version = "0.18.0"
 authors = []
 description = "Compass SDK"
 readme = "README.md"

tests/test_compass_client.py

@@ -166,3 +166,40 @@ def test_get_document_asset_image(requests_mock: Mocker):
     assert isinstance(asset, bytes)
     assert asset == b"test"
     assert content_type == "image/png"
+
+
+def test_direct_search_is_valid(requests_mock: Mocker):
+    # Register mock response for the direct_search endpoint
+    requests_mock.post(
+        "http://test.com/api/v1/indexes/test_index/_direct_search",
+        json={"hits": [], "scroll_id": "test_scroll_id"},
+    )
+
+    compass = CompassClient(index_url="http://test.com")
+    compass.direct_search(index_name="test_index", query={"match_all": {}})
+    assert requests_mock.request_history[0].method == "POST"
+    assert (
+        requests_mock.request_history[0].url
+        == "http://test.com/api/v1/indexes/test_index/_direct_search"
+    )
+    assert "query" in requests_mock.request_history[0].json()
+    assert "size" in requests_mock.request_history[0].json()
+
+
+def test_direct_search_scroll_is_valid(requests_mock: Mocker):
+    # Register mock response for the direct_search_scroll endpoint
+    requests_mock.post(
+        "http://test.com/api/v1/indexes/_direct_search/scroll",
+        json={"hits": [], "scroll_id": "test_scroll_id"},
+    )
+
+    compass = CompassClient(index_url="http://test.com")
+    compass.direct_search_scroll(scroll_id="test_scroll_id")
+    assert requests_mock.request_history[0].method == "POST"
+    assert (
+        requests_mock.request_history[0].url
+        == "http://test.com/api/v1/indexes/_direct_search/scroll"
+    )
+    request_body = requests_mock.request_history[0].json()
+    assert request_body["scroll_id"] == "test_scroll_id"
+    assert request_body["scroll"] == "1m"