Repository: HumanSignal/Adala
Pull request #343: feat: DIA-1969: Support MIG Image use case
Author: login: robot-ci-heartex
Created: 2025-02-27T12:54:09Z | Merged: 2025-03-17T13:10:38Z

<!-- Follow Merge -->
This PR was automatically generated via Follow Merge. [Workflow Run](https://github.com/HumanSignal/Adala/actions/runs/13900146210)
Please ensure that all linked upstream Pull Requests are merged before proceeding with this one.
- https://github.com/HumanSignal/label-studio-sdk/pull/423
<!-- Follow Merge -->



- **Added support for multiple images in vision models**: Introduced a new `MessageChunkType.IMAGE_URLS` type to handle lists of image URLs in prompts, extending the previous single-image support.
- **Implemented context window management**: Added functionality to ensure messages with images fit within the model's context window by detecting when token counts exceed limits and trimming messages accordingly.
- **Optimized message chunking**: Completely refactored `split_message_into_chunks()` function to handle different field types more efficiently, including the new multi-image capability.
- **Enhanced error handling**: Improved exception handling in LLM calls with better logging and proper token counting for error cases.
- **Added performance monitoring**: Implemented a `debug_time_it` decorator to measure and log execution times of key functions, helping identify performance bottlenecks.
- **Improved logging**: Added more detailed logging throughout the codebase, particularly for input field types and token counts.





Diff:
diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
index c30232eb..d57cd4c8 100644
--- a/adala/runtimes/_litellm.py
+++ b/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 --git a/adala/skills/collection/label_studio.py b/adala/skills/collection/label_studio.py
index 1e02596b..179b8990 100644
--- a/adala/skills/collection/label_studio.py
+++ b/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,
diff --git a/adala/utils/llm_utils.py b/adala/utils/llm_utils.py
index cb27d3e2..6329674a 100644
--- a/adala/utils/llm_utils.py
+++ b/adala/utils/llm_utils.py
@@ -2,15 +2,19 @@
 import logging
 import traceback
 import litellm
-from typing import Any, Dict, List, Type, Optional
-from pydantic import BaseModel
+from litellm import token_counter
+from collections import defaultdict
+from typing import Any, Dict, List, Type, Optional, Tuple
+from pydantic import BaseModel, Field
 from pydantic_core import to_jsonable_python
 from litellm.types.utils import Usage
+from litellm.utils import trim_messages
 from tenacity import Retrying, AsyncRetrying
 from instructor.exceptions import InstructorRetryException, IncompleteOutputException
 from instructor.client import Instructor, AsyncInstructor
 from adala.utils.parse import MessagesBuilder, MessageChunkType
 from adala.utils.exceptions import ConstrainedGenerationError
+from adala.utils.types import debug_time_it
 
 logger = logging.getLogger(__name__)
 
@@ -37,8 +41,8 @@ def _get_usage_dict(usage: Usage, model: str) -> Dict:
         data["_prompt_cost_usd"] = prompt_cost
         data["_completion_cost_usd"] = completion_cost
         data["_total_cost_usd"] = prompt_cost + completion_cost
-    except:
-        logger.exception(f"Failed to get cost for model {model}")
+    except Exception as e:
+        logger.exception(f"Failed to get cost for model {model}", exc_info=e)
         data["_prompt_cost_usd"] = None
         data["_completion_cost_usd"] = None
         data["_total_cost_usd"] = None
@@ -68,7 +72,11 @@ def _log_llm_exception(e) -> dict:
 
 
 def handle_llm_exception(
-    e: Exception, messages: List[Dict[str, str]], model: str, retries
+    e: Exception,
+    messages: List[Dict[str, str]],
+    model: str,
+    retries,
+    prompt_token_count: Optional[int] = None,
 ) -> tuple[Dict, Usage]:
     """Handle exceptions from LLM calls and return standardized error dict and usage stats.
 
@@ -93,11 +101,9 @@ def handle_llm_exception(
         # usage = e.total_usage
         # not available here, so have to approximate by hand, assuming the same error occurred each time
         n_attempts = retries.stop.max_attempt_number
-        # Note that the default model used in token_counter is gpt-3.5-turbo as of now - if model passed in
-        # does not match a mapped model, falls back to default
-        prompt_tokens = n_attempts * litellm.token_counter(
-            model=model, messages=messages[:-1]
-        )  # response is appended as the last message
+        if prompt_token_count is None:
+            prompt_token_count = token_counter(model=model, messages=messages[:-1])
+        prompt_tokens = n_attempts * prompt_token_count
         # TODO a pydantic validation error may be appended as the last message, don't know how to get the raw response in this case
         usage = Usage(
             prompt_tokens=prompt_tokens,
@@ -116,15 +122,46 @@ def handle_llm_exception(
     return _log_llm_exception(e), usage
 
 
+def _ensure_messages_fit_in_context_window(
+    messages: List[Dict[str, str]], model: str
+) -> Tuple[List[Dict[str, str]], int]:
+    """
+    Ensure that the messages fit in the context window of the model.
+    """
+    token_count = token_counter(model=model, messages=messages)
+    logger.debug(f"Prompt tokens count: {token_count}")
+
+    if model in litellm.model_cost:
+        # If we are able to identify the model context window, ensure the messages fit in it
+        max_tokens = litellm.model_cost[model].get(
+            "max_input_tokens", litellm.model_cost[model]["max_tokens"]
+        )
+        if token_count > max_tokens:
+            logger.info(
+                f"Prompt tokens count {token_count} exceeds max tokens {max_tokens} for model {model}. Trimming messages."
+            )
+            # TODO: in case it exceeds max tokens, content of the last message is truncated.
+            # to improve this, we implement:
+            # - UI-level warning for the user, use prediction_meta field for warnings as well as errors in future
+            # - sequential aggregation instead of trimming
+            # - potential v2 solution to downsample images instead of cutting them off (using quality=low instead of quality=auto in completion)
+            return trim_messages(messages, model=model), token_count
+    # in other cases, just return the original messages
+    return messages, token_count
+
+
+@debug_time_it
 def run_instructor_with_messages(
     client: Instructor,
     messages: List[Dict[str, Any]],
     response_model: Type[BaseModel],
     model: str,
+    canonical_model_provider_string: Optional[str] = None,
     max_tokens: Optional[int] = None,
     temperature: Optional[float] = None,
     seed: Optional[int] = None,
     retries: Optional[Retrying] = None,
+    ensure_messages_fit_in_context_window: bool = False,
     **kwargs,
 ) -> Dict[str, Any]:
     """
@@ -135,17 +172,24 @@ def run_instructor_with_messages(
         messages: The messages to send to the model
         response_model: The Pydantic model to validate the response against
         model: The model name to use
+        canonical_model_provider_string: The canonical model provider string to use
         max_tokens: Maximum tokens to generate
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
         retries: Retry policy to use
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion call
 
     Returns:
         Dict containing the parsed response and usage information
     """
     try:
-        # returns a pydantic model and completion info
+        prompt_token_count = None
+        if ensure_messages_fit_in_context_window:
+            messages, prompt_token_count = _ensure_messages_fit_in_context_window(
+                messages, canonical_model_provider_string or model
+            )
+
         response, completion = client.chat.completions.create_with_completion(
             messages=messages,
             response_model=response_model,
@@ -162,7 +206,9 @@ def run_instructor_with_messages(
         usage_model = completion.model
 
     except Exception as e:
-        dct, usage = handle_llm_exception(e, messages, model, retries)
+        dct, usage = handle_llm_exception(
+            e, messages, model, retries, prompt_token_count=prompt_token_count
+        )
         # With exceptions we don't have access to completion.model
         usage_model = model
 
@@ -172,15 +218,18 @@ def run_instructor_with_messages(
     return dct
 
 
+@debug_time_it
 async def arun_instructor_with_messages(
     client: AsyncInstructor,
     messages: List[Dict[str, Any]],
     response_model: Type[BaseModel],
     model: str,
+    canonical_model_provider_string: Optional[str] = None,
     max_tokens: Optional[int] = None,
     temperature: Optional[float] = None,
     seed: Optional[int] = None,
     retries: Optional[AsyncRetrying] = None,
+    ensure_messages_fit_in_context_window: bool = False,
     **kwargs,
 ) -> Dict[str, Any]:
     """
@@ -191,17 +240,24 @@ async def arun_instructor_with_messages(
         messages: The messages to send to the model
         response_model: The Pydantic model to validate the response against
         model: The model name to use
+        canonical_model_provider_string: The canonical model provider string to use
         max_tokens: Maximum tokens to generate
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
         retries: Retry policy to use
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion call
 
     Returns:
         Dict containing the parsed response and usage information
     """
     try:
-        # returns a pydantic model and completion info
+        prompt_token_count = None
+        if ensure_messages_fit_in_context_window:
+            messages, prompt_token_count = _ensure_messages_fit_in_context_window(
+                messages, canonical_model_provider_string or model
+            )
+
         response, completion = await client.chat.completions.create_with_completion(
             messages=messages,
             response_model=response_model,
@@ -218,7 +274,9 @@ async def arun_instructor_with_messages(
         usage_model = completion.model
 
     except Exception as e:
-        dct, usage = handle_llm_exception(e, messages, model, retries)
+        dct, usage = handle_llm_exception(
+            e, messages, model, retries, prompt_token_count=prompt_token_count
+        )
         # With exceptions we don't have access to completion.model
         usage_model = model
 
@@ -234,6 +292,7 @@ def run_instructor_with_payload(
     user_prompt_template: str,
     response_model: Type[BaseModel],
     model: str,
+    canonical_model_provider_string: Optional[str] = None,
     max_tokens: Optional[int] = None,
     temperature: Optional[float] = None,
     seed: Optional[int] = None,
@@ -243,6 +302,7 @@ def run_instructor_with_payload(
     input_field_types: Optional[Dict[str, MessageChunkType]] = None,
     split_into_chunks: bool = False,
     extra_fields: Optional[Dict[str, Any]] = None,
+    ensure_messages_fit_in_context_window: bool = False,
     **kwargs,
 ) -> Dict[str, Any]:
     """
@@ -254,6 +314,7 @@ def run_instructor_with_payload(
         user_prompt_template: The template to use for the user prompt
         response_model: The Pydantic model to validate the response against
         model: The model name to use
+        canonical_model_provider_string: The canonical model provider string to use
         max_tokens: Maximum tokens to generate
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
@@ -263,6 +324,7 @@ def run_instructor_with_payload(
         input_field_types: The types of the input fields
         split_into_chunks: Whether to split the user prompt into chunks
         extra_fields: Additional fields to send to the model
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion call
 
     Returns:
@@ -284,10 +346,12 @@ def run_instructor_with_payload(
         messages,
         response_model,
         model,
+        canonical_model_provider_string,
         max_tokens,
         temperature,
         seed,
         retries,
+        ensure_messages_fit_in_context_window=ensure_messages_fit_in_context_window,
         **kwargs,
     )
 
@@ -298,6 +362,7 @@ async def arun_instructor_with_payload(
     user_prompt_template: str,
     response_model: Type[BaseModel],
     model: str,
+    canonical_model_provider_string: Optional[str] = None,
     max_tokens: Optional[int] = None,
     temperature: Optional[float] = None,
     seed: Optional[int] = None,
@@ -307,6 +372,7 @@ async def arun_instructor_with_payload(
     input_field_types: Optional[Dict[str, MessageChunkType]] = None,
     extra_fields: Optional[Dict[str, Any]] = None,
     split_into_chunks: bool = False,
+    ensure_messages_fit_in_context_window: bool = False,
     **kwargs,
 ) -> Dict[str, Any]:
     """
@@ -327,6 +393,7 @@ async def arun_instructor_with_payload(
         input_field_types: The types of the input fields
         extra_fields: Additional fields to send to the model
         split_into_chunks: Whether to split the user prompt into chunks
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion call
 
     Returns:
@@ -348,10 +415,12 @@ async def arun_instructor_with_payload(
         messages,
         response_model,
         model,
+        canonical_model_provider_string,
         max_tokens,
         temperature,
         seed,
         retries,
+        ensure_messages_fit_in_context_window=ensure_messages_fit_in_context_window,
         **kwargs,
     )
 
@@ -362,15 +431,17 @@ def run_instructor_with_payloads(
     user_prompt_template: str,
     response_model: Type[BaseModel],
     model: str,
-    max_tokens: int,
-    temperature: float,
-    seed: Optional[int],
-    retries: Retrying,
+    canonical_model_provider_string: Optional[str] = None,
+    max_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    seed: Optional[int] = None,
+    retries: Optional[Retrying] = None,
     instructions_template: Optional[str] = None,
     instructions_first: bool = True,
     input_field_types: Optional[Dict[str, MessageChunkType]] = None,
     split_into_chunks: bool = False,
     extra_fields: Optional[Dict[str, Any]] = None,
+    ensure_messages_fit_in_context_window: bool = False,
     **kwargs,
 ) -> List[Dict[str, Any]]:
     """
@@ -383,6 +454,7 @@ def run_instructor_with_payloads(
         user_prompt_template: The template to use for the user prompt
         response_model: The Pydantic model to validate the responses against
         model: The model name to use
+        canonical_model_provider_string: The canonical model provider string to use
         max_tokens: Maximum tokens to generate
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
@@ -392,6 +464,7 @@ def run_instructor_with_payloads(
         input_field_types: The types of the input fields
         split_into_chunks: Whether to split the user prompt into chunks
         extra_fields: Additional fields to send to the model
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion calls
 
     Returns:
@@ -414,10 +487,12 @@ def run_instructor_with_payloads(
             messages,
             response_model,
             model,
+            canonical_model_provider_string,
             max_tokens,
             temperature,
             seed,
             retries,
+            ensure_messages_fit_in_context_window=ensure_messages_fit_in_context_window,
             **kwargs,
         )
         results.append(result)
@@ -431,6 +506,7 @@ async def arun_instructor_with_payloads(
     user_prompt_template: str,
     response_model: Type[BaseModel],
     model: str,
+    canonical_model_provider_string: Optional[str] = None,
     max_tokens: Optional[int] = None,
     temperature: Optional[float] = None,
     seed: Optional[int] = None,
@@ -440,6 +516,7 @@ async def arun_instructor_with_payloads(
     input_field_types: Optional[Dict[str, MessageChunkType]] = None,
     split_into_chunks: bool = False,
     extra_fields: Optional[Dict[str, Any]] = None,
+    ensure_messages_fit_in_context_window: bool = False,
     **kwargs,
 ) -> List[Dict[str, Any]]:
     """
@@ -451,6 +528,7 @@ async def arun_instructor_with_payloads(
         user_prompt_template: The template to use for the user prompt
         response_model: The Pydantic model to validate the responses against
         model: The model name to use
+        canonical_model_provider_string: The canonical model provider string to use
         max_tokens: Maximum tokens to generate
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
@@ -460,6 +538,7 @@ async def arun_instructor_with_payloads(
         input_field_types: The types of the input fields
         split_into_chunks: Whether to split the user prompt into chunks
         extra_fields: Additional fields to send to the model
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion calls
 
     Returns:
@@ -481,10 +560,12 @@ async def arun_instructor_with_payloads(
             messages_builder.get_messages(payload),
             response_model,
             model,
+            canonical_model_provider_string,
             max_tokens,
             temperature,
             seed,
             retries,
+            ensure_messages_fit_in_context_window=ensure_messages_fit_in_context_window,
             **kwargs,
         )
         for payload in payloads
diff --git a/adala/utils/parse.py b/adala/utils/parse.py
index 3a07dd3a..0a6f7880 100644
--- a/adala/utils/parse.py
+++ b/adala/utils/parse.py
@@ -104,7 +104,18 @@ class TemplateChunks(TypedDict):
 match_fields_regex = re.compile(r"(?<!\{)\{([a-zA-Z0-9_]+)\}(?!})")
 
 
-def parse_template(string, include_texts=True) -> List[TemplateChunks]:
+class MessageChunkType(Enum):
+    TEXT = "text"
+    IMAGE_URL = "image_url"
+    IMAGE_URLS = "image_urls"
+
+
+def parse_template(
+    string,
+    include_texts=True,
+    payload: Optional[Dict[str, Any]] = None,
+    input_field_types: Optional[Dict[str, MessageChunkType]] = None,
+) -> List[TemplateChunks]:
     """
     Parses a template string to extract output fields and the text between them.
 
@@ -125,34 +136,57 @@ def parse_template(string, include_texts=True) -> List[TemplateChunks]:
 
     chunks: List[TemplateChunks] = []
     last_index = 0
+    payload = payload or {}
+    input_field_types = input_field_types or {}
 
     for match in match_fields_regex.finditer(string):
         # for match in re.finditer(r'\{(.*?)\}', string):
         # Text before field
-        if last_index < match.start() and include_texts:
-            text = string[last_index : match.start()]
+        start = match.start()
+        if last_index < start and include_texts:
+            text = string[last_index:start]
             chunks.append(
                 {
                     "text": text,
                     "start": last_index,
-                    "end": match.start(),
+                    "end": start,
                     "type": "text",
+                    "data": None,
+                    "field_type": None,
                 }
             )
 
         # Field itself
-        field = match.group(1)
-        start = match.start()
         end = match.end()
-        chunks.append({"text": field, "start": start, "end": end, "type": "var"})
+        field = string[start:end].strip("{}")
+        # Extract the field name by removing the brackets
+        data = payload.get(field)
+        field_type = input_field_types.get(field, MessageChunkType.TEXT)
+        chunks.append(
+            {
+                "text": field,
+                "start": start,
+                "end": end,
+                "type": "var",
+                "data": data,
+                "field_type": field_type,
+            }
+        )
 
-        last_index = match.end()
+        last_index = end
 
     # Text after the last field
     if last_index < len(string) and include_texts:
         text = string[last_index:]
         chunks.append(
-            {"text": text, "start": last_index, "end": len(string), "type": "text"}
+            {
+                "text": text,
+                "start": last_index,
+                "end": len(string),
+                "type": "text",
+                "data": None,
+                "field_type": None,
+            }
         )
 
     return chunks
@@ -174,112 +208,104 @@ class ImageMessageChunk(TypedDict):
 Message = Union[str, List[MessageChunk]]
 
 
-class MessageChunkType(Enum):
-    TEXT = "text"
-    IMAGE_URL = "image_url"
-
-
 def split_message_into_chunks(
-    input_template: str, input_field_types: Dict[str, MessageChunkType], **input_fields
+    input_template: str, input_field_types: Dict[str, MessageChunkType], **payload
 ) -> List[MessageChunk]:
-    """Split a template string with field types into a list of message chunks.
-
-    Takes a template string with placeholders and splits it into chunks based on the field types,
-    preserving the text between placeholders.
+    """Split a template string into message chunks based on field types.
 
     Args:
-        input_template (str): Template string with placeholders, e.g. '{a} is a {b} is an {a}'
-        input_field_types (Dict[str, MessageChunkType]): Dict mapping field names to their types
-        **input_fields: Field values to substitute into template
+        input_template: Template string with placeholders like '{field_name}'
+        input_field_types: Mapping of field names to their chunk types
+        payload: Dictionary with values to substitute into the template instead of placeholders
 
     Returns:
-        List[Dict[str, str]]: List of message chunks with appropriate type and content.
-            Text chunks have format: {'type': 'text', 'text': str}
-            Image chunks have format: {'type': 'image_url', 'image_url': {'url': str}}
+        List of message chunks with appropriate type and content:
+        - Text chunks: {'type': 'text', 'text': str}
+        - Image chunks: {'type': 'image_url', 'image_url': {'url': str}}
 
     Example:
         >>> split_message_into_chunks(
-        ...     '{a} is a {b} is an {a}',
-        ...     {'a': MessageChunkType.TEXT, 'b': MessageChunkType.IMAGE_URL},
-        ...     a='the letter a',
-        ...     b='http://example.com/b.jpg'
+        ...     'Look at {image} and describe {text}',
+        ...     {'image': MessageChunkType.IMAGE_URL, 'text': MessageChunkType.TEXT},
+        ...     {'image': 'http://example.com/img.jpg', 'text': 'this content'}
         ... )
         [
-            {'type': 'text', 'text': 'the letter a is a '},
-            {'type': 'image_url', 'image_url': {'url': 'http://example.com/b.jpg'}},
-            {'type': 'text', 'text': ' is an the letter a'}
+            {'type': 'text', 'text': 'Look at '},
+            {'type': 'image_url', 'image_url': {'url': 'http://example.com/img.jpg'}},
+            {'type': 'text', 'text': ' and describe this content'}
         ]
     """
-    # Parse template to get field positions and surrounding text
-    parsed = parse_template(input_template)
-
-    def add_to_current_chunk(
-        current_chunk: Optional[MessageChunk], chunk: MessageChunk
-    ) -> MessageChunk:
-        if current_chunk:
-            current_chunk["text"] += chunk["text"]
-            return current_chunk
-        else:
-            return chunk
-
-    # Build chunks by iterating through parsed template parts
-    def build_chunks(
-        parsed: Iterable[TemplateChunks],
-    ) -> Generator[MessageChunk, None, None]:
-        current_chunk: Optional[MessageChunk] = None
-
-        for part in parsed:
-            if part["type"] == "text":
-                current_chunk = add_to_current_chunk(
-                    current_chunk, {"type": "text", "text": part["text"]}
-                )
-            elif part["type"] == "var":
-                field_value = part["text"]
-                try:
-                    field_type = input_field_types[field_value]
-                except KeyError:
-                    raise ValueError(
-                        f"Field {field_value} not found in input_field_types. Found fields: {input_field_types}"
-                    )
-                if field_type == MessageChunkType.TEXT:
-                    # try to substitute in variable and add to current chunk
-                    substituted_text = partial_str_format(
-                        f"{{{field_value}}}", **input_fields
-                    )
-                    if substituted_text != field_value:
-                        current_chunk = add_to_current_chunk(
-                            current_chunk, {"type": "text", "text": substituted_text}
-                        )
-                    else:
-                        # be permissive for unfound variables
-                        current_chunk = add_to_current_chunk(
-                            current_chunk,
-                            {"type": "text", "text": f"{{{field_value}}}"},
+    # Parse template to get chunks with field positions and types
+    parsed = parse_template(
+        input_template,
+        include_texts=True,
+        payload=payload,
+        input_field_types=input_field_types,
+    )
+
+    logger.debug(f"Parsed template: {parsed}")
+
+    result = []
+    current_text = ""
+
+    def _add_current_text_as_chunk():
+        # this function is used to flush `current_text` buffer into a text chunk, and start over
+        nonlocal current_text
+        if current_text:
+            result.append({"type": "text", "text": current_text})
+            current_text = ""
+
+    for part in parsed:
+        # iterate over parsed chunks - they already contains field types and placeholder values
+        if part["type"] == "text":
+            # each text chunk without placeholders is added to the current buffer and we continue
+            current_text += part["text"]
+        elif part["type"] == "var":
+            field_type = part["field_type"]
+            field_value = part["data"]
+            if field_value is None:
+                # if field value is not provided, it is assumed to be a text field
+                current_text += part["text"]
+            else:
+                match field_type:
+                    case MessageChunkType.TEXT:
+                        # For text fields, we don't break chunks and add text fields to current buffer
+                        current_text += (
+                            str(field_value)
+                            if field_value is not None
+                            else part["text"]
                         )
-                elif field_type == MessageChunkType.IMAGE_URL:
-                    substituted_text = partial_str_format(
-                        f"{{{field_value}}}", **input_fields
-                    )
-                    if substituted_text != field_value:
-                        # push current chunk, push image chunk, and start new chunk
-                        if current_chunk:
-                            yield current_chunk
-                        current_chunk = None
-                        yield {
-                            "type": "image_url",
-                            "image_url": {"url": input_fields[field_value]},
-                        }
-                    else:
-                        # be permissive for unfound variables
-                        current_chunk = add_to_current_chunk(
-                            current_chunk,
-                            {"type": "text", "text": f"{{{field_value}}}"},
+
+                    case MessageChunkType.IMAGE_URL:
+                        # Add remaining text as text chunk
+                        _add_current_text_as_chunk()
+                        # Add image URL as new image chunk
+                        result.append(
+                            {"type": "image_url", "image_url": {"url": field_value}}
                         )
 
-        if current_chunk:
-            yield current_chunk
+                    case MessageChunkType.IMAGE_URLS:
+                        assert isinstance(
+                            field_value, List
+                        ), "Image URLs must be a list"
+                        # Add remaining text as text chunk
+                        _add_current_text_as_chunk()
+                        # Add image URLs as new image chunks
+                        for url in field_value:
+                            result.append(
+                                {"type": "image_url", "image_url": {"url": url}}
+                            )
+
+                    case _:
+                        # Handle unknown field types as text
+                        current_text += part["text"]
+
+    # Add any remaining text
+    _add_current_text_as_chunk()
+
+    logger.debug(f"Result: {result}")
 
-    return list(build_chunks(parsed))
+    return result
 
 
 class MessagesBuilder(BaseModel):
@@ -296,8 +322,8 @@ def get_messages(self, payload: Dict[str, Any]):
                 lambda: MessageChunkType.TEXT
             )
             user_prompt = split_message_into_chunks(
-                self.user_prompt_template,
-                input_field_types,
+                input_template=self.user_prompt_template,
+                input_field_types=input_field_types,
                 **payload,
                 **self.extra_fields,
             )
diff --git a/adala/utils/types.py b/adala/utils/types.py
index 38db5eb7..c7f787ec 100644
--- a/adala/utils/types.py
+++ b/adala/utils/types.py
@@ -1,6 +1,13 @@
 from pydantic import BaseModel, Field, ConfigDict, field_validator
 from typing import List, Optional, Union
 from adala.utils.parse import parse_template
+import asyncio
+import time
+import logging
+import functools
+from typing import Any, Callable, TypeVar, cast
+
+logger = logging.getLogger(__name__)
 
 
 class BatchData(BaseModel):
@@ -22,3 +29,47 @@ class ErrorResponseModel(BaseModel):
         # guard against name collisions with other fields
         populate_by_name=False,
     )
+
+
+# Type variable for function return type
+T = TypeVar("T")
+
+
+def debug_time_it(func: Callable[..., T]) -> Callable[..., T]:
+    """
+    Decorator that measures the execution time of the wrapped function
+    and logs it at debug level. Works with both synchronous and asynchronous functions.
+
+    Args:
+        func: The function to be wrapped (sync or async)
+
+    Returns:
+        The wrapped function with time measurement
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        logger.debug(
+            f"Function '{func.__name__}' executed in {execution_time:.4f} seconds"
+        )
+        return result
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> T:
+        start_time = time.time()
+        result = await func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        logger.debug(
+            f"Async function '{func.__name__}' executed in {execution_time:.4f} seconds"
+        )
+        return result
+
+    if asyncio.iscoroutinefunction(func):
+        return cast(Callable[..., T], async_wrapper)
+    else:
+        return cast(Callable[..., T], wrapper)
diff --git a/poetry.lock b/poetry.lock
index 12c2b5da..3467cbcb 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "bc777b2985cbd1559869376d2403e6225fe6b475.zip", hash = "sha256:4d31c1798a1a1375f648d7afda20a524afcb4d782a5ee580503dc94c49e3449b"},
+    {file = "388257eeea2ed95ae769724bb8b6ad9c41511b96.zip", hash = "sha256:c86f91aebd989d762663aa49f4fff5dc128f168f70e4540b9d899e456a006f01"},
 ]
 
 [package.dependencies]
@@ -3722,7 +3722,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/bc777b2985cbd1559869376d2403e6225fe6b475.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/388257eeea2ed95ae769724bb8b6ad9c41511b96.zip"
 
 [[package]]
 name = "litellm"
@@ -8467,4 +8467,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "34f440804a475d14314eec103fbd2eda26506f7fa675ca6e6ae9dc040d6854c2"
+content-hash = "8930e9c676607922cdb1bfece4b688db95871e2d5b80c8bde2518d3311ad2675"
diff --git a/pyproject.toml b/pyproject.toml
index 7a133ed3..42a677de 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ celery = {version = "^5.3.6", extras = ["redis"]}
 kombu = ">=5.4.0rc2" # Pin version to fix https://github.com/celery/celery/issues/8030. TODO: remove when this fix will be included in celery
 uvicorn = "*"
 pydantic-settings = "^2.2.1"
-label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/bc777b2985cbd1559869376d2403e6225fe6b475.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/388257eeea2ed95ae769724bb8b6ad9c41511b96.zip"}
 kafka-python-ng = "^2.2.3"
 requests = "^2.32.0"
 # Using litellm from forked repo until vertex fix is released: https://github.com/BerriAI/litellm/issues/7904
diff --git a/tests/cassettes/test_label_studio_skill/test_label_studio_skill.yaml b/tests/cassettes/test_label_studio_skill/test_label_studio_skill_basic.yaml
similarity index 57%
rename from tests/cassettes/test_label_studio_skill/test_label_studio_skill.yaml
rename to tests/cassettes/test_label_studio_skill/test_label_studio_skill_basic.yaml
index dbebc690..b52c0d40 100644
--- a/tests/cassettes/test_label_studio_skill/test_label_studio_skill.yaml
+++ b/tests/cassettes/test_label_studio_skill/test_label_studio_skill_basic.yaml
@@ -16,7 +16,7 @@ interactions:
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.47.1
+      - OpenAI/Python 1.63.2
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -26,9 +26,13 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.47.1
+      - 1.63.2
       x-stainless-raw-response:
       - 'true'
+      x-stainless-read-timeout:
+      - '10.0'
+      x-stainless-retry-count:
+      - '0'
       x-stainless-runtime:
       - CPython
       x-stainless-runtime-version:
@@ -38,18 +42,20 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAAwAAAP//jFJNbxMxEL3vrxh84ZJFmw+aNJeKSiDChQtSoQitHHuya2p7LHs2Jaoi
-        8Tf4e/wS5E2aTUWRuFjyezNP783MQwEgjBZLEKqVrFyw5fXErS5w6t7dbj++ifLT9ub19INbzb+8
-        vf58K0a5g9bfUfFj1ytFLlhkQ/5Aq4iSMauO57OqqsaXi4uecKTR5rYmcDmj0hlvykk1mZXVvBwv
-        jt0tGYVJLOFrAQDw0L/Zp9f4QyyhGj0iDlOSDYrlqQhARLIZETIlk1h6FqOBVOQZfW999dKBJuMb
-        uEdrR8Ct9Hewo+4FvKd7kGvqOH+v4KaV/PvnrwTkMxDBGa+BScvd1bl4xE2XZA7oO2uP+P7k1lIT
-        Iq3TkT/hG+NNauuIMpHPzhJTED27LwC+9VPpngQVIZILXDPdoc+C49lBTgy7OCMvjyQTSzvg0+no
-        GbVaI0tj09lUhZKqRT10DiuQnTZ0RhRnmf8285z2Ibfxzf/ID4RSGBh1HSJqo54GHsoi5kv9V9lp
-        xr1hkTBujcKaDca8B40b2dnD/Yi0S4yu3hjfYAzRHI5oE+qqmk0W6/liIkWxL/4AAAD//wMANNbG
-        e1IDAAA=
+        H4sIAAAAAAAAA4xSwYrbMBC9+yumuvQSF9tJk91cSpeydFuWQil0oRSjSBNbqawR0rhtWPLvRXY2
+        zrK70IsO7808vTcz9xmAMFqsQahWsuq8za8ubr+9f3vVXX++u/7wpZCfaKdWd77Hr9/LWzFLHbTZ
+        oeKHrjeKOm+RDbmRVgElY1ItV4tyPq+Wl8uB6EijTW2N53xBeWecyauiWuTFKi8vjt0tGYVRrOFH
+        BgBwP7zJp9P4V6yhmD0gHcYoGxTrUxGACGQTImSMJrJ0LGYTqcgxusH6zesOdn1kkOADNUF2M9j0
+        DAlvMSBIpyGg1Htgghathz31r+Aj/QElHdzA+ENCgUnL/bvznwJu+yhTWtdbe8QPJ+uWGh9oE4/8
+        Cd8aZ2JbB5SRXLIZmbwY2EMG8HMYUf8otfCBOs810y90SbBcjHJiWsxEVtWRZGJpJ3y+nD2jVmtk
+        aWw8G7FQUrWop85pH7LXhs6I7CzzUzPPaY+5jWv+R34ilELPqGsfUBv1OPBUFjCd7UtlpxkPhkXE
+        8NsorNlgSHvQuJW9HY9JxH1k7OqtcQ0GH8x4UVtfF8vVfCUv58VSZIfsHwAAAP//AwDZWlUEXwMA
+        AA==
     headers:
+      CF-Cache-Status:
+      - DYNAMIC
       CF-RAY:
-      - 91497cde1d455bf0-LIS
+      - 91c864e2d94b193e-BCN
       Connection:
       - keep-alive
       Content-Encoding:
@@ -57,14 +63,14 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Wed, 19 Feb 2025 21:53:08 GMT
+      - Fri, 07 Mar 2025 07:31:36 GMT
       Server:
       - cloudflare
       Set-Cookie:
-      - __cf_bm=pdrPVnZs80kceMbFjP5hySilGriDjxO5ZV7rteMaOe8-1740001988-1.0.1.1-iD6jNgvCTil6rsU_6spFovSB0M4jV_wsXefMELr2VqhsCsypVsQDVSGcFc95vOcWOlCN2HD4zM0_wjHw2g.4_A;
-        path=/; expires=Wed, 19-Feb-25 22:23:08 GMT; domain=.api.openai.com; HttpOnly;
+      - __cf_bm=EWCL4yI.h.cjKyJAyYMqZDf5LIA16oiyMePtNXl6jGE-1741332696-1.0.1.1-3pg9E5nnOalXpiCSoMG7zuJBWXrHw2U9Zow_4hgJ2_gcioPQjgv8l88cmL0FBccjrtnJzVlVlt7dWDnpYjOVwp5uHL8MaSet8Q9Lr._QHJw;
+        path=/; expires=Fri, 07-Mar-25 08:01:36 GMT; domain=.api.openai.com; HttpOnly;
         Secure; SameSite=None
-      - _cfuvid=3936tT4iEjeNttD2XMU.Chdl_3WhiDD8G6V4bXj_vu0-1740001988047-0.0.1.1-604800000;
+      - _cfuvid=jGocxY0HpR2UOYQT1cR6yzZDBib5nqhIAPe0oxgS7rU-1741332696642-0.0.1.1-604800000;
         path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
       Transfer-Encoding:
       - chunked
@@ -74,12 +80,10 @@ interactions:
       - X-Request-ID
       alt-svc:
       - h3=":443"; ma=86400
-      cf-cache-status:
-      - DYNAMIC
       openai-organization:
       - heartex
       openai-processing-ms:
-      - '1584'
+      - '474'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -97,7 +101,7 @@ interactions:
       x-ratelimit-reset-tokens:
       - 0s
       x-request-id:
-      - req_a16fad289e23163891d6d35400521a48
+      - req_d143222aad7c6110a628bcceaf54448d
     status:
       code: 200
       message: OK
@@ -113,15 +117,16 @@ interactions:
       "description": "Correctly extracted `MyModel` with all the required parameters
       with correct types", "parameters": {"properties": {"title": {"anyOf": [{"type":
       "string"}, {"type": "null"}], "default": null, "title": "Title"}, "description":
-      {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "description":
-      "List of texts (one or more) for description", "title": "Description"}, "classification":
-      {"description": "Choices for title", "enum": ["Bug report", "Feature request",
-      "Question", "Other"], "title": "Classification", "type": "string"}, "rationale":
-      {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "description":
-      "List of texts (one or more) for title", "title": "Rationale"}, "evaluation":
-      {"description": "Rating for title (0 to 5)", "maximum": 5, "minimum": 0, "title":
-      "Evaluation", "type": "integer"}}, "required": ["classification", "evaluation"],
-      "type": "object"}}}]}'
+      {"anyOf": [{"type": "string"}, {"items": {"type": "string"}, "type": "array"},
+      {"type": "null"}], "default": null, "description": "Text or list of texts for
+      description", "title": "Description"}, "classification": {"description": "Choices
+      for title", "enum": ["Bug report", "Feature request", "Question", "Other"],
+      "title": "Classification", "type": "string"}, "rationale": {"anyOf": [{"type":
+      "string"}, {"items": {"type": "string"}, "type": "array"}, {"type": "null"}],
+      "default": null, "description": "Text or list of texts for title", "title":
+      "Rationale"}, "evaluation": {"description": "Rating for title (0 to 5)", "maximum":
+      5, "minimum": 0, "title": "Evaluation", "type": "integer"}}, "required": ["classification",
+      "evaluation"], "type": "object"}}}]}'
     headers:
       accept:
       - application/json
@@ -130,16 +135,16 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '1562'
+      - '1646'
       content-type:
       - application/json
       cookie:
-      - __cf_bm=pdrPVnZs80kceMbFjP5hySilGriDjxO5ZV7rteMaOe8-1740001988-1.0.1.1-iD6jNgvCTil6rsU_6spFovSB0M4jV_wsXefMELr2VqhsCsypVsQDVSGcFc95vOcWOlCN2HD4zM0_wjHw2g.4_A;
-        _cfuvid=3936tT4iEjeNttD2XMU.Chdl_3WhiDD8G6V4bXj_vu0-1740001988047-0.0.1.1-604800000
+      - __cf_bm=EWCL4yI.h.cjKyJAyYMqZDf5LIA16oiyMePtNXl6jGE-1741332696-1.0.1.1-3pg9E5nnOalXpiCSoMG7zuJBWXrHw2U9Zow_4hgJ2_gcioPQjgv8l88cmL0FBccjrtnJzVlVlt7dWDnpYjOVwp5uHL8MaSet8Q9Lr._QHJw;
+        _cfuvid=jGocxY0HpR2UOYQT1cR6yzZDBib5nqhIAPe0oxgS7rU-1741332696642-0.0.1.1-604800000
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.47.1
+      - OpenAI/Python 1.63.2
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -149,9 +154,13 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.47.1
+      - 1.63.2
       x-stainless-raw-response:
       - 'true'
+      x-stainless-read-timeout:
+      - '10.0'
+      x-stainless-retry-count:
+      - '0'
       x-stainless-runtime:
       - CPython
       x-stainless-runtime-version:
@@ -161,22 +170,21 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAAwAAAP//jFRRT9swEH7vrzj5ZS8tCgXW0LexiYEmxLQhpGmdIte5JMcc27MvsA7x
-        3yc7bdMCk5aHKLrv7r7PX+78OAIQVIo5CNVIVq3Tk7Npe5nTu/dXnw8/nK4+fbtu/9j85vyou7g9
-        vxDjWGGXd6h4U3WgbOs0MlnTw8qjZIxdD2fHWZYdnuZ5Alpboo5ltePJsZ20ZGgyzabHk2w2OczX
-        1Y0lhUHM4fsIAOAxvaNOU+JvMYdsvIm0GIKsUcy3SQDCWx0jQoZAgaVhMR5AZQ2jidJNp/UOwNbq
-        QkmtB+L+edz5HsySWhe/pvIjuVw9nGR3b6e3s6+/v9T56bXf4etbr1wSVHVGbU3awbfx+TMyAGFk
-        m2qvVlfJu/HzBOnrrkXDUbZ4XAgm1rgQ84W4BCXNGwZtazILMV6IEoPy5CLXywxgC9wgOC25sr5N
-        FUpHFytSclt01tXg0VnPKcMnRK45bxoECqFDUBql1ysgU8ZqDCDBebvU2MIDcZO4euLN+aUmXoGt
-        9mSM4aEh1QDFBsq2rTUQDY15y64GbiRDF9AHaOUK0CjbGUZ/AFHLzonBI5nKeoUBuKEAyxUElkym
-        ToRk5JKSArZRGJAZQ+jqGkPKkeBsHB2SGkqsUDFE0xqEsAqM7UGyA++l7jZmnTyJvd/1NHrt+8fO
-        FHqsuiD1ejzX8aftvGtbRw/Ds/EVFRkKTeFRhjRGIrB1PXfkSQyi21sV4bxtHRdsf6KJDacnJ30/
-        MazzgM42IFuWeogfHa23cb9fUSJLSru0XV8lVYPlUDqssexKsjvAaOfUL9W81rs/OZn6f9oPgFLo
-        GMvCeSxJ7Z94SPMYb7t/pW1dToJFQH9PCgsm9PFPlFjJTvd3kOgnpajI1Oidp3QRicoVWXY8zZez
-        fCrF6Gn0FwAA//8DAAwXP1CWBQAA
+        H4sIAAAAAAAAA4xUTW/bMAy951cQuuySFPlokja3dMOwdQiGDUUPXQZDkWmbnSxpEt3OCPLfB9lJ
+        nKYdMB8Mg4/vPZKmtO0BCErFAoQqJKvS6cHN1epuWd883C/p05/78mF19959lR+qL99+sxf9yLCb
+        R1R8YF0oWzqNTNa0sPIoGaPqaH45mkzGs+tZA5Q2RR1puePBpR2UZGgwHo4vB8P5YHS1ZxeWFAax
+        gB89AIBt8451mhT/iAUM+4dIiSHIHMXimAQgvNUxImQIFFgaFv0OVNYwmli6qbQ+AdhanSipdWfc
+        PtuT725YUutk6sNjsfx4jcvn2zoMb582D+H7yPgTv1a6dk1BWWXUcUgn+DG+ODMDEEaWDXdVr5rZ
+        9c8TpM+rEg3HssV2LZhY41os1uIzKGneMWibk1mL/lqkGJQnF71eZwBb4ALBacmZ9WXDUDpOMSMl
+        j6SbKgePznpuMnyDyL3nXYFAIVQIrdcGA0hw3m40lvBMXESznEwOZM4M+/BckCqAIoVRFYaU1Hs5
+        LiQD1y6GdA1k0lhTo76pcrAeSqkPg4TYTYH7vg5RqYnri6ZofJK6OrQ03YkXQ9313vr+ebIrHrMq
+        SL1fon18d9xKbfPYcThbMpGRoVAkHmVofrYIbF3rHX0aB1G9WGjhvC0dJ2x/oYmC4+m01RPdoevQ
+        6XwPsmWpu/hkNO6/oZekyJKajT8eMiVVgWlH7Q6brFKyJ0DvpOvX1byl3XZOJv8f+Q5QCh1jmjiP
+        KamXHXdpHuOd9K+045SbgkVA/0QKEyb08U+kmMlKtzeFCHVgLJOMTI7eeWquC5G5ZDibT+byejKc
+        id6u9xcAAP//AwCtSsqOPAUAAA==
     headers:
       CF-RAY:
-      - 91497ce9cb475bf0-LIS
+      - 91c864eb48a6193e-BCN
       Connection:
       - keep-alive
       Content-Encoding:
@@ -184,7 +192,7 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Wed, 19 Feb 2025 21:53:09 GMT
+      - Fri, 07 Mar 2025 07:31:40 GMT
       Server:
       - cloudflare
       Transfer-Encoding:
@@ -200,7 +208,7 @@ interactions:
       openai-organization:
       - heartex
       openai-processing-ms:
-      - '1443'
+      - '3642'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -218,7 +226,7 @@ interactions:
       x-ratelimit-reset-tokens:
       - 0s
       x-request-id:
-      - req_f226516bc99fc5fdc2297bb5bc5149d9
+      - req_e86560922a7bca1dc5b51c4edbee590d
     status:
       code: 200
       message: OK
@@ -234,15 +242,16 @@ interactions:
       "function": {"name": "MyModel", "description": "Correctly extracted `MyModel`
       with all the required parameters with correct types", "parameters": {"properties":
       {"title": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null,
-      "title": "Title"}, "description": {"anyOf": [{"type": "string"}, {"type": "null"}],
-      "default": null, "description": "List of texts (one or more) for description",
-      "title": "Description"}, "classification": {"description": "Choices for title",
-      "enum": ["Bug report", "Feature request", "Question", "Other"], "title": "Classification",
-      "type": "string"}, "rationale": {"anyOf": [{"type": "string"}, {"type": "null"}],
-      "default": null, "description": "List of texts (one or more) for title", "title":
-      "Rationale"}, "evaluation": {"description": "Rating for title (0 to 5)", "maximum":
-      5, "minimum": 0, "title": "Evaluation", "type": "integer"}}, "required": ["classification",
-      "evaluation"], "type": "object"}}}]}'
+      "title": "Title"}, "description": {"anyOf": [{"type": "string"}, {"items": {"type":
+      "string"}, "type": "array"}, {"type": "null"}], "default": null, "description":
+      "Text or list of texts for description", "title": "Description"}, "classification":
+      {"description": "Choices for title", "enum": ["Bug report", "Feature request",
+      "Question", "Other"], "title": "Classification", "type": "string"}, "rationale":
+      {"anyOf": [{"type": "string"}, {"items": {"type": "string"}, "type": "array"},
+      {"type": "null"}], "default": null, "description": "Text or list of texts for
+      title", "title": "Rationale"}, "evaluation": {"description": "Rating for title
+      (0 to 5)", "maximum": 5, "minimum": 0, "title": "Evaluation", "type": "integer"}},
+      "required": ["classification", "evaluation"], "type": "object"}}}]}'
     headers:
       accept:
       - application/json
@@ -251,16 +260,16 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '1598'
+      - '1682'
       content-type:
       - application/json
       cookie:
-      - __cf_bm=pdrPVnZs80kceMbFjP5hySilGriDjxO5ZV7rteMaOe8-1740001988-1.0.1.1-iD6jNgvCTil6rsU_6spFovSB0M4jV_wsXefMELr2VqhsCsypVsQDVSGcFc95vOcWOlCN2HD4zM0_wjHw2g.4_A;
-        _cfuvid=3936tT4iEjeNttD2XMU.Chdl_3WhiDD8G6V4bXj_vu0-1740001988047-0.0.1.1-604800000
+      - __cf_bm=EWCL4yI.h.cjKyJAyYMqZDf5LIA16oiyMePtNXl6jGE-1741332696-1.0.1.1-3pg9E5nnOalXpiCSoMG7zuJBWXrHw2U9Zow_4hgJ2_gcioPQjgv8l88cmL0FBccjrtnJzVlVlt7dWDnpYjOVwp5uHL8MaSet8Q9Lr._QHJw;
+        _cfuvid=jGocxY0HpR2UOYQT1cR6yzZDBib5nqhIAPe0oxgS7rU-1741332696642-0.0.1.1-604800000
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.47.1
+      - OpenAI/Python 1.63.2
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -270,9 +279,13 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.47.1
+      - 1.63.2
       x-stainless-raw-response:
       - 'true'
+      x-stainless-read-timeout:
+      - '10.0'
+      x-stainless-retry-count:
+      - '0'
       x-stainless-runtime:
       - CPython
       x-stainless-runtime-version:
@@ -282,22 +295,24 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAAwAAAP//jFTfb9s2EH73X3Hg0wrYgeK6seM3Z+iWbCi2YUGBti6EM3mSrqNIhqRs
-        p0H+94H0D9lZCkwPEnXf3XffHXl8GgAIVmIOQjYYZev06Gbc3l3/+s0sFurddjr7/ufVH8XPD/L2
-        0837yZ0Ypgi7+kYyHqIupG2dpsjW7GDpCSMl1svppCiKy+vZdQZaq0insNrF0cSOWjY8GhfjyaiY
-        ji5n++jGsqQg5vBlAADwlN9Jp1G0FXMohgdLSyFgTWJ+dAIQ3upkERgCh4gmimEPSmsimSTddFqf
-        ANFaXUrUuk+8e55O1n2zUOsyLD4uFjXdTW+v/vrt983b72jaz+vx+iTfjvrRZUFVZ+SxSSf40T5/
-        kQxAGGxz7IfHD7l3w5cO6OuuJROTbPG0FJGjpqWYL8XfnXPWRzC0gYo1QdIRlmK4FIqC9OxS0ux6
-        F2FjO61gRVCnzQOuYEMgs7Fz2qLKHAFslXngQlm5zWRSp05XLPHI9wth7DyBp4eOQsxuPsO4F3ff
-        EHAIXXof3NjUEBsCVIqTb8oV9kVU1gOelwI/ZQ1vhrBpWDbARiUNFABBUWBPOYpMg0aSgkObUXN8
-        BDa7XM7pvfILuG84AGquTYANxyalYYkaqvN6AsQGIyC3EC1w67xdE3SBPNDWkWcykgCNSr/pI9Hh
-        ijVHpnCRu0Fr1N2hYe+exdm2Pg9eW389Oa2eqi6g3h/jvf35OBfa1s7bVXhxzEXFhkNTesKQj5sI
-        0bpd7pQnZxDd2UgJ523rYhntP2QS4fhqsuMT/dj36Gy8B6ONqHv728nV8BW+UlFEzjN3HHOJsiHV
-        h/bjjp1iewIMTqr+r5rXuHeVs6n/D30PSEkukiqdJ8XyvOLezVO6FX/kduxyFiwC+TVLKiOTTzuh
-        qMJO7+4qER5DpLas2NTkned8YYnKlUUxGc9W09kYxeB58C8AAAD//wMAcf72EL4FAAA=
+        H4sIAAAAAAAAAwAAAP//jFTbahsxEH33Vwx6asE2vtVO/NgWQymBECdtoS6LLM3uKtFKqi5xjPG/
+        F2nXu3aaQvdBLHNmzpwZzejQAyCCkyUQVlLPKiMHH69u7lcPxfOtKc3tdjLfV+tvnxULD+tVKUk/
+        RujtIzJ/ihoyXRmJXmhVw8wi9RhZx4vZeDqdLEbjBFSao4xhhfGDmR5UQonBZDSZDUaLwfiqiS61
+        YOjIEn72AAAO6Yw6FccXsoRR/2Sp0DlaIFm2TgDEahkthDonnKfKk34HMq08qihdBSnPAK+1zBiV
+        sktcf4ez/65ZVMrsyU9n7lGtv1+b+Wg8n339cffk71afzvLV1HuTBOVBsbZJZ3hrX75KBkAUrVLs
+        zf4m9a7/2oHaIlSofJRNDhvihZe4IcsNWQdjtPWgcAe5kAhRh9uQ/oZwdMwKE5Mm1y8edjpIDluE
+        Il4eiBx2CCwZg5Ga8sThQOeJB4Zcs5dExmTsdC4YbflWSH2wCBZ/B3Q+udkE00bcfYkgnAvxPLkJ
+        VYAvESjnIvrGXK4pItcW6GUp8C5peN+HXSlYCULxqAEdUODohMUUlWKaDlMp/B6EqtMYIxvRQ7gv
+        hQMqRaEc7IQvkwfHXKhWCoX8sqxhqgufqQyn0j8cycUFHXtv/f86mzuLeXBUNgPZ2I/thEtdGKu3
+        7tXAkijMlZlF6tLgEOe1qXPHPCkDCRfLQYzVlfGZ10+oIuFkPqv5SLfAHbqYN6DXnsrOPp01+3fJ
+        l3H0VKTtaReWUVYi70K7xaWBC30G9M6q/lvNW9x15UIV/0PfAYyh8cgzY5ELdllx52Yxvm//cmu7
+        nAQTh/ZZMMy8QBtvgmNOg6xfHeL2zmOV5UIVaI0V6ekhuclG88V0Qa+noznpHXt/AAAA//8DAKiV
+        ZFyIBQAA
     headers:
+      CF-Cache-Status:
+      - DYNAMIC
       CF-RAY:
-      - 91497cf3e8b85bf0-LIS
+      - 91c86503fadb193e-BCN
       Connection:
       - keep-alive
       Content-Encoding:
@@ -305,7 +320,7 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Wed, 19 Feb 2025 21:53:15 GMT
+      - Fri, 07 Mar 2025 07:31:43 GMT
       Server:
       - cloudflare
       Transfer-Encoding:
@@ -316,12 +331,10 @@ interactions:
       - X-Request-ID
       alt-svc:
       - h3=":443"; ma=86400
-      cf-cache-status:
-      - DYNAMIC
       openai-organization:
       - heartex
       openai-processing-ms:
-      - '5402'
+      - '2389'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -339,7 +352,7 @@ interactions:
       x-ratelimit-reset-tokens:
       - 0s
       x-request-id:
-      - req_60a31a7131c4ac9f05ed02e403eb5302
+      - req_31397671a786d58cb8d6d080c8d49b15
     status:
       code: 200
       message: OK
diff --git a/tests/cassettes/test_label_studio_skill/test_label_studio_skill_partial_label_config.yaml b/tests/cassettes/test_label_studio_skill/test_label_studio_skill_partial_label_config.yaml
index b7fdf4ec..4e1a1e95 100644
--- a/tests/cassettes/test_label_studio_skill/test_label_studio_skill_partial_label_config.yaml
+++ b/tests/cassettes/test_label_studio_skill/test_label_studio_skill_partial_label_config.yaml
@@ -13,13 +13,10 @@ interactions:
       - '142'
       content-type:
       - application/json
-      cookie:
-      - __cf_bm=pdrPVnZs80kceMbFjP5hySilGriDjxO5ZV7rteMaOe8-1740001988-1.0.1.1-iD6jNgvCTil6rsU_6spFovSB0M4jV_wsXefMELr2VqhsCsypVsQDVSGcFc95vOcWOlCN2HD4zM0_wjHw2g.4_A;
-        _cfuvid=3936tT4iEjeNttD2XMU.Chdl_3WhiDD8G6V4bXj_vu0-1740001988047-0.0.1.1-604800000
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.47.1
+      - OpenAI/Python 1.63.2
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -29,9 +26,13 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.47.1
+      - 1.63.2
       x-stainless-raw-response:
       - 'true'
+      x-stainless-read-timeout:
+      - '10.0'
+      x-stainless-retry-count:
+      - '0'
       x-stainless-runtime:
       - CPython
       x-stainless-runtime-version:
@@ -41,18 +42,18 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAAwAAAP//jFJBbtswELzrFVteerEKyVabxJcAubQ5BcilRYtCoMmVxITisuSqjhEY
-        6Df6vb6koOxYCpoAvRDgzO5gZncfMwBhtFiDUJ1k1XubXy37a/5wG+TtrvtRfbXbdvuF/NXDXXlz
-        81EsUgdt7lDxU9c7Rb23yIbcgVYBJWNSLc+qoiiWxer9SPSk0aa21nNeUd4bZ/Jlsazy4iwvz4/d
-        HRmFUazhWwYA8Di+yafT+CDWUCyekB5jlC2K9akIQASyCREyRhNZOhaLiVTkGN1o/fptD5qMa2GL
-        1i6AO+nuYUfDG/hEW5AbGjh9L+FzJ/nPr98RyCUgQG+cBiYtd5dz8YDNEGUK6AZrj/j+5NZS6wNt
-        4pE/4Y1xJnZ1QBnJJWeRyYuR3WcA38epDM+CCh+o91wz3aNLgmV1kBPTLmbkxZFkYmknfLVavKBW
-        a2RpbJxNVSipOtRT57QCOWhDMyKbZf7XzEvah9zGtf8jPxFKoWfUtQ+ojXoeeCoLmC71tbLTjEfD
-        ImL4aRTWbDCkPWhs5GAP9yPiLjL2dWNci8EHcziixtflClFXjcJSZPvsLwAAAP//AwCeZnj5UgMA
+        H4sIAAAAAAAAA4xSTW/bMAy9+1dwuuwSD4kTJFkuw9Z9NAWGFT3sMhSGItG2OlkUJHprUOS/F7LT
+        OEU7YBcd3iOf3iP5kAEIo8UGhGokq9bb/NP6+8/L693Hi6uvyhdttQ331/ztx5ebm88XV2KSOmh3
+        h4qfut4par1FNuQGWgWUjEl1tlrM5vNiXax6oiWNNrXVnvMF5a1xJi+mxSKfrvLZ+tjdkFEYxQZ+
+        ZQAAD/2bfDqN92ID08kT0mKMskaxORUBiEA2IULGaCJLx2Iykooco+utb9+2cNdFBgk+UB1kO4Fd
+        x5DwBgOCdBoCSr0HJmjQethT9wYu6S8o6WALww8JBSYt9x/OfwpYdVGmtK6z9ogfTtYt1T7QLh75
+        E14ZZ2JTBpSRXLIZmbzo2UMGcNuPqHuWWvhAreeS6Te6JDhbDHJiXMxIFsWRZGJpR3y+nLyiVmpk
+        aWw8G7FQUjWox85xH7LThs6I7CzzSzOvaQ+5jav/R34klELPqEsfUBv1PPBYFjCd7b/KTjPuDYuI
+        4Y9RWLLBkPagsZKdHY5JxH1kbMvKuBqDD2a4qMqX0+VqvpLv59OlyA7ZIwAAAP//AwAbJADZXwMA
         AA==
     headers:
       CF-RAY:
-      - 91497e0d5edd5bd6-LIS
+      - 91c8681958a9193e-BCN
       Connection:
       - keep-alive
       Content-Encoding:
@@ -60,9 +61,15 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Wed, 19 Feb 2025 21:53:56 GMT
+      - Fri, 07 Mar 2025 07:33:47 GMT
       Server:
       - cloudflare
+      Set-Cookie:
+      - __cf_bm=yIR9IRtAoDtDYSRbRz65V1uVz0aOpmyCXJxUKHgFWbw-1741332827-1.0.1.1-XlDI1BLYNLz4NbsovQYC6lJe2DQFMVLbU.rIngLG4NuMuyS18B.QOZjy1miwZSbVq2qX5Q7k5ZK4AKKpALopTxZg3F4pvhu0bw2yBVWryyQ;
+        path=/; expires=Fri, 07-Mar-25 08:03:47 GMT; domain=.api.openai.com; HttpOnly;
+        Secure; SameSite=None
+      - _cfuvid=efhEBYVKW9HgP_4tPYEsUivokKWavkQchC9QguDremU-1741332827927-0.0.1.1-604800000;
+        path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
       Transfer-Encoding:
       - chunked
       X-Content-Type-Options:
@@ -76,7 +83,7 @@ interactions:
       openai-organization:
       - heartex
       openai-processing-ms:
-      - '1075'
+      - '434'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -88,13 +95,13 @@ interactions:
       x-ratelimit-remaining-requests:
       - '29999'
       x-ratelimit-remaining-tokens:
-      - '149999793'
+      - '149999794'
       x-ratelimit-reset-requests:
       - 2ms
       x-ratelimit-reset-tokens:
       - 0s
       x-request-id:
-      - req_25a51b32c9f2c719f83875515b352260
+      - req_a21db11835014a581fc78cdd7a9c704f
     status:
       code: 200
       message: OK
@@ -112,10 +119,10 @@ interactions:
       "Choices for [''title'']", "enum": ["Bug report", "Feature request", "Question",
       "Other"], "title": "Classification", "type": "string"}, "evaluation": {"description":
       "Rating for [''title''] (0 to 5)", "maximum": 5, "minimum": 0, "title": "Evaluation",
-      "type": "integer"}, "rationale": {"anyOf": [{"type": "string"}, {"type": "null"}],
-      "default": null, "description": "List of texts (one or more) for [''title'']",
-      "title": "Rationale"}}, "required": ["classification", "evaluation"], "type":
-      "object"}}}]}'
+      "type": "integer"}, "rationale": {"anyOf": [{"type": "string"}, {"items": {"type":
+      "string"}, "type": "array"}, {"type": "null"}], "default": null, "description":
+      "Text or list of texts for [''title'']", "title": "Rationale"}}, "required":
+      ["classification", "evaluation"], "type": "object"}}}]}'
     headers:
       accept:
       - application/json
@@ -124,16 +131,16 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '1310'
+      - '1352'
       content-type:
       - application/json
       cookie:
-      - __cf_bm=pdrPVnZs80kceMbFjP5hySilGriDjxO5ZV7rteMaOe8-1740001988-1.0.1.1-iD6jNgvCTil6rsU_6spFovSB0M4jV_wsXefMELr2VqhsCsypVsQDVSGcFc95vOcWOlCN2HD4zM0_wjHw2g.4_A;
-        _cfuvid=3936tT4iEjeNttD2XMU.Chdl_3WhiDD8G6V4bXj_vu0-1740001988047-0.0.1.1-604800000
+      - __cf_bm=yIR9IRtAoDtDYSRbRz65V1uVz0aOpmyCXJxUKHgFWbw-1741332827-1.0.1.1-XlDI1BLYNLz4NbsovQYC6lJe2DQFMVLbU.rIngLG4NuMuyS18B.QOZjy1miwZSbVq2qX5Q7k5ZK4AKKpALopTxZg3F4pvhu0bw2yBVWryyQ;
+        _cfuvid=efhEBYVKW9HgP_4tPYEsUivokKWavkQchC9QguDremU-1741332827927-0.0.1.1-604800000
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.47.1
+      - OpenAI/Python 1.63.2
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -143,9 +150,13 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.47.1
+      - 1.63.2
       x-stainless-raw-response:
       - 'true'
+      x-stainless-read-timeout:
+      - '10.0'
+      x-stainless-retry-count:
+      - '0'
       x-stainless-runtime:
       - CPython
       x-stainless-runtime-version:
@@ -155,22 +166,22 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAAwAAAP//jFTbattAEH33VwzzLAf5kpvfGkKhhAQaCoFURaxXI2mbvbGXuK7xv5eV
-        bEtJU6gexGrOnDPX1W4CgKLCFSBvWeDKyunNXH2JXx/ZtX1etE+feft7bTbV5u52/RQeMEsMs/5J
-        PBxZZ9woKykIo3uYO2KBkurscpnn+TxfXHSAMhXJRGtsmC7NVAktpvN8vpzml9PZ1YHdGsHJ4wq+
-        TwAAdt075akr+oUryLOjRZH3rCFcnZwA0BmZLMi8Fz4wHTAbQG50IJ1S11HKERCMkSVnUg6B+2c3
-        Og/NYlKWD7dXj/d3z8+f2MVt+7Ldztq8engiMYrXS29tl1AdNT81aYSf7Kt3wQBQM9Vx77f3Xe+y
-        9w7MNVGRDilt3BXIZaq7FpwlyQJXBd7EBhxZ40KBWYH0ymQ8oudZga77YJI6728tgfA+EnBJzMkt
-        CF0lOfLAwDqzlqRgI0IL0jSN0A0IHQyElsBKFmrjVAabVvAWRKJwo5TRkLoApoZ1bCC0LED05Dwo
-        tgXS3EQdyJ0VuMc3Fe4nH51/jAbnqI6eycNED/b9aUWkaVLO/t3EsRZa+LZ0xHzXefTB2D52itNF
-        wPhmu9A6o2wog3khnQTns+teD4cbMKCLIxhMYHLEOr/KPtArKwpMdOt32njOeEvVQB02n8VKmBEw
-        GVX9dzYfafeVC938j/wAcE42UFVaR5Xgbyse3BylH8S/3E5d7hJGT+5VcCqDIJcmUVHNouyvLfqt
-        D6TKWuiGnHWiu7tY23K2IKqWNacZTvaTPwAAAP//AwCgZ4QfyQQAAA==
+        H4sIAAAAAAAAA4xUTU/bQBC951eM5uwgJwYScuRCqzZSQbSHksoa1mN7Yb272l1DQ5T/Xq1NYkOp
+        VB+s9bx5bz7XuwkAygJXgKKmIBqrppfL9Y/Pj+b6U922X9eXj2r95erlJf0p7NPiGpPIMPcPLMKB
+        dSJMYxUHaXQPC8cUOKrOFqezLJsv58sOaEzBKtIqG6anZtpILafzdH46TRfT2fKVXRsp2OMK7iYA
+        ALvuHfPUBf/GFaTJwdKw91Qxro5OAOiMihYk76UPpAMmAyiMDqxj6rpVagQEY1QuSKkhcP/sRueh
+        WaRUnnn/XX+7uinXN3O6pZLLB5HWlI3i9dJb2yVUtlocmzTCj/bVu2AAqKnpuOvtuutd8t6BXNU2
+        rENMG3cbFCrWXUpBUXKDqw1ethU4tsaFDSYb5CdS7QE9Szboug9S3Hnf1gzS+5ZBKCantiB1EeXY
+        A4F15l5xA88y1KBMVUldgdTBQKgZrKJQGtck8FxLUYOMFGGaxmiIXQBTwn1bQagpAJUli+Ch9eyA
+        hGDvTza4xzcV7icfnX+NBue4bD2p14m+2vfHFVGmijn7dxPHUmrp69wx+a7z6IOxfewYp4uA7Zvt
+        QutMY0MezCPrKDifXfR6ONyAAc0OYDCB1Ih1tkw+0MsLDiS79TtuvCBRczFQh82ntpBmBExGVf+d
+        zUfafeVSV/8jPwBCsA1c5NZxIcXbigc3x/EH8S+3Y5e7hNGze5KC8yDZxUkUXFKr+muLfusDN3kp
+        dcXOOtndXSxtnp4vsgVdZOk5TvaTPwAAAP//AwDMOK6uyQQAAA==
     headers:
       CF-Cache-Status:
       - DYNAMIC
       CF-RAY:
-      - 91497e1628e25bd6-LIS
+      - 91c8681f2e5f193e-BCN
       Connection:
       - keep-alive
       Content-Encoding:
@@ -178,7 +189,7 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Wed, 19 Feb 2025 21:53:57 GMT
+      - Fri, 07 Mar 2025 07:33:48 GMT
       Server:
       - cloudflare
       Transfer-Encoding:
@@ -192,7 +203,7 @@ interactions:
       openai-organization:
       - heartex
       openai-processing-ms:
-      - '707'
+      - '761'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -210,7 +221,7 @@ interactions:
       x-ratelimit-reset-tokens:
       - 0s
       x-request-id:
-      - req_c0e90f6adb97b1275842e2622957ba77
+      - req_1baa8020a9af91058f329fcb44643439
     status:
       code: 200
       message: OK
@@ -229,9 +240,10 @@ interactions:
       report", "Feature request", "Question", "Other"], "title": "Classification",
       "type": "string"}, "evaluation": {"description": "Rating for [''title''] (0
       to 5)", "maximum": 5, "minimum": 0, "title": "Evaluation", "type": "integer"},
-      "rationale": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null,
-      "description": "List of texts (one or more) for [''title'']", "title": "Rationale"}},
-      "required": ["classification", "evaluation"], "type": "object"}}}]}'
+      "rationale": {"anyOf": [{"type": "string"}, {"items": {"type": "string"}, "type":
+      "array"}, {"type": "null"}], "default": null, "description": "Text or list of
+      texts for [''title'']", "title": "Rationale"}}, "required": ["classification",
+      "evaluation"], "type": "object"}}}]}'
     headers:
       accept:
       - application/json
@@ -240,16 +252,16 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '1346'
+      - '1388'
       content-type:
       - application/json
       cookie:
-      - __cf_bm=pdrPVnZs80kceMbFjP5hySilGriDjxO5ZV7rteMaOe8-1740001988-1.0.1.1-iD6jNgvCTil6rsU_6spFovSB0M4jV_wsXefMELr2VqhsCsypVsQDVSGcFc95vOcWOlCN2HD4zM0_wjHw2g.4_A;
-        _cfuvid=3936tT4iEjeNttD2XMU.Chdl_3WhiDD8G6V4bXj_vu0-1740001988047-0.0.1.1-604800000
+      - __cf_bm=yIR9IRtAoDtDYSRbRz65V1uVz0aOpmyCXJxUKHgFWbw-1741332827-1.0.1.1-XlDI1BLYNLz4NbsovQYC6lJe2DQFMVLbU.rIngLG4NuMuyS18B.QOZjy1miwZSbVq2qX5Q7k5ZK4AKKpALopTxZg3F4pvhu0bw2yBVWryyQ;
+        _cfuvid=efhEBYVKW9HgP_4tPYEsUivokKWavkQchC9QguDremU-1741332827927-0.0.1.1-604800000
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.47.1
+      - OpenAI/Python 1.63.2
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -259,9 +271,13 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.47.1
+      - 1.63.2
       x-stainless-raw-response:
       - 'true'
+      x-stainless-read-timeout:
+      - '10.0'
+      x-stainless-retry-count:
+      - '0'
       x-stainless-runtime:
       - CPython
       x-stainless-runtime-version:
@@ -271,21 +287,20 @@ interactions:
   response:
     body:
       string: !!binary |
-        H4sIAAAAAAAAAwAAAP//jFTfb9MwEH7vX3HyCyClVdp1dPRxG0gTDA3ohhBBkedckhuO7dnO1q7q
-        /46ctE06hkQerOS++7774busBwCMMjYHJkruRWXk8HRSXTws9cnN99Xi+l21eMpJ6dP3T2efzqqY
-        RYGhb+9Q+B1rJHRlJHrSqoWFRe4xqI5n0ziOJ/HRrAEqnaEMtML44VQPK1I0nMST6TCeDccnW3ap
-        SaBjc/g5AABYN2fIU2W4ZHOIo52lQud4gWy+dwJgVstgYdw5cp4rz6IOFFp5VCF1VUvZA7zWMhVc
-        yi5w+6x7712zuJTpR3/+zZTX519uqtnb5d3V/Y+r7P7rxedevFZ6ZZqE8lqJfZN6+N4+fxYMgCle
-        NdzL1WXTu+i5A7dFXaHyIW22TpiQoe6cBA+SCZsn7ANyX1sEi/c1Op+wKGH4wGW9czmOEmabDy6x
-        oSxKBHKuDueOR6oAVxujrYdcW+Cg8BFykgihQng9yrRYvongsSRRAqksJIEOOGToyGLLUoCq5Epg
-        yBpIgS8R3Mp5rF45ENzwW5LkCd0IFiU54JIK5eCRfBkCkeAS8sOSwAlU3JJ2o4Rt2EGPNoOX3n/1
-        rt5iXjsutzOxtW/2QyZ1Yay+dc9mhuWkyJWpRe6au2POa9PGDnGaCKw+mE9mrK6MT73+jSoITiYn
-        rR7rdqhDj7eTzrz2XPZYsx3rQC/N0HNqBni/M4KLErOO2u0OrzPSPWDQq/rvbF7SbisnVfyPfAcI
-        gcZjlhqLGYnDijs3i+EX8y+3fZebhJlD+0ACU09ow01kmPNatovP2tlKc1IFWmOp2X6Wm3R8hJhN
-        c4FjNtgM/gAAAP//AwByluPPCwUAAA==
+        H4sIAAAAAAAAAwAAAP//jFRda9swFH3Pr7jcpw2ckq8uaR5H1kHXjA3GRpmHUeRrW6ssKfpoWkL+
+        +5CdxG7XwfwgbJ17zj1X98r7AQCKHJeAvGKe10YO3y/W328uw2pz+/V+o1bbxXZ7t9o9fHj8UXyZ
+        YxIZevObuD+xLriujSQvtGphbol5iqrj+Ww8nU4Wk6sGqHVOMtJK44czPayFEsPJaDIbjubD8eLI
+        rrTg5HAJPwcAAPtmjT5VTo+4hFFy2qnJOVYSLs9BAGi1jDvInBPOM+Ux6UCulScVrasgZQ/wWsuM
+        Mym7xO2z7713h8WkzOSn3efVnfqotRJ+fKtXN2vJ7t11L18r/WQaQ0VQ/HxIPfy8v3yRDAAVqxvu
+        +mndnF3yMoDZMtSkfLSN+xS5jHUXgrMomeIyxWtiPlgCS9tAzqeYpEgPTIZTyGWSom0+mKSG8q0i
+        EM6FuJ54QpXggjHaeii0BQaKdlAISRArhDcXueaPbxPYVYJXIFQeTZADBjk5YallKSBVMcUpugZt
+        W5WjRaHAVwTuyXmqL1I84LN6D4PX3n/12mipCI7JY3+P+4fzwEhdGqs37kX/sRBKuCqzxFzTB3Re
+        mzZ3zNNkwPBs1tBYXRufeX1PKgpOJotWD7v70KGz6RH02jPZY83HySt6WU6eiWYYz/PPGa8o76jd
+        PWAhF7oHDHpV/+3mNe22cqHK/5HvAM7JeMozYykX/HnFXZil+Lv4V9j5lBvD6Mg+CE6ZF2RjJ3Iq
+        WJDtJcZ2MLJCqJKssaK5yViYbPRuPp2zq+noHQ4Ogz8AAAD//wMAWBCZL9cEAAA=
     headers:
       CF-RAY:
-      - 91497e1c38135bd6-LIS
+      - 91c86825dd8e193e-BCN
       Connection:
       - keep-alive
       Content-Encoding:
@@ -293,7 +308,7 @@ interactions:
       Content-Type:
       - application/json
       Date:
-      - Wed, 19 Feb 2025 21:53:58 GMT
+      - Fri, 07 Mar 2025 07:33:50 GMT
       Server:
       - cloudflare
       Transfer-Encoding:
@@ -309,7 +324,7 @@ interactions:
       openai-organization:
       - heartex
       openai-processing-ms:
-      - '853'
+      - '1060'
       openai-version:
       - '2020-10-01'
       strict-transport-security:
@@ -327,7 +342,7 @@ interactions:
       x-ratelimit-reset-tokens:
       - 0s
       x-request-id:
-      - req_20ea4e2c05e416eb5367cc2f1f0fa325
+      - req_503ca383dbbbccb07ae7c140f0020073
     status:
       code: 200
       message: OK
diff --git a/tests/test_label_studio_skill.py b/tests/test_label_studio_skill.py
index 9a63238a..032c8e89 100644
--- a/tests/test_label_studio_skill.py
+++ b/tests/test_label_studio_skill.py
@@ -7,14 +7,19 @@
 from label_studio_sdk.label_interface.objects import PredictionValue
 
 
-# NOTE: commented code can be used to recreate vcr cassettes, as they don't work with asyncio
+# NOTE: commented code can be used to recreate vcr cassettes, as they don't work with asyncio (yeah, frustrating - consider use other stub framework)
+# INSTRUCTIONS to recreate cassettes:
+# 1. Uncomment the lines in test you want to recreate
+# 2. Run the test with `pytest -vvv --record-mode=rewrite -k <test_name>`
+# 3. If you see mismatched asserts, verify the results and replace the test suite with the new results
+# 4. Don't forget to uncomment back the lines you commented out before recording the cassette
 
 
 @pytest.mark.asyncio
 @pytest.mark.vcr
-async def test_label_studio_skill():
+async def test_label_studio_skill_basic():
     # @pytest.mark.vcr
-    # def test_label_studio_skill():
+    # def test_label_studio_skill_basic():
 
     df = pd.DataFrame(
         [
@@ -76,8 +81,8 @@ async def test_label_studio_skill():
     assert predictions.classification.tolist() == ["Bug report", "Feature request"]
     assert predictions.evaluation.tolist() == [5, 5]
     assert predictions.rationale.tolist() == [
-        "The issue clearly indicates a problem with the login functionality of the platform, which is a common type of bug that users may encounter. The description reinforces this by stating the inability to log in, suggesting a potential defect in the system.",
-        "The issue is requesting the addition of support for a new file type (.docx), which indicates a desire for enhanced functionality in the application. This aligns with typical feature requests that aim to improve user experience and expand capabilities.",
+        "The issue describes a problem with logging into the platform, which is a technical issue that typically indicates a bug or malfunction in the login functionality.",
+        "The issue is requesting the addition of support for a new file type (.docx), which indicates a desire for new functionality in the application. This aligns with the definition of a feature request.",
     ]
 
 


Reviews:
user: login: matt-bernstein
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T16:43:13Z
user: login: matt-bernstein
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T16:48:35Z
user: login: matt-bernstein
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T16:54:54Z
user: login: matt-bernstein
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T16:58:27Z
user: login: niklub
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T14:19:48Z
user: login: niklub
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T14:19:55Z
user: login: niklub
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T14:21:34Z
user: login: niklub
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T14:35:31Z
user: login: niklub
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T14:36:00Z
user: login: matt-bernstein
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T16:20:06Z
user: login: matt-bernstein
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T18:50:08Z
user: login: matt-bernstein
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T18:50:20Z

Review comments:
diff_hunk: @@ -116,6 +122,30 @@ def handle_llm_exception(
     return _log_llm_exception(e), usage
 
 
+def _ensure_messages_fit_in_context_window(
+    messages: List[Dict[str, str]], model: str
path: adala/utils/llm_utils.py
user: login: matt-bernstein
body: merge in master and make sure `self.canonical_model_string` is being passed in here for model name
created_at: 2025-03-07T16:43:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -116,6 +122,30 @@ def handle_llm_exception(
     return _log_llm_exception(e), usage
 
 
+def _ensure_messages_fit_in_context_window(
+    messages: List[Dict[str, str]], model: str
+) -> Tuple[List[Dict[str, str]], int]:
+    """
+    Ensure that the messages fit in the context window of the model.
+    """
+    token_count = token_counter(model=model, messages=messages)
+    logger.debug(f"Prompt tokens count: {token_count}")
+
+    if model in litellm.model_cost:
+        # If we are able to identify the model context window, ensure the messages fit in it
+        max_tokens = litellm.model_cost[model].get(
+            "max_input_tokens", litellm.model_cost[model]["max_tokens"]
+        )
+        if token_count > max_tokens:
+            logger.info(
+                f"Prompt tokens count {token_count} exceeds max tokens {max_tokens} for model {model}. Trimming messages."
+            )
+            return trim_messages(messages, model=model), token_count
path: adala/utils/llm_utils.py
user: login: matt-bernstein
body: add comment, and warning in docs: this loses message content, does not send the entire message to the provider

potential v2 solution to downsample images instead of cutting them off (using quality=low instead of quality=auto in completion)
created_at: 2025-03-07T16:48:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -116,6 +122,30 @@ def handle_llm_exception(
     return _log_llm_exception(e), usage
 
 
+def _ensure_messages_fit_in_context_window(
+    messages: List[Dict[str, str]], model: str
+) -> Tuple[List[Dict[str, str]], int]:
+    """
+    Ensure that the messages fit in the context window of the model.
+    """
+    token_count = token_counter(model=model, messages=messages)
+    logger.debug(f"Prompt tokens count: {token_count}")
+
+    if model in litellm.model_cost:
+        # If we are able to identify the model context window, ensure the messages fit in it
+        max_tokens = litellm.model_cost[model].get(
+            "max_input_tokens", litellm.model_cost[model]["max_tokens"]
+        )
+        if token_count > max_tokens:
+            logger.info(
+                f"Prompt tokens count {token_count} exceeds max tokens {max_tokens} for model {model}. Trimming messages."
+            )
+            return trim_messages(messages, model=model), token_count
path: adala/utils/llm_utils.py
user: login: matt-bernstein
body: use `prediction_meta` field for warnings as well as errors in future
created_at: 2025-03-07T16:54:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -539,7 +539,14 @@ 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
path: adala/runtimes/_litellm.py
user: login: matt-bernstein
body: add comment - this may happen for any project type, just more likely for MIG, so this setting is a heuristic
created_at: 2025-03-07T16:58:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -116,6 +122,30 @@ def handle_llm_exception(
     return _log_llm_exception(e), usage
 
 
+def _ensure_messages_fit_in_context_window(
+    messages: List[Dict[str, str]], model: str
path: adala/utils/llm_utils.py
user: login: niklub
body: done
created_at: 2025-03-11T14:19:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -116,6 +122,30 @@ def handle_llm_exception(
     return _log_llm_exception(e), usage
 
 
+def _ensure_messages_fit_in_context_window(
+    messages: List[Dict[str, str]], model: str
+) -> Tuple[List[Dict[str, str]], int]:
+    """
+    Ensure that the messages fit in the context window of the model.
+    """
+    token_count = token_counter(model=model, messages=messages)
+    logger.debug(f"Prompt tokens count: {token_count}")
+
+    if model in litellm.model_cost:
+        # If we are able to identify the model context window, ensure the messages fit in it
+        max_tokens = litellm.model_cost[model].get(
+            "max_input_tokens", litellm.model_cost[model]["max_tokens"]
+        )
+        if token_count > max_tokens:
+            logger.info(
+                f"Prompt tokens count {token_count} exceeds max tokens {max_tokens} for model {model}. Trimming messages."
+            )
+            return trim_messages(messages, model=model), token_count
path: adala/utils/llm_utils.py
user: login: niklub
body: added comments
created_at: 2025-03-11T14:19:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -539,7 +539,14 @@ 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
path: adala/runtimes/_litellm.py
user: login: niklub
body: done
created_at: 2025-03-11T14:21:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -174,112 +208,104 @@ class ImageMessageChunk(TypedDict):
 Message = Union[str, List[MessageChunk]]
 
 
-class MessageChunkType(Enum):
-    TEXT = "text"
-    IMAGE_URL = "image_url"
-
-
 def split_message_into_chunks(
path: adala/utils/parse.py
user: login: niklub
body: Conducted some analysis, I propose not to roll back to the initial version.
Reason:
1. Performance: as the number of vars grows, the previous method shows degraded performance. The reason is due to repeated calls for partial_str_format() for each variable field, introducing linear complexity overhead
    ```
    split_message_into_chunks_v1: 331.67 ms avg
    split_message_into_chunks: 6.28 ms avg
    ```
2. More stable version: the previous version fails when encountering a non-existent field key
3. Control flow: `build_chunks` has multiple returns whereas new version is only one - bad practices according to single exit point pattern
4. Conciseness - less lines
created_at: 2025-03-11T14:35:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -174,112 +208,104 @@ class ImageMessageChunk(TypedDict):
 Message = Union[str, List[MessageChunk]]
 
 
-class MessageChunkType(Enum):
-    TEXT = "text"
-    IMAGE_URL = "image_url"
-
-
 def split_message_into_chunks(
path: adala/utils/parse.py
user: login: niklub
body: also added comments to improve readability
created_at: 2025-03-11T14:35:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,13 +175,19 @@ def run_instructor_with_messages(
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
         retries: Retry policy to use
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion call
 
     Returns:
         Dict containing the parsed response and usage information
     """
     try:
-        # returns a pydantic model and completion info
+        prompt_token_count = None
+        if ensure_messages_fit_in_context_window:
+            messages, prompt_token_count = _ensure_messages_fit_in_context_window(
+                messages, model
+            )
path: adala/utils/llm_utils.py
user: login: matt-bernstein
body: @niklub here is where `self.canonical_model_provider_string` should be used. here is the implementation: https://github.com/HumanSignal/Adala/blob/master/adala/utils/model_info_utils.py#L48

The regular self.model is the "endpoint" model name that the model provider understands. The canonical string is the "canonical" model name that litellm itself understands, where the model is listed with its supported parameters, cost per token, context window, etc in a big lookup table.
created_at: 2025-03-14T18:50:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,13 +175,19 @@ def run_instructor_with_messages(
         temperature: Temperature for sampling
         seed: Integer seed to reduce nondeterminism
         retries: Retry policy to use
+        ensure_messages_fit_in_context_window: Whether to ensure the messages fit in the context window (setting it to True will slow down the function)
         **kwargs: Additional arguments to pass to the completion call
 
     Returns:
         Dict containing the parsed response and usage information
     """
     try:
-        # returns a pydantic model and completion info
+        prompt_token_count = None
+        if ensure_messages_fit_in_context_window:
+            messages, prompt_token_count = _ensure_messages_fit_in_context_window(
+                messages, model
+            )
path: adala/utils/llm_utils.py
user: login: matt-bernstein
body: They are often different :upside_down_face: 
created_at: 2025-03-14T18:50:19Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov-commenter
created_at: 2025-02-27T12:57:17Z
body: ## [Codecov](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal) Report
:x: Patch coverage is `76.85185%` with `25 lines` in your changes missing coverage. Please review.
| [Files with missing lines](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal) | Patch % | Lines |
|---|---|---|
| [adala/utils/llm\_utils.py](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal#diff-YWRhbGEvdXRpbHMvbGxtX3V0aWxzLnB5) | 50.00% | [15 Missing :warning: ](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal) |
| [adala/utils/parse.py](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal#diff-YWRhbGEvdXRpbHMvcGFyc2UucHk=) | 82.97% | [8 Missing :warning: ](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal) |
| [adala/skills/collection/label\_studio.py](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal#diff-YWRhbGEvc2tpbGxzL2NvbGxlY3Rpb24vbGFiZWxfc3R1ZGlvLnB5) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/HumanSignal/Adala/pull/343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HumanSignal) |

:loudspeaker: Thoughts on this report? [Let us know!](https://github.com/codecov/feedback/issues/255)
author_association: NONE
user: login: niklub
created_at: 2025-03-01T10:39:01Z
body: /git merge
> [Workflow run](https://github.com/HumanSignal/Adala/actions/runs/13603953951)
> Successfully merged:  5 files changed, 100 insertions(+), 69 deletions(-)
author_association: CONTRIBUTOR
user: login: niklub
created_at: 2025-03-05T14:06:29Z
body: /git merge
> [Workflow run](https://github.com/HumanSignal/Adala/actions/runs/13677837044)
> Successfully merged:  3 files changed, 8 insertions(+), 7 deletions(-)
author_association: CONTRIBUTOR
user: login: niklub
created_at: 2025-03-13T11:53:55Z
body: /git merge
> [Workflow run](https://github.com/HumanSignal/Adala/actions/runs/13833966749)
> Successfully merged:  1 file changed, 1 insertion(+), 1 deletion(-)
author_association: CONTRIBUTOR