Added Staging bucket and update requirements to include reasoning engine

Author: holtskinner

generative_ai/gemini_reasoning_engine.py

@@ -17,16 +17,21 @@
 from vertexai.preview import reasoning_engines
 
 
-def create_reasoning_engine_basic(project_id: str) -> reasoning_engines.ReasoningEngine:
+def create_reasoning_engine_basic(
+    project_id: str, staging_bucket: str
+) -> reasoning_engines.ReasoningEngine:
 
     # [START generativeaionvertexai_create_reasoning_engine_basic]
     import vertexai
     from vertexai.preview import reasoning_engines
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "PROJECT_ID"
+    # staging_bucket = "gs://YOUR_BUCKET_NAME"
 
-    vertexai.init(project=project_id, location="us-central1")
+    vertexai.init(
+        project=project_id, location="us-central1", staging_bucket=staging_bucket
+    )
 
     class SimpleAdditionApp:
         def query(self, a: int, b: int):
@@ -60,7 +65,7 @@ def query(self, a: int, b: int):
 
 
 def create_reasoning_engine_advanced(
-    project_id: str, location: str
+    project_id: str, location: str, staging_bucket: str
 ) -> reasoning_engines.ReasoningEngine:
 
     # [START generativeaionvertexai_create_reasoning_engine_advanced]
@@ -71,6 +76,9 @@ def create_reasoning_engine_advanced(
     # TODO(developer): Update and un-comment below lines
     # project_id = "PROJECT_ID"
     # location = "us-central1"
+    # staging_bucket = "gs://YOUR_BUCKET_NAME"
+
+    vertexai.init(project=project_id, location=location, staging_bucket=staging_bucket)
 
     class LangchainApp:
         def __init__(self, project: str, location: str):
@@ -108,8 +116,6 @@ def query(self, question: str):
     app.set_up()
     print(app.query("What is Vertex AI?"))
 
-    vertexai.init(project=project_id, location=location)
-
     # Create a remote app with reasoning engine
     # This may take 1-2 minutes to finish because it builds a container and turn up HTTP servers.
     reasoning_engine = reasoning_engines.ReasoningEngine.create(

generative_ai/gemini_reasoning_engine_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import sys
 from typing import Generator
 
 import pytest
@@ -21,11 +22,15 @@
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 REGION = "us-central1"
+STAGING_BUCKET = "gs://ucaip-samples-us-central1/reasoning_engine"
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="Python 3.12 is not supported")
 @pytest.fixture(scope="module")
 def reasoning_engine_id() -> Generator[str, None, None]:
-    reasoning_engine = gemini_reasoning_engine.create_reasoning_engine_basic(PROJECT_ID)
+    reasoning_engine = gemini_reasoning_engine.create_reasoning_engine_basic(
+        PROJECT_ID, STAGING_BUCKET
+    )
     yield reasoning_engine.resource_name
     print("Deleting Reasoning Engine...")
     gemini_reasoning_engine.delete_reasoning_engine(
@@ -39,7 +44,7 @@ def test_create_reasoning_engine_basic(reasoning_engine_id: str) -> None:
 
 def test_create_reasoning_engine_advanced() -> None:
     reasoning_engine = gemini_reasoning_engine.create_reasoning_engine_advanced(
-        PROJECT_ID, REGION
+        PROJECT_ID, REGION, STAGING_BUCKET
     )
     assert reasoning_engine
     gemini_reasoning_engine.delete_reasoning_engine(

generative_ai/requirements.txt

@@ -2,6 +2,6 @@ pandas==1.3.5; python_version == '3.7'
 pandas==2.0.1; python_version > '3.7'
 pillow==9.5.0; python_version < '3.8'
 pillow==10.3.0; python_version >= '3.8'
-google-cloud-aiplatform[pipelines]==1.50.0
+google-cloud-aiplatform[pipelines,reasoningengine]==1.50.0
 google-auth==2.17.3
 anthropic[vertex]==0.25.6