diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py new file mode 100644 index 00000000000..480e1d6b8da --- /dev/null +++ b/litellm/integrations/opentelemetry.py @@ -0,0 +1,139 @@ +import os +from typing import Optional +from dataclasses import dataclass + +from litellm.integrations.custom_logger import CustomLogger + + +LITELLM_TRACER_NAME = "litellm" +LITELLM_RESOURCE = {"service.name": "litellm"} + + +@dataclass +class OpenTelemetryConfig: + from opentelemetry.sdk.trace.export import SpanExporter + + exporter: str | SpanExporter = "console" + endpoint: Optional[str] = None + bearer_token: Optional[str] = None + + @classmethod + def from_env(cls): + return cls( + exporter=os.getenv("OTEL_EXPORTER", "console"), + endpoint=os.getenv("OTEL_ENDPOINT"), + bearer_token=os.getenv("OTEL_BEARER_TOKEN"), + ) + + +class OpenTelemetry(CustomLogger): + def __init__(self, config=OpenTelemetryConfig.from_env()): + from opentelemetry import trace + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + + self.config = config + provider = TracerProvider(resource=Resource(attributes=LITELLM_RESOURCE)) + provider.add_span_processor(self._get_span_processor()) + + trace.set_tracer_provider(provider) + self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._handle_sucess(kwargs, response_obj, start_time, end_time) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._handle_failure(kwargs, response_obj, start_time, end_time) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._handle_sucess(kwargs, response_obj, start_time, end_time) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._handle_failure(kwargs, response_obj, start_time, end_time) + + def _handle_sucess(self, kwargs, response_obj, start_time, end_time): + from opentelemetry.trace import Status, StatusCode + + span = self.tracer.start_span( + name=self._get_span_name(kwargs), + start_time=self._to_ns(start_time), + context=self._get_span_context(kwargs), + ) + span.set_status(Status(StatusCode.OK)) + self.set_attributes(span, kwargs, response_obj) + span.end(end_time=self._to_ns(end_time)) + + def _handle_failure(self, kwargs, response_obj, start_time, end_time): + from opentelemetry.trace import Status, StatusCode + + span = self.tracer.start_span( + name=self._get_span_name(kwargs), + start_time=self._to_ns(start_time), + context=self._get_span_context(kwargs), + ) + span.set_status(Status(StatusCode.ERROR)) + self.set_attributes(span, kwargs, response_obj) + span.end(end_time=self._to_ns(end_time)) + + def set_attributes(self, span, kwargs, response_obj): + for key in ["model", "api_base", "api_version"]: + if key in kwargs: + span.set_attribute(key, kwargs[key]) + + def _to_ns(self, dt): + return int(dt.timestamp() * 1e9) + + def _get_span_name(self, kwargs): + return f"litellm-{kwargs.get('call_type', 'completion')}" + + def _get_span_context(self, kwargs): + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + + litellm_params = kwargs.get("litellm_params", {}) or {} + proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} + headers = proxy_server_request.get("headers", {}) or {} + traceparent = headers.get("traceparent", None) + + if traceparent is None: + return None + else: + carrier = {"traceparent": traceparent} + return TraceContextTextMapPropagator().extract(carrier=carrier) + + def _get_span_processor(self): + from opentelemetry.sdk.trace.export import ( + SpanExporter, + SimpleSpanProcessor, + BatchSpanProcessor, + ConsoleSpanExporter, + ) + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter as OTLPSpanExporterHTTP, + ) + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter as OTLPSpanExporterGRPC, + ) + + if isinstance(self.config.exporter, SpanExporter): + return SimpleSpanProcessor(self.config.exporter) + + if self.config.exporter == "console": + return BatchSpanProcessor(ConsoleSpanExporter()) + elif self.config.exporter == "otlp_http": + return BatchSpanProcessor( + OTLPSpanExporterHTTP( + endpoint=self.OTEL_ENDPOINT, + headers={"Authorization": f"Bearer {self.OTEL_BEARER_TOKEN}"}, + ) + ) + elif self.config.exporter == "otlp_grpc": + return BatchSpanProcessor( + OTLPSpanExporterGRPC( + endpoint=self.OTEL_ENDPOINT, + headers={"Authorization": f"Bearer {self.OTEL_BEARER_TOKEN}"}, + ) + ) + else: + return BatchSpanProcessor(ConsoleSpanExporter()) diff --git a/litellm/tests/test_opentelemetry.py b/litellm/tests/test_opentelemetry.py new file mode 100644 index 00000000000..b8d3379beab --- /dev/null +++ b/litellm/tests/test_opentelemetry.py @@ -0,0 +1,26 @@ +import asyncio +import litellm + +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +def test_otel_callback(): + exporter = InMemorySpanExporter() + + litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter=exporter))] + + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + ) + + asyncio.run( + litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + ) + ) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + 1 diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 77692e2ee92..114b96872f6 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -41,49 +41,39 @@ example_completion_result = { { "message": { "content": "Whispers of the wind carry dreams to me.", - "role": "assistant" + "role": "assistant", } } ], } example_embedding_result = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [ - -0.006929283495992422, - -0.005336422007530928, - -4.547132266452536e-05, - -0.024047505110502243, - -0.006929283495992422, - -0.005336422007530928, - -4.547132266452536e-05, - -0.024047505110502243, - -0.006929283495992422, - -0.005336422007530928, - -4.547132266452536e-05, - -0.024047505110502243, - ], - } - ], - "model": "text-embedding-3-small", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5 - } + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + ], + } + ], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, } example_image_generation_result = { - "created": 1589478378, - "data": [ - { - "url": "https://..." - }, - { - "url": "https://..." - } - ] + "created": 1589478378, + "data": [{"url": "https://..."}, {"url": "https://..."}], } @@ -185,7 +175,9 @@ def test_engines_model_chat_completions(mock_acompletion, client_no_auth): } print("testing proxy server with chat completions") - response = client_no_auth.post("/engines/gpt-3.5-turbo/chat/completions", json=test_data) + response = client_no_auth.post( + "/engines/gpt-3.5-turbo/chat/completions", json=test_data + ) mock_acompletion.assert_called_once_with( model="gpt-3.5-turbo", messages=[ @@ -249,7 +241,9 @@ def test_chat_completion_azure(mock_acompletion, client_no_auth): @mock_patch_acompletion() -def test_openai_deployments_model_chat_completions_azure(mock_acompletion, client_no_auth): +def test_openai_deployments_model_chat_completions_azure( + mock_acompletion, client_no_auth +): global headers try: # Your test data @@ -388,10 +382,10 @@ def test_img_gen(mock_aimage_generation, client_no_auth): response = client_no_auth.post("/v1/images/generations", json=test_data) mock_aimage_generation.assert_called_once_with( - model='dall-e-3', - prompt='A cute baby sea otter', + model="dall-e-3", + prompt="A cute baby sea otter", n=1, - size='1024x1024', + size="1024x1024", metadata=mock.ANY, proxy_server_request=mock.ANY, ) diff --git a/poetry.lock b/poetry.lock index 9c3e653c9af..ed4c4138d77 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2114,7 +2114,6 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, - {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -2122,15 +2121,8 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, - {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, - {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, - {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, - {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, - {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, - {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -2147,7 +2139,6 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, - {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -2155,7 +2146,6 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, - {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, diff --git a/requirements.txt b/requirements.txt index 3c5abae5a8d..23f72219855 100644 --- a/requirements.txt +++ b/requirements.txt @@ -28,6 +28,9 @@ python-multipart==0.0.9 # admin UI Pillow==10.3.0 azure-ai-contentsafety==1.0.0 # for azure content safety azure-identity==1.15.0 # for azure content safety +opentelemetry-api = "^1.25.0" +opentelemetry-sdk = "^1.25.0" +opentelemetry-exporter-otlp = "^1.25.0" ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.0 # for env