monocle-apptrace 0.5.3__py3-none-any.whl → 0.6.6__py3-none-any.whl

monocle_apptrace/instrumentation/metamodel/llamaindex/_helper.py

@@ -20,6 +20,11 @@ from monocle_apptrace.instrumentation.common.utils import (
 from monocle_apptrace.instrumentation.metamodel.finish_types import map_llamaindex_finish_reason_to_finish_type
 
 LLAMAINDEX_AGENT_NAME_KEY = "_active_agent_name"
+
+# Thread-local storage for current agent context
+import threading
+_thread_local = threading.local()
+
 logger = logging.getLogger(__name__)
 
 def get_status(result):
@@ -64,18 +69,18 @@ def get_tool_description(arguments):
         return ""
 
 def extract_tool_args(arguments):
-    tool_args = []
+    tool_args = {}
     if len(arguments['args']) > 1:
         for key, value in arguments['args'][2].items():
             # check if value is builtin type or a string
             if value is not None and isinstance(value, (str, int, float, bool)):
-                tool_args.append({key, value})
+                tool_args[key] = value
     else:
         for key, value in arguments['kwargs'].items():
             # check if value is builtin type or a string
             if value is not None and isinstance(value, (str, int, float, bool)):
-                tool_args.append({key, value})
-    return [get_json_dumps(tool_arg) for tool_arg in tool_args]
+                tool_args[key] = value
+    return get_json_dumps(tool_args)
 
 def extract_tool_response(response):
     if hasattr(response, 'raw_output'):
@@ -96,12 +101,27 @@ def get_agent_description(instance) -> str:
         return instance.description
     return ""
 
-def get_source_agent(parent_span:Span) -> str:
-    source_agent_name = parent_span.attributes.get(LLAMAINDEX_AGENT_NAME_KEY, "")
-    if source_agent_name == "" and parent_span.name.startswith("llama_index.core.agent.ReActAgent."):
-        # Fallback to the agent name from the parent span if not set
-        source_agent_name = "ReactAgent"
-    return source_agent_name
+def get_name(instance):
+    return instance.name if hasattr(instance, 'name') else ""
+
+def set_current_agent(agent_name: str):
+    """Set the current agent name in thread-local storage."""
+    _thread_local.current_agent = agent_name
+
+def get_current_agent() -> str:
+    """Get the current agent name from thread-local storage."""
+    return getattr(_thread_local, 'current_agent', '')
+
+def get_source_agent() -> str:
+    """Get the name of the agent that initiated the request."""
+    source_agent = get_value(LLAMAINDEX_AGENT_NAME_KEY)
+    if source_agent is not None and isinstance(source_agent,str) and source_agent != "":
+        return source_agent
+
+    source_agent = get_current_agent()
+    if source_agent:
+        return source_agent
+    return ""
 
 def get_target_agent(results) -> str:
     if hasattr(results, 'raw_input'):

monocle_apptrace/instrumentation/metamodel/llamaindex/entities/agent.py

@@ -111,7 +111,7 @@ TOOLS = {
               {
                 "_comment": "name of the agent",
                 "attribute": "name",
-                "accessor": lambda arguments: _helper.get_source_agent(arguments['parent_span'])
+                "accessor": lambda arguments: _helper.get_source_agent()
               },          
               {
                 "_comment": "agent type",
@@ -157,12 +157,13 @@ AGENT_DELEGATION = {
               {
                 "_comment": "name of the agent",
                 "attribute": "from_agent",
-                "accessor": lambda arguments: _helper.get_source_agent(arguments['parent_span'])
+                "accessor": lambda arguments: _helper.get_source_agent()
               },
               {
                 "_comment": "name of the agent called",
                 "attribute": "to_agent",
-                "accessor": lambda arguments: _helper.get_target_agent(arguments['result'])
+                "accessor": lambda arguments: _helper.get_target_agent(arguments['result']),
+                "phase": "post_execution"
               }
         ]
       ]

monocle_apptrace/instrumentation/metamodel/llamaindex/llamaindex_processor.py

@@ -2,7 +2,7 @@ from opentelemetry.context import attach, detach, get_current, get_value, set_va
 from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
 from monocle_apptrace.instrumentation.metamodel.llamaindex._helper import (
-    is_delegation_tool, LLAMAINDEX_AGENT_NAME_KEY, get_agent_name
+    is_delegation_tool, LLAMAINDEX_AGENT_NAME_KEY, get_agent_name, get_name, set_current_agent
 )
 from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.agent import (
     AGENT_DELEGATION
@@ -13,14 +13,14 @@ TOOL_INVOCATION_STARTED:str = "llamaindex.tool_invocation_started"
 class DelegationHandler(SpanHandler):
     # LlamaIndex uses an internal tool to initate delegation to other agents. The method is tool invoke() with tool name as `transfer_to_<agent_name>`.
     # Hence we usea different output processor for tool invoke() to format the span as agentic.delegation.
-    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None, is_post_exec:bool= False) -> bool:
         if is_delegation_tool(args, instance):
             agent_request_wrapper = to_wrap.copy()
             agent_request_wrapper["output_processor"] = AGENT_DELEGATION
         else:
             agent_request_wrapper = to_wrap
 
-        return super().hydrate_span(agent_request_wrapper, wrapped, instance, args, kwargs, result, span, parent_span, ex)
+        return super().hydrate_span(agent_request_wrapper, wrapped, instance, args, kwargs, result, span, parent_span, ex, is_post_exec)
 
 # There are two different APIs for tool calling FunctionTool.call() and AgentWorkflow.tool_call(). In case of single agent calling tool, only the FunctionTool.call() is used. In case of multi agent case,
 # the AgentWorkflow.tool_call() is used which inturn calls FunctionTool.call(). We can't entirely rely on the FunctionTool.call() to extract tool span details, especially the agent delegation details are not available there.
@@ -29,7 +29,10 @@ class LlamaIndexToolHandler(DelegationHandler):
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         cur_context = get_current()
         cur_context = set_value(TOOL_INVOCATION_STARTED, True, cur_context)
-        return attach(cur_context)
+        current_agent = get_value(LLAMAINDEX_AGENT_NAME_KEY)
+        if current_agent is not None:
+            cur_context = set_value(LLAMAINDEX_AGENT_NAME_KEY, current_agent, cur_context)
+        return attach(cur_context), None
 
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token=None):
         if token:
@@ -44,8 +47,13 @@ class LlamaIndexSingleAgenttToolHandlerWrapper(DelegationHandler):
 class LlamaIndexAgentHandler(SpanHandler):
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         cur_context = get_current()
+        agent_name = get_name(instance)
+
+        # Set both OpenTelemetry context and thread-local storage
+        set_current_agent(agent_name)
+        cur_context = set_value(LLAMAINDEX_AGENT_NAME_KEY, agent_name, cur_context)
         cur_context = set_value(AGENT_PREFIX_KEY, "handoff", cur_context)
-        return attach(cur_context)
+        return attach(cur_context), None
 
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token=None):
         if token:
@@ -53,8 +61,8 @@ class LlamaIndexAgentHandler(SpanHandler):
     
     # LlamaIndex uses direct OpenAI call for agent inferences. Given that the workflow type is set to llamaindex, the openAI inference does not record the input/output events.
     # To avoid this, we set the workflow type to generic for agent inference spans so we can capture the prompts and responses.
-    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
-        retval = super().hydrate_span(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex)
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None, is_post_exec:bool= False) -> bool:
+        retval = super().hydrate_span(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex, is_post_exec)
         if SpanHandler.is_root_span(parent_span):
             span.set_attribute(LLAMAINDEX_AGENT_NAME_KEY, "")
         else:

monocle_apptrace/instrumentation/metamodel/llamaindex/methods.py

@@ -89,17 +89,10 @@ LLAMAINDEX_METHODS = [
     {
         "package": "llama_index.core.agent",
         "object": "ReActAgent",
-        "method": "chat",
+        "method": "run",
         "wrapper_method": task_wrapper,
         "output_processor": AGENT
     },
-    {
-        "package": "llama_index.core.agent",
-        "object": "ReActAgent",
-        "method": "achat",
-        "wrapper_method": atask_wrapper,
-        "output_processor": AGENT
-    },
     {
         "package": "llama_index.core.agent.workflow.function_agent",
         "object": "FunctionAgent",

monocle_apptrace/instrumentation/metamodel/mcp/_helper.py

@@ -22,7 +22,7 @@ def get_output_text(arguments):
         for tool in arguments["result"].tools:
             if hasattr(tool, "name"):
                 tools.append(tool.name)
-        return tools
+        return ", ".join(tools)  
     if (
         "result" in arguments
         and hasattr(arguments["result"], "content")
@@ -32,12 +32,12 @@ def get_output_text(arguments):
         for content in arguments["result"].content:
             if hasattr(content, "text"):
                 ret_val.append(content.text)
-        return ret_val
+        return " ".join(ret_val)
 
 
 def get_name(arguments):
     """Get the name of the tool from the instance."""
-    if 'parent_span' in arguments:
+    if 'parent_span' in arguments and arguments['parent_span']:
         arguments['parent_span'].set_attribute("is_mcp", True)
     args = arguments["args"]
     if (
@@ -63,17 +63,18 @@ def get_params_arguments(arguments):
 
     args = arguments["args"]
     if (
-        args
+       args
         and hasattr(args[0], "root")
         and hasattr(args[0].root, "params")
         and hasattr(args[0].root.params, "arguments")
     ):
         # If the first argument has a root with params and arguments, return those arguments
         try:
-            return json.dumps(args[0].root.params.arguments)
+            return [json.dumps(args[0].root.params.arguments)]
         except (TypeError, ValueError) as e:
             logger.error(f"Error serializing arguments: {e}")
-            return str(args[0].root.params.arguments)
+            return [str(args[0].root.params.arguments)]
+    return []
 
 
 def get_url(arguments):

monocle_apptrace/instrumentation/metamodel/mistral/_helper.py

@@ -15,16 +15,43 @@ from monocle_apptrace.instrumentation.common.utils import (
     try_option,
     get_exception_message,
 )
-from monocle_apptrace.instrumentation.metamodel.finish_types import map_anthropic_finish_reason_to_finish_type
+from monocle_apptrace.instrumentation.metamodel.finish_types import map_mistral_finish_reason_to_finish_type
 from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY, INFERENCE_AGENT_DELEGATION, INFERENCE_TURN_END, INFERENCE_TOOL_CALL
 
-
 logger = logging.getLogger(__name__)
 
+
 def extract_provider_name(instance):
     provider_url: Option[str] = try_option(getattr, instance._client.base_url, 'host')
     return provider_url.unwrap_or(None)
 
+def update_input_span_events(kwargs):
+    """Extract embedding input for spans"""
+    if "inputs" in kwargs and isinstance(kwargs["inputs"], list):
+        # Join multiple strings into one
+        return " | ".join(kwargs["inputs"])
+    elif "inputs" in kwargs and isinstance(kwargs["inputs"], str):
+        return kwargs["inputs"]
+    return ""
+
+def update_output_span_events(results):
+    """Extract embedding output for spans"""
+    try:
+        if hasattr(results, "data") and isinstance(results.data, list):
+            embeddings = results.data
+            # just return the indices, not full vectors
+            embedding_summaries = [
+                f"index={e.index}, dim={len(e.embedding)}"
+                for e in embeddings
+            ]
+            output = "\n".join(embedding_summaries)
+            if len(output) > 200:
+                output = output[:200] + "..."
+            return output
+    except Exception as e:
+        logger.warning("Error in update_output_span_events: %s", str(e))
+    return ""
+
 def extract_inference_endpoint(instance):
     inference_endpoint: Option[str] = try_option(getattr, instance._client, 'base_url').map(str)
     if inference_endpoint.is_none() and "meta" in instance.client.__dict__:
@@ -32,16 +59,18 @@ def extract_inference_endpoint(instance):
 
     return inference_endpoint.unwrap_or(extract_provider_name(instance))
 
-def dummy_method(arguents):
+
+def dummy_method(arguments):
     pass
 
+
 def extract_messages(kwargs):
     """Extract system and user messages"""
     try:
         messages = []
         if "system" in kwargs and isinstance(kwargs["system"], str):
             messages.append({"system": kwargs["system"]})
-        if 'messages' in kwargs and len(kwargs['messages']) >0:
+        if 'messages' in kwargs and kwargs['messages']:
             for msg in kwargs['messages']:
                 if msg.get('content') and msg.get('role'):
                     messages.append({msg['role']: msg['content']})
@@ -50,6 +79,7 @@ def extract_messages(kwargs):
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
         return []
 
+
 def get_exception_status_code(arguments):
     exc = arguments.get("exception")
     if exc is not None and hasattr(exc, "status_code"):
@@ -73,7 +103,7 @@ def extract_assistant_message(arguments):
     Returns a JSON string like {"assistant": "<text>"}.
     """
     try:
-        result = arguments.get("result")
+        result = arguments.get("result") if isinstance(arguments, dict) else arguments
         if result is None:
             return ""
 
@@ -86,9 +116,10 @@ def extract_assistant_message(arguments):
         if isinstance(result, list):
             content = []
             for chunk in result:
-                # Each chunk may have delta attribute
-                if hasattr(chunk, "delta") and hasattr(chunk.delta, "content"):
-                    content.append(chunk.delta.content or "")
+                if hasattr(chunk, "data") and hasattr(chunk.data, "choices") and chunk.data.choices:
+                    choice = chunk.data.choices[0]
+                    if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
+                        content.append(choice.delta.content or "")
             return get_json_dumps({"assistant": "".join(content)})
 
         return ""
@@ -98,77 +129,95 @@ def extract_assistant_message(arguments):
         return ""
 
 
-
-def update_span_from_llm_response(response):
+'''def update_span_from_llm_response(response):
     meta_dict = {}
     if response is not None and hasattr(response, "usage"):
-        if hasattr(response, "usage") and response.usage is not None:
-            token_usage = response.usage
-        else:
-            response_metadata = response.response_metadata
-            token_usage = response_metadata.get("token_usage")
+        token_usage = getattr(response, "usage", None) or getattr(response, "response_metadata", {}).get("token_usage")
         if token_usage is not None:
             meta_dict.update({"completion_tokens": getattr(response.usage, "output_tokens", 0)})
             meta_dict.update({"prompt_tokens": getattr(response.usage, "input_tokens", 0)})
-            meta_dict.update({"total_tokens": getattr(response.usage, "input_tokens", 0)+getattr(response.usage, "output_tokens", 0)})
-    return meta_dict
+            meta_dict.update({"total_tokens": getattr(response.usage, "input_tokens", 0) + getattr(response.usage, "output_tokens", 0)})
+    return meta_dict'''
+
+def update_span_from_llm_response(result, include_token_counts=False):
+    tokens = {
+        "completion_tokens": getattr(result, "completion_tokens", 0),
+        "prompt_tokens": getattr(result, "prompt_tokens", 0),
+        "total_tokens": getattr(result, "total_tokens", 0),
+    } if include_token_counts else {}
+    # Add other metadata fields like finish_reason, etc.
+    return {**tokens}
+
 
 def extract_finish_reason(arguments):
-    """Extract stop_reason from Anthropic response (Claude)."""
+    """
+    Extract stop_reason from a Mistral response or stream chunks.
+    Works for both streaming (list of chunks) and full responses.
+    """
     try:
-        # Arguments may be a dict with 'result' or just the response object
         response = arguments.get("result") if isinstance(arguments, dict) else arguments
-        if response is not None and hasattr(response, "stop_reason"):
+        if response is None:
+            return None
+
+        # Handle full response: single object with stop_reason
+        if hasattr(response, "stop_reason") and response.stop_reason:
             return response.stop_reason
+
+        # Handle streaming: list of chunks, last chunk may have finish_reason
+        if isinstance(response, list):
+            for chunk in reversed(response):
+                if hasattr(chunk, "data") and hasattr(chunk.data, "choices") and chunk.data.choices:
+                    fr = getattr(chunk.data.choices[0], "finish_reason", None)
+                    if fr is not None:
+                        return fr
+
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_finish_reason: %s", str(e))
         return None
+
     return None
 
+
 def map_finish_reason_to_finish_type(finish_reason):
-    """Map Anthropic stop_reason to finish_type, similar to OpenAI mapping."""
-    return map_anthropic_finish_reason_to_finish_type(finish_reason)
+    """Map Mistral stop_reason to finish_type, similar to OpenAI mapping."""
+    return map_mistral_finish_reason_to_finish_type(finish_reason)
+
 
 def agent_inference_type(arguments):
-    """Extract agent inference type from Anthropic response"""
+    """Extract agent inference type from Mistral response"""
     try:
         status = get_status_code(arguments)
-        if status == 'success' or status == 'completed':
-            response = arguments["result"]
-            
+        if status in ('success', 'completed'):
+            response = arguments.get("result")
+            if response is None:
+                return INFERENCE_TURN_END
+
             # Check if stop_reason indicates tool use
-            if hasattr(response, "stop_reason") and response.stop_reason == "tool_use":
-                # Check if this is agent delegation by looking at tool names
-                if hasattr(response, "content") and response.content:
-                    agent_prefix = get_value(AGENT_PREFIX_KEY)
-                    for content_block in response.content:
-                        if (hasattr(content_block, "type") and 
-                            content_block.type == "tool_use" and
-                            hasattr(content_block, "name")):
-                            tool_name = content_block.name
-                            if agent_prefix and tool_name.startswith(agent_prefix):
-                                return INFERENCE_AGENT_DELEGATION
-                    # If we found tool use but no agent delegation, it's a regular tool call
-                    return INFERENCE_TOOL_CALL
-            
+            stop_reason = getattr(response, "stop_reason", None)
+            if stop_reason == "tool_use" and hasattr(response, "content") and response.content:
+                agent_prefix = get_value(AGENT_PREFIX_KEY)
+                for content_block in response.content:
+                    if getattr(content_block, "type", None) == "tool_use" and hasattr(content_block, "name"):
+                        if agent_prefix and content_block.name.startswith(agent_prefix):
+                            return INFERENCE_AGENT_DELEGATION
+                return INFERENCE_TOOL_CALL
+
             # Fallback: check the extracted message for tool content
             assistant_message = extract_assistant_message(arguments)
             if assistant_message:
                 try:
                     message = json.loads(assistant_message)
-                    if message and isinstance(message, dict):
-                        assistant_content = message.get("assistant", "")
-                        if assistant_content:
-                            agent_prefix = get_value(AGENT_PREFIX_KEY)
-                            if agent_prefix and agent_prefix in assistant_content:
-                                return INFERENCE_AGENT_DELEGATION
+                    assistant_content = message.get("assistant", "") if isinstance(message, dict) else ""
+                    agent_prefix = get_value(AGENT_PREFIX_KEY)
+                    if agent_prefix and agent_prefix in assistant_content:
+                        return INFERENCE_AGENT_DELEGATION
                 except (json.JSONDecodeError, TypeError):
-                    # If JSON parsing fails, fall back to string analysis
                     agent_prefix = get_value(AGENT_PREFIX_KEY)
                     if agent_prefix and agent_prefix in assistant_message:
                         return INFERENCE_AGENT_DELEGATION
-        
+
         return INFERENCE_TURN_END
+
     except Exception as e:
         logger.warning("Warning: Error occurred in agent_inference_type: %s", str(e))
-        return INFERENCE_TURN_END
+        return INFERENCE_TURN_END

monocle_apptrace/instrumentation/metamodel/mistral/entities/inference.py

@@ -4,13 +4,13 @@ from monocle_apptrace.instrumentation.common.utils import get_error_message, res
 
 MISTRAL_INFERENCE = {
     "type": SPAN_TYPES.INFERENCE,
+    "subtype": lambda arguments: _helper.agent_inference_type(arguments),
     "attributes": [
         [
             {
                 "_comment": "provider type ,name , deployment , inference_endpoint",
                 "attribute": "type",
                 "accessor": lambda arguments: 'inference.mistral'
-
             },
             {
                 "attribute": "provider_name",
@@ -52,9 +52,16 @@ MISTRAL_INFERENCE = {
                     "accessor": lambda arguments: get_error_message(arguments)
                 },
                 {
-                    "_comment": "this is result from LLM",
+                    "_comment": "this is result from LLM, works for streaming and non-streaming",
                     "attribute": "response",
-                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments)
+                    "accessor": lambda arguments: (
+                        # Handle streaming: combine chunks if result is iterable and doesn't have 'choices'
+                        _helper.extract_assistant_message(
+                            {"result": list(arguments["result"])}
+                            if hasattr(arguments.get("result"), "__iter__") and not hasattr(arguments.get("result"), "choices")
+                            else arguments
+                        )
+                    )
                 }
             ]
         },
@@ -62,8 +69,11 @@ MISTRAL_INFERENCE = {
             "name": "metadata",
             "attributes": [
                 {
-                    "_comment": "this is metadata usage from LLM",
-                    "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'])
+                    "_comment": "this is metadata usage from LLM, includes token counts",
+                    "accessor": lambda arguments: _helper.update_span_from_llm_response(
+                        arguments.get("result"),
+                        include_token_counts=True  # new flag for streaming handling
+                    )
                 },
                 {
                     "_comment": "finish reason from Anthropic response",
@@ -74,10 +84,6 @@ MISTRAL_INFERENCE = {
                     "_comment": "finish type mapped from finish reason",
                     "attribute": "finish_type",
                     "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(_helper.extract_finish_reason(arguments))
-                },
-                {
-                    "attribute": "inference_sub_type",
-                    "accessor": lambda arguments: _helper.agent_inference_type(arguments)
                 }
             ]
         }

monocle_apptrace/instrumentation/metamodel/mistral/entities/retrieval.py

@@ -0,0 +1,41 @@
+from monocle_apptrace.instrumentation.metamodel.mistral import _helper
+from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
+
+MISTRAL_RETRIEVAL = {
+    "type": "embedding",
+    "attributes": [
+        [
+            {
+                "_comment": "LLM Model",
+                "attribute": "name",
+                "accessor": lambda arguments: resolve_from_alias(arguments['kwargs'], ['model'])
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: 'model.embedding.' + resolve_from_alias(arguments['kwargs'], ['model'])
+            }
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "embedding input",
+                    "attribute": "input",
+                    "accessor": lambda arguments: _helper.update_input_span_events(arguments["kwargs"])
+                }
+            ]
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "embedding output summary",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.update_output_span_events(arguments["result"])
+                }
+            ]
+        }
+    ]
+}

monocle_apptrace/instrumentation/metamodel/mistral/methods.py

@@ -1,5 +1,6 @@
 from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, atask_wrapper
 from monocle_apptrace.instrumentation.metamodel.mistral.entities.inference import MISTRAL_INFERENCE
+from monocle_apptrace.instrumentation.metamodel.mistral.entities.retrieval import MISTRAL_RETRIEVAL
 
 MISTRAL_METHODS = [
     {
@@ -33,7 +34,15 @@ MISTRAL_METHODS = [
         "span_handler": "non_framework_handler",
         "wrapper_method": atask_wrapper,
         "output_processor": MISTRAL_INFERENCE,
-    }
+    },
+    {
+        "package": "mistralai.embeddings",    # where Embeddings is defined
+        "object": "Embeddings",               # sync embeddings client
+        "method": "create",                   # sync create
+        "span_handler": "non_framework_handler",
+        "wrapper_method": task_wrapper,
+        "output_processor": MISTRAL_RETRIEVAL
+    },
 ]