fix(ollama): instrument embed API

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

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

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

packages/opentelemetry-instrumentation-ollama/tests/test_embeddings.py

@@ -1,11 +1,254 @@
+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,
 )
 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]]}
+
+
+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]},
+        },
+    )
+
+
 @pytest.mark.vcr
 def test_ollama_embeddings_legacy(
     instrument_legacy, ollama_client, span_exporter, log_exporter