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