From b31f57179d8036dcb8357089dee65291c59c48e9 Mon Sep 17 00:00:00 2001 From: Aditya Jha Date: Mon, 7 Sep 2026 16:33:12 +0530 Subject: [PATCH 1/5] Propagate OpenTelemetry context on inference API calls. Flow, Query Tuning, and Observability go through SingleStoreChatFactory without Analyst's trace hooks, so UMG never gets session/turn baggage and per-question cost rows stay empty. Co-authored-by: Cursor --- singlestoredb/ai/chat.py | 36 ++++++++++++++++++-- singlestoredb/tests/test_ai_chat.py | 52 +++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) create mode 100644 singlestoredb/tests/test_ai_chat.py diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py index 6636fe8d5..5e66d942b 100644 --- a/singlestoredb/ai/chat.py +++ b/singlestoredb/ai/chat.py @@ -9,6 +9,30 @@ from singlestoredb import manage_workspaces from singlestoredb.management.inference_api import InferenceAPIInfo + +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) + + try: from langchain_openai import ChatOpenAI except ImportError: @@ -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,15 @@ 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) + if http_client is None: + http_client = httpx.Client(timeout=httpx.Timeout(None)) + _attach_otel_request_hook(http_client) + openai_kwargs['http_client'] = http_client + if http_async_client is None: + http_async_client = httpx.AsyncClient(timeout=httpx.Timeout(None)) + _attach_otel_request_hook(http_async_client) + openai_kwargs['http_async_client'] = http_async_client return ChatOpenAI( **openai_kwargs, **kwargs, diff --git a/singlestoredb/tests/test_ai_chat.py b/singlestoredb/tests/test_ai_chat.py new file mode 100644 index 000000000..93cfaef70 --- /dev/null +++ b/singlestoredb/tests/test_ai_chat.py @@ -0,0 +1,52 @@ +#!/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') + 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() + + +if __name__ == '__main__': + unittest.main() From 4e6f3ea06bc446e5afac63363646cb39463998dc Mon Sep 17 00:00:00 2001 From: Aditya Jha Date: Mon, 7 Sep 2026 20:18:07 +0530 Subject: [PATCH 2/5] Move helper functions below top-level imports to fix flake8 E402 --- singlestoredb/ai/chat.py | 42 ++++++++++++++++++++-------------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py index 5e66d942b..8f5481a1a 100644 --- a/singlestoredb/ai/chat.py +++ b/singlestoredb/ai/chat.py @@ -10,6 +10,27 @@ from singlestoredb.management.inference_api import InferenceAPIInfo +try: + from langchain_openai import ChatOpenAI +except ImportError: + raise ImportError( + 'Could not import langchain_openai python package. ' + 'Please install it with `pip install langchain_openai`.', + ) + +try: + from langchain_aws import ChatBedrockConverse +except ImportError: + raise ImportError( + 'Could not import langchain-aws python package. ' + 'Please install it with `pip install langchain-aws`.', + ) + +import boto3 +from botocore import UNSIGNED +from botocore.config import Config + + def _inject_otel_headers(headers: Any) -> None: try: from opentelemetry.propagate import inject as otel_inject @@ -33,27 +54,6 @@ def _attach_otel_request_hook( hooks.append(_httpx_inject_otel) -try: - from langchain_openai import ChatOpenAI -except ImportError: - raise ImportError( - 'Could not import langchain_openai python package. ' - 'Please install it with `pip install langchain_openai`.', - ) - -try: - from langchain_aws import ChatBedrockConverse -except ImportError: - raise ImportError( - 'Could not import langchain-aws python package. ' - 'Please install it with `pip install langchain-aws`.', - ) - -import boto3 -from botocore import UNSIGNED -from botocore.config import Config - - def SingleStoreChatFactory( model_name: str, api_key: Optional[str] = None, From 1d3ecf7eee5b269c21be65013a2076432a2e2e7b Mon Sep 17 00:00:00 2001 From: Aditya Jha Date: Mon, 7 Sep 2026 21:01:17 +0530 Subject: [PATCH 3/5] Default OpenAI httpx clients to 600s timeout and inherit caller timeout --- singlestoredb/ai/chat.py | 8 ++++++-- singlestoredb/tests/test_ai_chat.py | 7 +++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py index 8f5481a1a..520064e5a 100644 --- a/singlestoredb/ai/chat.py +++ b/singlestoredb/ai/chat.py @@ -187,12 +187,16 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: streaming=streaming, ) http_async_client = kwargs.pop('http_async_client', None) + default_timeout = httpx.Timeout(600.0) if http_client is None: - http_client = httpx.Client(timeout=httpx.Timeout(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: - http_async_client = httpx.AsyncClient(timeout=httpx.Timeout(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( diff --git a/singlestoredb/tests/test_ai_chat.py b/singlestoredb/tests/test_ai_chat.py index 93cfaef70..fbe464437 100644 --- a/singlestoredb/tests/test_ai_chat.py +++ b/singlestoredb/tests/test_ai_chat.py @@ -47,6 +47,13 @@ def test_httpx_hook_appends_once(self): finally: client.close() + def test_openai_default_timeout_is_600(self): + client = httpx.Client(timeout=httpx.Timeout(600.0)) + try: + self.assertEqual(client.timeout.read, 600.0) + finally: + client.close() + if __name__ == '__main__': unittest.main() From f3a4109f3c83fd510e000dc593a1077dd5beb874 Mon Sep 17 00:00:00 2001 From: Aditya Jha Date: Mon, 7 Sep 2026 21:14:21 +0530 Subject: [PATCH 4/5] Test SingleStoreChatFactory timeout defaults and inheritance --- singlestoredb/tests/test_ai_chat.py | 48 ++++++++++++++++++++++++++--- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/singlestoredb/tests/test_ai_chat.py b/singlestoredb/tests/test_ai_chat.py index fbe464437..31eeabc55 100644 --- a/singlestoredb/tests/test_ai_chat.py +++ b/singlestoredb/tests/test_ai_chat.py @@ -47,12 +47,52 @@ def test_httpx_hook_appends_once(self): finally: client.close() - def test_openai_default_timeout_is_600(self): - client = httpx.Client(timeout=httpx.Timeout(600.0)) + @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.assertEqual(client.timeout.read, 600.0) + self.assertIsNotNone(http_client) + self.assertEqual(http_client.timeout.read, 600.0) + self.assertIsNotNone(http_async_client) + self.assertEqual(http_async_client.timeout.read, 600.0) finally: - client.close() + 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__': From cbe5f25b519d984bc4678b0650dc0487ed76d37c Mon Sep 17 00:00:00 2001 From: Aditya Jha Date: Tue, 8 Sep 2026 19:03:20 +0530 Subject: [PATCH 5/5] copilot review comments --- singlestoredb/ai/chat.py | 2 +- singlestoredb/tests/test_ai_chat.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py index 520064e5a..453a5f74b 100644 --- a/singlestoredb/ai/chat.py +++ b/singlestoredb/ai/chat.py @@ -187,7 +187,7 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: streaming=streaming, ) http_async_client = kwargs.pop('http_async_client', None) - default_timeout = httpx.Timeout(600.0) + 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) diff --git a/singlestoredb/tests/test_ai_chat.py b/singlestoredb/tests/test_ai_chat.py index 31eeabc55..6658ce89a 100644 --- a/singlestoredb/tests/test_ai_chat.py +++ b/singlestoredb/tests/test_ai_chat.py @@ -25,6 +25,7 @@ def fake_inject(target): 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', { @@ -61,8 +62,10 @@ def test_openai_default_timeout_is_600(self, mock_chat_openai): 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()