feat: Add the system tests for Langgraph prebuilt template

PiperOrigin-RevId: 734299775

Author: shawn-yang-google

tests/system/vertexai/test_reasoning_engines.py

@@ -106,3 +106,74 @@ def runnable_builder(**kwargs):
  bucket = client.bucket(shared_state["staging_bucket_name"])
  assert bucket.exists()
  assert bucket.get_blob(f"test-gcs-dir-name/{_BLOB_FILENAME}").exists()
+
+ def test_langgraph_template(self, shared_state):
+ super().setup_method()
+ credentials, _ = auth.default(
+ scopes=["https://www.googleapis.com/auth/cloud-platform"]
+ )
+ vertexai.init(
+ project=e2e_base._PROJECT,
+ location=e2e_base._LOCATION,
+ staging_bucket=f"gs://{shared_state['staging_bucket_name']}",
+ credentials=credentials,
+ )
+
+ # System tests are currently affected by contamination in the Gemini
+ # model and ToolConfig test fixture.
+ # To eliminate false positives, we are mocking the runnnable builder to
+ # make the system tests hermetic.
+ # This change will be reverted once the the test fixture is corrected.
+ class LanggraphAgentNoDependencies:
+ """LangGraph Agent with no dependencies."""
+
+ def invoke(self, input, **kwargs) -> str:
+ return "Testing langgraph agent with no dependencies."
+
+ def runnable_builder(**kwargs):
+ """Creates a LangGraph Runnable."""
+ return LanggraphAgentNoDependencies()
+
+ # Test prebuilt langgraph_template
+ created_app = reasoning_engines.ReasoningEngine.create(
+ reasoning_engines.LanggraphAgent(
+ model="gemini-1.5-pro-preview-0409",
+ runnable_builder=runnable_builder,
+ ),
+ requirements=["google-cloud-aiplatform[reasoningengine,langchain]"],
+ display_name="test-display-name",
+ description="test-description",
+ gcs_dir_name="test-gcs-dir-name",
+ )
+ shared_state.setdefault("resources", [])
+ shared_state["resources"].append(created_app) # Deletion at teardown.
+ got_app = reasoning_engines.ReasoningEngine(created_app.resource_name)
+
+ # Test resource attributes
+ assert isinstance(created_app.resource_name, str)
+ assert got_app.resource_name == created_app.resource_name
+ assert got_app.gca_resource.name == got_app.resource_name
+ assert got_app.gca_resource.display_name == "test-display-name"
+ assert got_app.gca_resource.description == "test-description"
+
+ # Test operation schemas
+ assert got_app.operation_schemas() == created_app.operation_schemas()
+
+ # Test query response
+ # (Wrap in a try-except block because of non-determinism from Gemini.)
+ try:
+ response = created_app.query(input="hello")
+ assert response.get("input") == "hello"
+ response = got_app.query(input="hello")
+ assert response.get("input") == "hello"
+ except exceptions.FailedPrecondition:
+ traceback.print_exc()
+ except Exception:
+ traceback.print_exc()
+
+ # Test GCS Bucket subdirectory creation
+ # Original: https://github.com/googleapis/python-aiplatform/issues/3650
+ client = storage.Client(project=e2e_base._PROJECT)
+ bucket = client.bucket(shared_state["staging_bucket_name"])
+ assert bucket.exists()
+ assert bucket.get_blob(f"test-gcs-dir-name/{_BLOB_FILENAME}").exists()