diff --git a/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/__init__.py b/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/__init__.py index 76336b94f2..6f5742611d 100644 --- a/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/__init__.py +++ b/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/__init__.py @@ -57,6 +57,10 @@ "method": "chat", "span_name": "ollama.chat", }, + { + "method": "embed", + "span_name": "ollama.embeddings", + }, { "method": "embeddings", "span_name": "ollama.embeddings", @@ -266,7 +270,7 @@ def _llm_request_type_by_method(method_name): return LLMRequestTypeValues.CHAT elif method_name == "generate": return LLMRequestTypeValues.COMPLETION - elif method_name == "embeddings": + elif method_name in ("embed", "embeddings"): return LLMRequestTypeValues.EMBEDDING else: return LLMRequestTypeValues.UNKNOWN @@ -541,7 +545,7 @@ def _instrument(self, **kwargs): except (ImportError, AttributeError): # _copy_messages not available in older versions, skip it pass - # instrument all llm methods (generate/chat/embeddings) via _request dispatch wrapper + # instrument all llm methods (generate/chat/embed/embeddings) via _request dispatch wrapper wrap_function_wrapper( "ollama._client", "Client._request", diff --git a/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/event_emitter.py b/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/event_emitter.py index 32c4c35586..1e37763b20 100644 --- a/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/event_emitter.py +++ b/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/event_emitter.py @@ -66,12 +66,17 @@ def emit_message_events(llm_request_type, args, kwargs, event_logger): MessageEvent(content=content, role=role, tool_calls=tool_calls), event_logger, ) - elif ( - llm_request_type == LLMRequestTypeValues.COMPLETION - or LLMRequestTypeValues.EMBEDDING - ): + elif llm_request_type == LLMRequestTypeValues.COMPLETION: prompt = json_data.get("prompt", "") emit_event(MessageEvent(content=prompt, role="user"), event_logger) + elif llm_request_type == LLMRequestTypeValues.EMBEDDING: + prompt = json_data.get("prompt", "") + prompt = json_data.get("input", prompt) + if isinstance(prompt, (list, tuple)): + for prompt_content in prompt: + emit_event(MessageEvent(content=prompt_content, role="user"), event_logger) + else: + emit_event(MessageEvent(content=prompt, role="user"), event_logger) else: raise ValueError( "It wasn't possible to emit the input events due to an unknown llm_request_type." @@ -104,14 +109,22 @@ def emit_choice_events(llm_request_type, response: dict, event_logger): event_logger, ) elif llm_request_type == LLMRequestTypeValues.EMBEDDING: - emit_event( - ChoiceEvent( - index=0, - message={"content": response.get("embedding"), "role": "assistant"}, - finish_reason="unknown", - ), - event_logger, - ) + embedding = response.get("embedding") + if embedding is not None: + embeddings = [embedding] + else: + embeddings = response.get("embeddings") + if not isinstance(embeddings, (list, tuple)): + embeddings = [embeddings] + for index, embedding in enumerate(embeddings): + emit_event( + ChoiceEvent( + index=index, + message={"content": embedding, "role": "assistant"}, + finish_reason="unknown", + ), + event_logger, + ) else: raise ValueError( "It wasn't possible to emit the choice events due to an unknown llm_request_type." diff --git a/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/span_utils.py b/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/span_utils.py index e108f04c21..45fcbdffde 100644 --- a/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/span_utils.py +++ b/packages/opentelemetry-instrumentation-ollama/opentelemetry/instrumentation/ollama/span_utils.py @@ -1,4 +1,5 @@ import json +from collections.abc import Sequence from opentelemetry.instrumentation.ollama.utils import dont_throw, should_send_prompts from opentelemetry.semconv._incubating.attributes import ( @@ -17,6 +18,24 @@ def _set_span_attribute(span, name, value): return +def _set_prompt_attributes(span, prompt): + if isinstance(prompt, Sequence) and not isinstance(prompt, (str, bytes)): + for index, prompt_content in enumerate(prompt): + _set_span_attribute( + span, + f"{GenAIAttributes.GEN_AI_PROMPT}.{index}.role", + "user", + ) + _set_span_attribute( + span, + f"{GenAIAttributes.GEN_AI_PROMPT}.{index}.content", + prompt_content, + ) + else: + _set_span_attribute(span, f"{GenAIAttributes.GEN_AI_PROMPT}.0.role", "user") + _set_span_attribute(span, f"{GenAIAttributes.GEN_AI_PROMPT}.0.content", prompt) + + @dont_throw def set_input_attributes(span, llm_request_type, kwargs): if not span.is_recording(): @@ -40,11 +59,13 @@ def set_input_attributes(span, llm_request_type, kwargs): _set_prompts(span, json_data.get("messages")) if json_data.get("tools"): set_tools_attributes(span, json_data.get("tools")) + elif llm_request_type == LLMRequestTypeValues.EMBEDDING: + prompt = json_data.get("input") + if prompt is None: + prompt = json_data.get("prompt") + _set_prompt_attributes(span, prompt) else: - _set_span_attribute(span, f"{GenAIAttributes.GEN_AI_PROMPT}.0.role", "user") - _set_span_attribute( - span, f"{GenAIAttributes.GEN_AI_PROMPT}.0.content", json_data.get("prompt") - ) + _set_prompt_attributes(span, json_data.get("prompt")) @dont_throw diff --git a/packages/opentelemetry-instrumentation-ollama/tests/test_embeddings.py b/packages/opentelemetry-instrumentation-ollama/tests/test_embeddings.py index 4deba2e207..f029ea2671 100644 --- a/packages/opentelemetry-instrumentation-ollama/tests/test_embeddings.py +++ b/packages/opentelemetry-instrumentation-ollama/tests/test_embeddings.py @@ -1,4 +1,9 @@ +import importlib + +import ollama import pytest +from opentelemetry.instrumentation.ollama import OllamaInstrumentor +from opentelemetry.instrumentation.ollama.utils import TRACELOOP_TRACE_CONTENT from opentelemetry.sdk._logs import ReadableLogRecord from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, @@ -6,6 +11,300 @@ from opentelemetry.semconv_ai import SpanAttributes +EMBED_RESPONSE = {"embeddings": [[0.1, 0.2, 0.3]]} +BATCH_EMBED_RESPONSE = {"embeddings": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]} +LEGACY_EMBEDDING_RESPONSE = {"embedding": [0.1, 0.2, 0.3]} + + +def _mock_ollama_requests(monkeypatch, response=EMBED_RESPONSE): + client_module = importlib.import_module("ollama._client") + calls = [] + + def request(self, cls, *args, stream=False, **kwargs): + calls.append({"args": args, "kwargs": kwargs}) + return cls(**response) + + async def async_request(self, cls, *args, stream=False, **kwargs): + calls.append({"args": args, "kwargs": kwargs}) + return cls(**response) + + monkeypatch.setattr(client_module.Client, "_request", request) + monkeypatch.setattr(client_module.AsyncClient, "_request", async_request) + + return calls + + +def _instrument_ollama( + tracer_provider, meter_provider, logger_provider=None, use_legacy_attributes=True +): + instrumentor = OllamaInstrumentor(use_legacy_attributes=use_legacy_attributes) + instrument_kwargs = { + "tracer_provider": tracer_provider, + "meter_provider": meter_provider, + } + if logger_provider: + instrument_kwargs["logger_provider"] = logger_provider + instrumentor.instrument(**instrument_kwargs) + return instrumentor + + +def _assert_embed_request(calls, expected_input): + assert len(calls) == 1 + assert calls[0]["args"][1] == "/api/embed" + assert calls[0]["kwargs"]["json"]["input"] == expected_input + + +def _assert_embed_span(ollama_span, prompt_content=None): + assert ollama_span.name == "ollama.embeddings" + assert ollama_span.attributes.get(f"{GenAIAttributes.GEN_AI_SYSTEM}") == "Ollama" + assert ( + ollama_span.attributes.get(f"{SpanAttributes.LLM_REQUEST_TYPE}") + == "embedding" + ) + assert not ollama_span.attributes.get(f"{SpanAttributes.LLM_IS_STREAMING}") + assert ( + ollama_span.attributes.get(f"{GenAIAttributes.GEN_AI_REQUEST_MODEL}") + == "nomic-embed-text" + ) + if prompt_content is not None: + assert ( + ollama_span.attributes.get( + f"{GenAIAttributes.GEN_AI_PROMPT}.0.content" + ) + == prompt_content + ) + + +def test_ollama_embed_legacy( + monkeypatch, tracer_provider, meter_provider, span_exporter, log_exporter +): + calls = _mock_ollama_requests(monkeypatch) + instrumentor = _instrument_ollama(tracer_provider, meter_provider) + + try: + response = ollama.Client().embed( + model="nomic-embed-text", input="OpenTelemetry" + ) + finally: + instrumentor.uninstrument() + + assert response.embeddings == EMBED_RESPONSE["embeddings"] + _assert_embed_request(calls, "OpenTelemetry") + + spans = span_exporter.get_finished_spans() + ollama_span = spans[0] + _assert_embed_span(ollama_span, "OpenTelemetry") + + logs = log_exporter.get_finished_logs() + assert ( + len(logs) == 0 + ), "Assert that it doesn't emit logs when use_legacy_attributes is True" + + +def test_ollama_embed_multiple_inputs_legacy( + monkeypatch, tracer_provider, meter_provider, span_exporter +): + inputs = ["first text", "second text"] + calls = _mock_ollama_requests(monkeypatch) + instrumentor = _instrument_ollama(tracer_provider, meter_provider) + + try: + ollama.Client().embed(model="nomic-embed-text", input=inputs) + finally: + instrumentor.uninstrument() + + _assert_embed_request(calls, inputs) + + spans = span_exporter.get_finished_spans() + ollama_span = spans[0] + assert ollama_span.name == "ollama.embeddings" + assert ( + ollama_span.attributes.get(f"{SpanAttributes.LLM_REQUEST_TYPE}") + == "embedding" + ) + assert ( + ollama_span.attributes.get(f"{GenAIAttributes.GEN_AI_REQUEST_MODEL}") + == "nomic-embed-text" + ) + assert ( + ollama_span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.0.content") + == "first text" + ) + assert ( + ollama_span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.1.content") + == "second text" + ) + + +@pytest.mark.asyncio +async def test_ollama_async_embed_with_events_with_content( + monkeypatch, + tracer_provider, + logger_provider, + meter_provider, + span_exporter, + log_exporter, +): + monkeypatch.setenv(TRACELOOP_TRACE_CONTENT, "True") + calls = _mock_ollama_requests(monkeypatch) + instrumentor = _instrument_ollama( + tracer_provider, + meter_provider, + logger_provider=logger_provider, + use_legacy_attributes=False, + ) + + try: + response = await ollama.AsyncClient().embed( + model="nomic-embed-text", input="OpenTelemetry" + ) + finally: + instrumentor.uninstrument() + + assert response.embeddings == EMBED_RESPONSE["embeddings"] + _assert_embed_request(calls, "OpenTelemetry") + + spans = span_exporter.get_finished_spans() + ollama_span = spans[0] + _assert_embed_span(ollama_span) + + logs = log_exporter.get_finished_logs() + assert len(logs) == 2 + assert_message_in_logs( + logs[0], "gen_ai.user.message", {"content": "OpenTelemetry"} + ) + assert_message_in_logs( + logs[1], + "gen_ai.choice", + { + "index": 0, + "finish_reason": "unknown", + "message": {"content": EMBED_RESPONSE["embeddings"][0]}, + }, + ) + + +def test_ollama_embed_multiple_inputs_with_events_with_content( + monkeypatch, + tracer_provider, + logger_provider, + meter_provider, + span_exporter, + log_exporter, +): + monkeypatch.setenv(TRACELOOP_TRACE_CONTENT, "True") + inputs = ["first text", "second text"] + calls = _mock_ollama_requests(monkeypatch, response=BATCH_EMBED_RESPONSE) + instrumentor = _instrument_ollama( + tracer_provider, + meter_provider, + logger_provider=logger_provider, + use_legacy_attributes=False, + ) + + try: + response = ollama.Client().embed(model="nomic-embed-text", input=inputs) + finally: + instrumentor.uninstrument() + + assert response.embeddings == BATCH_EMBED_RESPONSE["embeddings"] + _assert_embed_request(calls, inputs) + + spans = span_exporter.get_finished_spans() + ollama_span = spans[0] + _assert_embed_span(ollama_span) + + logs = log_exporter.get_finished_logs() + assert len(logs) == 4 + + user_message_logs = [ + log for log in logs if log.log_record.event_name == "gen_ai.user.message" + ] + choice_logs = [log for log in logs if log.log_record.event_name == "gen_ai.choice"] + assert len(user_message_logs) == 2 + assert len(choice_logs) == 2 + + assert_message_in_logs( + user_message_logs[0], "gen_ai.user.message", {"content": inputs[0]} + ) + assert_message_in_logs( + user_message_logs[1], "gen_ai.user.message", {"content": inputs[1]} + ) + assert_message_in_logs( + choice_logs[0], + "gen_ai.choice", + { + "index": 0, + "finish_reason": "unknown", + "message": {"content": BATCH_EMBED_RESPONSE["embeddings"][0]}, + }, + ) + assert_message_in_logs( + choice_logs[1], + "gen_ai.choice", + { + "index": 1, + "finish_reason": "unknown", + "message": {"content": BATCH_EMBED_RESPONSE["embeddings"][1]}, + }, + ) + + +def test_ollama_embeddings_legacy_response_with_events_with_content( + monkeypatch, + tracer_provider, + logger_provider, + meter_provider, + span_exporter, + log_exporter, +): + monkeypatch.setenv(TRACELOOP_TRACE_CONTENT, "True") + calls = _mock_ollama_requests(monkeypatch, response=LEGACY_EMBEDDING_RESPONSE) + instrumentor = _instrument_ollama( + tracer_provider, + meter_provider, + logger_provider=logger_provider, + use_legacy_attributes=False, + ) + + try: + response = ollama.Client().embeddings( + model="nomic-embed-text", prompt="OpenTelemetry" + ) + finally: + instrumentor.uninstrument() + + assert response.embedding == LEGACY_EMBEDDING_RESPONSE["embedding"] + assert len(calls) == 1 + assert calls[0]["args"][1] == "/api/embeddings" + assert calls[0]["kwargs"]["json"]["prompt"] == "OpenTelemetry" + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + ollama_span = spans[0] + _assert_embed_span(ollama_span) + + logs = log_exporter.get_finished_logs() + user_message_logs = [ + log for log in logs if log.log_record.event_name == "gen_ai.user.message" + ] + choice_logs = [log for log in logs if log.log_record.event_name == "gen_ai.choice"] + assert len(user_message_logs) == 1 + assert len(choice_logs) == 1 + assert_message_in_logs( + user_message_logs[0], "gen_ai.user.message", {"content": "OpenTelemetry"} + ) + assert_message_in_logs( + choice_logs[0], + "gen_ai.choice", + { + "index": 0, + "finish_reason": "unknown", + "message": {"content": LEGACY_EMBEDDING_RESPONSE["embedding"]}, + }, + ) + + @pytest.mark.vcr def test_ollama_embeddings_legacy( instrument_legacy, ollama_client, span_exporter, log_exporter