diff --git a/pyproject.toml b/pyproject.toml index b086b68a..2629f827 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "ujson>=5.10.0", "boto3==1.38.0", "setuptools", + "Deprecated==1.2.18", ] requires-python = ">=3.9" @@ -47,7 +48,7 @@ dev = [ "qdrant-client", "graphlit-client", "python-dotenv", - "pinecone", + "pinecone>=3.1.0,<=6.0.2", "langchain", "langchain-community", "langchain-openai", diff --git a/src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py b/src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py index bdbe8fdb..70e0a156 100644 --- a/src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py +++ b/src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py @@ -5,6 +5,10 @@ "METHOD": "aws_bedrock.invoke_model", "ENDPOINT": "/invoke-model", }, + "INVOKE_MODEL_WITH_RESPONSE_STREAM": { + "METHOD": "aws_bedrock.invoke_model_with_response_stream", + "ENDPOINT": "/invoke-model-with-response-stream", + }, "CONVERSE": { "METHOD": AWSBedrockMethods.CONVERSE.value, "ENDPOINT": "/converse", diff --git a/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py b/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py index 2f354d86..bb02259d 100644 --- a/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py +++ b/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py @@ -15,8 +15,10 @@ """ import json +import io from wrapt import ObjectProxy +from itertools import tee from .stream_body_wrapper import BufferedStreamBody from functools import wraps from langtrace.trace_attributes import ( @@ -43,6 +45,7 @@ set_span_attributes, set_usage_attributes, ) +from langtrace_python_sdk.utils import set_event_prompt def converse_stream(original_method, version, tracer): @@ -104,7 +107,7 @@ def traced_method(wrapped, instance, args, kwargs): def patch_converse_stream(original_method, tracer, version): def traced_method(*args, **kwargs): modelId = kwargs.get("modelId") - (vendor, _) = modelId.split(".") + vendor, _ = parse_vendor_and_model_name_from_model_id(modelId) input_content = [ { "role": message.get("role", "user"), @@ -128,7 +131,9 @@ def traced_method(*args, **kwargs): response = original_method(*args, **kwargs) if span.is_recording(): - set_span_streaming_response(span, response) + stream1, stream2 = tee(response["stream"]) + set_span_streaming_response(span, stream1) + response["stream"] = stream2 return response return traced_method @@ -137,7 +142,7 @@ def traced_method(*args, **kwargs): def patch_converse(original_method, tracer, version): def traced_method(*args, **kwargs): modelId = kwargs.get("modelId") - (vendor, _) = modelId.split(".") + vendor, _ = parse_vendor_and_model_name_from_model_id(modelId) input_content = [ { "role": message.get("role", "user"), @@ -167,12 +172,29 @@ def traced_method(*args, **kwargs): return traced_method +def parse_vendor_and_model_name_from_model_id(model_id): + if model_id.startswith("arn:aws:bedrock:"): + # This needs to be in one of the following forms: + # arn:aws:bedrock:region:account-id:foundation-model/vendor.model-name + # arn:aws:bedrock:region:account-id:custom-model/vendor.model-name/model-id + parts = model_id.split("/") + identifiers = parts[1].split(".") + return identifiers[0], identifiers[1] + parts = model_id.split(".") + if len(parts) == 1: + return parts[0], parts[0] + else: + return parts[-2], parts[-1] + + def patch_invoke_model(original_method, tracer, version): def traced_method(*args, **kwargs): modelId = kwargs.get("modelId") - (vendor, _) = modelId.split(".") + vendor, _ = parse_vendor_and_model_name_from_model_id(modelId) span_attributes = { **get_langtrace_attributes(version, vendor, vendor_type="framework"), + SpanAttributes.LLM_PATH: APIS["INVOKE_MODEL"]["ENDPOINT"], + SpanAttributes.LLM_IS_STREAMING: False, **get_extra_attributes(), } with tracer.start_as_current_span( @@ -193,9 +215,11 @@ def patch_invoke_model_with_response_stream(original_method, tracer, version): @wraps(original_method) def traced_method(*args, **kwargs): modelId = kwargs.get("modelId") - (vendor, _) = modelId.split(".") + vendor, _ = parse_vendor_and_model_name_from_model_id(modelId) span_attributes = { **get_langtrace_attributes(version, vendor, vendor_type="framework"), + SpanAttributes.LLM_PATH: APIS["INVOKE_MODEL_WITH_RESPONSE_STREAM"]["ENDPOINT"], + SpanAttributes.LLM_IS_STREAMING: True, **get_extra_attributes(), } span = tracer.start_span( @@ -217,7 +241,7 @@ def handle_streaming_call(span, kwargs, response): def stream_finished(response_body): request_body = json.loads(kwargs.get("body")) - (vendor, model) = kwargs.get("modelId").split(".") + vendor, model = parse_vendor_and_model_name_from_model_id(kwargs.get("modelId")) set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model) set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model) @@ -241,18 +265,22 @@ def stream_finished(response_body): def handle_call(span, kwargs, response): modelId = kwargs.get("modelId") - (vendor, model_name) = modelId.split(".") + vendor, model_name = parse_vendor_and_model_name_from_model_id(modelId) + read_response_body = response.get("body").read() + request_body = json.loads(kwargs.get("body")) + response_body = json.loads(read_response_body) response["body"] = BufferedStreamBody( - response["body"]._raw_stream, response["body"]._content_length + io.BytesIO(read_response_body), len(read_response_body) ) - request_body = json.loads(kwargs.get("body")) - response_body = json.loads(response.get("body").read()) set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, modelId) set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, modelId) if vendor == "amazon": - set_amazon_attributes(span, request_body, response_body) + if model_name.startswith("titan-embed-text"): + set_amazon_embedding_attributes(span, request_body, response_body) + else: + set_amazon_attributes(span, request_body, response_body) if vendor == "anthropic": if "prompt" in request_body: @@ -356,6 +384,27 @@ def set_amazon_attributes(span, request_body, response_body): set_event_completion(span, completions) +def set_amazon_embedding_attributes(span, request_body, response_body): + input_text = request_body.get("inputText") + set_event_prompt(span, input_text) + + embeddings = response_body.get("embedding", []) + input_tokens = response_body.get("inputTextTokenCount") + set_usage_attributes( + span, + { + "input_tokens": input_tokens, + "output": len(embeddings), + }, + ) + set_span_attribute( + span, SpanAttributes.LLM_REQUEST_MODEL, request_body.get("modelId") + ) + set_span_attribute( + span, SpanAttributes.LLM_RESPONSE_MODEL, request_body.get("modelId") + ) + + def set_anthropic_completions_attributes(span, request_body, response_body): set_span_attribute( span, @@ -442,10 +491,10 @@ def _set_response_attributes(span, kwargs, result): ) -def set_span_streaming_response(span, response): +def set_span_streaming_response(span, response_stream): streaming_response = "" role = None - for event in response["stream"]: + for event in response_stream: if "messageStart" in event: role = event["messageStart"]["role"] elif "contentBlockDelta" in event: @@ -475,13 +524,15 @@ def __init__( stream_done_callback=None, ): super().__init__(response) - self._stream_done_callback = stream_done_callback self._accumulating_body = {"generation": ""} + self.last_chunk = None def __iter__(self): for event in self.__wrapped__: + # Process the event self._process_event(event) + # Yield the original event immediately yield event def _process_event(self, event): @@ -496,7 +547,11 @@ def _process_event(self, event): self._stream_done_callback(decoded_chunk) return if "generation" in decoded_chunk: - self._accumulating_body["generation"] += decoded_chunk.get("generation") + generation = decoded_chunk.get("generation") + if self.last_chunk == generation: + return + self.last_chunk = generation + self._accumulating_body["generation"] += generation if type == "message_start": self._accumulating_body = decoded_chunk.get("message") @@ -505,9 +560,11 @@ def _process_event(self, event): decoded_chunk.get("content_block") ) elif type == "content_block_delta": - self._accumulating_body["content"][-1]["text"] += decoded_chunk.get( - "delta" - ).get("text") + text = decoded_chunk.get("delta").get("text") + if self.last_chunk == text: + return + self.last_chunk = text + self._accumulating_body["content"][-1]["text"] += text elif self.has_finished(type, decoded_chunk): self._accumulating_body["invocation_metrics"] = decoded_chunk.get( diff --git a/src/langtrace_python_sdk/instrumentation/pinecone/instrumentation.py b/src/langtrace_python_sdk/instrumentation/pinecone/instrumentation.py index 7608e901..8ad3f27b 100644 --- a/src/langtrace_python_sdk/instrumentation/pinecone/instrumentation.py +++ b/src/langtrace_python_sdk/instrumentation/pinecone/instrumentation.py @@ -33,7 +33,7 @@ class PineconeInstrumentation(BaseInstrumentor): The PineconeInstrumentation class represents the Pinecone instrumentation""" def instrumentation_dependencies(self) -> Collection[str]: - return ["pinecone >= 3.1.0"] + return ["pinecone >= 3.1.0", "pinecone <= 6.0.2"] def _instrument(self, **kwargs): tracer_provider = kwargs.get("tracer_provider") diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py index 91d803c2..32ffaf06 100644 --- a/src/langtrace_python_sdk/version.py +++ b/src/langtrace_python_sdk/version.py @@ -1 +1 @@ -__version__ = "3.8.18" +__version__ = "3.8.21" diff --git a/src/tests/aws_bedrock/cassettes/test_chat_completion.yaml b/src/tests/aws_bedrock/cassettes/test_chat_completion.yaml new file mode 100644 index 00000000..3ec871bb --- /dev/null +++ b/src/tests/aws_bedrock/cassettes/test_chat_completion.yaml @@ -0,0 +1,45 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": "Say this is a test three times"}], + "anthropic_version": "bedrock-2023-05-31", "max_tokens": 100}' + headers: + Accept: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + Content-Length: + - '139' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS4zOC4xOCBtZC9Cb3RvY29yZSMxLjM4LjE4IHVhLzIuMSBvcy9tYWNvcyMyNC40LjAg + bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjEzLjEgbWQvcHlpbXBsI0NQeXRob24gbS9aLGIg + Y2ZnL3JldHJ5LW1vZGUjc3RhbmRhcmQgQm90b2NvcmUvMS4zOC4xOA== + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-3-7-sonnet-20250219-v1%3A0/invoke + response: + body: + string: '{"id":"msg_bdrk_01NJB1bDTLkFh6pgfoAD5hkb","type":"message","role":"assistant","model":"claude-3-7-sonnet-20250219","content":[{"type":"text","text":"This + is a test.\nThis is a test.\nThis is a test."}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":14,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":20}}' + headers: + Connection: + - keep-alive + Content-Length: + - '355' + Content-Type: + - application/json + Date: + - Mon, 19 May 2025 16:42:05 GMT + X-Amzn-Bedrock-Input-Token-Count: + - '14' + X-Amzn-Bedrock-Invocation-Latency: + - '926' + X-Amzn-Bedrock-Output-Token-Count: + - '20' + x-amzn-RequestId: + - c0a92363-ec28-4a8b-9c09-571131d946b0 + status: + code: 200 + message: OK +version: 1 diff --git a/src/tests/aws_bedrock/cassettes/test_generate_embedding.yaml b/src/tests/aws_bedrock/cassettes/test_generate_embedding.yaml new file mode 100644 index 00000000..1deaaba1 --- /dev/null +++ b/src/tests/aws_bedrock/cassettes/test_generate_embedding.yaml @@ -0,0 +1,41 @@ +interactions: +- request: + body: '{"inputText": "Say this is a test three times"}' + headers: + Accept: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + Content-Length: + - '47' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS4zOC4xOCBtZC9Cb3RvY29yZSMxLjM4LjE4IHVhLzIuMSBvcy9tYWNvcyMyNC40LjAg + bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjEzLjEgbWQvcHlpbXBsI0NQeXRob24gbS9aLGIg + Y2ZnL3JldHJ5LW1vZGUjc3RhbmRhcmQgQm90b2NvcmUvMS4zOC4xOA== + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v1/invoke + response: + body: + string: '{"embedding":[0.53515625,-0.64453125,-0.51953125,-0.1787109375,0.5234375,0.498046875,0.019287109375,-2.2172927856445312E-5,-0.1796875,1.220703125E-4,-0.54296875,-0.078125,0.07177734375,0.0966796875,-0.380859375,-0.0030517578125,-0.197265625,-0.283203125,-0.107421875,-0.03759765625,-0.765625,-0.0947265625,0.06298828125,0.73828125,0.263671875,0.00750732421875,0.234375,-0.423828125,0.023681640625,-0.076171875,-0.3984375,0.486328125,0.396484375,-1.140625,0.451171875,-0.1669921875,-0.0888671875,0.1474609375,0.5859375,-0.10888671875,0.076171875,0.3359375,0.34375,0.52734375,-0.1416015625,0.26171875,0.099609375,-0.375,-0.2373046875,0.490234375,-0.19140625,0.318359375,0.1396484375,-0.35546875,0.3671875,0.2578125,0.443359375,0.361328125,0.48046875,-0.361328125,0.0169677734375,-0.11767578125,-0.294921875,-0.72265625,-0.0771484375,0.1640625,-0.08154296875,-0.333984375,-0.61328125,0.00750732421875,-0.1220703125,-0.546875,-0.291015625,0.138671875,-0.01806640625,0.37109375,-0.00848388671875,0.08935546875,-0.06005859375,0.54296875,-0.216796875,-0.42578125,-0.0299072265625,0.412109375,-0.5390625,0.09765625,-0.55859375,0.5703125,1.354217529296875E-4,-0.50390625,-0.2021484375,0.88671875,0.2099609375,-0.044677734375,0.1669921875,0.43359375,-0.3984375,0.875,0.6015625,-0.15625,-0.228515625,0.1591796875,0.134765625,0.7734375,-0.138671875,0.29296875,-0.232421875,0.365234375,0.11572265625,0.50390625,-0.004486083984375,-0.279296875,0.6171875,-1.03125,-0.169921875,-0.4609375,0.08984375,-0.01336669921875,-0.1962890625,-0.3828125,0.53125,0.0830078125,0.353515625,-0.1650390625,-0.00799560546875,-0.6015625,0.142578125,0.2158203125,-0.046875,0.423828125,0.2451171875,0.33203125,-0.494140625,-0.4140625,0.1650390625,0.1533203125,-0.53515625,0.023193359375,0.275390625,0.09716796875,-0.18359375,-0.083984375,0.71875,0.2138671875,0.21875,0.026611328125,-0.1318359375,0.2109375,-0.93359375,-0.005706787109375,0.263671875,-0.419921875,-0.34375,-0.7578125,0.1396484375,-0.40234375,0.103515625,0.435546875,-0.625,0.150390625,0.3203125,0.6328125,-0.82421875,-0.59375,0.1591796875,-0.46875,0.08349609375,0.146484375,-0.1845703125,0.171875,0.421875,-0.34765625,0.37109375,0.140625,-0.04638671875,-0.03466796875,0.490234375,0.061767578125,-0.2255859375,-0.236328125,-0.234375,0.220703125,-0.016845703125,0.002716064453125,1.2578125,0.16015625,0.12109375,0.46484375,-0.19921875,-0.00909423828125,0.287109375,0.578125,-0.47265625,-0.353515625,-0.365234375,0.41796875,-0.248046875,-0.357421875,-0.56640625,-0.294921875,-0.353515625,0.1982421875,0.45703125,-0.06689453125,0.162109375,0.388671875,0.29296875,-0.051025390625,0.0732421875,0.30078125,0.11181640625,0.6484375,-0.349609375,1.0703125,0.5625,0.408203125,0.359375,0.220703125,-0.06982421875,-0.78125,-0.1962890625,-0.369140625,-0.016357421875,-0.287109375,-0.25390625,0.3671875,-0.004791259765625,-0.2275390625,0.40625,0.734375,-0.392578125,0.06298828125,-0.75390625,0.49609375,-0.01422119140625,-0.4296875,-0.609375,0.31640625,0.41015625,0.275390625,0.146484375,-0.39453125,-0.515625,-0.6953125,-0.267578125,-0.5703125,0.76171875,-0.271484375,0.28125,-0.140625,0.0810546875,0.71484375,-0.34375,-0.2001953125,-0.365234375,-0.0361328125,-0.62890625,0.39453125,0.60546875,-0.201171875,0.2216796875,-0.416015625,0.1865234375,0.330078125,0.8671875,-0.337890625,0.421875,-0.30859375,-0.030029296875,-0.28125,1.354217529296875E-4,-0.11669921875,0.3671875,0.16015625,-0.1728515625,-0.06298828125,-0.251953125,0.0322265625,0.66015625,-0.53125,0.205078125,0.375,0.54296875,0.3203125,-0.212890625,-0.478515625,-0.263671875,0.20703125,0.201171875,-0.1357421875,-0.59375,-0.1787109375,0.54296875,0.20703125,0.07275390625,0.203125,-0.1943359375,0.443359375,0.318359375,0.609375,-0.181640625,-0.1728515625,0.375,-1.2421875,0.1953125,-0.244140625,0.103515625,-0.1591796875,0.1455078125,-0.03466796875,-0.154296875,0.20703125,-0.341796875,0.283203125,-0.1611328125,0.328125,-0.2255859375,0.486328125,-0.83203125,-0.185546875,0.203125,-0.357421875,-0.357421875,-0.4453125,0.271484375,-0.00823974609375,0.298828125,-0.1826171875,0.2578125,-0.515625,-0.71875,-1.03125,-0.068359375,0.1875,1.0390625,0.07373046875,0.1572265625,0.45703125,-0.337890625,-0.296875,-0.0078125,-0.19921875,-0.0634765625,-0.58984375,-0.4453125,-0.08935546875,0.002227783203125,0.1416015625,-0.27734375,0.1396484375,0.10498046875,-0.185546875,0.2890625,-0.66015625,0.400390625,-0.0908203125,0.703125,0.193359375,-0.376953125,-0.06591796875,-0.09912109375,-0.0888671875,-0.0038909912109375,0.1103515625,-0.51953125,-0.279296875,-0.53125,-0.361328125,0.0294189453125,0.7265625,-0.09521484375,0.62890625,-0.224609375,-0.140625,0.1259765625,0.322265625,0.6015625,-0.396484375,-0.451171875,-0.384765625,-0.134765625,-0.54296875,-0.11865234375,0.275390625,-0.26171875,0.796875,-0.302734375,-0.8828125,0.78515625,-0.333984375,0.1611328125,0.37109375,0.234375,0.042724609375,0.53515625,0.2333984375,-0.061279296875,-0.10107421875,0.16796875,0.18359375,-0.005950927734375,0.435546875,-0.0927734375,-0.275390625,0.05029296875,1.2421875,-0.169921875,-0.05322265625,-0.3125,-0.2216796875,-0.21484375,-0.1806640625,-0.58984375,0.55859375,-0.390625,0.140625,0.32421875,-0.1484375,0.06884765625,-0.2138671875,0.2578125,0.255859375,0.26953125,0.107421875,0.1474609375,0.578125,-0.10791015625,0.294921875,-0.1962890625,0.404296875,0.10693359375,1.1484375,-0.65234375,0.365234375,0.39453125,0.5625,0.169921875,0.1875,0.58203125,-0.103515625,-0.72265625,-0.71484375,0.310546875,0.1572265625,-0.11083984375,-0.55078125,-0.1044921875,-0.28125,-1.0546875,-0.357421875,0.1728515625,0.40234375,0.150390625,-0.376953125,-0.33203125,0.67578125,-0.267578125,0.020751953125,0.38671875,0.2353515625,0.67578125,0.302734375,-0.185546875,0.265625,0.076171875,0.2314453125,1.4765625,-0.287109375,0.1015625,-0.42578125,-0.1728515625,-0.007049560546875,-0.08935546875,0.65234375,-0.26171875,0.197265625,-0.427734375,-0.95703125,-0.34765625,-0.5703125,0.1259765625,-0.11279296875,-0.453125,0.30859375,-0.10009765625,0.1337890625,0.07666015625,-0.19921875,0.07421875,0.56640625,0.71484375,0.640625,-0.482421875,-0.06640625,-0.2353515625,0.224609375,-4.00543212890625E-4,-1.0625,0.470703125,0.26953125,0.2353515625,0.2021484375,-0.470703125,0.35546875,0.01611328125,0.462890625,-0.265625,0.03857421875,-0.6796875,-0.11962890625,0.0272216796875,-0.30859375,0.1279296875,-0.2421875,0.330078125,0.5546875,-0.023193359375,-0.80859375,-0.2490234375,0.267578125,0.2216796875,0.2138671875,0.125,0.69921875,-0.10986328125,0.91015625,-0.28125,-0.005859375,-0.3828125,0.1923828125,-0.251953125,-0.625,-0.298828125,-0.224609375,-0.333984375,0.62890625,-0.109375,0.2490234375,0.11865234375,-0.022216796875,0.8203125,-0.291015625,0.0015869140625,0.5390625,-0.2353515625,0.310546875,0.07275390625,0.109375,0.04248046875,1.1328125,1.3203125,0.0498046875,-0.1826171875,0.306640625,0.2001953125,0.48046875,0.1943359375,-0.039306640625,0.26171875,-0.455078125,-0.0693359375,-0.2314453125,0.00958251953125,-0.283203125,0.0380859375,0.232421875,-0.06640625,-0.34375,-0.328125,0.09375,0.70703125,-0.1376953125,0.1416015625,-0.35546875,-0.00421142578125,1.328125,-0.33203125,-0.703125,-0.34765625,-0.59765625,-0.373046875,0.00714111328125,0.80078125,-0.06298828125,0.0361328125,-0.032958984375,-0.0771484375,-0.640625,0.1689453125,-0.37109375,-0.68359375,-0.1337890625,-0.255859375,0.439453125,-0.8984375,-0.11328125,0.31640625,-0.466796875,-0.345703125,-0.349609375,0.828125,-0.40234375,0.8359375,-0.00433349609375,-0.35546875,0.5546875,-0.54296875,0.61328125,-0.14453125,0.033935546875,0.53515625,0.205078125,-0.00799560546875,-0.26953125,-0.49609375,-0.2060546875,0.6015625,-0.01397705078125,-0.57421875,-0.26171875,0.06298828125,0.609375,-0.9375,-0.171875,-0.0859375,0.0419921875,0.6328125,-0.6796875,0.10791015625,0.609375,-0.2255859375,0.7265625,-0.193359375,-0.01708984375,-0.353515625,-1.392364501953125E-4,0.578125,-1.046875,-0.028564453125,-0.1650390625,0.419921875,0.01153564453125,0.047607421875,-1.1328125,0.46875,0.201171875,0.044677734375,-0.06640625,-0.0654296875,-0.142578125,-0.67578125,0.11376953125,-0.00360107421875,0.052490234375,-0.038330078125,-0.22265625,0.35546875,0.07080078125,-0.24609375,-0.1787109375,-0.1630859375,-0.70703125,-0.2021484375,0.36328125,-0.1630859375,0.099609375,-0.453125,0.0966796875,-0.7421875,0.435546875,0.4609375,-0.265625,0.298828125,0.05859375,-0.4765625,-0.85546875,-0.0703125,0.703125,0.010986328125,-0.294921875,0.34375,-0.091796875,-0.5,-0.8671875,-0.296875,-0.5859375,0.1328125,0.0947265625,0.13671875,0.07666015625,-0.48046875,0.271484375,0.314453125,0.478515625,0.13671875,0.328125,0.21875,0.1123046875,0.2490234375,-0.369140625,-0.050537109375,-0.126953125,-0.353515625,0.1669921875,-0.345703125,0.30859375,0.5546875,0.244140625,-0.0966796875,0.349609375,0.3828125,0.1796875,0.37109375,0.15625,-0.287109375,0.0908203125,-0.361328125,0.0908203125,-0.11181640625,-0.1552734375,0.08056640625,0.04052734375,0.59375,-0.08837890625,-0.1435546875,-0.4453125,0.77734375,0.515625,0.546875,-0.40625,-0.021728515625,-0.5078125,1.0234375,-0.072265625,-0.216796875,0.11767578125,-0.2890625,-0.390625,-0.019775390625,0.70703125,0.3515625,-0.484375,0.54296875,0.01446533203125,-0.11669921875,-0.279296875,-0.25390625,-0.6796875,0.80859375,-0.16015625,-0.1884765625,0.431640625,-0.09912109375,-0.1943359375,-0.0849609375,-0.0703125,0.345703125,0.349609375,-0.08349609375,-0.3515625,0.1650390625,0.5546875,-0.498046875,0.302734375,-0.140625,-0.12060546875,0.5390625,0.3046875,-0.02734375,0.08837890625,0.25,-0.12353515625,0.375,0.11474609375,0.228515625,-0.003631591796875,1.28125,0.66796875,-0.466796875,0.453125,0.71484375,0.205078125,0.01318359375,0.462890625,-0.99609375,-0.2353515625,0.76171875,-0.1533203125,0.2578125,0.55859375,-1.0078125,-0.041015625,0.16015625,0.1279296875,0.2265625,0.49609375,-0.07080078125,-0.498046875,0.337890625,-0.59375,-0.267578125,-0.2294921875,0.14453125,0.453125,-0.96875,0.010986328125,0.37890625,1.046875,-0.032470703125,0.10791015625,-0.1650390625,-0.1416015625,0.150390625,-0.3203125,-0.11328125,0.251953125,-0.265625,-0.79296875,0.427734375,-0.177734375,-0.55859375,0.55859375,-0.40234375,-0.400390625,-0.091796875,0.00244140625,-0.09716796875,0.0157470703125,0.41796875,-0.306640625,-0.56640625,0.55078125,0.12451171875,-0.1220703125,-0.451171875,-0.318359375,0.0302734375,-0.31640625,0.150390625,-0.54296875,-0.2734375,-0.28515625,0.80859375,-0.103515625,-0.0341796875,-0.47265625,-0.0966796875,0.03369140625,0.390625,0.2392578125,1.5,-0.0240478515625,-0.7734375,-0.60546875,0.1298828125,-0.1728515625,-0.0224609375,-0.3046875,-0.1298828125,-0.40234375,0.119140625,-0.55078125,0.12353515625,0.1650390625,0.25390625,-0.07666015625,0.294921875,-0.2265625,-0.2197265625,0.10693359375,0.173828125,-0.392578125,0.236328125,-0.177734375,0.62109375,0.1845703125,0.80078125,0.2314453125,-0.138671875,-0.22265625,0.0556640625,-0.427734375,-1.0078125,-0.201171875,0.130859375,-0.052978515625,0.107421875,0.390625,-0.3203125,-0.1357421875,0.1064453125,-0.609375,-0.396484375,0.26171875,-0.28125,0.189453125,-0.388671875,0.4140625,-0.58984375,-0.002960205078125,-0.369140625,0.7265625,0.1181640625,-0.439453125,-0.5078125,-0.142578125,0.1552734375,0.1904296875,0.28515625,-0.1328125,-0.177734375,-0.072265625,0.9609375,0.4453125,-0.62109375,0.88671875,0.12353515625,0.010986328125,-0.31640625,-0.1865234375,0.10205078125,-0.169921875,0.56640625,-0.357421875,0.14453125,-0.02685546875,-0.4765625,0.130859375,0.271484375,-0.0654296875,-1.03125,-0.5859375,-0.11181640625,-0.0859375,0.1435546875,0.4375,0.458984375,-0.2158203125,-0.28515625,-0.59375,-0.1279296875,2.9921531677246094E-5,0.00994873046875,0.73046875,0.2119140625,-0.76171875,-0.2060546875,0.008056640625,-0.09130859375,0.0228271484375,0.15234375,0.045654296875,0.64453125,0.0012359619140625,0.1474609375,0.345703125,0.671875,-0.2265625,-0.58203125,0.59375,0.04931640625,0.435546875,-0.3046875,-0.68359375,-0.494140625,0.412109375,0.0703125,0.169921875,1.0625,0.31640625,-0.9921875,-0.068359375,-0.349609375,-0.314453125,-0.1904296875,0.310546875,0.1748046875,0.49609375,0.84765625,-0.423828125,0.061767578125,-0.2451171875,0.0908203125,-0.408203125,-0.703125,0.8046875,-0.349609375,0.3046875,-0.54296875,-0.5234375,-0.3359375,-0.06640625,-0.275390625,0.1474609375,0.8828125,-0.55078125,0.10791015625,-0.103515625,1.0390625,-0.51953125,-0.42578125,0.77734375,-0.216796875,0.5546875,0.11669921875,-0.443359375,-0.380859375,-0.2109375,0.6171875,-0.83984375,0.5234375,0.337890625,0.01068115234375,0.053466796875,5.07354736328125E-4,-0.68359375,-0.1044921875,-0.07861328125,-0.388671875,-0.0712890625,-0.578125,-0.1796875,-0.609375,-0.228515625,0.2578125,0.98046875,-0.21484375,0.9609375,-0.427734375,-0.640625,0.173828125,0.33203125,-0.0439453125,0.05224609375,0.047607421875,-0.9609375,0.330078125,0.0888671875,-0.36328125,0.275390625,-0.119140625,0.263671875,-0.134765625,-0.10791015625,-0.578125,-0.076171875,0.66015625,0.0732421875,0.00140380859375,0.06201171875,0.00860595703125,-0.431640625,0.404296875,0.353515625,0.1728515625,0.2138671875,-0.2470703125,0.578125,-0.0225830078125,-0.515625,0.10107421875,0.375,0.33203125,0.53125,0.1826171875,-0.32421875,0.2373046875,-0.3046875,0.94140625,0.5625,-0.232421875,-0.1884765625,-0.004150390625,0.310546875,-0.2119140625,0.466796875,-0.5,-0.2158203125,0.171875,0.57421875,0.1611328125,-0.56640625,-0.00139617919921875,-0.4765625,0.5,-0.038330078125,0.99609375,0.06787109375,-0.0262451171875,0.06982421875,0.12158203125,-0.11767578125,0.14453125,-0.2119140625,0.8671875,0.5859375,0.412109375,-0.220703125,-0.71484375,-0.07763671875,0.25390625,0.75390625,0.3671875,-1.59375,0.326171875,0.1044921875,0.0029754638671875,-0.0264892578125,-0.119140625,-0.25390625,-0.345703125,0.244140625,0.302734375,0.23828125,0.390625,0.65625,0.43359375,0.0361328125,-0.0556640625,0.0703125,0.060302734375,-0.2890625,-0.25390625,0.34375,-0.8046875,0.169921875,0.59375,-0.33984375,0.2275390625,-0.01092529296875,-0.45703125,-0.6015625,0.059814453125,0.85546875,0.038330078125,-0.263671875,0.1484375,0.494140625,0.07275390625,-0.01190185546875,0.921875,-0.1591796875,0.7109375,-0.19140625,-0.275390625,-0.369140625,-0.201171875,-0.70703125,-0.41015625,0.408203125,-0.1435546875,-0.302734375,-1.046875,-0.255859375,0.010498046875,0.875,0.4375,-0.25390625,-0.70703125,0.10546875,0.0947265625,-0.302734375,0.15234375,0.63671875,-0.023681640625,0.345703125,0.640625,-0.09130859375,-0.50390625,0.173828125,0.0027923583984375,0.3359375,-0.7421875,0.1865234375,0.10400390625,-0.61328125,0.5234375,0.271484375,-0.41796875,-0.15625,-0.341796875,0.318359375,-0.3515625,-0.1708984375,-0.1435546875,-0.2490234375,0.095703125,0.38671875,0.10791015625,-0.1416015625,-0.27734375,0.1591796875,-0.61328125,1.4921875,0.2451171875,0.00750732421875,-0.055419921875,-0.09033203125,-0.07470703125,-0.1103515625,-0.322265625,-0.181640625,0.65234375,-0.427734375,0.033935546875,-0.85546875,-0.47265625,-0.0103759765625,0.796875,-0.10546875,-0.44140625,0.5703125,0.002044677734375,0.337890625,0.44140625,0.2060546875,-0.0283203125,0.01495361328125,0.41015625,-0.703125,-0.2421875,-0.388671875,-0.6875,-0.26171875,-0.1337890625,0.32421875,0.63671875,0.1884765625,0.1982421875,0.1474609375,-0.083984375,-0.46875,-0.58984375,-0.8359375,0.41796875,0.1376953125,0.03955078125,-0.291015625,-0.361328125,-0.78515625,-0.337890625,-0.123046875,0.24609375,-0.125,0.39453125,-0.70703125,0.302734375,-0.00714111328125,-0.68359375,-0.419921875,0.373046875,0.177734375,0.130859375,-0.42578125,0.45703125,-0.94921875,0.287109375,-0.6171875,-0.1826171875,-0.40234375,-0.40625,-0.228515625,0.1396484375,-0.1083984375,-0.255859375,-0.65625,-0.9296875,-0.365234375,-0.14453125,6.37054443359375E-4,0.271484375,0.6796875,0.1064453125,0.376953125,-0.4140625,0.1484375,0.051025390625,-0.2197265625,-0.283203125,-0.42578125,-0.2255859375,0.5546875,-0.1328125,0.1796875,-0.1650390625,-0.0732421875,-0.11669921875,0.173828125,0.1982421875,-0.1689453125,0.369140625,-1.3515625,-0.25390625,-0.310546875,-0.38671875,-1.1015625,-0.083984375,0.12451171875,-0.111328125,-0.458984375,-0.30859375,0.154296875,-0.2021484375,0.6015625,0.11376953125,-0.5703125,0.41015625,-0.4375,0.66015625,-0.37890625,-0.421875,-0.26953125,0.5546875,-0.00506591796875,0.6484375,0.038330078125,-0.1484375,0.0272216796875,-0.4921875,0.2470703125,0.267578125,-0.251953125,-0.369140625,0.042236328125,-0.150390625,-0.1728515625,0.466796875,-0.341796875,0.443359375,0.6015625,-0.154296875,0.07861328125,0.330078125,-0.671875,0.0233154296875,-0.625,1.28125,-0.25390625,0.6484375,-0.1689453125,0.36328125,-0.078125,0.4140625,0.71875,-0.09716796875,-0.0419921875,-0.68359375,-0.228515625,0.64453125,-0.201171875,0.10302734375,0.232421875,-0.20703125,-0.443359375,0.36328125,-0.1748046875,0.34765625,0.09912109375,0.07861328125,0.197265625,-0.455078125,0.1259765625,-0.1845703125,-0.470703125,0.984375,0.0184326171875,-0.01080322265625,-0.150390625,0.259765625,0.31640625,0.46875,0.09814453125,0.68359375,-0.322265625,-0.1416015625,-0.6953125,0.142578125,-0.1787109375,0.314453125,0.09619140625,0.53515625,0.49609375,-0.1298828125,-1.2109375,0.1611328125,0.625,-0.169921875,-0.2275390625,0.06298828125,-0.77734375,-0.5,-0.2373046875,-0.1787109375,0.32421875,0.47265625,0.40234375,0.345703125,0.099609375,0.45703125,-0.09765625,0.240234375,-0.2373046875,-0.333984375,0.62890625,0.1630859375,0.1279296875,-0.25390625,0.42578125,0.009521484375,0.470703125,-0.7421875,0.357421875,-0.16015625,0.0264892578125,-0.2451171875,-0.318359375,-0.06201171875,-0.326171875,-0.0859375,-0.3046875,0.52734375,0.1357421875,0.376953125,-0.042236328125,0.57421875,-0.0751953125,0.796875,-0.4609375,0.53125,-0.71484375,0.5078125,-0.6875,0.6796875,-0.09716796875,-0.6875,-0.2734375,-0.1103515625,-0.55078125,0.054443359375,-0.3359375,0.08447265625,-0.26953125,0.18359375,0.75390625,0.06689453125,-0.1611328125,0.259765625,-1.109375,-0.1708984375,-0.421875,-0.408203125,0.0022125244140625,0.671875,-0.150390625,0.45703125,0.6484375,0.28515625,-0.11572265625,-0.55078125,0.625,0.11328125,-0.1435546875,0.23828125,0.060546875,-0.2451171875,-0.07763671875,0.53125,-0.16796875,0.12451171875,-0.0625,0.265625,-0.2578125,0.09912109375,0.2431640625,0.23046875,0.357421875,0.21875,-0.326171875,0.244140625,-0.17578125,-0.359375,-0.32421875,-0.453125,-0.7734375,0.0289306640625,-0.0830078125,0.5,-0.50390625,-0.79296875,-0.042236328125,-0.1240234375,-0.33984375,0.5625,1.515625,0.228515625,0.357421875,-0.546875,-0.515625,0.62109375,-0.53125,-0.3203125,0.2119140625,-0.6640625,-0.1328125,0.08056640625,-1.1953125,0.373046875,0.0791015625,-0.28125,0.0296630859375,0.2353515625,-0.330078125,0.2275390625,0.036865234375,0.330078125,-0.9140625,0.4375,0.349609375,-0.0703125,-0.072265625,0.251953125,0.69140625,0.1328125,0.0016326904296875,-0.421875,0.326171875,-0.51953125,0.3671875,0.1982421875,-0.01336669921875,0.05908203125,-0.05908203125,-0.1240234375,0.01220703125,-0.6328125,0.77734375],"inputTextTokenCount":7}' + headers: + Connection: + - keep-alive + Content-Length: + - '18891' + Content-Type: + - application/json + Date: + - Mon, 19 May 2025 16:42:06 GMT + X-Amzn-Bedrock-Input-Token-Count: + - '7' + X-Amzn-Bedrock-Invocation-Latency: + - '68' + x-amzn-RequestId: + - 1b70cfa3-5b0a-41b1-acfb-271252abc5be + status: + code: 200 + message: OK +version: 1 diff --git a/src/tests/aws_bedrock/conftest.py b/src/tests/aws_bedrock/conftest.py new file mode 100644 index 00000000..5bfc8067 --- /dev/null +++ b/src/tests/aws_bedrock/conftest.py @@ -0,0 +1,48 @@ +"""Unit tests configuration module.""" + +import pytest +import os + +from boto3.session import Session +from botocore.config import Config + +from langtrace_python_sdk.instrumentation.aws_bedrock.instrumentation import ( + AWSBedrockInstrumentation, +) + + +@pytest.fixture(autouse=True) +def environment(): + if not os.getenv("AWS_ACCESS_KEY_ID"): + os.environ["AWS_ACCESS_KEY_ID"] = "test_aws_access_key_id" + if not os.getenv("AWS_SECRET_ACCESS_KEY"): + os.environ["AWS_SECRET_ACCESS_KEY"] = "test_aws_secret_access_key" + + +@pytest.fixture +def aws_bedrock_client(): + bedrock_config = Config( + region_name="us-east-1", + connect_timeout=300, + read_timeout=300, + retries={"total_max_attempts": 2, "mode": "standard"}, + ) + return Session().client("bedrock-runtime", config=bedrock_config) + + +@pytest.fixture(scope="module") +def vcr_config(): + return { + "filter_headers": [ + "authorization", + "X-Amz-Date", + "X-Amz-Security-Token", + "amz-sdk-invocation-id", + "amz-sdk-request", + ] + } + + +@pytest.fixture(scope="session", autouse=True) +def instrument(): + AWSBedrockInstrumentation().instrument() diff --git a/src/tests/aws_bedrock/test_invoke_model.py b/src/tests/aws_bedrock/test_invoke_model.py new file mode 100644 index 00000000..8abb4318 --- /dev/null +++ b/src/tests/aws_bedrock/test_invoke_model.py @@ -0,0 +1,135 @@ +import pytest +import json +from tests.utils import ( + assert_completion_in_events, + assert_prompt_in_events, + assert_token_count, +) +from importlib_metadata import version as v + +from langtrace.trace_attributes import SpanAttributes +from langtrace_python_sdk.constants.instrumentation.aws_bedrock import APIS + +ANTHROPIC_VERSION = "bedrock-2023-05-31" + + +@pytest.mark.vcr() +def test_chat_completion(exporter, aws_bedrock_client): + model_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0" + messages_value = [{"role": "user", "content": "Say this is a test three times"}] + + kwargs = { + "modelId": model_id, + "accept": "application/json", + "contentType": "application/json", + "body": json.dumps( + { + "messages": messages_value, + "anthropic_version": ANTHROPIC_VERSION, + "max_tokens": 100, + } + ), + } + + aws_bedrock_client.invoke_model(**kwargs) + spans = exporter.get_finished_spans() + completion_span = spans[-1] + assert completion_span.name == "aws_bedrock.invoke_model" + + attributes = completion_span.attributes + + assert attributes.get(SpanAttributes.LANGTRACE_SDK_NAME) == "langtrace-python-sdk" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "anthropic" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_TYPE) == "framework" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_VERSION) == v("boto3") + assert attributes.get(SpanAttributes.LANGTRACE_VERSION) == v("langtrace-python-sdk") + assert attributes.get(SpanAttributes.LLM_PATH) == APIS["INVOKE_MODEL"]["ENDPOINT"] + assert attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) == model_id + assert attributes.get(SpanAttributes.LLM_IS_STREAMING) is False + assert_prompt_in_events(completion_span.events) + assert_completion_in_events(completion_span.events) + assert_token_count(attributes) + + +@pytest.mark.skip(reason="Skipping streaming test due to no streaming support in vcrpy") +def test_chat_completion_streaming(exporter, aws_bedrock_client): + model_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0" + messages_value = [{"role": "user", "content": "Say this is a test three times"}] + + kwargs = { + "modelId": model_id, + "accept": "application/json", + "contentType": "application/json", + "body": json.dumps( + { + "messages": messages_value, + "anthropic_version": ANTHROPIC_VERSION, + "max_tokens": 100, + } + ), + } + + response = aws_bedrock_client.invoke_model_with_response_stream(**kwargs) + chunk_count = 0 + + for chunk in response["body"]: + if chunk: + chunk_count += 1 + + spans = exporter.get_finished_spans() + streaming_span = spans[-1] + assert streaming_span.name == "aws_bedrock.invoke_model_with_response_stream" + + attributes = streaming_span.attributes + + assert attributes.get(SpanAttributes.LANGTRACE_SDK_NAME) == "langtrace-python-sdk" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "anthropic" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_TYPE) == "framework" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_VERSION) == v("boto3") + assert attributes.get(SpanAttributes.LANGTRACE_VERSION) == v("langtrace-python-sdk") + assert ( + attributes.get(SpanAttributes.LLM_PATH) + == APIS["INVOKE_MODEL_WITH_RESPONSE_STREAM"]["ENDPOINT"] + ) + assert ( + attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) + == "claude-3-7-sonnet-20250219-v1:0" + ) + assert attributes.get(SpanAttributes.LLM_IS_STREAMING) is True + assert_prompt_in_events(streaming_span.events) + assert_completion_in_events(streaming_span.events) + assert_token_count(attributes) + + +@pytest.mark.vcr() +def test_generate_embedding(exporter, aws_bedrock_client): + model_id = "amazon.titan-embed-text-v1" + + kwargs = { + "modelId": model_id, + "accept": "application/json", + "contentType": "application/json", + "body": json.dumps( + { + "inputText": "Say this is a test three times", + } + ), + } + + aws_bedrock_client.invoke_model(**kwargs) + spans = exporter.get_finished_spans() + completion_span = spans[-1] + assert completion_span.name == "aws_bedrock.invoke_model" + + attributes = completion_span.attributes + + assert attributes.get(SpanAttributes.LANGTRACE_SDK_NAME) == "langtrace-python-sdk" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_NAME) == "amazon" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_TYPE) == "framework" + assert attributes.get(SpanAttributes.LANGTRACE_SERVICE_VERSION) == v("boto3") + assert attributes.get(SpanAttributes.LANGTRACE_VERSION) == v("langtrace-python-sdk") + assert attributes.get(SpanAttributes.LLM_PATH) == APIS["INVOKE_MODEL"]["ENDPOINT"] + assert attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) == model_id + assert attributes.get(SpanAttributes.LLM_IS_STREAMING) is False + assert_prompt_in_events(completion_span.events) + assert_token_count(attributes) diff --git a/src/tests/aws_bedrock/test_model_id_parsing.py b/src/tests/aws_bedrock/test_model_id_parsing.py new file mode 100644 index 00000000..832e32eb --- /dev/null +++ b/src/tests/aws_bedrock/test_model_id_parsing.py @@ -0,0 +1,34 @@ +import pytest +from langtrace_python_sdk.instrumentation.aws_bedrock.patch import ( + parse_vendor_and_model_name_from_model_id, +) + + +def test_model_id_parsing(): + model_id = "anthropic.claude-3-opus-20240229" + vendor, model_name = parse_vendor_and_model_name_from_model_id(model_id) + assert vendor == "anthropic" + assert model_name == "claude-3-opus-20240229" + + +def test_model_id_parsing_cross_region_inference(): + model_id = "us.anthropic.claude-3-opus-20240229" + vendor, model_name = parse_vendor_and_model_name_from_model_id(model_id) + assert vendor == "anthropic" + assert model_name == "claude-3-opus-20240229" + + +def test_model_id_parsing_arn_custom_model_inference(): + model_id = ( + "arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.my-model/abc123" + ) + vendor, model_name = parse_vendor_and_model_name_from_model_id(model_id) + assert vendor == "amazon" + assert model_name == "my-model" + + +def test_model_id_parsing_arn_foundation_model_inference(): + model_id = "arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3-opus-20240229" + vendor, model_name = parse_vendor_and_model_name_from_model_id(model_id) + assert vendor == "anthropic" + assert model_name == "claude-3-opus-20240229"