Add memory bank to modules/agent-engine (#3844)

* feat(agent-engine): add support for memory bank configuration

* refactor(agent-engine): remove source_path_override and revert to standard source_path

* provider version upgrade from 7.17 to 7.27

* docs(agent-engine): fix README validation and update tables

---------

Co-authored-by: Hemanand <hemr@google.com>

modules/agent-engine/README.md

@@ -16,6 +16,7 @@ The module creates Agent Engine and related dependencies.
 - [Private networking: setup PSC-I](#private-networking-setup-psc-i)
 - [Specify an encryption key](#specify-an-encryption-key)
 - [Define environment variables and use secrets](#define-environment-variables-and-use-secrets)
+- [Memory Bank](#memory-bank)
 - [Getting values from context](#getting-values-from-context)
 - [Variables](#variables)
 - [Outputs](#outputs)
@@ -344,25 +345,63 @@ module "agent_engine" {
 # tftest inventory=environment.yaml
 ```
 
+## Memory Bank
+
+You can optionally configure a Memory Bank to provide long-term persistent memory for your agent.
+
+```hcl
+module "agent_engine" {
+  source     = "./fabric/modules/agent-engine"
+  name       = "my-agent"
+  project_id = var.project_id
+  region     = var.region
+
+  agent_engine_config = {
+    agent_framework = "google-adk"
+  }
+
+  deployment_files = {
+    source_config = {
+      source_path = "assets/src/source.tar.gz"
+    }
+  }
+
+  memory_bank_config = {
+    disable_memory_revisions = false
+    generation_config = {
+      model = "projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-001"
+    }
+    similarity_search_config = {
+      embedding_model = "projects/my-project/locations/us-central1/publishers/google/models/text-embedding-005"
+    }
+    ttl_config = {
+      default_ttl = "2592000s" # 30 days
+    }
+  }
+}
+# tftest skip
+```
+
 ## Getting values from context
 
-The module allows you to dynamically reference context values for resources created outside this module, through the `context` variable. This includes the definition of custom roles, iam_principals, locations, kms_keys and project ids.
+The module allows you to dynamically reference context values for resources created outside this module, through the `context` variable. This includes the definition of custom roles, iam_principals, locations, kms_keys, models and project ids.
 <!-- BEGIN TFDOC -->
 ## Variables
 
 | name | description | type | required | default |
 |---|---|:---:|:---:|:---:|
 | [agent_engine_config](variables.tf#L17) | The agent configuration. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> | ✓ |  |
-| [name](variables.tf#L122) | The name of the agent. | <code>string</code> | ✓ |  |
-| [project_id](variables.tf#L141) | The id of the project where to deploy the agent. | <code>string</code> | ✓ |  |
-| [region](variables.tf#L147) | The region where to deploy the agent. | <code>string</code> | ✓ |  |
+| [name](variables.tf#L146) | The name of the agent. | <code>string</code> | ✓ |  |
+| [project_id](variables.tf#L165) | The id of the project where to deploy the agent. | <code>string</code> | ✓ |  |
+| [region](variables.tf#L171) | The region where to deploy the agent. | <code>string</code> | ✓ |  |
 | [bucket_config](variables.tf#L40) | The GCS bucket configuration. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>&#123;&#125;</code> |
 | [context](variables.tf#L52) | Context-specific interpolations. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>&#123;&#125;</code> |
-| [deployment_files](variables.tf#L65) | The to source files path and names. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>&#123;&#8230;&#125;</code> |
-| [description](variables.tf#L102) | The Agent Engine description. | <code>string</code> |  | <code>&#34;Terraform managed.&#34;</code> |
-| [encryption_key](variables.tf#L109) | The full resource name of the Cloud KMS CryptoKey. | <code>string</code> |  | <code>null</code> |
-| [managed](variables.tf#L115) | Whether the Terraform module should control the code updates. | <code>bool</code> |  | <code>true</code> |
-| [networking_config](variables.tf#L128) | Networking configuration. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>null</code> |
+| [deployment_files](variables.tf#L66) | The to source files path and names. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>&#123;&#8230;&#125;</code> |
+| [description](variables.tf#L103) | The Agent Engine description. | <code>string</code> |  | <code>&#34;Terraform managed.&#34;</code> |
+| [encryption_key](variables.tf#L110) | The full resource name of the Cloud KMS CryptoKey. | <code>string</code> |  | <code>null</code> |
+| [managed](variables.tf#L116) | Whether the Terraform module should control the code updates. | <code>bool</code> |  | <code>true</code> |
+| [memory_bank_config](variables.tf#L123) | Configuration for the memory bank. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>null</code> |
+| [networking_config](variables.tf#L152) | Networking configuration. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>null</code> |
 | [service_account_config](variables-serviceaccount.tf#L18) | Service account configurations. | <code>object&#40;&#123;&#8230;&#125;&#41;</code> |  | <code>&#123;&#125;</code> |
 
 ## Outputs

modules/agent-engine/agent-managed.tf

@@ -15,6 +15,7 @@
  */
 
 resource "google_vertex_ai_reasoning_engine" "managed" {
+  provider     = google-beta
   count        = var.managed ? 1 : 0
   display_name = var.name
   project      = local.project_id
@@ -147,4 +148,69 @@ resource "google_vertex_ai_reasoning_engine" "managed" {
       }
     }
   }
+
+  dynamic "context_spec" {
+    for_each = var.memory_bank_config != null ? { 1 = 1 } : {}
+
+    content {
+      memory_bank_config {
+        disable_memory_revisions = var.memory_bank_config.disable_memory_revisions
+
+        dynamic "generation_config" {
+          for_each = (
+            var.memory_bank_config.generation_config != null ? { 1 = 1 } : {}
+          )
+          content {
+            model = lookup(
+              local.ctx.models,
+              var.memory_bank_config.generation_config.model,
+              var.memory_bank_config.generation_config.model
+            )
+          }
+        }
+
+        dynamic "similarity_search_config" {
+          for_each = (
+            var.memory_bank_config.similarity_search_config != null ? { 1 = 1 } : {}
+          )
+          content {
+            embedding_model = lookup(
+              local.ctx.models,
+              var.memory_bank_config.similarity_search_config.embedding_model,
+              var.memory_bank_config.similarity_search_config.embedding_model
+            )
+          }
+        }
+
+        dynamic "ttl_config" {
+          for_each = (
+            var.memory_bank_config.ttl_config != null ? { 1 = 1 } : {}
+          )
+          content {
+            default_ttl                 = var.memory_bank_config.ttl_config.default_ttl
+            memory_revision_default_ttl = var.memory_bank_config.ttl_config.memory_revision_default_ttl
+
+            dynamic "granular_ttl_config" {
+              for_each = (
+                var.memory_bank_config.ttl_config.granular_ttl_config != null
+                ? { 1 = 1 }
+                : {}
+              )
+              content {
+                create_ttl = (
+                  var.memory_bank_config.ttl_config.granular_ttl_config.create_ttl
+                )
+                generate_created_ttl = (
+                  var.memory_bank_config.ttl_config.granular_ttl_config.generate_created_ttl
+                )
+                generate_updated_ttl = (
+                  var.memory_bank_config.ttl_config.granular_ttl_config.generate_updated_ttl
+                )
+              }
+            }
+          }
+        }
+      }
+    }
+  }
 }

modules/agent-engine/agent-unmanaged.tf

@@ -15,6 +15,7 @@
  */
 
 resource "google_vertex_ai_reasoning_engine" "unmanaged" {
+  provider     = google-beta
   count        = var.managed ? 0 : 1
   display_name = var.name
   project      = local.project_id
@@ -148,6 +149,71 @@ resource "google_vertex_ai_reasoning_engine" "unmanaged" {
     }
   }
 
+  dynamic "context_spec" {
+    for_each = var.memory_bank_config != null ? { 1 = 1 } : {}
+
+    content {
+      memory_bank_config {
+        disable_memory_revisions = var.memory_bank_config.disable_memory_revisions
+
+        dynamic "generation_config" {
+          for_each = (
+            var.memory_bank_config.generation_config != null ? { 1 = 1 } : {}
+          )
+          content {
+            model = lookup(
+              local.ctx.models,
+              var.memory_bank_config.generation_config.model,
+              var.memory_bank_config.generation_config.model
+            )
+          }
+        }
+
+        dynamic "similarity_search_config" {
+          for_each = (
+            var.memory_bank_config.similarity_search_config != null ? { 1 = 1 } : {}
+          )
+          content {
+            embedding_model = lookup(
+              local.ctx.models,
+              var.memory_bank_config.similarity_search_config.embedding_model,
+              var.memory_bank_config.similarity_search_config.embedding_model
+            )
+          }
+        }
+
+        dynamic "ttl_config" {
+          for_each = (
+            var.memory_bank_config.ttl_config != null ? { 1 = 1 } : {}
+          )
+          content {
+            default_ttl                 = var.memory_bank_config.ttl_config.default_ttl
+            memory_revision_default_ttl = var.memory_bank_config.ttl_config.memory_revision_default_ttl
+
+            dynamic "granular_ttl_config" {
+              for_each = (
+                var.memory_bank_config.ttl_config.granular_ttl_config != null
+                ? { 1 = 1 }
+                : {}
+              )
+              content {
+                create_ttl = (
+                  var.memory_bank_config.ttl_config.granular_ttl_config.create_ttl
+                )
+                generate_created_ttl = (
+                  var.memory_bank_config.ttl_config.granular_ttl_config.generate_created_ttl
+                )
+                generate_updated_ttl = (
+                  var.memory_bank_config.ttl_config.granular_ttl_config.generate_updated_ttl
+                )
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   lifecycle {
     ignore_changes = [
       spec[0].package_spec,

modules/agent-engine/variables.tf

@@ -56,6 +56,7 @@ variable "context" {
     iam_principals = optional(map(string), {})
     locations      = optional(map(string), {})
     kms_keys       = optional(map(string), {})
+    models         = optional(map(string), {})
     project_ids    = optional(map(string), {})
   })
   nullable = false
@@ -119,6 +120,29 @@ variable "managed" {
   default     = true
 }
 
+variable "memory_bank_config" {
+  description = "Configuration for the memory bank."
+  type = object({
+    disable_memory_revisions = optional(bool)
+    generation_config = optional(object({
+      model = string
+    }))
+    similarity_search_config = optional(object({
+      embedding_model = string
+    }))
+    ttl_config = optional(object({
+      default_ttl                 = optional(string)
+      memory_revision_default_ttl = optional(string)
+      granular_ttl_config = optional(object({
+        create_ttl           = optional(string)
+        generate_created_ttl = optional(string)
+        generate_updated_ttl = optional(string)
+      }))
+    }))
+  })
+  default = null
+}
+
 variable "name" {
   description = "The name of the agent."
   type        = string

modules/agent-engine/versions.tf

@@ -19,11 +19,11 @@ terraform {
   required_providers {
     google = {
       source  = "hashicorp/google"
-      version = ">= 7.17.0, < 8.0.0" # tftest
+      version = ">= 7.27.0, < 8.0.0" # tftest
     }
     google-beta = {
       source  = "hashicorp/google-beta"
-      version = ">= 7.17.0, < 8.0.0" # tftest
+      version = ">= 7.27.0, < 8.0.0" # tftest
     }
   }
   provider_meta "google" {

modules/agent-engine/versions.tofu

@@ -19,11 +19,11 @@ terraform {
   required_providers {
     google = {
       source  = "hashicorp/google"
-      version = ">= 7.17.0, < 8.0.0" # tftest
+      version = ">= 7.27.0, < 8.0.0" # tftest
     }
     google-beta = {
       source  = "hashicorp/google-beta"
-      version = ">= 7.17.0, < 8.0.0" # tftest
+      version = ">= 7.27.0, < 8.0.0" # tftest
     }
   }
   provider_meta "google" {