feat: DIA-1969: Support MIG Image use case (#343)

Co-authored-by: niklub <[email protected]>
Co-authored-by: nik <[email protected]>
Co-authored-by: niklub <[email protected]>

Author: 3 people

Diff for: adala/runtimes/_litellm.py

@@ -539,7 +539,16 @@ async def batch_to_batch(
             )
 
         extra_fields = extra_fields or {}
+        input_field_types = input_field_types or {}
         records = batch.to_dict(orient="records")
+        # in multi-image cases, the number of tokens can be too large for the context window
+        # so we need to split the payloads into chunks
+        # we use this heuristic for MIG projects as they more likely to have multi-image inputs
+        # for other data types, we skip checking the context window as it will be slower
+        ensure_messages_fit_in_context_window = any(
+            input_field_types.get(field) == MessageChunkType.IMAGE_URLS
+            for field in input_field_types
+        )
 
         df_data = await arun_instructor_with_payloads(
             client=self.client,
@@ -556,6 +565,7 @@ async def batch_to_batch(
             instructions_first=instructions_first,
             instructions_template=instructions_template,
             extra_fields=extra_fields,
+            ensure_messages_fit_in_context_window=ensure_messages_fit_in_context_window,
             **self.model_extra,
         )
 

Diff for: adala/skills/collection/label_studio.py

@@ -163,7 +163,15 @@ async def aapply(
                             f"Image tag {tag.name} has multiple variables: {variables}. Cannot mark these variables as image inputs."
                         )
                         continue
-                    input_field_types[variables[0]] = MessageChunkType.IMAGE_URL
+                    input_field_types[variables[0]] = (
+                        MessageChunkType.IMAGE_URLS
+                        if tag.attr.get("valueList")
+                        else MessageChunkType.IMAGE_URL
+                    )
+
+                logger.debug(
+                    f"Using VisionRuntime with input field types: {input_field_types}"
+                )
                 output = await runtime.batch_to_batch(
                     input,
                     input_template=self.input_template,