diff --git a/llama-index-core/llama_index/core/instrumentation/dispatcher.py b/llama-index-core/llama_index/core/instrumentation/dispatcher.py
index 10d8cecb896897bd5c317ad8d100ae130c1acd2e..727d7cba827ce434448bceba5fb6103b4ca60ee6 100644
--- a/llama-index-core/llama_index/core/instrumentation/dispatcher.py
+++ b/llama-index-core/llama_index/core/instrumentation/dispatcher.py
@@ -1,5 +1,4 @@
from typing import Any, List, Optional, Dict
-import functools
import inspect
import uuid
from llama_index.core.bridge.pydantic import BaseModel, Field
@@ -9,6 +8,7 @@ from llama_index.core.instrumentation.span_handlers import (
NullSpanHandler,
)
from llama_index.core.instrumentation.events.base import BaseEvent
+import wrapt
class Dispatcher(BaseModel):
@@ -58,68 +58,119 @@ class Dispatcher(BaseModel):
else:
c = c.parent
- def span_enter(self, *args, id: str, **kwargs) -> None:
- """Send notice to handlers that a span with id has started."""
+ def span_enter(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ **kwargs: Any,
+ ) -> None:
+ """Send notice to handlers that a span with id_ has started."""
c = self
while c:
for h in c.span_handlers:
- h.span_enter(*args, id=id, **kwargs)
+ h.span_enter(
+ *args,
+ id_=id_,
+ bound_args=bound_args,
+ instance=instance,
+ **kwargs,
+ )
if not c.propagate:
c = None
else:
c = c.parent
- def span_drop(self, *args, id: str, err: Optional[Exception], **kwargs) -> None:
- """Send notice to handlers that a span with id is being dropped."""
+ def span_drop(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ err: Optional[BaseException] = None,
+ **kwargs: Any,
+ ) -> None:
+ """Send notice to handlers that a span with id_ is being dropped."""
c = self
while c:
for h in c.span_handlers:
- h.span_drop(*args, id=id, err=err, **kwargs)
+ h.span_drop(
+ *args,
+ id_=id_,
+ bound_args=bound_args,
+ instance=instance,
+ err=err,
+ **kwargs,
+ )
if not c.propagate:
c = None
else:
c = c.parent
- def span_exit(self, *args, id: str, result: Optional[Any] = None, **kwargs) -> None:
- """Send notice to handlers that a span with id is exiting."""
+ def span_exit(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ result: Optional[Any] = None,
+ **kwargs: Any,
+ ) -> None:
+ """Send notice to handlers that a span with id_ is exiting."""
c = self
while c:
for h in c.span_handlers:
- h.span_exit(*args, id=id, result=result, **kwargs)
+ h.span_exit(
+ *args,
+ id_=id_,
+ bound_args=bound_args,
+ instance=instance,
+ result=result,
+ **kwargs,
+ )
if not c.propagate:
c = None
else:
c = c.parent
def span(self, func):
- @functools.wraps(func)
- def wrapper(*args, **kwargs):
- id = f"{func.__qualname__}-{uuid.uuid4()}"
- self.span_enter(*args, id=id, **kwargs)
+ @wrapt.decorator
+ def wrapper(func, instance, args, kwargs):
+ bound_args = inspect.signature(func).bind(*args, **kwargs)
+ id_ = f"{func.__qualname__}-{uuid.uuid4()}"
+ self.span_enter(id_=id_, bound_args=bound_args, instance=instance)
try:
result = func(*args, **kwargs)
- except Exception as e:
- self.span_drop(*args, id=id, err=e, **kwargs)
+ except BaseException as e:
+ self.span_drop(id_=id_, bound_args=bound_args, instance=instance, err=e)
+ raise
else:
- self.span_exit(*args, id=id, result=result, **kwargs)
+ self.span_exit(
+ id_=id_, bound_args=bound_args, instance=instance, result=result
+ )
return result
- @functools.wraps(func)
- async def async_wrapper(*args, **kwargs):
- id = f"{func.__qualname__}-{uuid.uuid4()}"
- self.span_enter(*args, id=id, **kwargs)
+ @wrapt.decorator
+ async def async_wrapper(func, instance, args, kwargs):
+ bound_args = inspect.signature(func).bind(*args, **kwargs)
+ id_ = f"{func.__qualname__}-{uuid.uuid4()}"
+ self.span_enter(id_=id_, bound_args=bound_args, instance=instance)
try:
result = await func(*args, **kwargs)
- except Exception as e:
- self.span_drop(*args, id=id, err=e, **kwargs)
+ except BaseException as e:
+ self.span_drop(id_=id_, bound_args=bound_args, instance=instance, err=e)
+ raise
else:
- self.span_exit(*args, id=id, result=result, **kwargs)
+ self.span_exit(
+ id_=id_, bound_args=bound_args, instance=instance, result=result
+ )
return result
if inspect.iscoroutinefunction(func):
- return async_wrapper
+ return async_wrapper(func)
else:
- return wrapper
+ return wrapper(func)
@property
def log_name(self) -> str:
diff --git a/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py b/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py
index a54a8d7b30b79906ddcf4834718d0368ff924a68..78beb7c2956b29da9625f771725b3b528dd8483b 100644
--- a/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py
+++ b/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py
@@ -1,3 +1,4 @@
+import inspect
from abc import abstractmethod
from typing import Any, Dict, Generic, Optional, TypeVar
@@ -22,52 +23,100 @@ class BaseSpanHandler(BaseModel, Generic[T]):
"""Class name."""
return "BaseSpanHandler"
- def span_enter(self, *args, id: str, **kwargs) -> None:
+ def span_enter(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ **kwargs: Any,
+ ) -> None:
"""Logic for entering a span."""
- if id in self.open_spans:
+ if id_ in self.open_spans:
pass # should probably raise an error here
else:
# TODO: thread safe?
span = self.new_span(
- *args, id=id, parent_span_id=self.current_span_id, **kwargs
+ id_=id_,
+ bound_args=bound_args,
+ instance=instance,
+ parent_span_id=self.current_span_id,
)
if span:
- self.open_spans[id] = span
- self.current_span_id = id
+ self.open_spans[id_] = span
+ self.current_span_id = id_
- def span_exit(self, *args, id: str, result: Optional[Any] = None, **kwargs) -> None:
+ def span_exit(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ result: Optional[Any] = None,
+ **kwargs: Any,
+ ) -> None:
"""Logic for exiting a span."""
- span = self.prepare_to_exit_span(*args, id=id, result=result, **kwargs)
+ span = self.prepare_to_exit_span(
+ id_=id_, bound_args=bound_args, instance=instance, result=result
+ )
if span:
- if self.current_span_id == id:
- self.current_span_id = self.open_spans[id].parent_id
- del self.open_spans[id]
+ if self.current_span_id == id_:
+ self.current_span_id = self.open_spans[id_].parent_id
+ del self.open_spans[id_]
- def span_drop(self, *args, id: str, err: Optional[Exception], **kwargs) -> None:
+ def span_drop(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ err: Optional[BaseException] = None,
+ **kwargs: Any,
+ ) -> None:
"""Logic for dropping a span i.e. early exit."""
- span = self.prepare_to_drop_span(*args, id=id, err=err, **kwargs)
+ span = self.prepare_to_drop_span(
+ id_=id_, bound_args=bound_args, instance=instance, err=err
+ )
if span:
- if self.current_span_id == id:
- self.current_span_id = self.open_spans[id].parent_id
- del self.open_spans[id]
+ if self.current_span_id == id_:
+ self.current_span_id = self.open_spans[id_].parent_id
+ del self.open_spans[id_]
@abstractmethod
def new_span(
- self, *args, id: str, parent_span_id: Optional[str], **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ parent_span_id: Optional[str] = None,
+ **kwargs: Any,
) -> Optional[T]:
"""Create a span."""
...
@abstractmethod
def prepare_to_exit_span(
- self, *args, id: str, result: Optional[Any] = None, **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ result: Optional[Any] = None,
+ **kwargs: Any,
) -> Optional[T]:
"""Logic for preparing to exit a span."""
...
@abstractmethod
def prepare_to_drop_span(
- self, *args, id: str, err: Optional[Exception], **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ err: Optional[BaseException] = None,
+ **kwargs: Any,
) -> Optional[T]:
"""Logic for preparing to drop a span."""
...
diff --git a/llama-index-core/llama_index/core/instrumentation/span_handlers/null.py b/llama-index-core/llama_index/core/instrumentation/span_handlers/null.py
index fccc12c42c943a44567c78796073693964a4d005..a4c6000ebcd0ac271d517759fb566976662002b5 100644
--- a/llama-index-core/llama_index/core/instrumentation/span_handlers/null.py
+++ b/llama-index-core/llama_index/core/instrumentation/span_handlers/null.py
@@ -1,3 +1,4 @@
+import inspect
from typing import Optional, Any
from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler
from llama_index.core.instrumentation.span.base import BaseSpan
@@ -9,26 +10,61 @@ class NullSpanHandler(BaseSpanHandler[BaseSpan]):
"""Class name."""
return "NullSpanHandler"
- def span_enter(self, *args, id: str, **kwargs) -> None:
+ def span_enter(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ **kwargs: Any
+ ) -> None:
"""Logic for entering a span."""
return
- def span_exit(self, *args, id: str, result: Optional[Any], **kwargs) -> None:
+ def span_exit(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ result: Optional[Any] = None,
+ **kwargs: Any
+ ) -> None:
"""Logic for exiting a span."""
return
- def new_span(self, *args, id: str, parent_span_id: Optional[str], **kwargs) -> None:
+ def new_span(
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ parent_span_id: Optional[str] = None,
+ **kwargs: Any
+ ) -> None:
"""Create a span."""
return
def prepare_to_exit_span(
- self, *args, id: str, result: Optional[Any] = None, **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ result: Optional[Any] = None,
+ **kwargs: Any
) -> None:
"""Logic for exiting a span."""
return
def prepare_to_drop_span(
- self, *args, id: str, err: Optional[Exception], **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ err: Optional[BaseException] = None,
+ **kwargs: Any
) -> None:
"""Logic for droppping a span."""
if err:
diff --git a/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py b/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
index 96c54b2fb2982f8389b5b0294d9f3868ffee7358..18f9bddc4f27d1254e712728084181718b9969bf 100644
--- a/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
+++ b/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
@@ -1,3 +1,4 @@
+import inspect
from typing import Any, cast, List, Optional, TYPE_CHECKING
from llama_index.core.bridge.pydantic import Field
from llama_index.core.instrumentation.span.simple import SimpleSpan
@@ -21,16 +22,28 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
return "SimpleSpanHandler"
def new_span(
- self, *args, id: str, parent_span_id: Optional[str], **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ parent_span_id: Optional[str] = None,
+ **kwargs: Any,
) -> SimpleSpan:
"""Create a span."""
- return SimpleSpan(id_=id, parent_id=parent_span_id)
+ return SimpleSpan(id_=id_, parent_id=parent_span_id)
def prepare_to_exit_span(
- self, *args, id: str, result: Optional[Any] = None, **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ result: Optional[Any] = None,
+ **kwargs: Any,
) -> SimpleSpan:
"""Logic for preparing to drop a span."""
- span = self.open_spans[id]
+ span = self.open_spans[id_]
span = cast(SimpleSpan, span)
span.end_time = datetime.now()
span.duration = (span.end_time - span.start_time).total_seconds()
@@ -38,11 +51,17 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
return span
def prepare_to_drop_span(
- self, *args, id: str, err: Optional[Exception], **kwargs
+ self,
+ *args: Any,
+ id_: str,
+ bound_args: inspect.BoundArguments,
+ instance: Optional[Any] = None,
+ err: Optional[BaseException] = None,
+ **kwargs: Any,
) -> SimpleSpan:
"""Logic for droppping a span."""
- if id in self.open_spans:
- return self.open_spans[id]
+ if id_ in self.open_spans:
+ return self.open_spans[id_]
return None
def _get_trace_trees(self) -> List["Tree"]:
diff --git a/llama-index-core/llama_index/core/llms/llm.py b/llama-index-core/llama_index/core/llms/llm.py
index 7f3e8ae623fe58844aaac34ce61a9b6e55c42537..678d9dfcea89b500c1c5a1c9a334578d5e95dea3 100644
--- a/llama-index-core/llama_index/core/llms/llm.py
+++ b/llama-index-core/llama_index/core/llms/llm.py
@@ -286,6 +286,7 @@ class LLM(BaseLLM):
# -- Structured outputs --
+ @dispatcher.span
def structured_predict(
self,
output_cls: BaseModel,
@@ -331,6 +332,7 @@ class LLM(BaseLLM):
return program(**prompt_args)
+ @dispatcher.span
async def astructured_predict(
self,
output_cls: BaseModel,
@@ -418,6 +420,7 @@ class LLM(BaseLLM):
dispatcher.event(LLMPredictEndEvent())
return self._parse_output(output)
+ @dispatcher.span
def stream(
self,
prompt: BasePromptTemplate,
@@ -500,6 +503,7 @@ class LLM(BaseLLM):
dispatcher.event(LLMPredictEndEvent())
return self._parse_output(output)
+ @dispatcher.span
async def astream(
self,
prompt: BasePromptTemplate,
@@ -544,6 +548,7 @@ class LLM(BaseLLM):
return stream_tokens
+ @dispatcher.span
def predict_and_call(
self,
tools: List["BaseTool"],
@@ -595,6 +600,7 @@ class LLM(BaseLLM):
return output
+ @dispatcher.span
async def apredict_and_call(
self,
tools: List["BaseTool"],
diff --git a/llama-index-core/pyproject.toml b/llama-index-core/pyproject.toml
index 50974ee5da5c3249f23b1c57b9cf2104d9dc33f4..e91bf83c830120a5cff0bdd6484edb31154547b3 100644
--- a/llama-index-core/pyproject.toml
+++ b/llama-index-core/pyproject.toml
@@ -84,6 +84,7 @@ tqdm = "^4.66.1"
pillow = ">=9.0.0"
PyYAML = ">=6.0.1"
llamaindex-py-client = "^0.1.15"
+wrapt = "*"
[tool.poetry.extras]
gradientai = [
diff --git a/llama-index-core/tests/instrumentation/test_dispatcher.py b/llama-index-core/tests/instrumentation/test_dispatcher.py
index 4805a959bb09fdc7db9d346c47cb27bc59998f15..c256719b4b1f857cead0b68396510f0bbce76d07 100644
--- a/llama-index-core/tests/instrumentation/test_dispatcher.py
+++ b/llama-index-core/tests/instrumentation/test_dispatcher.py
@@ -1,3 +1,6 @@
+import inspect
+from asyncio import CancelledError
+
import pytest
import llama_index.core.instrumentation as instrument
from llama_index.core.instrumentation.dispatcher import Dispatcher
@@ -5,6 +8,9 @@ from unittest.mock import patch, MagicMock
dispatcher = instrument.get_dispatcher("test")
+value_error = ValueError("value error")
+cancelled_error = CancelledError("cancelled error")
+
@dispatcher.span
def func(*args, a, b=3, **kwargs):
@@ -16,6 +22,34 @@ async def async_func(*args, a, b=3, **kwargs):
return a + b
+@dispatcher.span
+def func_exc(*args, a, b=3, c=4, **kwargs):
+ raise value_error
+
+
+@dispatcher.span
+async def async_func_exc(*args, a, b=3, c=4, **kwargs):
+ raise cancelled_error
+
+
+class _TestObject:
+ @dispatcher.span
+ def func(self, *args, a, b=3, **kwargs):
+ return a + b
+
+ @dispatcher.span
+ async def async_func(self, *args, a, b=3, **kwargs):
+ return a + b
+
+ @dispatcher.span
+ def func_exc(self, *args, a, b=3, c=4, **kwargs):
+ raise value_error
+
+ @dispatcher.span
+ async def async_func_exc(self, *args, a, b=3, c=4, **kwargs):
+ raise cancelled_error
+
+
@patch.object(Dispatcher, "span_exit")
@patch.object(Dispatcher, "span_enter")
@patch("llama_index.core.instrumentation.dispatcher.uuid")
@@ -29,60 +63,129 @@ def test_dispatcher_span_args(mock_uuid, mock_span_enter, mock_span_exit):
# assert
# span_enter
span_id = f"{func.__qualname__}-mock"
+ bound_args = inspect.signature(func).bind(1, 2, a=3, c=5)
+ mock_span_enter.assert_called_once()
+ args, kwargs = mock_span_enter.call_args
+ assert args == ()
+ assert kwargs == {"id_": span_id, "bound_args": bound_args, "instance": None}
+
+ # span_exit
+ args, kwargs = mock_span_exit.call_args
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": None,
+ "result": result,
+ }
+
+
+@patch.object(Dispatcher, "span_exit")
+@patch.object(Dispatcher, "span_enter")
+@patch("llama_index.core.instrumentation.dispatcher.uuid")
+def test_dispatcher_span_args_with_instance(mock_uuid, mock_span_enter, mock_span_exit):
+ # arrange
+ mock_uuid.uuid4.return_value = "mock"
+
+ # act
+ instance = _TestObject()
+ result = instance.func(1, 2, a=3, c=5)
+
+ # assert
+ # span_enter
+ span_id = f"{instance.func.__qualname__}-mock"
+ bound_args = inspect.signature(instance.func).bind(1, 2, a=3, c=5)
mock_span_enter.assert_called_once()
args, kwargs = mock_span_enter.call_args
- assert args == (1, 2)
- assert kwargs == {"id": span_id, "a": 3, "c": 5}
+ assert args == ()
+ assert kwargs == {"id_": span_id, "bound_args": bound_args, "instance": instance}
# span_exit
args, kwargs = mock_span_exit.call_args
- assert args == (1, 2)
- assert kwargs == {"id": span_id, "a": 3, "c": 5, "result": result}
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": instance,
+ "result": result,
+ }
@patch.object(Dispatcher, "span_exit")
@patch.object(Dispatcher, "span_drop")
@patch.object(Dispatcher, "span_enter")
@patch("llama_index.core.instrumentation.dispatcher.uuid")
-@patch(f"{__name__}.func")
def test_dispatcher_span_drop_args(
- mock_func: MagicMock,
mock_uuid: MagicMock,
mock_span_enter: MagicMock,
mock_span_drop: MagicMock,
mock_span_exit: MagicMock,
):
# arrange
- class CustomException(Exception):
- pass
+ mock_uuid.uuid4.return_value = "mock"
+
+ with pytest.raises(ValueError):
+ # act
+ _ = func_exc(7, a=3, b=5, c=2, d=5)
+
+ # assert
+ # span_enter
+ mock_span_enter.assert_called_once()
+
+ # span_drop
+ mock_span_drop.assert_called_once()
+ span_id = f"{func_exc.__qualname__}-mock"
+ bound_args = inspect.signature(func_exc).bind(7, a=3, b=5, c=2, d=5)
+ args, kwargs = mock_span_drop.call_args
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": None,
+ "err": value_error,
+ }
+
+ # span_exit
+ mock_span_exit.assert_not_called()
+
+@patch.object(Dispatcher, "span_exit")
+@patch.object(Dispatcher, "span_drop")
+@patch.object(Dispatcher, "span_enter")
+@patch("llama_index.core.instrumentation.dispatcher.uuid")
+def test_dispatcher_span_drop_args(
+ mock_uuid: MagicMock,
+ mock_span_enter: MagicMock,
+ mock_span_drop: MagicMock,
+ mock_span_exit: MagicMock,
+):
+ # arrange
mock_uuid.uuid4.return_value = "mock"
- mock_func.side_effect = CustomException
- with pytest.raises(CustomException):
+ with pytest.raises(ValueError):
# act
- result = func(7, a=3, b=5, c=2, d=5)
-
- # assert
- # span_enter
- mock_span_enter.assert_called_once()
-
- # span_drop
- mock_span_drop.assert_called_once()
- span_id = f"{func.__qualname__}-mock"
- args, kwargs = mock_span_exit.call_args
- assert args == (7,)
- assert kwargs == {
- "id": span_id,
- "a": 3,
- "b": 5,
- "c": 2,
- "d": 2,
- "err": CustomException,
- }
-
- # span_exit
- mock_span_exit.assert_not_called()
+ instance = _TestObject()
+ _ = instance.func_exc(7, a=3, b=5, c=2, d=5)
+
+ # assert
+ # span_enter
+ mock_span_enter.assert_called_once()
+
+ # span_drop
+ mock_span_drop.assert_called_once()
+ span_id = f"{instance.func_exc.__qualname__}-mock"
+ bound_args = inspect.signature(instance.func_exc).bind(7, a=3, b=5, c=2, d=5)
+ args, kwargs = mock_span_drop.call_args
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": instance,
+ "err": value_error,
+ }
+
+ # span_exit
+ mock_span_exit.assert_not_called()
@pytest.mark.asyncio()
@@ -99,15 +202,55 @@ async def test_dispatcher_async_span_args(mock_uuid, mock_span_enter, mock_span_
# assert
# span_enter
span_id = f"{async_func.__qualname__}-mock"
+ bound_args = inspect.signature(async_func).bind(1, 2, a=3, c=5)
mock_span_enter.assert_called_once()
args, kwargs = mock_span_enter.call_args
- assert args == (1, 2)
- assert kwargs == {"id": span_id, "a": 3, "c": 5}
+ assert args == ()
+ assert kwargs == {"id_": span_id, "bound_args": bound_args, "instance": None}
# span_exit
args, kwargs = mock_span_exit.call_args
- assert args == (1, 2)
- assert kwargs == {"id": span_id, "a": 3, "c": 5, "result": result}
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": None,
+ "result": result,
+ }
+
+
+@pytest.mark.asyncio()
+@patch.object(Dispatcher, "span_exit")
+@patch.object(Dispatcher, "span_enter")
+@patch("llama_index.core.instrumentation.dispatcher.uuid")
+async def test_dispatcher_async_span_args_with_instance(
+ mock_uuid, mock_span_enter, mock_span_exit
+):
+ # arrange
+ mock_uuid.uuid4.return_value = "mock"
+
+ # act
+ instance = _TestObject()
+ result = await instance.async_func(1, 2, a=3, c=5)
+
+ # assert
+ # span_enter
+ span_id = f"{instance.async_func.__qualname__}-mock"
+ bound_args = inspect.signature(instance.async_func).bind(1, 2, a=3, c=5)
+ mock_span_enter.assert_called_once()
+ args, kwargs = mock_span_enter.call_args
+ assert args == ()
+ assert kwargs == {"id_": span_id, "bound_args": bound_args, "instance": instance}
+
+ # span_exit
+ args, kwargs = mock_span_exit.call_args
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": instance,
+ "result": result,
+ }
@pytest.mark.asyncio()
@@ -115,42 +258,75 @@ async def test_dispatcher_async_span_args(mock_uuid, mock_span_enter, mock_span_
@patch.object(Dispatcher, "span_drop")
@patch.object(Dispatcher, "span_enter")
@patch("llama_index.core.instrumentation.dispatcher.uuid")
-@patch(f"{__name__}.async_func")
-async def test_dispatcher_aysnc_span_drop_args(
- mock_func: MagicMock,
+async def test_dispatcher_async_span_drop_args(
mock_uuid: MagicMock,
mock_span_enter: MagicMock,
mock_span_drop: MagicMock,
mock_span_exit: MagicMock,
):
# arrange
- class CustomException(Exception):
- pass
+ mock_uuid.uuid4.return_value = "mock"
+
+ with pytest.raises(CancelledError):
+ # act
+ _ = await async_func_exc(7, a=3, b=5, c=2, d=5)
+ # assert
+ # span_enter
+ mock_span_enter.assert_called_once()
+
+ # span_drop
+ mock_span_drop.assert_called_once()
+ span_id = f"{async_func_exc.__qualname__}-mock"
+ bound_args = inspect.signature(async_func_exc).bind(7, a=3, b=5, c=2, d=5)
+ args, kwargs = mock_span_drop.call_args
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": None,
+ "err": cancelled_error,
+ }
+
+ # span_exit
+ mock_span_exit.assert_not_called()
+
+
+@pytest.mark.asyncio()
+@patch.object(Dispatcher, "span_exit")
+@patch.object(Dispatcher, "span_drop")
+@patch.object(Dispatcher, "span_enter")
+@patch("llama_index.core.instrumentation.dispatcher.uuid")
+async def test_dispatcher_async_span_drop_args_with_instance(
+ mock_uuid: MagicMock,
+ mock_span_enter: MagicMock,
+ mock_span_drop: MagicMock,
+ mock_span_exit: MagicMock,
+):
+ # arrange
mock_uuid.uuid4.return_value = "mock"
- mock_func.side_effect = CustomException
- with pytest.raises(CustomException):
+ with pytest.raises(CancelledError):
# act
- result = await async_func(7, a=3, b=5, c=2, d=5)
-
- # assert
- # span_enter
- mock_span_enter.assert_called_once()
-
- # span_drop
- mock_span_drop.assert_called_once()
- span_id = f"{func.__qualname__}-mock"
- args, kwargs = mock_span_exit.call_args
- assert args == (7,)
- assert kwargs == {
- "id": span_id,
- "a": 3,
- "b": 5,
- "c": 2,
- "d": 2,
- "err": CustomException,
- }
-
- # span_exit
- mock_span_exit.assert_not_called()
+ instance = _TestObject()
+ _ = await instance.async_func_exc(7, a=3, b=5, c=2, d=5)
+
+ # assert
+ # span_enter
+ mock_span_enter.assert_called_once()
+
+ # span_drop
+ mock_span_drop.assert_called_once()
+ span_id = f"{instance.async_func_exc.__qualname__}-mock"
+ bound_args = inspect.signature(instance.async_func_exc).bind(7, a=3, b=5, c=2, d=5)
+ args, kwargs = mock_span_drop.call_args
+ assert args == ()
+ assert kwargs == {
+ "id_": span_id,
+ "bound_args": bound_args,
+ "instance": instance,
+ "err": cancelled_error,
+ }
+
+ # span_exit
+ mock_span_exit.assert_not_called()