nat.llm.utils.thinking#

Attributes#

`ModelType`
`MessagesType`
`logger`

Classes#

`FunctionArgumentWrapper`	Wrapper for the arguments and keyword arguments of a function.
`BaseThinkingInjector`	Base class for thinking injectors.

Functions#

`_make_thinking_decorator`(injector)
`patch_with_thinking`(→ ModelType)	Patch the given object with a decorator that injects a system prompt into the supplied messages.

Module Contents#

ModelType#

MessagesType#

logger#

class FunctionArgumentWrapper(*args: Any, **kwargs: Any)#

Wrapper for the arguments and keyword arguments of a function.

The arguments and keyword arguments are stored in the args and kwargs attributes, respectively.

Initialize the FunctionArgumentWrapper.

Args:: args: The arguments to the function. kwargs: The keyword arguments to the function.

args = ()#

kwargs#

class BaseThinkingInjector#

Base class for thinking injectors.

Args:: system_prompt: The system prompt to inject. function_names: The function names to inject the system prompt into.

system_prompt: str#

function_names: list[str]#

abstractmethod inject(*args, **kwargs) → FunctionArgumentWrapper#

Inject the system prompt into the arguments.

Args:: args: The arguments to inject the system prompt into. kwargs: The keyword arguments to inject the system prompt into.
Returns:: FunctionArgumentWrapper: An object that contains the transformed args and kwargs.

_make_thinking_decorator(injector: BaseThinkingInjector)#

patch_with_thinking( obj: ModelType, injector: BaseThinkingInjector, ) → ModelType#

Patch the given object with a decorator that injects a system prompt into the supplied messages. There is an assumption that the first non-object argument is the messages.

Args:

obj: The object to patch. injector: The injector to use.

Returns:

The patched object.

Examples:

>>> from nat.llm.utils.thinking import BaseThinkingInjector
>>> from nat.llm.utils.thinking import FunctionArgumentWrapper
>>> from nat.llm.utils.thinking import patch_with_thinking
>>>
>>> class MockClass:
...     def sync_method(self, *args, **kwargs):
...         return (args, kwargs)
...
>>> mock_obj_1 = MockClass()
>>> class AddThinking(BaseThinkingInjector):
...     def inject(self, x: str, *args, **kwargs) -> FunctionArgumentWrapper:
...         return FunctionArgumentWrapper(("thinking " + x), *args, **kwargs)
>>>
>>> patched_obj = patch_with_thinking(mock_obj_1, AddThinking(
...     system_prompt="thinking",
...     function_names=["sync_method"],
... ))
>>> patched_obj.sync_method("test", 1, 2, 3, foo="bar")
(('thinking test', 1, 2, 3), {'foo': 'bar'})
>>>
>>> mock_obj_2 = MockClass()
>>> class AddThinkingWithArgs(BaseThinkingInjector):
...     def inject(self, *args, **kwargs) -> FunctionArgumentWrapper:
...         return FunctionArgumentWrapper("thinking", *args, **kwargs)
>>>
>>> patched_obj = patch_with_thinking(mock_obj_2, AddThinkingWithArgs(
...     system_prompt="thinking",
...     function_names=["sync_method"],
... ))
>>> patched_obj.sync_method("test", 1, 2, 3, foo="bar")
(('thinking', 'test', 1, 2, 3), {'foo': 'bar'})
>>>
>>> mock_obj_3 = MockClass()
>>> class AddThinkingWithKwargs(BaseThinkingInjector):
...     def inject(self, *args, **kwargs) -> FunctionArgumentWrapper:
...         return FunctionArgumentWrapper(*args, thinking=True, **kwargs)
>>>
>>> patched_obj = patch_with_thinking(mock_obj_3, AddThinkingWithKwargs(
...     system_prompt="thinking",
...     function_names=["sync_method"],
... ))
>>> patched_obj.sync_method("test", 1, 2, 3, foo="bar")
(('test', 1, 2, 3), {'thinking': True, 'foo': 'bar'})