fix: support conditional imports of http client libraries (#194)

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)

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

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")

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]