diff --git a/mindee/input/inference_parameters.py b/mindee/input/inference_parameters.py index fd9f0dc6..0df0495a 100644 --- a/mindee/input/inference_parameters.py +++ b/mindee/input/inference_parameters.py @@ -29,3 +29,5 @@ class InferenceParameters: """Options for polling. Set only if having timeout issues.""" close_file: bool = True """Whether to close the file after parsing.""" + text_context: Optional[str] = None + """Additional text context used by the model during inference. Not recommended, for specific use only.""" diff --git a/mindee/mindee_http/mindee_api_v2.py b/mindee/mindee_http/mindee_api_v2.py index 940836c4..8a7260fd 100644 --- a/mindee/mindee_http/mindee_api_v2.py +++ b/mindee/mindee_http/mindee_api_v2.py @@ -94,6 +94,8 @@ def req_post_inference_enqueue( data["webhook_ids"] = params.webhook_ids if params.alias and len(params.alias): data["alias"] = params.alias + if params.text_context and (params.text_context): + data["text_context"] = params.text_context if isinstance(input_source, LocalInputSource): files = {"file": input_source.read_contents(params.close_file)} diff --git a/tests/v2/test_client.py b/tests/v2/test_client.py index fb70db24..7ae18db4 100644 --- a/tests/v2/test_client.py +++ b/tests/v2/test_client.py @@ -129,7 +129,8 @@ def test_enqueue_and_parse_path_with_env_token(custom_base_url_client): ) with pytest.raises(MindeeHTTPErrorV2): custom_base_url_client.enqueue_and_get_inference( - input_doc, InferenceParameters("dummy-model") + input_doc, + InferenceParameters("dummy-model", text_context="ignore this message"), ) @@ -176,7 +177,7 @@ def test_error_handling(custom_base_url_client): @pytest.mark.v2 -def test_enqueue(custom_base_url_client): +def test_queue_get(custom_base_url_client): response = custom_base_url_client.get_job("12345678-1234-1234-1234-123456789ABC") assert isinstance(response, JobResponse) assert isinstance(response.job, Job) diff --git a/tests/v2/test_client_integration.py b/tests/v2/test_client_integration.py index 173ad627..3330f7fb 100644 --- a/tests/v2/test_client_integration.py +++ b/tests/v2/test_client_integration.py @@ -131,6 +131,7 @@ def test_parse_file_filled_single_page_must_succeed( polygon=None, confidence=None, alias="py_integration_filled_single", + text_context="this is an invoice.", ) response: InferenceResponse = v2_client.enqueue_and_get_inference( @@ -172,7 +173,9 @@ def test_invalid_uuid_must_throw_error(v2_client: ClientV2) -> None: input_path: Path = FILE_TYPES_DIR / "pdf" / "blank_1.pdf" input_source = PathInput(input_path) - params = InferenceParameters(model_id="INVALID MODEL ID") + params = InferenceParameters( + model_id="INVALID MODEL ID", text_context="ignore this message" + ) with pytest.raises(MindeeHTTPErrorV2) as exc_info: v2_client.enqueue_inference(input_source, params)