diff --git a/generative_ai/gemini_reasoning_engine.py b/generative_ai/gemini_reasoning_engine.py new file mode 100644 index 00000000000..b25dd2c73cf --- /dev/null +++ b/generative_ai/gemini_reasoning_engine.py @@ -0,0 +1,204 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, List, Union + +from vertexai.preview import reasoning_engines + + +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", staging_bucket=staging_bucket + ) + + class SimpleAdditionApp: + def query(self, a: int, b: int) -> str: + """Query the application. + + Args: + a: The first input number + b: The second input number + + Returns: + int: The additional result. + """ + + return f"{int(a)} + {int(b)} is {int(a + b)}" + + # Locally test + app = SimpleAdditionApp() + app.query(a=1, b=2) + + # Create a remote app with reasoning engine. + # This may take 1-2 minutes to finish. + reasoning_engine = reasoning_engines.ReasoningEngine.create( + SimpleAdditionApp(), + display_name="Demo Addition App", + description="A simple demo addition app", + requirements=[], + extra_packages=[], + ) + # [END generativeaionvertexai_create_reasoning_engine_basic] + return reasoning_engine + + +def create_reasoning_engine_advanced( + project_id: str, location: str, staging_bucket: str +) -> reasoning_engines.ReasoningEngine: + # [START generativeaionvertexai_create_reasoning_engine_advanced] + + from typing import List + + import vertexai + from vertexai.preview import reasoning_engines + + # 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) -> None: + self.project_id = project + self.location = location + + def set_up(self) -> None: + from langchain_core.prompts import ChatPromptTemplate + from langchain_google_vertexai import ChatVertexAI + + system = ( + "You are a helpful assistant that answers questions " + "about Google Cloud." + ) + human = "{text}" + prompt = ChatPromptTemplate.from_messages( + [("system", system), ("human", human)] + ) + chat = ChatVertexAI(project=self.project_id, location=self.location) + self.chain = prompt | chat + + def query(self, question: str) -> Union[str, List[Union[str, Dict]]]: + """Query the application. + + Args: + question: The user prompt. + + Returns: + str: The LLM response. + """ + return self.chain.invoke({"text": question}).content + + # Locally test + app = LangchainApp(project=project_id, location=location) + app.set_up() + print(app.query("What is Vertex AI?")) + + # 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( + LangchainApp(project=project_id, location=location), + requirements=[ + "google-cloud-aiplatform==1.50.0", + "langchain-google-vertexai", + "langchain-core", + ], + display_name="Demo LangChain App", + description="This is a simple LangChain app.", + # sys_version="3.10", # Optional + extra_packages=[], + ) + # [END generativeaionvertexai_create_reasoning_engine_advanced] + return reasoning_engine + + +def query_reasoning_engine(project_id: str, reasoning_engine_id: str) -> object: + # [START generativeaionvertexai_query_reasoning_engine] + import vertexai + from vertexai.preview import reasoning_engines + + # TODO(developer): Update and un-comment below lines + # project_id = "PROJECT_ID" + # reasoning_engine_id = "REASONING_ENGINE_ID" + + vertexai.init(project=project_id, location="us-central1") + reasoning_engine = reasoning_engines.ReasoningEngine(reasoning_engine_id) + + # Replace with kwargs for `.query()` method. + response = reasoning_engine.query(a=1, b=2) + print(response) + # [END generativeaionvertexai_query_reasoning_engine] + return response + + +def list_reasoning_engines(project_id: str) -> List[reasoning_engines.ReasoningEngine]: + # [START generativeaionvertexai_list_reasoning_engines] + import vertexai + from vertexai.preview import reasoning_engines + + # TODO(developer): Update and un-comment below lines + # project_id = "PROJECT_ID" + + vertexai.init(project=project_id, location="us-central1") + + reasoning_engine_list = reasoning_engines.ReasoningEngine.list() + print(reasoning_engine_list) + # [END generativeaionvertexai_list_reasoning_engines] + return reasoning_engine_list + + +def get_reasoning_engine( + project_id: str, reasoning_engine_id: str +) -> reasoning_engines.ReasoningEngine: + # [START generativeaionvertexai_get_reasoning_engine] + import vertexai + from vertexai.preview import reasoning_engines + + # TODO(developer): Update and un-comment below lines + # project_id = "PROJECT_ID" + # reasoning_engine_id = "REASONING_ENGINE_ID" + + vertexai.init(project=project_id, location="us-central1") + + reasoning_engine = reasoning_engines.ReasoningEngine(reasoning_engine_id) + print(reasoning_engine) + # [END generativeaionvertexai_get_reasoning_engine] + return reasoning_engine + + +def delete_reasoning_engine(project_id: str, reasoning_engine_id: str) -> None: + # [START generativeaionvertexai_delete_reasoning_engine] + import vertexai + from vertexai.preview import reasoning_engines + + # TODO(developer): Update and un-comment below lines + # project_id = "PROJECT_ID" + # reasoning_engine_id = "REASONING_ENGINE_ID" + + vertexai.init(project=project_id, location="us-central1") + + reasoning_engine = reasoning_engines.ReasoningEngine(reasoning_engine_id) + reasoning_engine.delete() + # [END generativeaionvertexai_delete_reasoning_engine] diff --git a/generative_ai/gemini_reasoning_engine_test.py b/generative_ai/gemini_reasoning_engine_test.py new file mode 100644 index 00000000000..243eecffc73 --- /dev/null +++ b/generative_ai/gemini_reasoning_engine_test.py @@ -0,0 +1,74 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +from typing import Generator + +import pytest + +import gemini_reasoning_engine + +PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") +REGION = "us-central1" +STAGING_BUCKET = "gs://ucaip-samples-us-central1" + + +@pytest.fixture(scope="module") +def reasoning_engine_id() -> Generator[str, None, None]: + 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( + PROJECT_ID, reasoning_engine.resource_name + ) + + +@pytest.mark.skipif( + sys.version_info >= (3, 12), reason="requires Python version lower than 3.12" +) +def test_create_reasoning_engine_basic(reasoning_engine_id: str) -> None: + assert reasoning_engine_id + + +def test_create_reasoning_engine_advanced() -> None: + reasoning_engine = gemini_reasoning_engine.create_reasoning_engine_advanced( + PROJECT_ID, REGION, STAGING_BUCKET + ) + assert reasoning_engine + gemini_reasoning_engine.delete_reasoning_engine( + PROJECT_ID, reasoning_engine.resource_name + ) + + +def test_query_reasoning_engine(reasoning_engine_id: str) -> None: + response = gemini_reasoning_engine.query_reasoning_engine( + PROJECT_ID, reasoning_engine_id + ) + assert response + assert response == "1 + 2 is 3" + + +def test_list_reasoning_engines() -> None: + response = gemini_reasoning_engine.list_reasoning_engines(PROJECT_ID) + assert response + + +def test_get_reasoning_engine(reasoning_engine_id: str) -> None: + response = gemini_reasoning_engine.get_reasoning_engine( + PROJECT_ID, reasoning_engine_id + ) + assert response diff --git a/generative_ai/requirements-test.txt b/generative_ai/requirements-test.txt index fceefa0d35e..5867597a7da 100644 --- a/generative_ai/requirements-test.txt +++ b/generative_ai/requirements-test.txt @@ -1,4 +1,4 @@ backoff==2.2.1 -google-api-core==2.17.1 +google-api-core==2.19.0 pytest==8.1.1 pytest-asyncio==0.23.6 diff --git a/generative_ai/requirements.txt b/generative_ai/requirements.txt index 709fb2fa04a..e92c5e171cb 100644 --- a/generative_ai/requirements.txt +++ b/generative_ai/requirements.txt @@ -2,6 +2,8 @@ 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-auth==2.17.3 +google-cloud-aiplatform[pipelines,reasoningengine]==1.50.0 +google-auth==2.29.0 anthropic[vertex]==0.25.6 +langchain-core==0.1.52 +langchain-google-vertexai==1.0.3