From 03798115ab4264d3ca0f7433a4bccc6b007f72ac Mon Sep 17 00:00:00 2001 From: sebastianMindee Date: Wed, 16 Apr 2025 16:08:36 +0200 Subject: [PATCH] :sparkles: add support for rag param in workflow executions --- mindee/client.py | 2 +- mindee/input/workflow_options.py | 4 ++++ mindee/mindee_http/workflow_endpoint.py | 2 ++ tests/workflows/test_workflow_integration.py | 2 +- 4 files changed, 8 insertions(+), 2 deletions(-) diff --git a/mindee/client.py b/mindee/client.py index 85efa354..11c61af0 100644 --- a/mindee/client.py +++ b/mindee/client.py @@ -263,7 +263,7 @@ def execute_workflow( if not options: options = WorkflowOptions( - alias=None, priority=None, full_text=False, public_url=None + alias=None, priority=None, full_text=False, public_url=None, rag=False ) return self._send_to_workflow(GeneratedV1, input_source, workflow_id, options) diff --git a/mindee/input/workflow_options.py b/mindee/input/workflow_options.py index 32b56d45..8e368921 100644 --- a/mindee/input/workflow_options.py +++ b/mindee/input/workflow_options.py @@ -14,6 +14,8 @@ class WorkflowOptions: """Whether to include the full OCR text response in compatible APIs.""" public_url: Optional[str] """A unique, encrypted URL for accessing the document validation interface without requiring authentication.""" + rag: bool + """Whether to enable Retrieval-Augmented Generation.""" def __init__( self, @@ -21,8 +23,10 @@ def __init__( priority: Optional[ExecutionPriority] = None, full_text: Optional[bool] = False, public_url: Optional[str] = None, + rag: Optional[bool] = False, ): self.alias = alias self.priority = priority self.full_text = full_text if full_text else False self.public_url = public_url + self.rag = rag if rag else False diff --git a/mindee/mindee_http/workflow_endpoint.py b/mindee/mindee_http/workflow_endpoint.py index 4fe26d87..7cd4cd2b 100644 --- a/mindee/mindee_http/workflow_endpoint.py +++ b/mindee/mindee_http/workflow_endpoint.py @@ -45,6 +45,8 @@ def workflow_execution_post( params = {} if options.full_text: params["full_text_ocr"] = "true" + if options.rag: + params["rag"] = "true" if isinstance(input_source, UrlInputSource): data["document"] = input_source.url diff --git a/tests/workflows/test_workflow_integration.py b/tests/workflows/test_workflow_integration.py index 3d963393..9acbec39 100644 --- a/tests/workflows/test_workflow_integration.py +++ b/tests/workflows/test_workflow_integration.py @@ -30,7 +30,7 @@ def test_workflow(mindee_client: Client, workflow_id: str, input_path: str): current_date_time = datetime.now().strftime("%Y-%m-%d-%H:%M:%S") alias = f"python-{current_date_time}" priority = ExecutionPriority.LOW - options = WorkflowOptions(alias=alias, priority=priority) + options = WorkflowOptions(alias=alias, priority=priority, rag=True) response = mindee_client.execute_workflow(input_source, workflow_id, options)