Use graphql query builder

Author: NarekA

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
     "tenacity>=8.0.1",
     "watchfiles>=0.19.0,<0.20",
     "truss-transfer>=0.0.37,<0.0.40",
+    "gql-query-builder (>=0.1.7,<0.2.0)",
 ]
 
 [project.urls]

truss/remote/baseten/api.py

@@ -1,8 +1,10 @@
+import json
 import logging
 from enum import Enum
 from typing import Any, Dict, List, Mapping, Optional
 
 import requests
+from gql_query_builder import GqlQuery
 from pydantic import BaseModel, Field
 
 from truss.remote.baseten import custom_types as b10_types
@@ -299,37 +301,39 @@ def deploy_chain_atomic(
         chain_name: Optional[str] = None,
         environment: Optional[str] = None,
         is_draft: bool = False,
+        original_source_artifact_s3_key: Optional[str] = None,
+        allow_truss_download: Optional[bool] = True,
     ):
-        entrypoint_str = _chainlet_data_atomic_to_graphql_mutation(entrypoint)
-
-        dependencies_str = ", ".join(
-            [
+        if allow_truss_download is None:
+            allow_truss_download = True
+
+        # Build the mutation parameters
+        mutation_params = {
+            "chain_id": chain_id,
+            "chain_name": chain_name,
+            "environment": environment,
+            "original_source_artifact_s3_key": original_source_artifact_s3_key,
+            "allow_truss_download": "false" if allow_truss_download is False else None,
+            "is_draft": is_draft,
+            "entrypoint": _chainlet_data_atomic_to_graphql_mutation(entrypoint),
+            "dependencies": [
                 _chainlet_data_atomic_to_graphql_mutation(dependency)
                 for dependency in dependencies
-            ]
-        )
+            ],
+            "truss_user_env": "$trussUserEnv",
+        }
+        mutation_params = {
+            str(k): v for k, v in mutation_params.items() if v is not None
+        }
 
-        query_string = f"""
-            mutation ($trussUserEnv: String) {{
-                deploy_chain_atomic(
-                    {f'chain_id: "{chain_id}"' if chain_id else ""}
-                    {f'chain_name: "{chain_name}"' if chain_name else ""}
-                    {f'environment: "{environment}"' if environment else ""}
-                    is_draft: {str(is_draft).lower()}
-                    entrypoint: {entrypoint_str}
-                    dependencies: [{dependencies_str}]
-                    truss_user_env: $trussUserEnv
-                ) {{
-                    chain_deployment {{
-                        id
-                        chain {{
-                            id
-                            hostname
-                        }}
-                    }}
-                }}
-            }}
-        """
+        gql = GqlQuery()
+        gql.operation(
+            "mutation",
+            "deploy_chain_atomic",
+            mutation_params,
+            ["chain_deployment { id chain { id hostname } }"],
+        )
+        query_string = gql.generate()
 
         resp = self._post_graphql_query(
             query_string, variables={"trussUserEnv": truss_user_env.json()}

truss/tests/remote/baseten/test_api.py

@@ -353,8 +353,8 @@ def test_deploy_chain_deployment(mock_post, baseten_api):
 
     gql_mutation = mock_post.call_args[1]["json"]["query"]
 
-    assert 'environment: "production"' in gql_mutation
-    assert 'chain_id: "chain_id"' in gql_mutation
+    assert "environment: production" in gql_mutation
+    assert "chain_id: chain_id" in gql_mutation
     assert "dependencies:" in gql_mutation
     assert "entrypoint:" in gql_mutation
 
@@ -378,8 +378,8 @@ def test_deploy_chain_deployment_with_gitinfo(mock_post, baseten_api):
 
     gql_mutation = mock_post.call_args[1]["json"]["query"]
 
-    assert 'environment: "production"' in gql_mutation
-    assert 'chain_id: "chain_id"' in gql_mutation
+    assert "environment: production" in gql_mutation
+    assert "chain_id: chain_id" in gql_mutation
     assert "dependencies:" in gql_mutation
     assert "entrypoint:" in gql_mutation
 
@@ -402,12 +402,70 @@ def test_deploy_chain_deployment_no_environment(mock_post, baseten_api):
 
     gql_mutation = mock_post.call_args[1]["json"]["query"]
 
-    assert 'chain_id: "chain_id"' in gql_mutation
+    assert "chain_id: chain_id" in gql_mutation
     assert "environment" not in gql_mutation
     assert "dependencies:" in gql_mutation
     assert "entrypoint:" in gql_mutation
 
 
+@mock.patch("requests.post", return_value=mock_deploy_chain_deployment_response())
+def test_deploy_chain_deployment_with_dependencies(mock_post, baseten_api):
+    dependencies = [
+        ChainletDataAtomic(
+            name="dependency-1",
+            oracle=OracleData(
+                model_name="dep-model-1",
+                s3_key="dep-s3-key-1",
+                encoded_config_str="dep-encoded-config-str-1",
+            ),
+        ),
+        ChainletDataAtomic(
+            name="dependency-2",
+            oracle=OracleData(
+                model_name="dep-model-2",
+                s3_key="dep-s3-key-2",
+                encoded_config_str="dep-encoded-config-str-2",
+            ),
+        ),
+    ]
+
+    baseten_api.deploy_chain_atomic(
+        environment="production",
+        chain_id="chain_id",
+        dependencies=dependencies,
+        entrypoint=ChainletDataAtomic(
+            name="chainlet-1",
+            oracle=OracleData(
+                model_name="model-1",
+                s3_key="s3-key-1",
+                encoded_config_str="encoded-config-str-1",
+            ),
+        ),
+        truss_user_env=b10_types.TrussUserEnv.collect(),
+    )
+
+    gql_mutation = mock_post.call_args[1]["json"]["query"]
+
+    # Single regex to check all assertions
+    import re
+
+    pattern = (
+        r"(?=.*environment: production)"
+        r"(?=.*chain_id: chain_id)"
+        r"(?=.*dependencies:)"
+        r"(?=.*entrypoint:)"
+        r'(?=.*name: "dependency-1")'
+        r'(?=.*name: "dependency-2")'
+        r'(?=.*model_name: "dep-model-1")'
+        r'(?=.*model_name: "dep-model-2")'
+        r'(?=.*s3_key: "dep-s3-key-1")'
+        r'(?=.*s3_key: "dep-s3-key-2")'
+    )
+    assert re.search(pattern, gql_mutation), (
+        f"GraphQL mutation does not contain all expected elements: {gql_mutation}"
+    )
+
+
 @mock.patch("requests.post", return_value=mock_upsert_training_project_response())
 def test_upsert_training_project(mock_post, baseten_api):
     baseten_api.upsert_training_project(