diff --git a/llama-index-core/llama_index/core/base/llms/base.py b/llama-index-core/llama_index/core/base/llms/base.py index fa7c816980b21da10f146553ddfadbb9e1c2956d..aa4946fda5bbb7440a876f5a0429d9346efaff6b 100644 --- a/llama-index-core/llama_index/core/base/llms/base.py +++ b/llama-index-core/llama_index/core/base/llms/base.py @@ -4,6 +4,7 @@ from typing import ( Sequence, ) +from llama_index.core import instrumentation from llama_index.core.base.llms.types import ( ChatMessage, ChatResponse, @@ -253,3 +254,23 @@ class BaseLLM(ChainableMixin, BaseComponent): print(response.text, end="", flush=True) ``` """ + + def __init_subclass__(cls, **kwargs) -> None: + """ + Decorate the abstract methods' implementations for each subclass. + `__init_subclass__` is analogous to `__init__` because classes are also objects. + """ + super().__init_subclass__(**kwargs) + dispatcher = instrumentation.get_dispatcher(cls.__module__) + for attr in ( + "chat", + "complete", + "stream_chat", + "stream_complete", + "achat", + "acomplete", + "astream_chat", + "astream_complete", + ): + if callable(method := cls.__dict__.get(attr)): + setattr(cls, attr, dispatcher.span(method)) diff --git a/llama-index-core/llama_index/core/llms/callbacks.py b/llama-index-core/llama_index/core/llms/callbacks.py index 0e9ad8b837ea6fd0170cb9141b94cd81b95fa6d6..44d7368f288d216d1cbc2b6763c27689a7907435 100644 --- a/llama-index-core/llama_index/core/llms/callbacks.py +++ b/llama-index-core/llama_index/core/llms/callbacks.py @@ -209,6 +209,22 @@ def llm_chat_callback() -> Callable: if not is_wrapped: f.__wrapped__ = True # type: ignore + # Update the wrapper function to look like the wrapped function. + # See e.g. https://github.com/python/cpython/blob/0abf997e75bd3a8b76d920d33cc64d5e6c2d380f/Lib/functools.py#L57 + for attr in ( + "__module__", + "__name__", + "__qualname__", + "__doc__", + "__annotations__", + "__type_params__", + ): + if v := getattr(f, attr, None): + setattr(async_dummy_wrapper, attr, v) + setattr(wrapped_async_llm_chat, attr, v) + setattr(dummy_wrapper, attr, v) + setattr(wrapped_llm_chat, attr, v) + if asyncio.iscoroutinefunction(f): if is_wrapped: return async_dummy_wrapper @@ -394,6 +410,22 @@ def llm_completion_callback() -> Callable: if not is_wrapped: f.__wrapped__ = True # type: ignore + # Update the wrapper function to look like the wrapped function. + # See e.g. https://github.com/python/cpython/blob/0abf997e75bd3a8b76d920d33cc64d5e6c2d380f/Lib/functools.py#L57 + for attr in ( + "__module__", + "__name__", + "__qualname__", + "__doc__", + "__annotations__", + "__type_params__", + ): + if v := getattr(f, attr, None): + setattr(async_dummy_wrapper, attr, v) + setattr(wrapped_async_llm_predict, attr, v) + setattr(dummy_wrapper, attr, v) + setattr(wrapped_llm_predict, attr, v) + if asyncio.iscoroutinefunction(f): if is_wrapped: return async_dummy_wrapper