diff --git a/langfuse/openai.py b/langfuse/openai.py index 32637569b..0b77e1729 100644 --- a/langfuse/openai.py +++ b/langfuse/openai.py @@ -23,6 +23,8 @@ See docs for more details: https://langfuse.com/docs/integrations/openai """ +import base64 +import binascii import json import os import types @@ -963,8 +965,16 @@ def _get_langfuse_data_from_default_response( if hasattr(first_embedding, "embedding") else first_embedding.get("embedding", []) ) + if isinstance(embedding_vector, str): + try: + # Base64 embeddings contain one float32 per dimension. + dimensions = len(base64.b64decode(embedding_vector)) // 4 + except (binascii.Error, ValueError): + dimensions = len(embedding_vector) + else: + dimensions = len(embedding_vector) if embedding_vector else 0 completion = { - "dimensions": len(embedding_vector) if embedding_vector else 0, + "dimensions": dimensions, "count": len(data), } diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index a2a4f7f4f..ae1bc1c58 100644 --- a/tests/unit/test_openai.py +++ b/tests/unit/test_openai.py @@ -3,6 +3,7 @@ from types import SimpleNamespace from unittest.mock import patch +import httpx import pytest from openai.types.chat import ChatCompletion from openai.types.responses import ParsedResponseOutputMessage, ParsedResponseOutputText @@ -1015,6 +1016,53 @@ def test_embedding_exports_dimensions_and_count( } +@pytest.mark.parametrize( + "encoded_embedding, dimensions", + [ + ("AAAAAA==", 1), + ("AAAAAAAAAAA=", 2), + ("AAAAAAAAAAAAAAAA", 3), + pytest.param("!invalid!", 9, id="invalid-base64"), + pytest.param("\u2603", 1, id="non-ascii"), + ], +) +def test_base64_embedding_exports_vector_dimensions( + langfuse_memory_client, get_span, json_attr, encoded_embedding, dimensions +): + + def respond(request): + return httpx.Response( + 200, + json={ + "model": "text-embedding-3-small", + "object": "list", + "data": [ + {"index": 0, "object": "embedding", "embedding": encoded_embedding} + ], + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + + with lf_openai.OpenAI( + api_key="test", + http_client=httpx.Client(transport=httpx.MockTransport(respond)), + ) as openai_client: + result = openai_client.embeddings.create( + name="unit-openai-base64-embedding", + model="text-embedding-3-small", + input="hello world", + encoding_format="base64", + ) + + assert result.data[0].embedding == encoded_embedding + langfuse_memory_client.flush() + span = get_span("unit-openai-base64-embedding") + assert json_attr(span, LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT) == { + "dimensions": dimensions, + "count": 1, + } + + def _make_chat_response(**extra_response_fields): return SimpleNamespace( model="gpt-4o-mini",