Propagate OpenTelemetry context on inference API calls MCDB-100371

singlestoredb/ai/chat.py

@@ -9,6 +9,7 @@
 from singlestoredb import manage_workspaces
 from singlestoredb.management.inference_api import InferenceAPIInfo
 
+
 try:
     from langchain_openai import ChatOpenAI
 except ImportError:
@@ -30,6 +31,29 @@
 from botocore.config import Config
 
 
+def _inject_otel_headers(headers: Any) -> None:
+    try:
+        from opentelemetry.propagate import inject as otel_inject
+    except ImportError:
+        return
+    try:
+        otel_inject(headers)
+    except Exception:
+        return
+
+
+def _httpx_inject_otel(request: httpx.Request) -> None:
+    _inject_otel_headers(request.headers)
+
+
+def _attach_otel_request_hook(
+    client: Union[httpx.Client, httpx.AsyncClient],
+) -> None:
+    hooks = client.event_hooks.setdefault('request', [])
+    if _httpx_inject_otel not in hooks:
+        hooks.append(_httpx_inject_otel)
+
+
 def SingleStoreChatFactory(
     model_name: str,
     api_key: Optional[str] = None,
@@ -119,6 +143,7 @@ def _inject_headers(request: Any, **_ignored: Any) -> None:
                 obo_val = obo_token_getter()
                 if obo_val:
                     request.headers['X-S2-OBO'] = obo_val
+            _inject_otel_headers(request.headers)
             request.headers.pop('X-Amz-Date', None)
             request.headers.pop('X-Amz-Security-Token', None)
 
@@ -161,8 +186,19 @@ def _inject_headers(request: Any, **_ignored: Any) -> None:
         model=model_name,
         streaming=streaming,
     )
-    if http_client is not None:
-        openai_kwargs['http_client'] = http_client
+    http_async_client = kwargs.pop('http_async_client', None)
+    default_timeout = httpx.Timeout(timeout=600.0, connect=5.0)
+    if http_client is None:
+        http_client = httpx.Client(timeout=default_timeout)
+    _attach_otel_request_hook(http_client)
+    openai_kwargs['http_client'] = http_client
+    if http_async_client is None:
+        async_timeout = (
+            http_client.timeout if http_client is not None else default_timeout
+        )
+        http_async_client = httpx.AsyncClient(timeout=async_timeout)
+    _attach_otel_request_hook(http_async_client)
+    openai_kwargs['http_async_client'] = http_async_client
     return ChatOpenAI(
         **openai_kwargs,
         **kwargs,

singlestoredb/tests/test_ai_chat.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# type: ignore
+"""SingleStoreChatFactory OTel header injection tests."""
+import unittest
+from types import ModuleType
+from unittest.mock import patch
+
+try:
+    import httpx
+    from singlestoredb.ai import chat as chat_mod
+except ImportError:
+    httpx = None
+    chat_mod = None
+
+
+@unittest.skipIf(chat_mod is None, 'singlestoredb.ai.chat dependencies missing')
+class TestInjectOtelHeaders(unittest.TestCase):
+
+    def test_calls_otel_inject(self):
+        headers = {}
+
+        def fake_inject(target):
+            target['baggage'] = 'session=abc,turn=def'
+
+        fake_propagate = ModuleType('opentelemetry.propagate')
+        fake_propagate.inject = fake_inject
+        fake_otel = ModuleType('opentelemetry')
+        fake_otel.__path__ = []  # mark as package for submodules
+        with patch.dict(
+            'sys.modules',
+            {
+                'opentelemetry': fake_otel,
+                'opentelemetry.propagate': fake_propagate,
+            },
+        ):
+            chat_mod._inject_otel_headers(headers)
+        self.assertEqual(headers['baggage'], 'session=abc,turn=def')
+
+    def test_httpx_hook_appends_once(self):
+        client = httpx.Client()
+        try:
+            chat_mod._attach_otel_request_hook(client)
+            chat_mod._attach_otel_request_hook(client)
+            self.assertEqual(
+                client.event_hooks['request'].count(chat_mod._httpx_inject_otel),
+                1,
+            )
+        finally:
+            client.close()
+
+    @patch('singlestoredb.ai.chat.ChatOpenAI')
+    def test_openai_default_timeout_is_600(self, mock_chat_openai):
+        chat_mod.SingleStoreChatFactory(
+            model_name='gpt-4o',
+            base_url='https://example.com',
+            hosting_platform='OpenAI',
+        )
+        self.assertTrue(mock_chat_openai.called)
+        _, kwargs = mock_chat_openai.call_args
+        http_client = kwargs.get('http_client')
+        http_async_client = kwargs.get('http_async_client')
+        try:
+            self.assertIsNotNone(http_client)
+            self.assertEqual(http_client.timeout.read, 600.0)
+            self.assertEqual(http_client.timeout.connect, 5.0)
+            self.assertIsNotNone(http_async_client)
+            self.assertEqual(http_async_client.timeout.read, 600.0)
+            self.assertEqual(http_async_client.timeout.connect, 5.0)
+        finally:
+            if http_client:
+                http_client.close()
+            if http_async_client:
+                import asyncio
+                asyncio.run(http_async_client.aclose())
+
+    @patch('singlestoredb.ai.chat.ChatOpenAI')
+    def test_openai_inherits_custom_client_timeout(self, mock_chat_openai):
+        custom_sync = httpx.Client(timeout=httpx.Timeout(42.0))
+        http_async_client = None
+        try:
+            chat_mod.SingleStoreChatFactory(
+                model_name='gpt-4o',
+                base_url='https://example.com',
+                hosting_platform='OpenAI',
+                http_client=custom_sync,
+            )
+            self.assertTrue(mock_chat_openai.called)
+            _, kwargs = mock_chat_openai.call_args
+            http_client = kwargs.get('http_client')
+            http_async_client = kwargs.get('http_async_client')
+            self.assertIs(http_client, custom_sync)
+            self.assertIsNotNone(http_async_client)
+            self.assertEqual(http_async_client.timeout.read, 42.0)
+        finally:
+            custom_sync.close()
+            if http_async_client:
+                import asyncio
+                asyncio.run(http_async_client.aclose())
+
+
+if __name__ == '__main__':
+    unittest.main()