From 004bb0d464a3d159dde7a24b8b84e8e0e3325bb8 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Tue, 11 Feb 2025 17:13:38 -0500
Subject: [PATCH 01/13] feat: DIA-1868: Azure AI Studio support in Prompts


From 3d4e88f502dd945f34632564a6e42b0253bbe1e8 Mon Sep 17 00:00:00 2001
From: fern-api <115122769+fern-api[bot]@users.noreply.github.com>
Date: Fri, 7 Feb 2025 16:52:45 +0000
Subject: [PATCH 02/13] [submodules] Bump HumanSignal/label-studio-sdk version

Workflow run: https://github.com/HumanSignal/Adala/actions/runs/13204307644
---
 poetry.lock    | 6 +++---
 pyproject.toml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index e6816d43..1aab6404 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "2547755090341885fedf17181f72290a0b48034a.zip", hash = "sha256:5dc5480b2c1755e33a2756cc5056391357f36d5c10e55822a16ece1db7511439"},
+    {file = "82951d8e0d652f071eb607df4e699365c8d89534.zip", hash = "sha256:edf73dd21cd9f63d3908d768a02e9a99fc75d474a1856745911947fdb4de9529"},
 ]
 
 [package.dependencies]
@@ -3721,7 +3721,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/2547755090341885fedf17181f72290a0b48034a.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/82951d8e0d652f071eb607df4e699365c8d89534.zip"
 
 [[package]]
 name = "litellm"
@@ -8486,4 +8486,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "df6cceebe54905fb3fe5720cd3109383d03f8a9a65e934dd6abeb839e859e750"
+content-hash = "a2a5ecc5add41eb47de3cd17ed4c062c9f0e0111138881586cb4b89437f511a9"
diff --git a/pyproject.toml b/pyproject.toml
index 75ecdc7e..6bf7367e 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/2547755090341885fedf17181f72290a0b48034a.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/82951d8e0d652f071eb607df4e699365c8d89534.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

From a0a12e26bf1f7a55f0492267e2753cc2bb1a3175 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Mon, 10 Feb 2025 17:13:17 +0000
Subject: [PATCH 03/13] [submodules] Bump HumanSignal/label-studio-sdk version

Workflow run: https://github.com/HumanSignal/Adala/actions/runs/13246481736
---
 poetry.lock    | 6 +++---
 pyproject.toml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 1aab6404..04280ccd 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "82951d8e0d652f071eb607df4e699365c8d89534.zip", hash = "sha256:edf73dd21cd9f63d3908d768a02e9a99fc75d474a1856745911947fdb4de9529"},
+    {file = "c3464d51e546db3a5acf2399d8277db7f76cc79e.zip", hash = "sha256:2eff83e45dd0a58e648744646d6a1863d9fc11cfb77fd01fca03712e3b08a589"},
 ]
 
 [package.dependencies]
@@ -3721,7 +3721,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/82951d8e0d652f071eb607df4e699365c8d89534.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/c3464d51e546db3a5acf2399d8277db7f76cc79e.zip"
 
 [[package]]
 name = "litellm"
@@ -8486,4 +8486,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "a2a5ecc5add41eb47de3cd17ed4c062c9f0e0111138881586cb4b89437f511a9"
+content-hash = "03ad7c6eb0ca626950ace8b20b08a10eb35700ffb59b5dbc8dbda82ddc82a6e3"
diff --git a/pyproject.toml b/pyproject.toml
index 6bf7367e..0be88c05 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/82951d8e0d652f071eb607df4e699365c8d89534.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/c3464d51e546db3a5acf2399d8277db7f76cc79e.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

From 4a2000069ac8cb71c590843a9a21498379018f68 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Tue, 11 Feb 2025 17:12:13 -0500
Subject: [PATCH 04/13] add azure ai foundry

---
 server/app.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/server/app.py b/server/app.py
index 19b7f42c..6ba779ca 100644
--- a/server/app.py
+++ b/server/app.py
@@ -246,6 +246,7 @@ async def submit_batch(batch: BatchData):
 
 @app.post("/validate-connection", response_model=Response[ValidateConnectionResponse])
 async def validate_connection(request: ValidateConnectionRequest):
+    # TODO: move this logic to LSE, this is the last place Adala needs to be updated when adding a provider connection
     multi_model_provider_test_models = {
         "openai": "gpt-4o-mini",
         "vertexai": "vertex_ai/gemini-1.5-flash",
@@ -290,6 +291,9 @@ async def validate_connection(request: ValidateConnectionRequest):
         if provider.lower() == "azureopenai":
             model = "azure/" + request.deployment_name
             model_extra = {"base_url": request.endpoint}
+        elif provider.lower() == "azureaifoundry":
+            model = "azure_ai/" + request.deployment_name
+            model_extra = {"base_url": request.endpoint}
         elif provider.lower() == "custom":
             model = "openai/" + request.deployment_name
             model_extra = (
@@ -364,6 +368,8 @@ async def estimate_cost(
     agent = request.agent
     provider = request.provider
     runtime = agent.get_runtime()
+    with open("cost_estimate.log", "w") as f:
+        json.dump(request.model_dump(), f, indent=4)
 
     try:
         cost_estimates = []

From f66e9664a8ded85d7fb577efff6f282b505e611d Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Wed, 12 Feb 2025 15:04:02 -0500
Subject: [PATCH 05/13] remove debug

---
 server/app.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/server/app.py b/server/app.py
index 6ba779ca..a7e369a4 100644
--- a/server/app.py
+++ b/server/app.py
@@ -368,8 +368,6 @@ async def estimate_cost(
     agent = request.agent
     provider = request.provider
     runtime = agent.get_runtime()
-    with open("cost_estimate.log", "w") as f:
-        json.dump(request.model_dump(), f, indent=4)
 
     try:
         cost_estimates = []

From 67049810127130688477a33140fcfc0e8db349be Mon Sep 17 00:00:00 2001
From: matt-bernstein <matt-bernstein@users.noreply.github.com>
Date: Wed, 12 Feb 2025 21:24:08 +0000
Subject: [PATCH 06/13] Sync Follow Merge dependencies

Workflow run: https://github.com/HumanSignal/Adala/actions/runs/13295140355
---
 poetry.lock    | 6 +++---
 pyproject.toml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 04280ccd..18d7023a 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "c3464d51e546db3a5acf2399d8277db7f76cc79e.zip", hash = "sha256:2eff83e45dd0a58e648744646d6a1863d9fc11cfb77fd01fca03712e3b08a589"},
+    {file = "f4e56f0104013f532e6c5817300cb9c5a6188b9d.zip", hash = "sha256:27ec4b60e4c8c9660bc07488df89d33eccb052ef811a1b40a9f325b74687a41f"},
 ]
 
 [package.dependencies]
@@ -3721,7 +3721,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/c3464d51e546db3a5acf2399d8277db7f76cc79e.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/f4e56f0104013f532e6c5817300cb9c5a6188b9d.zip"
 
 [[package]]
 name = "litellm"
@@ -8486,4 +8486,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "03ad7c6eb0ca626950ace8b20b08a10eb35700ffb59b5dbc8dbda82ddc82a6e3"
+content-hash = "f6c2d817ea638d3ea3caa69fd600ed56b4c2689270f94dae87e278eae9c982be"
diff --git a/pyproject.toml b/pyproject.toml
index 0be88c05..a585b99d 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/c3464d51e546db3a5acf2399d8277db7f76cc79e.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/f4e56f0104013f532e6c5817300cb9c5a6188b9d.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

From a56c87ed33cc9c1fc1022d81d8b7445d34791260 Mon Sep 17 00:00:00 2001
From: matt-bernstein <matt-bernstein@users.noreply.github.com>
Date: Wed, 12 Feb 2025 21:26:19 +0000
Subject: [PATCH 07/13] Sync Follow Merge dependencies

Workflow run: https://github.com/HumanSignal/Adala/actions/runs/13295173444
---
 poetry.lock    | 6 +++---
 pyproject.toml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 18d7023a..fe71c3dc 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "f4e56f0104013f532e6c5817300cb9c5a6188b9d.zip", hash = "sha256:27ec4b60e4c8c9660bc07488df89d33eccb052ef811a1b40a9f325b74687a41f"},
+    {file = "c39030dda4f5320ec9af1b15918b85d4c274f4c1.zip", hash = "sha256:94b69c0d84d420d1ccda4b14c5bcb6b52ba00983caa8227b13e550e0475808cf"},
 ]
 
 [package.dependencies]
@@ -3721,7 +3721,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/f4e56f0104013f532e6c5817300cb9c5a6188b9d.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/c39030dda4f5320ec9af1b15918b85d4c274f4c1.zip"
 
 [[package]]
 name = "litellm"
@@ -8486,4 +8486,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "f6c2d817ea638d3ea3caa69fd600ed56b4c2689270f94dae87e278eae9c982be"
+content-hash = "0d185cbfa5bb89a02eaff1b25a26ebd6148342799aa2397295e760ddcb3ebfa4"
diff --git a/pyproject.toml b/pyproject.toml
index a585b99d..f6dcd4ec 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/f4e56f0104013f532e6c5817300cb9c5a6188b9d.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/c39030dda4f5320ec9af1b15918b85d4c274f4c1.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

From e3cdee91b68a41fcbdca658e8d3ed279e9768839 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Thu, 13 Feb 2025 06:42:58 -0500
Subject: [PATCH 08/13] better error handling for cost estimation

---
 adala/runtimes/_litellm.py | 36 +++++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
index 0defec7d..7c833621 100644
--- a/adala/runtimes/_litellm.py
+++ b/adala/runtimes/_litellm.py
@@ -163,8 +163,13 @@ def normalize_litellm_model_and_provider(model_name: str, provider: str):
     if "/" in model_name:
         model_name = model_name.split("/", 1)[1]
     provider = provider.lower()
+    # TODO: move this logic to LSE, this is the last place Adala needs to be updated when adding a provider connection
     if provider == "vertexai":
         provider = "vertex_ai"
+    if provider == "azureopenai":
+        provider = "azure"
+    if provider == "azureaifoundry":
+        provider = "azure_ai"
 
     return model_name, provider
 
@@ -600,11 +605,32 @@ def _estimate_cost(
     ):
         prompt_tokens = cls._get_prompt_tokens(user_prompt, model, output_fields)
         completion_tokens = cls._get_completion_tokens(model, output_fields, provider)
-        prompt_cost, completion_cost = litellm.cost_per_token(
-            model=model,
-            prompt_tokens=prompt_tokens,
-            completion_tokens=completion_tokens,
-        )
+        # amazingly, litellm.cost_per_token refers to a hardcoded dictionary litellm.model_cost which is case-sensitive with inconsistent casing.....
+        # Example: 'azure_ai/deepseek-r1' vs 'azure_ai/Llama-3.3-70B-Instruct'
+        # so we ctually have no way of determining the correct casing or reliably fixing it.
+        # we can at least try all-loweercase.
+        candidate_model_names = [model, model.lower()]
+        # ...and Azure AI Foundry openai models are not listed there, but under Azure OpenAI
+        if model.startswith("azure_ai/"):
+            candidate_model_names.append(model.replace("azure_ai/", "azure/"))
+            candidate_model_names.append(model.replace("azure_ai/", "azure/").lower())
+        for candidate_model_name in candidate_model_names:
+            try:
+                prompt_cost, completion_cost = litellm.cost_per_token(
+                    model=model,
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                )
+                break
+            except Exception as e:
+                # it also doesn't have a type to catch:
+                # Exception("This model isn't mapped yet. model=azure_ai/deepseek-R1, custom_llm_provider=azure_ai. Add it here - https://github.com/ BerriAI/litellm/blob/main/model_prices_and_context_window.json.")
+                if "model isn't mapped" in str(e):
+                    prompt_cost, completion_cost = None, None
+                raise e
+        if prompt_cost is None or completion_cost is None:
+            raise ValueError(f"Model {model} for provider {provider} not found.")
+
         total_cost = prompt_cost + completion_cost
 
         return prompt_cost, completion_cost, total_cost

From d792385ec901e484d7cb2ea3cf92ed366eea9c16 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Thu, 13 Feb 2025 16:15:24 -0500
Subject: [PATCH 09/13] working azure cost estimations

---
 adala/runtimes/_litellm.py | 61 ++++++++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 16 deletions(-)

diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
index 7c833621..1a3f419c 100644
--- a/adala/runtimes/_litellm.py
+++ b/adala/runtimes/_litellm.py
@@ -583,17 +583,26 @@ def _get_prompt_tokens(string: str, model: str, output_fields: List[str]) -> int
 
     @staticmethod
     def _get_completion_tokens(
-        model: str, output_fields: Optional[List[str]], provider: str
+        candidate_model_names: List[str],
+        output_fields: Optional[List[str]],
+        provider: str,
     ) -> int:
-        model, provider = normalize_litellm_model_and_provider(model, provider)
-        max_tokens = litellm.get_model_info(
-            model=model, custom_llm_provider=provider
-        ).get("max_tokens", None)
-        if not max_tokens:
-            raise ValueError
-        # extremely rough heuristic, from testing on some anecdotal examples
-        n_outputs = len(output_fields) if output_fields else 1
-        return min(max_tokens, 4 * n_outputs)
+        for model in candidate_model_names:
+            model, provider = normalize_litellm_model_and_provider(model, provider)
+            try:
+                max_tokens = litellm.get_model_info(
+                    model=model, custom_llm_provider=provider
+                ).get("max_tokens", None)
+            except Exception as e:
+                if "model isn't mapped" in str(e):
+                    continue
+                else:
+                    raise e
+            if not max_tokens:
+                raise ValueError
+            # extremely rough heuristic, from testing on some anecdotal examples
+            n_outputs = len(output_fields) if output_fields else 1
+            return min(max_tokens, 4 * n_outputs)
 
     @classmethod
     def _estimate_cost(
@@ -604,20 +613,24 @@ def _estimate_cost(
         provider: str,
     ):
         prompt_tokens = cls._get_prompt_tokens(user_prompt, model, output_fields)
-        completion_tokens = cls._get_completion_tokens(model, output_fields, provider)
         # amazingly, litellm.cost_per_token refers to a hardcoded dictionary litellm.model_cost which is case-sensitive with inconsistent casing.....
         # Example: 'azure_ai/deepseek-r1' vs 'azure_ai/Llama-3.3-70B-Instruct'
-        # so we ctually have no way of determining the correct casing or reliably fixing it.
-        # we can at least try all-loweercase.
+        # so we have no way of determining the correct casing or reliably fixing it.
+        # we can at least try all-lowercase.
         candidate_model_names = [model, model.lower()]
         # ...and Azure AI Foundry openai models are not listed there, but under Azure OpenAI
         if model.startswith("azure_ai/"):
             candidate_model_names.append(model.replace("azure_ai/", "azure/"))
             candidate_model_names.append(model.replace("azure_ai/", "azure/").lower())
+
+        completion_tokens = cls._get_completion_tokens(
+            candidate_model_names.copy(), output_fields, provider
+        )
+
         for candidate_model_name in candidate_model_names:
             try:
                 prompt_cost, completion_cost = litellm.cost_per_token(
-                    model=model,
+                    model=candidate_model_name,
                     prompt_tokens=prompt_tokens,
                     completion_tokens=completion_tokens,
                 )
@@ -627,7 +640,8 @@ def _estimate_cost(
                 # Exception("This model isn't mapped yet. model=azure_ai/deepseek-R1, custom_llm_provider=azure_ai. Add it here - https://github.com/ BerriAI/litellm/blob/main/model_prices_and_context_window.json.")
                 if "model isn't mapped" in str(e):
                     prompt_cost, completion_cost = None, None
-                raise e
+                else:
+                    raise e
         if prompt_cost is None or completion_cost is None:
             raise ValueError(f"Model {model} for provider {provider} not found.")
 
@@ -650,10 +664,25 @@ def get_cost_estimate(
             cumulative_prompt_cost = 0
             cumulative_completion_cost = 0
             cumulative_total_cost = 0
+            # for azure, we need the canonical model name, not the deployment name
+            if self.model.startswith("azure/"):
+                messages = [{"role": "user", "content": "Hey, how's it going?"}]
+                response = litellm.completion(
+                    messages=messages,
+                    model=self.model,
+                    max_tokens=10,
+                    temperature=self.temperature,
+                    seed=self.seed,
+                    # extra inference params passed to this runtime
+                    **self.model_extra,
+                )
+                model = "azure/" + response.model
+            else:
+                model = self.model
             for user_prompt in user_prompts:
                 prompt_cost, completion_cost, total_cost = self._estimate_cost(
                     user_prompt=user_prompt,
-                    model=self.model,
+                    model=model,
                     output_fields=output_fields,
                     provider=provider,
                 )

From e902a22f60d86caad3be94267277e2b90ccdc0fe Mon Sep 17 00:00:00 2001
From: matt-bernstein <matt-bernstein@users.noreply.github.com>
Date: Thu, 13 Feb 2025 21:30:10 +0000
Subject: [PATCH 10/13] Sync Follow Merge dependencies

Workflow run: https://github.com/HumanSignal/Adala/actions/runs/13317432634
---
 poetry.lock    | 6 +++---
 pyproject.toml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index fe71c3dc..766437b1 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "c39030dda4f5320ec9af1b15918b85d4c274f4c1.zip", hash = "sha256:94b69c0d84d420d1ccda4b14c5bcb6b52ba00983caa8227b13e550e0475808cf"},
+    {file = "73e406721b62513d2b7c3962d8af9fd072caeff0.zip", hash = "sha256:82a4744ae81d5889a2814d40fd66deeb7faf462a92af4070fa6f74fda4bdb019"},
 ]
 
 [package.dependencies]
@@ -3721,7 +3721,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/c39030dda4f5320ec9af1b15918b85d4c274f4c1.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/73e406721b62513d2b7c3962d8af9fd072caeff0.zip"
 
 [[package]]
 name = "litellm"
@@ -8486,4 +8486,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "0d185cbfa5bb89a02eaff1b25a26ebd6148342799aa2397295e760ddcb3ebfa4"
+content-hash = "848feec2dc0ec8f8e84f0f2243a9705677b767051f62dae384dd0ce23fe0e8e9"
diff --git a/pyproject.toml b/pyproject.toml
index f6dcd4ec..95e65c7a 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/c39030dda4f5320ec9af1b15918b85d4c274f4c1.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/73e406721b62513d2b7c3962d8af9fd072caeff0.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

From bdaced5b49d0a17a286ad44b39373e4d20ef3f82 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Thu, 13 Feb 2025 17:00:29 -0500
Subject: [PATCH 11/13] remve debug

---
 adala/runtimes/_litellm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
index 1a3f419c..d1a2b437 100644
--- a/adala/runtimes/_litellm.py
+++ b/adala/runtimes/_litellm.py
@@ -624,7 +624,7 @@ def _estimate_cost(
             candidate_model_names.append(model.replace("azure_ai/", "azure/").lower())
 
         completion_tokens = cls._get_completion_tokens(
-            candidate_model_names.copy(), output_fields, provider
+            candidate_model_names, output_fields, provider
         )
 
         for candidate_model_name in candidate_model_names:

From a23182e00122605ea11211f4f8063f151c552688 Mon Sep 17 00:00:00 2001
From: Matt Bernstein <matt@humansignal.com>
Date: Fri, 14 Feb 2025 10:46:12 -0500
Subject: [PATCH 12/13] fix completion tokens loop

---
 adala/runtimes/_litellm.py | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
index d1a2b437..edce85d6 100644
--- a/adala/runtimes/_litellm.py
+++ b/adala/runtimes/_litellm.py
@@ -587,22 +587,21 @@ def _get_completion_tokens(
         output_fields: Optional[List[str]],
         provider: str,
     ) -> int:
+        max_tokens = None
         for model in candidate_model_names:
-            model, provider = normalize_litellm_model_and_provider(model, provider)
             try:
-                max_tokens = litellm.get_model_info(
-                    model=model, custom_llm_provider=provider
-                ).get("max_tokens", None)
+                max_tokens = litellm.get_model_info(model=model).get("max_tokens", None)
+                break
             except Exception as e:
                 if "model isn't mapped" in str(e):
                     continue
                 else:
                     raise e
-            if not max_tokens:
-                raise ValueError
-            # extremely rough heuristic, from testing on some anecdotal examples
-            n_outputs = len(output_fields) if output_fields else 1
-            return min(max_tokens, 4 * n_outputs)
+        if not max_tokens:
+            raise ValueError
+        # extremely rough heuristic, from testing on some anecdotal examples
+        n_outputs = len(output_fields) if output_fields else 1
+        return min(max_tokens, 4 * n_outputs)
 
     @classmethod
     def _estimate_cost(
@@ -622,11 +621,13 @@ def _estimate_cost(
         if model.startswith("azure_ai/"):
             candidate_model_names.append(model.replace("azure_ai/", "azure/"))
             candidate_model_names.append(model.replace("azure_ai/", "azure/").lower())
+        candidate_model_names = list(set(candidate_model_names))
 
         completion_tokens = cls._get_completion_tokens(
             candidate_model_names, output_fields, provider
         )
 
+        prompt_cost, completion_cost = None, None
         for candidate_model_name in candidate_model_names:
             try:
                 prompt_cost, completion_cost = litellm.cost_per_token(
@@ -634,14 +635,14 @@ def _estimate_cost(
                     prompt_tokens=prompt_tokens,
                     completion_tokens=completion_tokens,
                 )
-                break
             except Exception as e:
                 # it also doesn't have a type to catch:
                 # Exception("This model isn't mapped yet. model=azure_ai/deepseek-R1, custom_llm_provider=azure_ai. Add it here - https://github.com/ BerriAI/litellm/blob/main/model_prices_and_context_window.json.")
                 if "model isn't mapped" in str(e):
-                    prompt_cost, completion_cost = None, None
-                else:
-                    raise e
+                    pass
+            if prompt_cost is not None and completion_cost is not None:
+                break
+
         if prompt_cost is None or completion_cost is None:
             raise ValueError(f"Model {model} for provider {provider} not found.")
 

From c7551f9ab5f55498fe443924e17f5256e1285a06 Mon Sep 17 00:00:00 2001
From: matt-bernstein <matt-bernstein@users.noreply.github.com>
Date: Wed, 19 Feb 2025 15:06:56 +0000
Subject: [PATCH 13/13] Sync Follow Merge dependencies

Workflow run: https://github.com/HumanSignal/Adala/actions/runs/13415740545
---
 poetry.lock    | 6 +++---
 pyproject.toml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 766437b1..185d4951 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3696,7 +3696,7 @@ optional = false
 python-versions = ">=3.9,<4"
 groups = ["main"]
 files = [
-    {file = "73e406721b62513d2b7c3962d8af9fd072caeff0.zip", hash = "sha256:82a4744ae81d5889a2814d40fd66deeb7faf462a92af4070fa6f74fda4bdb019"},
+    {file = "51f3da6b9636b081ed172d725c2ad91184a0c459.zip", hash = "sha256:1cf344e74307e5cbea6734f830b1225477a2a32b04735a4b0365ecc0e26432cb"},
 ]
 
 [package.dependencies]
@@ -3721,7 +3721,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/73e406721b62513d2b7c3962d8af9fd072caeff0.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/51f3da6b9636b081ed172d725c2ad91184a0c459.zip"
 
 [[package]]
 name = "litellm"
@@ -8486,4 +8486,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.9.7 || >3.9.7,<3.13"
-content-hash = "848feec2dc0ec8f8e84f0f2243a9705677b767051f62dae384dd0ce23fe0e8e9"
+content-hash = "7099169a2eecf504631f11dd37ab49a9ea3703370f7eb1acd6b8ec558cacf70c"
diff --git a/pyproject.toml b/pyproject.toml
index 95e65c7a..fe37261e 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/73e406721b62513d2b7c3962d8af9fd072caeff0.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/51f3da6b9636b081ed172d725c2ad91184a0c459.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