diff --git a/ai/watsonxai-generation/watsonx-ai.graphql b/ai/watsonxai-generation/watsonx-ai.graphql
index c9a760d..04d48d4 100644
--- a/ai/watsonxai-generation/watsonx-ai.graphql
+++ b/ai/watsonxai-generation/watsonx-ai.graphql
@@ -1,34 +1,36 @@
 extend type Query {
   """
-  Infers the next tokens for a given deployed model from the `input` text.
+  Internal helper that accepts a Secret token and returns it as JSON.
+  This allows the Secret type to be converted to JSON for injection.
+  """
+  _token_to_json(token: Secret!): JSON
+    @value(
+      script: {
+        src: """
+        {"token": token }
+        """
+        language: JSONATA
+      }
+    )
+
   """
-  wxai_generation(
-    input: String
-    model_id: String! = "ibm/granite-3-8b-instruct"
-    parameters: WXAI_GenerationParameters
-  ): WXAI_GenerationResponse
+  Extract IBM IAM token and inject it into wxai_generation for authentication.
+  This makes the access_token available as an expansion variable.
+  """
+  _inject_ibm_iam_token: JSON
+    @inject(on: [{ expose: true, types: "Query", fields: "wxai_generation" }])
     @sequence(
       steps: [
         { query: "ibm_iam_token" }
-        {
-          query: "_wxai_generation"
-          arguments: [
-            { name: "token", field: "§0" }
-            { name: "input", argument: "input" }
-            { name: "model_id", argument: "model_id" }
-            { name: "parameters", argument: "parameters" }
-          ]
-        }
+        { query: "_token_to_json", arguments: [{ name: "token", field: "§0" }] }
       ]
     )
-
   """
-  Calls the generation endpoint.
+  Infers the next tokens for a given deployed model from the `input` text.
   """
-  _wxai_generation(
-    token: Secret!
+  wxai_generation(
     input: String
-    model_id: String
+    model_id: String = "ibm/granite-3-8b-instruct"
     parameters: WXAI_GenerationParameters
   ): WXAI_GenerationResponse
     @rest(