feat: add support for labels on generate content requests

.github/workflows/ci.yml

@@ -51,8 +51,18 @@ jobs:
   build:
     needs: test
     runs-on: ubuntu-latest
+    if: github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev'
     steps:
-      - name: Build & Release
-        uses: FullFact/ff_release@v2
+      - name: Bump version and push tag
+        id: tag_version
+        uses: mathieudutour/[email protected]
         with:
-          docker_build: false
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          release_branches: main
+          pre_release_branches: dev
+      - name: Create a GitHub release
+        uses: ncipollo/release-action@v1
+        with:
+          tag: ${{ steps.tag_version.outputs.new_tag }}
+          name: Release ${{ steps.tag_version.outputs.new_tag }}
+          body: ${{ steps.tag_version.outputs.changelog }}

src/genai_utils/gemini.py

@@ -37,6 +37,12 @@
     "top_p": 1,
 }
 
+DEFAULT_LABELS = {
+    key.removeprefix("GENAI_LABEL_").lower(): value
+    for key, value in os.environ.items()
+    if key.startswith("GENAI_LABEL_")
+}
+
 
 class GeminiError(Exception):
     """
@@ -180,7 +186,11 @@ def add_citations(response: types.GenerateContentResponse) -> str:
             links.append(follow_redirect(url))
 
     # Sort supports by end_index in descending order to avoid shifting issues when inserting.
-    sorted_supports: list[types.GroundingSupport] = sorted(supports, key=lambda s: s.segment.end_index, reverse=True)  # type: ignore
+    sorted_supports: list[types.GroundingSupport] = sorted(
+        supports,
+        key=lambda s: s.segment.end_index,  # type: ignore
+        reverse=True,
+    )
 
     for support in sorted_supports:
         segment = support.segment
@@ -202,6 +212,52 @@ def add_citations(response: types.GenerateContentResponse) -> str:
     return text
 
 
+def validate_labels(labels: dict[str, str]) -> None:
+    """
+    Validates labels for GCP requirements.
+
+    GCP label requirements:
+    - Keys must start with a lowercase letter
+    - Keys and values can only contain lowercase letters, numbers, hyphens, and underscores
+    - Keys and values must be max 63 characters
+    - Keys cannot be empty
+
+    Raises:
+        GeminiError: If labels don't meet GCP requirements
+    """
+    label_pattern = re.compile(r"^[a-z0-9_-]{1,63}$")
+    key_start_pattern = re.compile(r"^[a-z]")
+
+    for key, value in labels.items():
+        if not key:
+            raise GeminiError("Label keys cannot be empty")
+
+        if len(key) > 63:
+            raise GeminiError(
+                f"Label key '{key}' exceeds 63 characters (length: {len(key)})"
+            )
+
+        if len(value) > 63:
+            raise GeminiError(
+                f"Label value for key '{key}' exceeds 63 characters (length: {len(value)})"
+            )
+
+        if not key_start_pattern.match(key):
+            raise GeminiError(f"Label key '{key}' must start with a lowercase letter")
+
+        if not label_pattern.match(key):
+            raise GeminiError(
+                f"Label key '{key}' contains invalid characters. "
+                "Only lowercase letters, numbers, hyphens, and underscores are allowed"
+            )
+
+        if not label_pattern.match(value):
+            raise GeminiError(
+                f"Label value '{value}' for key '{key}' contains invalid characters. "
+                "Only lowercase letters, numbers, hyphens, and underscores are allowed"
+            )
+
+
 def check_grounding_ran(response: types.GenerateContentResponse) -> bool:
     """
     Checks if grounding ran and logs some metadata about the grounding.
@@ -246,6 +302,7 @@ def run_prompt(
     model_config: ModelConfig | None = None,
     use_grounding: bool = False,
     inline_citations: bool = False,
+    labels: dict[str, str] = {},
 ) -> str:
     """
     Runs a prompt through the model.
@@ -298,6 +355,10 @@ class Movie(BaseModel):
         Whether output should include citations inline with the text.
         These citations will be links to be used as evidence.
         This is only possible if grounding is set to true.
+    labels: dict[str, str]
+        Optional labels to attach to the API call for tracking and monitoring purposes.
+        Labels are key-value pairs that can be used to organize and filter requests
+        in Google Cloud logs and metrics.
 
     Returns
     -------
@@ -341,6 +402,8 @@ class Movie(BaseModel):
 
     if inline_citations and not use_grounding:
         raise GeminiError("Inline citations only work if `use_grounding = True`")
+    merged_labels = DEFAULT_LABELS | labels
+    validate_labels(merged_labels)
 
     response = client.models.generate_content(
         model=model_config.model_name,
@@ -349,6 +412,7 @@ class Movie(BaseModel):
             system_instruction=system_instruction,
             safety_settings=safety_settings,
             **built_gen_config,
+            labels=merged_labels,
         ),
     )
 

tests/genai_utils/test_labels.py

@@ -0,0 +1,189 @@
+import os
+from unittest.mock import Mock, patch
+
+import pytest
+from google.genai import Client
+from google.genai.models import Models
+
+from genai_utils.gemini import GeminiError, ModelConfig, run_prompt, validate_labels
+
+
+class DummyResponse:
+    candidates = "yes!"
+    text = "response!"
+
+
+def get_dummy():
+    return DummyResponse()
+
+
+def test_validate_labels_valid():
+    """Test that valid labels pass validation"""
+    valid_labels = {
+        "team": "ai",
+        "project": "genai-utils",
+        "environment": "production",
+        "version": "1-2-3",
+        "my_label": "my_value",
+    }
+    # Should not raise any exception
+    validate_labels(valid_labels)
+
+
+def test_validate_labels_empty_key():
+    """Test that empty keys are rejected"""
+    with pytest.raises(GeminiError, match="cannot be empty"):
+        validate_labels({"": "value"})
+
+
+def test_validate_labels_key_too_long():
+    """Test that keys exceeding 63 characters are rejected"""
+    long_key = "a" * 64
+    with pytest.raises(GeminiError, match="exceeds 63 characters"):
+        validate_labels({long_key: "value"})
+
+
+def test_validate_labels_value_too_long():
+    """Test that values exceeding 63 characters are rejected"""
+    long_value = "a" * 64
+    with pytest.raises(GeminiError, match="exceeds 63 characters"):
+        validate_labels({"key": long_value})
+
+
+def test_validate_labels_key_starts_with_number():
+    """Test that keys starting with numbers are rejected"""
+    with pytest.raises(GeminiError, match="must start with a lowercase letter"):
+        validate_labels({"1key": "value"})
+
+
+def test_validate_labels_key_starts_with_uppercase():
+    """Test that keys starting with uppercase are rejected"""
+    with pytest.raises(GeminiError, match="must start with a lowercase letter"):
+        validate_labels({"Key": "value"})
+
+
+@pytest.mark.parametrize(
+    "invalid_key", ["key@value", "key.name", "key$", "key with space", "key/name"]
+)
+def test_validate_labels_key_invalid_characters(invalid_key):
+    """Test that keys with invalid characters are rejected"""
+    with pytest.raises(GeminiError, match="contains invalid characters"):
+        validate_labels({invalid_key: "value"})
+
+
+@pytest.mark.parametrize(
+    "invalid_value", ["value@", "value.txt", "value$", "value with space", "value/"]
+)
+def test_validate_labels_value_invalid_characters(invalid_value):
+    """Test that values with invalid characters are rejected"""
+    with pytest.raises(GeminiError, match="contains invalid characters"):
+        validate_labels({"key": invalid_value})
+
+
+def test_validate_labels_max_length_valid():
+    """Test that keys and values at exactly 63 characters are valid"""
+    max_key = "a" * 63
+    max_value = "b" * 63
+    # Should not raise any exception
+    validate_labels({max_key: max_value})
+
+
+def test_validate_labels_valid_special_chars():
+    """Test that valid special characters (hyphens, underscores) are accepted"""
+    valid_labels = {
+        "my-key": "my-value",
+        "my_key": "my_value",
+        "my-key_name": "my-value_123",
+        "key123": "value456",
+    }
+    # Should not raise any exception
+    validate_labels(valid_labels)
+
+
+@patch("genai_utils.gemini.genai.Client")
+def test_run_prompt_with_valid_labels(mock_client):
+    """Test that run_prompt accepts and uses valid labels"""
+    client = Mock(Client)
+    models = Mock(Models)
+
+    models.generate_content.return_value = get_dummy()
+    client.models = models
+    mock_client.return_value = client
+
+    labels = {"team": "ai", "project": "test"}
+
+    run_prompt(
+        "test prompt",
+        labels=labels,
+        model_config=ModelConfig(
+            project="project", location="location", model_name="model"
+        ),
+    )
+
+    # Verify the call was made
+    assert models.generate_content.called
+    call_kwargs = models.generate_content.call_args[1]
+    assert "config" in call_kwargs
+    assert call_kwargs["config"].labels == labels
+
+
+@patch("genai_utils.gemini.genai.Client")
+def test_run_prompt_with_invalid_labels(mock_client):
+    """Test that run_prompt rejects invalid labels"""
+    client = Mock(Client)
+    models = Mock(Models)
+
+    models.generate_content.return_value = get_dummy()
+    client.models = models
+    mock_client.return_value = client
+
+    invalid_labels = {"Invalid": "value"}  # uppercase key
+
+    with pytest.raises(GeminiError, match="must start with a lowercase letter"):
+        run_prompt(
+            "test prompt",
+            labels=invalid_labels,
+            model_config=ModelConfig(
+                project="project", location="location", model_name="model"
+            ),
+        )
+
+
+@patch("genai_utils.gemini.genai.Client")
+@patch.dict(os.environ, {"GENAI_LABEL_TEAM": "ai", "GENAI_LABEL_ENV": "test"})
+def test_run_prompt_merges_env_labels(mock_client):
+    """Test that run_prompt merges environment labels with request labels"""
+    # Need to reload the module to pick up the new environment variables
+    import importlib
+
+    import genai_utils.gemini
+
+    importlib.reload(genai_utils.gemini)
+
+    client = Mock(Client)
+    models = Mock(Models)
+
+    models.generate_content.return_value = get_dummy()
+    client.models = models
+    mock_client.return_value = client
+
+    request_labels = {"project": "test"}
+
+    genai_utils.gemini.run_prompt(
+        "test prompt",
+        labels=request_labels,
+        model_config=ModelConfig(
+            project="project", location="location", model_name="model"
+        ),
+    )
+
+    # Verify the call was made with merged labels
+    assert models.generate_content.called
+    call_kwargs = models.generate_content.call_args[1]
+    assert "config" in call_kwargs
+
+    # Should contain both env labels (team, env) and request label (project)
+    merged_labels = call_kwargs["config"].labels
+    assert "team" in merged_labels
+    assert "env" in merged_labels
+    assert "project" in merged_labels