From 59c0a353559d3376970ae198ef3041e7a7dba23f Mon Sep 17 00:00:00 2001 From: Andrew Liu <159852527+aliu39@users.noreply.github.com> Date: Mon, 5 May 2025 16:05:20 -0700 Subject: [PATCH 1/4] Update vertex provider code and test --- src/sentry/llm/providers/vertex.py | 52 +++++++++++++++------------- tests/sentry/llm/test_vertex.py | 55 ++++++++++++++++++++++-------- 2 files changed, 69 insertions(+), 38 deletions(-) diff --git a/src/sentry/llm/providers/vertex.py b/src/sentry/llm/providers/vertex.py index 0cccb41af7fd06..e6cc6164441901 100644 --- a/src/sentry/llm/providers/vertex.py +++ b/src/sentry/llm/providers/vertex.py @@ -2,11 +2,12 @@ import google.auth import google.auth.transport.requests -import requests +from google import genai +from google.genai.types import GenerateContentConfig, HttpOptions from sentry.llm.exceptions import VertexRequestFailed from sentry.llm.providers.base import LlmModelBase -from sentry.llm.types import UseCaseConfig +from sentry.llm.types import ProviderConfig, UseCaseConfig logger = logging.getLogger(__name__) @@ -20,6 +21,15 @@ class VertexProvider(LlmModelBase): candidate_count = 1 # we only want one candidate returned at the moment top_p = 1 # TODO: make this configurable? + def __init__(self, provider_config: ProviderConfig) -> None: + super().__init__(provider_config) + self.client = genai.Client( + vertexai=True, + project=provider_config["options"]["project"], + location=provider_config["options"]["location"], + http_options=HttpOptions(api_version="v1"), + ) + def _complete_prompt( self, *, @@ -30,35 +40,29 @@ def _complete_prompt( max_output_tokens: int, ) -> str | None: + model = usecase_config["options"]["model"] content = f"{prompt} {message}" if prompt else message + generate_config = GenerateContentConfig( + candidate_count=self.candidate_count, + max_output_tokens=max_output_tokens, + temperature=temperature, + top_p=self.top_p, + ) - payload = { - "instances": [{"content": content}], - "parameters": { - "candidateCount": self.candidate_count, - "maxOutputTokens": max_output_tokens, - "temperature": temperature, - "topP": self.top_p, - }, - } - - headers = { - "Authorization": f"Bearer {self._get_access_token()}", - "Content-Type": "application/json", - } - vertex_url = self.provider_config["options"]["url"] - vertex_url += usecase_config["options"]["model"] + ":predict" - - response = requests.post(vertex_url, headers=headers, json=payload) + response = self.client.models.generate_content( + model=model, + contents=content, + config=generate_config, + ) if response.status_code != 200: logger.error( - "Request failed with status code and response text.", - extra={"status_code": response.status_code, "response_text": response.text}, + "Vertex request failed.", + extra={"status_code": response.status_code}, ) - raise VertexRequestFailed(f"Response {response.status_code}: {response.text}") + raise VertexRequestFailed(f"Response {response.status_code}") - return response.json()["predictions"][0]["content"] + return response.text def _get_access_token(self) -> str: # https://stackoverflow.com/questions/53472429/how-to-get-a-gcp-bearer-token-programmatically-with-python diff --git a/tests/sentry/llm/test_vertex.py b/tests/sentry/llm/test_vertex.py index 09f43cf85e15ec..5b1bd8165d900d 100644 --- a/tests/sentry/llm/test_vertex.py +++ b/tests/sentry/llm/test_vertex.py @@ -1,29 +1,53 @@ -from unittest.mock import patch +from contextlib import contextmanager +from unittest.mock import Mock, patch from sentry.llm.usecases import LLMUseCase, complete_prompt -def test_complete_prompt(set_sentry_option): +@contextmanager +def mock_options(set_sentry_option): with ( set_sentry_option( "llm.provider.options", - {"vertex": {"models": ["vertex-1.0"], "options": {"url": "fake_url"}}}, + { + "vertex": { + "models": ["vertex-1.0"], + "options": {"project": "my-gcp-project", "location": "us-central1"}, + } + }, ), set_sentry_option( "llm.usecases.options", {"example": {"provider": "vertex", "options": {"model": "vertex-1.0"}}}, ), + ): + yield + + +class MockGenaiClient: + def __init__(self, mock_generate_content): + self.models = type( + "obj", + (object,), + {"generate_content": mock_generate_content}, + )() + + +def test_complete_prompt(set_sentry_option): + + mock_generate_content = Mock( + return_value=type( + "obj", + (object,), + {"status_code": 200, "text": "hello world"}, + )() + ) + + with ( + mock_options(set_sentry_option), patch( - "sentry.llm.providers.vertex.VertexProvider._get_access_token", - return_value="fake_token", - ), - patch( - "requests.post", - return_value=type( - "obj", - (object,), - {"status_code": 200, "json": lambda x: {"predictions": [{"content": ""}]}}, - )(), + "sentry.llm.providers.vertex.genai.Client", + return_value=MockGenaiClient(mock_generate_content), ), ): res = complete_prompt( @@ -33,4 +57,7 @@ def test_complete_prompt(set_sentry_option): temperature=0.0, max_output_tokens=1024, ) - assert res == "" + + assert res == "hello world" + assert mock_generate_content.call_count == 1 + assert mock_generate_content.call_args[1]["model"] == "vertex-1.0" From fe464b8ba7dc5e7e02a54d61702c499a85e66fd8 Mon Sep 17 00:00:00 2001 From: Andrew Liu <159852527+aliu39@users.noreply.github.com> Date: Mon, 5 May 2025 16:21:37 -0700 Subject: [PATCH 2/4] Add error test --- tests/sentry/llm/test_vertex.py | 44 ++++++++++++++++++++++++--------- 1 file changed, 33 insertions(+), 11 deletions(-) diff --git a/tests/sentry/llm/test_vertex.py b/tests/sentry/llm/test_vertex.py index 5b1bd8165d900d..d87c65c0be84cf 100644 --- a/tests/sentry/llm/test_vertex.py +++ b/tests/sentry/llm/test_vertex.py @@ -1,10 +1,12 @@ -from contextlib import contextmanager from unittest.mock import Mock, patch -from sentry.llm.usecases import LLMUseCase, complete_prompt +import pytest +from sentry.llm.exceptions import VertexRequestFailed +from sentry.llm.usecases import LLMUseCase, complete_prompt, llm_provider_backends -@contextmanager + +@pytest.fixture def mock_options(set_sentry_option): with ( set_sentry_option( @@ -33,8 +35,8 @@ def __init__(self, mock_generate_content): )() -def test_complete_prompt(set_sentry_option): - +def test_complete_prompt(mock_options): + llm_provider_backends.clear() mock_generate_content = Mock( return_value=type( "obj", @@ -43,12 +45,9 @@ def test_complete_prompt(set_sentry_option): )() ) - with ( - mock_options(set_sentry_option), - patch( - "sentry.llm.providers.vertex.genai.Client", - return_value=MockGenaiClient(mock_generate_content), - ), + with patch( + "sentry.llm.providers.vertex.genai.Client", + return_value=MockGenaiClient(mock_generate_content), ): res = complete_prompt( usecase=LLMUseCase.EXAMPLE, @@ -61,3 +60,26 @@ def test_complete_prompt(set_sentry_option): assert res == "hello world" assert mock_generate_content.call_count == 1 assert mock_generate_content.call_args[1]["model"] == "vertex-1.0" + + +def test_complete_prompt_error(mock_options): + llm_provider_backends.clear() + mock_generate_content = Mock( + return_value=type( + "obj", + (object,), + {"status_code": 400}, + )() + ) + + with patch( + "sentry.llm.providers.vertex.genai.Client", + return_value=MockGenaiClient(mock_generate_content), + ): + with pytest.raises(VertexRequestFailed): + complete_prompt( + usecase=LLMUseCase.EXAMPLE, + message="message here", + temperature=0.0, + max_output_tokens=1024, + ) From 3372dba7b6d459e8d442867c8ff4b0ca8d5fe533 Mon Sep 17 00:00:00 2001 From: Andrew Liu <159852527+aliu39@users.noreply.github.com> Date: Mon, 5 May 2025 16:35:16 -0700 Subject: [PATCH 3/4] Instantiate new client for each prompt --- src/sentry/llm/providers/vertex.py | 23 ++++++++++++----------- tests/sentry/llm/test_vertex.py | 4 ++-- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/sentry/llm/providers/vertex.py b/src/sentry/llm/providers/vertex.py index e6cc6164441901..39bf579fb2ae80 100644 --- a/src/sentry/llm/providers/vertex.py +++ b/src/sentry/llm/providers/vertex.py @@ -7,7 +7,7 @@ from sentry.llm.exceptions import VertexRequestFailed from sentry.llm.providers.base import LlmModelBase -from sentry.llm.types import ProviderConfig, UseCaseConfig +from sentry.llm.types import UseCaseConfig logger = logging.getLogger(__name__) @@ -21,15 +21,6 @@ class VertexProvider(LlmModelBase): candidate_count = 1 # we only want one candidate returned at the moment top_p = 1 # TODO: make this configurable? - def __init__(self, provider_config: ProviderConfig) -> None: - super().__init__(provider_config) - self.client = genai.Client( - vertexai=True, - project=provider_config["options"]["project"], - location=provider_config["options"]["location"], - http_options=HttpOptions(api_version="v1"), - ) - def _complete_prompt( self, *, @@ -49,7 +40,8 @@ def _complete_prompt( top_p=self.top_p, ) - response = self.client.models.generate_content( + client = self._create_genai_client() + response = client.models.generate_content( model=model, contents=content, config=generate_config, @@ -64,6 +56,15 @@ def _complete_prompt( return response.text + # Separate method to allow mocking + def _create_genai_client(self): + return genai.Client( + vertexai=True, + project=self.provider_config["options"]["project"], + location=self.provider_config["options"]["location"], + http_options=HttpOptions(api_version="v1"), + ) + def _get_access_token(self) -> str: # https://stackoverflow.com/questions/53472429/how-to-get-a-gcp-bearer-token-programmatically-with-python diff --git a/tests/sentry/llm/test_vertex.py b/tests/sentry/llm/test_vertex.py index d87c65c0be84cf..c544d194bbd5cd 100644 --- a/tests/sentry/llm/test_vertex.py +++ b/tests/sentry/llm/test_vertex.py @@ -46,7 +46,7 @@ def test_complete_prompt(mock_options): ) with patch( - "sentry.llm.providers.vertex.genai.Client", + "sentry.llm.providers.vertex.VertexProvider._create_genai_client", return_value=MockGenaiClient(mock_generate_content), ): res = complete_prompt( @@ -73,7 +73,7 @@ def test_complete_prompt_error(mock_options): ) with patch( - "sentry.llm.providers.vertex.genai.Client", + "sentry.llm.providers.vertex.VertexProvider._create_genai_client", return_value=MockGenaiClient(mock_generate_content), ): with pytest.raises(VertexRequestFailed): From f019b80482315fd46b633f69addbffb7c404db58 Mon Sep 17 00:00:00 2001 From: Andrew Liu <159852527+aliu39@users.noreply.github.com> Date: Mon, 5 May 2025 17:13:48 -0700 Subject: [PATCH 4/4] Update option names --- src/sentry/llm/providers/vertex.py | 4 ++-- tests/sentry/llm/test_vertex.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/sentry/llm/providers/vertex.py b/src/sentry/llm/providers/vertex.py index 39bf579fb2ae80..73733c8a57aeea 100644 --- a/src/sentry/llm/providers/vertex.py +++ b/src/sentry/llm/providers/vertex.py @@ -60,8 +60,8 @@ def _complete_prompt( def _create_genai_client(self): return genai.Client( vertexai=True, - project=self.provider_config["options"]["project"], - location=self.provider_config["options"]["location"], + project=self.provider_config["options"]["gcp_project"], + location=self.provider_config["options"]["gcp_location"], http_options=HttpOptions(api_version="v1"), ) diff --git a/tests/sentry/llm/test_vertex.py b/tests/sentry/llm/test_vertex.py index c544d194bbd5cd..133ef06a980e00 100644 --- a/tests/sentry/llm/test_vertex.py +++ b/tests/sentry/llm/test_vertex.py @@ -14,7 +14,7 @@ def mock_options(set_sentry_option): { "vertex": { "models": ["vertex-1.0"], - "options": {"project": "my-gcp-project", "location": "us-central1"}, + "options": {"gcp_project": "my-gcp-project", "gcp_location": "us-central1"}, } }, ),