from __future__ import annotations import json from unittest.mock import MagicMock, PropertyMock, patch import pytest from vibe.core.config import ProviderConfig from vibe.core.llm.backend.vertex import ( VertexAnthropicAdapter, VertexCredentials, build_vertex_base_url, build_vertex_endpoint, ) from vibe.core.types import AvailableFunction, AvailableTool, LLMMessage, Role @pytest.fixture def adapter(): adapter = VertexAnthropicAdapter() with patch.object( VertexCredentials, "access_token", new_callable=PropertyMock, return_value="fake-token", ): yield adapter @pytest.fixture def provider(): return ProviderConfig( name="vertex", api_base="", project_id="test-project", region="us-central1", api_style="vertex-anthropic", ) class TestBuildVertexEndpoint: def test_non_streaming(self): endpoint = build_vertex_endpoint( "us-central1", "my-project", "claude-3-5-sonnet" ) assert endpoint == ( "/v1/projects/my-project/locations/us-central1/" "publishers/anthropic/models/claude-3-5-sonnet:rawPredict" ) def test_streaming(self): endpoint = build_vertex_endpoint( "us-central1", "my-project", "claude-3-5-sonnet", streaming=True ) assert endpoint == ( "/v1/projects/my-project/locations/us-central1/" "publishers/anthropic/models/claude-3-5-sonnet:streamRawPredict" ) def test_base_url(self): base = build_vertex_base_url("us-central1") assert base == "https://us-central1-aiplatform.googleapis.com" def test_global_endpoint(self): endpoint = build_vertex_endpoint("global", "my-project", "claude-3-5-sonnet") assert endpoint == ( "/v1/projects/my-project/locations/global/" "publishers/anthropic/models/claude-3-5-sonnet:rawPredict" ) def test_global_base_url(self): base = build_vertex_base_url("global") assert base == "https://aiplatform.googleapis.com" class TestPrepareRequest: def test_basic_request(self, adapter, provider): messages = [LLMMessage(role=Role.user, content="Hello")] req = adapter.prepare_request( model_name="claude-3-5-sonnet", messages=messages, temperature=0.5, tools=None, max_tokens=1024, tool_choice=None, enable_streaming=False, provider=provider, ) payload = json.loads(req.body) assert payload["anthropic_version"] == "vertex-2023-10-16" assert "model" not in payload assert payload["max_tokens"] == 1024 assert "temperature" not in payload assert req.headers["Authorization"] == "Bearer fake-token" assert req.headers["anthropic-beta"] == adapter.BETA_FEATURES assert "rawPredict" in req.endpoint assert "streamRawPredict" not in req.endpoint assert req.base_url == "https://us-central1-aiplatform.googleapis.com" def test_streaming_request(self, adapter, provider): messages = [LLMMessage(role=Role.user, content="Hello")] req = adapter.prepare_request( model_name="claude-3-5-sonnet", messages=messages, temperature=0.5, tools=None, max_tokens=1024, tool_choice=None, enable_streaming=True, provider=provider, ) payload = json.loads(req.body) assert payload.get("stream") is True assert "streamRawPredict" in req.endpoint def test_no_beta_features_for_vertex(self, adapter, provider): """Vertex AI doesn't support the same beta features as direct Anthropic API.""" messages = [LLMMessage(role=Role.user, content="Hello")] req = adapter.prepare_request( model_name="claude-3-5-sonnet", messages=messages, temperature=0.5, tools=None, max_tokens=1024, tool_choice=None, enable_streaming=False, provider=provider, ) # Vertex AI doesn't support prompt-caching or other beta features assert req.headers.get("anthropic-beta", "") == "" def test_with_extended_thinking(self, adapter, provider): messages = [LLMMessage(role=Role.user, content="Hello")] req = adapter.prepare_request( model_name="claude-3-5-sonnet", messages=messages, temperature=0.5, tools=None, max_tokens=1024, tool_choice=None, enable_streaming=False, provider=provider, thinking="medium", ) payload = json.loads(req.body) assert payload["thinking"] == {"type": "adaptive", "display": "summarized"} assert payload["output_config"] == {"effort": "medium"} assert payload["max_tokens"] == 1024 assert "temperature" not in payload def test_with_tools(self, adapter, provider): messages = [LLMMessage(role=Role.user, content="Hello")] tools = [ AvailableTool( function=AvailableFunction( name="test_tool", description="A test tool", parameters={"type": "object", "properties": {}}, ) ) ] req = adapter.prepare_request( model_name="claude-3-5-sonnet", messages=messages, temperature=0.5, tools=tools, max_tokens=1024, tool_choice=None, enable_streaming=False, provider=provider, ) payload = json.loads(req.body) assert len(payload["tools"]) == 1 assert payload["tools"][0]["name"] == "test_tool" def test_missing_project_id(self, adapter): provider = ProviderConfig( name="vertex", api_base="", region="us-central1", api_style="vertex-anthropic", ) with pytest.raises(ValueError, match="project_id"): adapter.prepare_request( model_name="claude-3-5-sonnet", messages=[LLMMessage(role=Role.user, content="Hello")], temperature=0.5, tools=None, max_tokens=1024, tool_choice=None, enable_streaming=False, provider=provider, ) def test_missing_region(self, adapter): provider = ProviderConfig( name="vertex", api_base="", project_id="test-project", api_style="vertex-anthropic", ) with pytest.raises(ValueError, match="region"): adapter.prepare_request( model_name="claude-3-5-sonnet", messages=[LLMMessage(role=Role.user, content="Hello")], temperature=0.5, tools=None, max_tokens=1024, tool_choice=None, enable_streaming=False, provider=provider, ) def test_default_max_tokens(self, adapter, provider): messages = [LLMMessage(role=Role.user, content="Hello")] req = adapter.prepare_request( model_name="claude-3-5-sonnet", messages=messages, temperature=0.5, tools=None, max_tokens=None, tool_choice=None, enable_streaming=False, provider=provider, ) payload = json.loads(req.body) assert payload["max_tokens"] == adapter.DEFAULT_MAX_TOKENS class TestParseFullResponse: def test_simple_text_response(self, adapter, provider): data = { "content": [{"type": "text", "text": "Hello there!"}], "usage": {"input_tokens": 10, "output_tokens": 5}, } chunk = adapter.parse_response(data, provider) assert chunk.message.content == "Hello there!" assert chunk.usage.prompt_tokens == 10 assert chunk.usage.completion_tokens == 5 def test_response_with_tool_calls(self, adapter, provider): data = { "content": [ {"type": "text", "text": "Let me help."}, { "type": "tool_use", "id": "tool_123", "name": "search", "input": {"query": "test"}, }, ], "usage": {"input_tokens": 20, "output_tokens": 15}, } chunk = adapter.parse_response(data, provider) assert chunk.message.content == "Let me help." assert len(chunk.message.tool_calls) == 1 assert chunk.message.tool_calls[0].id == "tool_123" assert chunk.message.tool_calls[0].function.name == "search" assert json.loads(chunk.message.tool_calls[0].function.arguments) == { "query": "test" } def test_response_with_thinking(self, adapter, provider): data = { "content": [ { "type": "thinking", "thinking": "Let me think...", "signature": "sig123", }, {"type": "text", "text": "Here's my answer."}, ], "usage": {"input_tokens": 30, "output_tokens": 20}, } chunk = adapter.parse_response(data, provider) assert chunk.message.content == "Here's my answer." assert chunk.message.reasoning_content == "Let me think..." assert chunk.message.reasoning_signature == "sig123" def test_response_with_cache_tokens(self, adapter, provider): data = { "content": [{"type": "text", "text": "Hello"}], "usage": { "input_tokens": 10, "cache_creation_input_tokens": 5, "cache_read_input_tokens": 3, "output_tokens": 7, }, } chunk = adapter.parse_response(data, provider) assert chunk.usage.prompt_tokens == 18 assert chunk.usage.completion_tokens == 7 def test_response_with_redacted_thinking(self, adapter, provider): data = { "content": [ {"type": "redacted_thinking", "data": "redacted_data_here"}, {"type": "text", "text": "Answer."}, ], "usage": {"input_tokens": 10, "output_tokens": 5}, } chunk = adapter.parse_response(data, provider) assert chunk.message.content == "Answer." assert chunk.message.reasoning_content is None def test_response_empty_usage(self, adapter, provider): data = {"content": [{"type": "text", "text": "Hello"}], "usage": {}} chunk = adapter.parse_response(data, provider) assert chunk.usage.prompt_tokens == 0 assert chunk.usage.completion_tokens == 0 class TestStreamingEvents: def test_message_start(self, adapter, provider): data = { "type": "message_start", "message": { "usage": { "input_tokens": 100, "cache_creation_input_tokens": 20, "cache_read_input_tokens": 10, } }, } chunk = adapter.parse_response(data, provider) assert chunk.usage is not None assert chunk.usage.prompt_tokens == 130 assert chunk.usage.completion_tokens == 0 def test_message_start_without_usage(self, adapter, provider): data = {"type": "message_start", "message": {}} chunk = adapter.parse_response(data, provider) assert chunk.message.role == Role.assistant def test_content_block_start_tool_use(self, adapter, provider): data = { "type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "tool_abc", "name": "search"}, } chunk = adapter.parse_response(data, provider) assert chunk.message.tool_calls is not None assert len(chunk.message.tool_calls) == 1 assert chunk.message.tool_calls[0].id == "tool_abc" assert chunk.message.tool_calls[0].function.name == "search" assert chunk.message.tool_calls[0].index == 0 def test_content_block_start_thinking(self, adapter, provider): data = { "type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}, } chunk = adapter.parse_response(data, provider) assert chunk.message.reasoning_content is not None def test_content_block_start_redacted_thinking(self, adapter, provider): data = { "type": "content_block_start", "index": 0, "content_block": {"type": "redacted_thinking", "data": "abc"}, } chunk = adapter.parse_response(data, provider) assert chunk.message.content is None assert chunk.message.reasoning_content is None def test_content_block_delta_text(self, adapter, provider): data = { "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}, } chunk = adapter.parse_response(data, provider) assert chunk.message.content == "Hello" def test_content_block_delta_thinking(self, adapter, provider): data = { "type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "I think..."}, } chunk = adapter.parse_response(data, provider) assert chunk.message.reasoning_content == "I think..." def test_content_block_delta_input_json(self, adapter, provider): data = { "type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '{"key":'}, } chunk = adapter.parse_response(data, provider) assert chunk.message.tool_calls is not None assert chunk.message.tool_calls[0].function.arguments == '{"key":' def test_content_block_stop(self, adapter, provider): data = {"type": "content_block_stop", "index": 0} chunk = adapter.parse_response(data, provider) assert chunk.message.content is None assert chunk.message.reasoning_content is None def test_message_delta_with_usage(self, adapter, provider): data = {"type": "message_delta", "usage": {"output_tokens": 42}} chunk = adapter.parse_response(data, provider) assert chunk.usage is not None assert chunk.usage.completion_tokens == 42 assert chunk.usage.prompt_tokens == 0 def test_message_delta_without_usage(self, adapter, provider): data = {"type": "message_delta", "usage": {}} chunk = adapter.parse_response(data, provider) assert chunk.message.role == Role.assistant def test_unknown_event_returns_empty_chunk(self, adapter, provider): data = {"type": "ping"} chunk = adapter.parse_response(data, provider) assert chunk.message.role == Role.assistant assert chunk.message.content is None def test_signature_delta(self, adapter, provider): data = { "type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "sig_abc"}, } chunk = adapter.parse_response(data, provider) assert chunk.message.reasoning_signature == "sig_abc" def test_message_start_resets_state(self, adapter, provider): adapter._current_index = 5 data = {"type": "message_start", "message": {"usage": {"input_tokens": 10}}} adapter.parse_response(data, provider) assert adapter._current_index == 0 def test_full_streaming_sequence(self, adapter, provider): chunks = [] # message_start chunks.append( adapter.parse_response( {"type": "message_start", "message": {"usage": {"input_tokens": 50}}}, provider, ) ) assert chunks[-1].usage.prompt_tokens == 50 # thinking block adapter.parse_response( { "type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}, }, provider, ) chunks.append( adapter.parse_response( { "type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "Analyzing..."}, }, provider, ) ) assert chunks[-1].message.reasoning_content == "Analyzing..." adapter.parse_response({"type": "content_block_stop", "index": 0}, provider) # text block chunks.append( adapter.parse_response( { "type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Here's the result."}, }, provider, ) ) assert chunks[-1].message.content == "Here's the result." # tool use chunks.append( adapter.parse_response( { "type": "content_block_start", "index": 2, "content_block": { "type": "tool_use", "id": "tool_1", "name": "search", }, }, provider, ) ) assert chunks[-1].message.tool_calls[0].function.name == "search" # message_delta with final usage chunks.append( adapter.parse_response( {"type": "message_delta", "usage": {"output_tokens": 100}}, provider ) ) assert chunks[-1].usage.completion_tokens == 100 class TestHelperMethods: def test_has_thinking_content_true(self, adapter): messages = [ {"role": "user", "content": "Hello"}, { "role": "assistant", "content": [ {"type": "thinking", "thinking": "Let me think..."}, {"type": "text", "text": "Answer"}, ], }, ] assert adapter._has_thinking_content(messages) is True def test_has_thinking_content_false(self, adapter): messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Just text"}, ] assert adapter._has_thinking_content(messages) is False def test_has_thinking_content_empty(self, adapter): assert adapter._has_thinking_content([]) is False def test_has_thinking_content_non_list_content(self, adapter): messages = [ {"role": "assistant", "content": [{"type": "text", "text": "no thinking"}]} ] assert adapter._has_thinking_content(messages) is False def test_add_cache_control_to_last_user_message(self, adapter): messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] adapter._add_cache_control_to_last_user_message(messages) assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral"} def test_add_cache_control_skips_non_user(self, adapter): messages = [ {"role": "assistant", "content": [{"type": "text", "text": "Hello"}]} ] adapter._add_cache_control_to_last_user_message(messages) assert "cache_control" not in messages[0]["content"][0] def test_add_cache_control_skips_string_content(self, adapter): messages = [{"role": "user", "content": "Hello"}] adapter._add_cache_control_to_last_user_message(messages) assert messages[0]["content"] == "Hello" def test_add_cache_control_tool_result(self, adapter): messages = [ { "role": "user", "content": [ {"type": "tool_result", "tool_use_id": "123", "content": "result"} ], } ] adapter._add_cache_control_to_last_user_message(messages) assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral"} def test_add_cache_control_empty_messages(self, adapter): messages: list[dict] = [] adapter._add_cache_control_to_last_user_message(messages) assert messages == [] class TestVertexCredentials: def _make_creds( self, *, valid: bool = True, token: str | None = "tok" ) -> MagicMock: creds = MagicMock() creds.valid = valid creds.token = token return creds @patch("vibe.core.llm.backend.vertex.google.auth.default") def test_initializes_credentials_on_first_access(self, mock_default: MagicMock): creds = self._make_creds() mock_default.return_value = (creds, "project") vc = VertexCredentials() token = vc.access_token assert token == "tok" mock_default.assert_called_once() @patch("vibe.core.llm.backend.vertex.google.auth.default") def test_caches_credentials_across_calls(self, mock_default: MagicMock): creds = self._make_creds() mock_default.return_value = (creds, "project") vc = VertexCredentials() _ = vc.access_token _ = vc.access_token _ = vc.access_token mock_default.assert_called_once() @patch("vibe.core.llm.backend.vertex.google.auth.default") def test_refreshes_when_token_invalid(self, mock_default: MagicMock): creds = self._make_creds(valid=False) mock_default.return_value = (creds, "project") vc = VertexCredentials() _ = vc.access_token creds.refresh.assert_called_once() @patch("vibe.core.llm.backend.vertex.google.auth.default") def test_skips_refresh_when_token_valid(self, mock_default: MagicMock): creds = self._make_creds(valid=True) mock_default.return_value = (creds, "project") vc = VertexCredentials() _ = vc.access_token creds.refresh.assert_not_called() @patch("vibe.core.llm.backend.vertex.google.auth.default") def test_raises_when_token_is_none(self, mock_default: MagicMock): creds = self._make_creds(valid=True, token=None) mock_default.return_value = (creds, "project") vc = VertexCredentials() with pytest.raises(RuntimeError, match="did not produce a token"): _ = vc.access_token