From 10e26ebb1df42334b33f41d174dbc0f1e6bc6cad Mon Sep 17 00:00:00 2001 From: Anton Dubovik Date: Thu, 28 Nov 2024 09:33:47 +0000 Subject: [PATCH] fix: support conditional imports of http client libraries (#194) --- aidial_sdk/header_propagator.py | 34 +++++++++++++++++++++----------- aidial_sdk/telemetry/init.py | 35 +++++++++++++++++++++++++-------- noxfile.py | 6 ++++-- poetry.lock | 2 +- pyproject.toml | 8 ++++---- 5 files changed, 58 insertions(+), 27 deletions(-) diff --git a/aidial_sdk/header_propagator.py b/aidial_sdk/header_propagator.py index 77a4b4e..f4ee3d9 100644 --- a/aidial_sdk/header_propagator.py +++ b/aidial_sdk/header_propagator.py @@ -2,9 +2,6 @@ from contextvars import ContextVar from typing import MutableMapping, Optional -import aiohttp -import httpx -import requests import wrapt from fastapi import FastAPI from starlette.types import ASGIApp, Receive, Scope, Send @@ -59,9 +56,21 @@ def _instrument_fast_api(self, app: FastAPI): app.add_middleware(FastAPIMiddleware, api_key=self._api_key) def _instrument_aiohttp(self): + try: + import aiohttp + except ImportError: + return + + async def _on_request_start( + session: aiohttp.ClientSession, + trace_config_ctx: types.SimpleNamespace, + params: aiohttp.TraceRequestStartParams, + ): + self._modify_headers(str(params.url), params.headers) + def instrumented_init(wrapped, instance, args, kwargs): trace_config = aiohttp.TraceConfig() - trace_config.on_request_start.append(self._on_aiohttp_request_start) + trace_config.on_request_start.append(_on_request_start) trace_configs = list(kwargs.get("trace_configs") or []) trace_configs.append(trace_config) @@ -73,15 +82,12 @@ def instrumented_init(wrapped, instance, args, kwargs): aiohttp.ClientSession, "__init__", instrumented_init ) - async def _on_aiohttp_request_start( - self, - session: aiohttp.ClientSession, - trace_config_ctx: types.SimpleNamespace, - params: aiohttp.TraceRequestStartParams, - ): - self._modify_headers(str(params.url), params.headers) - def _instrument_requests(self): + try: + import requests + except ImportError: + return + def instrumented_send(wrapped, instance, args, kwargs): request: requests.PreparedRequest = args[0] self._modify_headers(request.url or "", request.headers) @@ -90,6 +96,10 @@ def instrumented_send(wrapped, instance, args, kwargs): wrapt.wrap_function_wrapper(requests.Session, "send", instrumented_send) def _instrument_httpx(self): + try: + import httpx + except ImportError: + return def instrumented_build_request(wrapped, instance, args, kwargs): request: httpx.Request = wrapped(*args, **kwargs) diff --git a/aidial_sdk/telemetry/init.py b/aidial_sdk/telemetry/init.py index f4693f2..5fa7552 100644 --- a/aidial_sdk/telemetry/init.py +++ b/aidial_sdk/telemetry/init.py @@ -10,13 +10,8 @@ OTLPSpanExporter, ) from opentelemetry.exporter.prometheus import PrometheusMetricReader -from opentelemetry.instrumentation.aiohttp_client import ( - AioHttpClientInstrumentor, -) from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor -from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor from opentelemetry.instrumentation.logging import LoggingInstrumentor -from opentelemetry.instrumentation.requests import RequestsInstrumentor from opentelemetry.instrumentation.system_metrics import ( SystemMetricsInstrumentor, ) @@ -57,10 +52,34 @@ def init_telemetry( set_tracer_provider(tracer_provider) - RequestsInstrumentor().instrument() - AioHttpClientInstrumentor().instrument() + try: + from opentelemetry.instrumentation.requests import ( + RequestsInstrumentor, + ) + + RequestsInstrumentor().instrument() + except ImportError: + pass + + try: + from opentelemetry.instrumentation.aiohttp_client import ( + AioHttpClientInstrumentor, + ) + + AioHttpClientInstrumentor().instrument() + except ImportError: + pass + URLLibInstrumentor().instrument() - HTTPXClientInstrumentor().instrument() + + try: + from opentelemetry.instrumentation.httpx import ( + HTTPXClientInstrumentor, + ) + + HTTPXClientInstrumentor().instrument() + except ImportError: + pass if config.tracing.logging: # Setting the root logger format in order to include diff --git a/noxfile.py b/noxfile.py index fb9e6b6..be8a81c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -40,6 +40,8 @@ def format(session: nox.Session): def test(session: nox.Session, pydantic: str, httpx: str) -> None: """Runs tests""" session.run("poetry", "install", external=True) - session.install(f"pydantic=={pydantic}") - session.install(f"httpx=={httpx}") + session.install( + f"pydantic=={pydantic}", + f"httpx=={httpx}", + ) session.run("pytest") diff --git a/poetry.lock b/poetry.lock index abd1cbf..ce048bf 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2235,4 +2235,4 @@ telemetry = ["opentelemetry-api", "opentelemetry-exporter-otlp-proto-grpc", "ope [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "8eeb4f517f3b2bf01030560cabf84d3430d58baf31d2d9ef1602f20a3573d5b8" +content-hash = "e3e7990015104bea82768a07c91a20eb0e7a2cd5503e2ffd6f449fea52ed3b9d" diff --git a/pyproject.toml b/pyproject.toml index 9a4888c..2625c01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,10 +13,7 @@ python = ">=3.8.1,<4.0" fastapi = ">=0.51,<1.0" uvicorn = ">=0.19,<1.0" pydantic = ">=1.10,<3" -requests = "^2.19" -wrapt = "^1.14" -aiohttp = "^3.8.3" -httpx = ">=0.25.0,<1.0" +wrapt = ">=1.10,<2" # Telemetry extras opentelemetry-sdk = {version = "^1.20.0", optional = true} @@ -53,8 +50,11 @@ pytest = "^8.2" pytest-asyncio = "^0.24.0" nox = "^2023.4.22" pillow = "^10.2.0" +httpx = "^0.25.0" respx = "^0.21.1" +aiohttp = "^3.8.3" aioresponses = "^0.7.6" +requests = "^2.19" responses = "^0.25.3" [tool.poetry.group.lint.dependencies]