remove bedrock invoke model unit tests since we are now using upstream extension

Author: yiyuan-he

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py

@@ -1,16 +1,12 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-import json
-import math
 import os
 from importlib.metadata import PackageNotFoundError
-from io import BytesIO
 from typing import Any, Dict
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
 import gevent.monkey
-from botocore.response import StreamingBody
 
 from amazon.opentelemetry.distro.patches._instrumentation_patch import (
     AWS_GEVENT_PATCH_MODULES,
@@ -213,84 +209,9 @@ def _test_patched_botocore_instrumentation(self):
         bedrock_agent_runtime_sucess_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent-runtime")
         self.assertEqual(len(bedrock_agent_runtime_sucess_attributes), 0)
 
-        # BedrockRuntime - Amazon Titan
+        # BedrockRuntime
         self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS)
 
-        self._test_patched_bedrock_runtime_invoke_model(
-            model_id="amazon.titan-embed-text-v1",
-            max_tokens=512,
-            temperature=0.9,
-            top_p=0.75,
-            finish_reason="FINISH",
-            input_tokens=123,
-            output_tokens=456,
-        )
-
-        self._test_patched_bedrock_runtime_invoke_model(
-            model_id="amazon.nova-pro-v1:0",
-            max_tokens=500,
-            temperature=0.9,
-            top_p=0.7,
-            finish_reason="FINISH",
-            input_tokens=123,
-            output_tokens=456,
-        )
-
-        # BedrockRuntime - Anthropic Claude
-        self._test_patched_bedrock_runtime_invoke_model(
-            model_id="anthropic.claude-v2:1",
-            max_tokens=512,
-            temperature=0.5,
-            top_p=0.999,
-            finish_reason="end_turn",
-            input_tokens=23,
-            output_tokens=36,
-        )
-
-        # BedrockRuntime - Meta LLama
-        self._test_patched_bedrock_runtime_invoke_model(
-            model_id="meta.llama2-13b-chat-v1",
-            max_tokens=512,
-            temperature=0.5,
-            top_p=0.9,
-            finish_reason="stop",
-            input_tokens=31,
-            output_tokens=36,
-        )
-
-        # BedrockRuntime - Cohere Command-r
-        cohere_input = "Hello, world"
-        cohere_output = "Goodbye, world"
-
-        self._test_patched_bedrock_runtime_invoke_model(
-            model_id="cohere.command-r-v1:0",
-            max_tokens=512,
-            temperature=0.5,
-            top_p=0.75,
-            finish_reason="COMPLETE",
-            input_tokens=math.ceil(len(cohere_input) / 6),
-            output_tokens=math.ceil(len(cohere_output) / 6),
-            input_prompt=cohere_input,
-            output_prompt=cohere_output,
-        )
-
-        # BedrockRuntime - Mistral
-        msg = "Hello World"
-        mistral_input = f"<s>[INST] {msg} [/INST]"
-        mistral_output = "Goodbye, World"
-
-        self._test_patched_bedrock_runtime_invoke_model(
-            model_id="mistral.mistral-7b-instruct-v0:2",
-            max_tokens=512,
-            temperature=0.5,
-            top_p=0.9,
-            finish_reason="stop",
-            input_tokens=math.ceil(len(mistral_input) / 6),
-            output_tokens=math.ceil(len(mistral_output) / 6),
-            input_prompt=mistral_input,
-            output_prompt=mistral_output,
-        )
-
         # SecretsManager
         self.assertTrue("secretsmanager" in _KNOWN_EXTENSIONS)
         secretsmanager_attributes: Dict[str, str] = _do_extract_secretsmanager_attributes()
@@ -374,130 +295,6 @@ def _test_patched_bedrock_instrumentation(self):
         self.assertEqual(len(bedrock_sucess_attributes), 1)
         self.assertEqual(bedrock_sucess_attributes["aws.bedrock.guardrail.id"], _BEDROCK_GUARDRAIL_ID)
 
-    def _test_patched_bedrock_runtime_invoke_model(self, **args):
-        model_id = args.get("model_id", None)
-        max_tokens = args.get("max_tokens", None)
-        temperature = args.get("temperature", None)
-        top_p = args.get("top_p", None)
-        finish_reason = args.get("finish_reason", None)
-        input_tokens = args.get("input_tokens", None)
-        output_tokens = args.get("output_tokens", None)
-        input_prompt = args.get("input_prompt", None)
-        output_prompt = args.get("output_prompt", None)
-
-        def get_model_response_request():
-            request_body = {}
-            response_body = {}
-
-            if "amazon.titan" in model_id:
-                request_body = {
-                    "textGenerationConfig": {
-                        "maxTokenCount": max_tokens,
-                        "temperature": temperature,
-                        "topP": top_p,
-                    }
-                }
-
-                response_body = {
-                    "inputTextTokenCount": input_tokens,
-                    "results": [
-                        {
-                            "tokenCount": output_tokens,
-                            "outputText": "testing",
-                            "completionReason": finish_reason,
-                        }
-                    ],
-                }
-
-            if "amazon.nova" in model_id:
-                request_body = {
-                    "inferenceConfig": {
-                        "max_new_tokens": max_tokens,
-                        "temperature": temperature,
-                        "topP": top_p,
-                    }
-                }
-
-                response_body = {
-                    "output": {"message": {"content": [{"text": ""}], "role": "assistant"}},
-                    "stopReason": finish_reason,
-                    "usage": {"inputTokens": input_tokens, "outputTokens": output_tokens},
-                }
-
-            if "anthropic.claude" in model_id:
-                request_body = {
-                    "anthropic_version": "bedrock-2023-05-31",
-                    "max_tokens": max_tokens,
-                    "temperature": temperature,
-                    "top_p": top_p,
-                }
-
-                response_body = {
-                    "stop_reason": finish_reason,
-                    "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
-                }
-
-            if "meta.llama" in model_id:
-                request_body = {
-                    "max_gen_len": max_tokens,
-                    "temperature": temperature,
-                    "top_p": top_p,
-                }
-
-                response_body = {
-                    "prompt_token_count": input_tokens,
-                    "generation_token_count": output_tokens,
-                    "stop_reason": finish_reason,
-                }
-
-            if "cohere.command" in model_id:
-                request_body = {
-                    "message": input_prompt,
-                    "max_tokens": max_tokens,
-                    "temperature": temperature,
-                    "p": top_p,
-                }
-
-                response_body = {
-                    "text": output_prompt,
-                    "finish_reason": finish_reason,
-                }
-
-            if "mistral" in model_id:
-                request_body = {
-                    "prompt": input_prompt,
-                    "max_tokens": max_tokens,
-                    "temperature": temperature,
-                    "top_p": top_p,
-                }
-
-                response_body = {"outputs": [{"text": output_prompt, "stop_reason": finish_reason}]}
-
-            json_bytes = json.dumps(response_body).encode("utf-8")
-
-            return json.dumps(request_body), StreamingBody(BytesIO(json_bytes), len(json_bytes))
-
-        request_body, response_body = get_model_response_request()
-
-        bedrock_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock(
-            "bedrock-runtime", operation="InvokeModel", model_id=model_id, request_body=request_body
-        )
-        bedrock_runtime_success_attributes: Dict[str, str] = _do_on_success_bedrock(
-            "bedrock-runtime", operation="InvokeModel", model_id=model_id, streaming_body=response_body
-        )
-
-        bedrock_runtime_attributes.update(bedrock_runtime_success_attributes)
-
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.system"], _GEN_AI_SYSTEM)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.request.model"], model_id)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.request.max_tokens"], max_tokens)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.request.temperature"], temperature)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.request.top_p"], top_p)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.usage.input_tokens"], input_tokens)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.usage.output_tokens"], output_tokens)
-        self.assertEqual(bedrock_runtime_attributes["gen_ai.response.finish_reasons"], [finish_reason])
-
     def _test_patched_bedrock_agent_instrumentation(self):
         """For bedrock-agent service, both extract_attributes and on_success provides attributes,
         the attributes depend on the API being invoked."""