From 751d474aa66dc3ab404d463acbb2fb7fee066dfd Mon Sep 17 00:00:00 2001
From: matt-bernstein <60152561+matt-bernstein@users.noreply.github.com>
Date: Wed, 29 Jan 2025 17:36:14 -0500
Subject: [PATCH] feat: DIA-1716: Google AI Studio (Gemini) support in Prompts
 (#6890)

Co-authored-by: hakan458 <hakan@heartex.com>
Co-authored-by: fern-api <115122769+fern-api[bot]@users.noreply.github.com>
Co-authored-by: matt-bernstein <matt-bernstein@users.noreply.github.com>
Co-authored-by: robot-ci-heartex <87703623+robot-ci-heartex@users.noreply.github.com>
---
 ..._alter_modelproviderconnection_provider.py | 31 +++++++++++++++++++
 label_studio/ml_model_providers/models.py     |  1 +
 ...4_alter_thirdpartymodelversion_provider.py | 29 +++++++++++++++++
 poetry.lock                                   |  6 ++--
 pyproject.toml                                |  2 +-
 5 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 label_studio/ml_model_providers/migrations/0007_alter_modelproviderconnection_provider.py
 create mode 100644 label_studio/ml_models/migrations/0014_alter_thirdpartymodelversion_provider.py

diff --git a/label_studio/ml_model_providers/migrations/0007_alter_modelproviderconnection_provider.py b/label_studio/ml_model_providers/migrations/0007_alter_modelproviderconnection_provider.py
new file mode 100644
index 000000000000..911f764f3b9f
--- /dev/null
+++ b/label_studio/ml_model_providers/migrations/0007_alter_modelproviderconnection_provider.py
@@ -0,0 +1,31 @@
+# Generated by Django 5.1.4 on 2025-01-10 18:38
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        (
+            "ml_model_providers",
+            "0006_modelproviderconnection_google_application_credentials_and_more",
+        ),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="modelproviderconnection",
+            name="provider",
+            field=models.CharField(
+                choices=[
+                    ("OpenAI", "OpenAI"),
+                    ("AzureOpenAI", "AzureOpenAI"),
+                    ("VertexAI", "VertexAI"),
+                    ("Gemini", "Gemini"),
+                    ("Custom", "Custom"),
+                ],
+                default="OpenAI",
+                max_length=255,
+            ),
+        ),
+    ]
diff --git a/label_studio/ml_model_providers/models.py b/label_studio/ml_model_providers/models.py
index d008f955588c..3828bdc45b81 100644
--- a/label_studio/ml_model_providers/models.py
+++ b/label_studio/ml_model_providers/models.py
@@ -12,6 +12,7 @@ class ModelProviders(models.TextChoices):
     OPENAI = 'OpenAI', _('OpenAI')
     AZURE_OPENAI = 'AzureOpenAI', _('AzureOpenAI')
     VERTEX_AI = 'VertexAI', _('VertexAI')
+    GEMINI = 'Gemini', _('Gemini')
     CUSTOM = 'Custom', _('Custom')
 
 
diff --git a/label_studio/ml_models/migrations/0014_alter_thirdpartymodelversion_provider.py b/label_studio/ml_models/migrations/0014_alter_thirdpartymodelversion_provider.py
new file mode 100644
index 000000000000..0e5d86397233
--- /dev/null
+++ b/label_studio/ml_models/migrations/0014_alter_thirdpartymodelversion_provider.py
@@ -0,0 +1,29 @@
+# Generated by Django 5.1.4 on 2025-01-10 18:38
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("ml_models", "0013_alter_thirdpartymodelversion_provider"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="thirdpartymodelversion",
+            name="provider",
+            field=models.CharField(
+                choices=[
+                    ("OpenAI", "OpenAI"),
+                    ("AzureOpenAI", "AzureOpenAI"),
+                    ("VertexAI", "VertexAI"),
+                    ("Gemini", "Gemini"),
+                    ("Custom", "Custom"),
+                ],
+                default="OpenAI",
+                help_text="The model provider to use e.g. OpenAI",
+                max_length=255,
+            ),
+        ),
+    ]
diff --git a/poetry.lock b/poetry.lock
index da5cf2e0872a..b5d7ffec5843 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2140,7 +2140,7 @@ python-versions = ">=3.9,<4"
 groups = ["main"]
 markers = "python_version >= \"3.12\" or python_version <= \"3.11\""
 files = [
-    {file = "4a9715c6b0b619371e89c09ea8d1c86ce5c880df.zip", hash = "sha256:04ee2d775c6ee4048f2306416c67d4f6b94837346a35de5b2aaa13e65ee4acf8"},
+    {file = "510e71ab176061588a0f7cfec92e43c207bb5d9f.zip", hash = "sha256:0e04dadc233210140e2d38f241c962f998d54ba7f19f58d686d8eff748434333"},
 ]
 
 [package.dependencies]
@@ -2165,7 +2165,7 @@ xmljson = "0.2.1"
 
 [package.source]
 type = "url"
-url = "https://github.com/HumanSignal/label-studio-sdk/archive/4a9715c6b0b619371e89c09ea8d1c86ce5c880df.zip"
+url = "https://github.com/HumanSignal/label-studio-sdk/archive/510e71ab176061588a0f7cfec92e43c207bb5d9f.zip"
 
 [[package]]
 name = "launchdarkly-server-sdk"
@@ -5034,4 +5034,4 @@ uwsgi = ["pyuwsgi", "uwsgitop"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.10,<4"
-content-hash = "d8f2e14335e74a3defac358d5716eabe073e79ecf365d439acac6574e93f33fd"
+content-hash = "9687682016b7a7d1eea5ff7274ca99dde9a378ef65d673799b6b38c0d3e5afe6"
diff --git a/pyproject.toml b/pyproject.toml
index 4839650960bf..55c9d9f558c6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -208,7 +208,7 @@ django-migration-linter = "^5.1.0"
 setuptools = ">=75.4.0"
 
 # Humansignal repo dependencies
-label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/4a9715c6b0b619371e89c09ea8d1c86ce5c880df.zip"}
+label-studio-sdk = {url = "https://github.com/HumanSignal/label-studio-sdk/archive/510e71ab176061588a0f7cfec92e43c207bb5d9f.zip"}
 
 [tool.poetry.group.test.dependencies]
 pytest = "7.2.2"