Merge pull request #24 from tjmlabs/filter

Jonathan-Adly · web-flow · commit 0459e0a7c50d · 2024-11-14T19:05:39.000-05:00
Add filter method to the sdk
diff --git a/colivara_py/client.py b/colivara_py/client.py
@@ -1,27 +1,16 @@
-import os
-import requests
-from typing import Optional, Dict, Any, List, Union
-from .models import (
-    CollectionIn,
-    CollectionOut,
-    GenericError,
-    GenericMessage,
-    PatchCollectionIn,
-    DocumentIn,
-    DocumentOut,
-    DocumentInPatch,
-    QueryIn,
-    QueryOut,
-    QueryFilter,
-    FileOut,
-    EmbeddingsOut,
-    TaskEnum,
-    EmbeddingsIn,
-)
 import base64
+import os
 from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import requests
 from pydantic import ValidationError
 
+from .models import (CollectionIn, CollectionOut, DocumentIn, DocumentInPatch,
+                     DocumentOut, EmbeddingsIn, EmbeddingsOut, FileOut,
+                     GenericError, GenericMessage, PatchCollectionIn,
+                     QueryFilter, QueryIn, QueryOut, TaskEnum)
+
 
 class ColiVara:
     def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None):
@@ -479,6 +468,76 @@ def search(
         else:
             response.raise_for_status()
 
+    def filter(
+        self,
+        query_filter: Dict[str, Any],
+        expand: Optional[str] = None,
+    ) -> List[Union[DocumentOut, CollectionOut]]:
+        """
+        Filter for documents and collections that meet the criteria of the filter.
+
+        Args:
+            query_filter (Dict[str, Any]): A dictionary specifying the filter criteria.
+                The filter can be used to narrow down the search based on specific criteria.
+                The dictionary should contain the following keys:
+                - "on": "document" or "collection"
+                - "key": str or List[str]
+                - "value": Optional[Union[str, int, float, bool]]
+                - "lookup": One of "key_lookup", "contains", "contained_by", "has_key", "has_keys", "has_any_keys"
+            expand (Optional[str]): A comma-separated list of fields to expand in the response.
+                Currently, only "pages" is supported, the document's pages will be included if provided.
+
+
+        Returns:
+            DocumentOut: The retrieved documents with their details.
+            CollectionOut: The retrieved collections with their details.
+
+        Raises:
+            ValueError: If the query_filter is invalid.
+            requests.HTTPError: If the API request fails.
+
+        Example:
+            # Simple filter
+            results = client.filter({
+                "on": "document",
+                "key": "category",
+                "value": "AI",
+                "lookup": "contains"
+            })
+
+            # Filter with a list of keys
+            results = client.filter({
+                "on": "collection",
+                "key": ["tag1", "tag2"],
+                "lookup": "has_keys"
+            })
+        """
+
+        request_url = f"{self.base_url}/v1/filter/"
+
+        try:
+            filter_obj = QueryFilter(**query_filter)
+            payload = filter_obj.model_dump()
+        except ValidationError as e:
+            raise ValueError(f"Invalid query_filter: {str(e)}")
+
+        params = {"expand": expand}
+
+        response = requests.post(
+            request_url, json=payload, params=params, headers=self.headers
+        )
+
+        if response.status_code == 200:
+            if query_filter["on"] == "document":
+                return [DocumentOut(**doc) for doc in response.json()]
+            else:
+                return [CollectionOut(**col) for col in response.json()]
+        elif response.status_code == 503:
+            error = GenericError(**response.json())
+            raise ValueError(f"Service unavailable: {error.detail}")
+        else:
+            response.raise_for_status()
+
     def file_to_imgbase64(self, file_path: str) -> List[FileOut]:
         """
         Converts a file to a list of base64 encoded images.
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -1,24 +1,17 @@
+import base64
 import os
+from pathlib import Path
+
 import pytest
-import base64
-from colivara_py import ColiVara, AsyncColiVara
-from colivara_py.models import (
-    CollectionOut,
-    DocumentOut,
-    QueryOut,
-    PageOutQuery,
-    FileOut,
-    EmbeddingsOut,
-    PatchCollectionIn,
-    DocumentIn,
-    DocumentInPatch,
-    QueryFilter,
-    GenericMessage,
-)
 import responses
-from requests.exceptions import HTTPError
 from pydantic import ValidationError
-from pathlib import Path
+from requests.exceptions import HTTPError
+
+from colivara_py import AsyncColiVara, ColiVara
+from colivara_py.models import (CollectionOut, DocumentIn, DocumentInPatch,
+                                DocumentOut, EmbeddingsOut, FileOut,
+                                GenericMessage, PageOutQuery,
+                                PatchCollectionIn, QueryFilter, QueryOut)
 
 
 def test_colivara_init_no_api_key():
@@ -988,6 +981,113 @@ def test_search_with_filter(api_key):
     assert result.results[0].document_metadata["category"] == "AI"
 
 
+@responses.activate
+def test_filter_documents(api_key):
+    os.environ["COLIVARA_API_KEY"] = api_key
+    base_url = "https://api.test.com"
+    client = ColiVara(base_url=base_url)
+
+    expected_out = [
+        {
+            "id": 1,
+            "name": "Test Document Fixture",
+            "metadata": {"important": True},
+            "url": "https://www.example.com",
+            "num_pages": 1,
+            "collection_name": "Test Collection Fixture",
+        }
+    ]
+
+    responses.add(
+        responses.POST, f"{client.base_url}/v1/filter/", json=expected_out, status=200
+    )
+
+    query_filter = {
+        "on": "document",
+        "key": "important",
+        "value": True,
+    }
+    result = client.filter(query_filter=query_filter)
+    assert isinstance(result, list)
+    assert len(result) == 1
+
+
+@responses.activate
+def test_filter_documents_expand(api_key):
+    os.environ["COLIVARA_API_KEY"] = api_key
+    base_url = "https://api.test.com"
+    client = ColiVara(base_url=base_url)
+
+    expected_out = [
+        {
+            "id": 1,
+            "name": "Test Document Fixture",
+            "metadata": {"important": True},
+            "url": "https://www.example.com",
+            "num_pages": 1,
+            "collection_name": "Test Collection Fixture",
+            "pages": [
+                {
+                    "document_name": "Test Document Fixture",
+                    "img_base64": "base64_string",
+                    "page_number": 1,
+                }
+            ],
+        }
+    ]
+
+    responses.add(
+        responses.POST,
+        f"{client.base_url}/v1/filter/?expand=pages",
+        json=expected_out,
+        status=200,
+    )
+
+    query_filter = {
+        "on": "document",
+        "key": "important",
+        "value": True,
+    }
+    result = client.filter(query_filter=query_filter, expand="pages")
+    assert isinstance(result, list)
+    assert len(result) == 1
+
+
+@responses.activate
+def test_filter_collections(api_key):
+    os.environ["COLIVARA_API_KEY"] = api_key
+    base_url = "https://api.test.com"
+    client = ColiVara(base_url=base_url)
+
+    expected_out = [
+        {
+            "id": 1,
+            "name": "test_collection",
+            "metadata": {"description": "A test collection"},
+            "num_documents": 2,
+        },
+        {
+            "id": 2,
+            "name": "another_test_collection",
+            "metadata": {"description": "Another test collection"},
+            "num_documents": 3,
+        },
+    ]
+
+    responses.add(
+        responses.POST, f"{client.base_url}/v1/filter/", json=expected_out, status=200
+    )
+
+    query_filter = {
+        "on": "collection",
+        "key": "important",
+        "value": True,
+    }
+    result = client.filter(query_filter=query_filter, expand="pages")
+    assert isinstance(result, list)
+    assert len(result) == 2
+
+
 @responses.activate
 def test_search_service_unavailable(api_key):
     os.environ["COLIVARA_API_KEY"] = api_key
@@ -1006,6 +1106,31 @@ def test_search_service_unavailable(api_key):
     assert "Service unavailable" in str(exc_info.value)
 
 
+@responses.activate
+def test_filter_service_unavailable(api_key):
+    os.environ["COLIVARA_API_KEY"] = api_key
+    base_url = "https://api.test.com"
+    client = ColiVara(base_url=base_url)
+
+    error_response = {"detail": "Service is temporarily unavailable"}
+
+    responses.add(
+        responses.POST, f"{client.base_url}/v1/filter/", json=error_response, status=503
+    )
+
+    with pytest.raises(ValueError) as exc_info:
+        client.filter(
+            query_filter={
+                "on": "document",
+                "key": "category",
+                "value": "AI",
+                "lookup": "contains",
+            }
+        )
+
+    assert "Service unavailable" in str(exc_info.value)
+
+
 @responses.activate
 def test_search_invalid_filter(api_key):
     os.environ["COLIVARA_API_KEY"] = api_key
@@ -1018,6 +1143,18 @@ def test_search_invalid_filter(api_key):
     assert "Invalid query_filter" in str(exc_info.value)
 
 
+@responses.activate
+def test_filter_invalid_filter(api_key):
+    os.environ["COLIVARA_API_KEY"] = api_key
+    base_url = "https://api.test.com"
+    client = ColiVara(base_url=base_url)
+
+    with pytest.raises(ValueError) as exc_info:
+        client.filter(query_filter={"invalid": "filter"})
+
+    assert "Invalid query_filter" in str(exc_info.value)
+
+
 @responses.activate
 def test_search_http_error(api_key):
     os.environ["COLIVARA_API_KEY"] = api_key
@@ -1033,6 +1170,28 @@ def test_search_http_error(api_key):
         client.search("what is 1+1?")
 
 
+@responses.activate
+def test_filter_http_error(api_key):
+    os.environ["COLIVARA_API_KEY"] = api_key
+    base_url = "https://api.test.com"
+    client = ColiVara(base_url=base_url)
+    responses.add(
+        responses.POST,
+        f"{client.base_url}/v1/filter/",
+        json={"error": "Internal Server Error"},
+        status=500,
+    )
+    with pytest.raises(HTTPError):
+        client.filter(
+            query_filter={
+                "on": "document",
+                "key": "category",
+                "value": "AI",
+                "lookup": "contains",
+            }
+        )
+
+
 @pytest.fixture
 def test_file_path(tmp_path):
     file_content = b"Test file content"
