wehub-resource-sync
97e91a83f3
Ruff / Ruff (push) Has been cancelled
Test / Core Tests (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.10) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.11) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.12) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.13) (push) Has been cancelled
Test / Offline Coverage Tests (Python 3.9) (push) Has been cancelled
Test / Full Coverage (Python 3.11) (push) Has been cancelled
Test / Core Provider Tests (OpenAI) (push) Has been cancelled
Test / Core Provider Tests (Anthropic) (push) Has been cancelled
Test / Core Provider Tests (Google) (push) Has been cancelled
Test / Core Provider Tests (Other) (push) Has been cancelled
Test / Anthropic Tests (push) Has been cancelled
Test / Gemini Tests (push) Has been cancelled
Test / Google GenAI Tests (push) Has been cancelled
Test / Vertex AI Tests (push) Has been cancelled
Test / OpenAI Tests (push) Has been cancelled
Test / Writer Tests (push) Has been cancelled
Test / Auto Client Tests (push) Has been cancelled
ty / type-check (push) Has been cancelled
366 lines
12 KiB
Python
366 lines
12 KiB
Python
from __future__ import annotations
|
|
from enum import Enum
|
|
from collections import defaultdict
|
|
from typing import Any, Literal, TypeVar, Protocol, Union
|
|
|
|
import inspect
|
|
import traceback
|
|
import warnings
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class HookName(Enum):
|
|
COMPLETION_KWARGS = "completion:kwargs"
|
|
COMPLETION_RESPONSE = "completion:response"
|
|
COMPLETION_ERROR = "completion:error"
|
|
COMPLETION_LAST_ATTEMPT = "completion:last_attempt"
|
|
PARSE_ERROR = "parse:error"
|
|
|
|
|
|
# Handler protocol types for type safety
|
|
class CompletionKwargsHandler(Protocol):
|
|
"""Protocol for completion kwargs handlers."""
|
|
|
|
def __call__(self, *args: Any, **kwargs: Any) -> None: ...
|
|
|
|
|
|
class CompletionResponseHandler(Protocol):
|
|
"""Protocol for completion response handlers."""
|
|
|
|
def __call__(self, response: Any) -> None: ...
|
|
|
|
|
|
class CompletionErrorHandler(Protocol):
|
|
"""Protocol for completion error and last attempt handlers."""
|
|
|
|
def __call__(
|
|
self,
|
|
error: Exception,
|
|
*,
|
|
attempt_number: int = ...,
|
|
max_attempts: int | None = ...,
|
|
is_last_attempt: bool = ...,
|
|
) -> None: ...
|
|
|
|
|
|
class ParseErrorHandler(Protocol):
|
|
"""Protocol for parse error handlers."""
|
|
|
|
def __call__(self, error: Exception, **kwargs: Any) -> None: ...
|
|
|
|
|
|
# Type alias for hook name parameter
|
|
HookNameType = Union[
|
|
HookName,
|
|
Literal[
|
|
"completion:kwargs",
|
|
"completion:response",
|
|
"completion:error",
|
|
"completion:last_attempt",
|
|
"parse:error",
|
|
],
|
|
]
|
|
|
|
# Type alias for all handler types
|
|
HandlerType = Union[
|
|
CompletionKwargsHandler,
|
|
CompletionResponseHandler,
|
|
CompletionErrorHandler,
|
|
ParseErrorHandler,
|
|
]
|
|
|
|
|
|
class Hooks:
|
|
"""
|
|
Hooks class for handling and emitting events related to completion processes.
|
|
|
|
This class provides a mechanism to register event handlers and emit events
|
|
for various stages of the completion process.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the hooks container."""
|
|
self._handlers: defaultdict[HookName, list[HandlerType]] = defaultdict(list)
|
|
|
|
def on(
|
|
self,
|
|
hook_name: HookNameType,
|
|
handler: HandlerType,
|
|
) -> None:
|
|
"""
|
|
Register an event handler for a specific event.
|
|
|
|
This method allows you to attach a handler function to a specific event.
|
|
When the event is emitted, all registered handlers for that event will be called.
|
|
|
|
Args:
|
|
hook_name: The event to listen for. This can be either a HookName enum
|
|
value or a string representation of the event name.
|
|
handler: The function to be called when the event is emitted.
|
|
|
|
Raises:
|
|
ValueError: If the hook_name is not a valid HookName enum or string representation.
|
|
|
|
Example:
|
|
>>> def on_completion_kwargs(*args: Any, **kwargs: Any) -> None:
|
|
... print(f"Completion kwargs: {args}, {kwargs}")
|
|
>>> hooks = Hooks()
|
|
>>> hooks.on(HookName.COMPLETION_KWARGS, on_completion_kwargs)
|
|
>>> hooks.emit_completion_arguments(model="gpt-3.5-turbo", temperature=0.7)
|
|
Completion kwargs: (), {'model': 'gpt-3.5-turbo', 'temperature': 0.7}
|
|
"""
|
|
hook_name = self.get_hook_name(hook_name)
|
|
self._handlers[hook_name].append(handler)
|
|
|
|
def get_hook_name(self, hook_name: HookNameType) -> HookName:
|
|
"""
|
|
Convert a string hook name to its corresponding enum value.
|
|
|
|
Args:
|
|
hook_name: Either a HookName enum value or string representation.
|
|
|
|
Returns:
|
|
The corresponding HookName enum value.
|
|
|
|
Raises:
|
|
ValueError: If the string doesn't match any HookName enum value.
|
|
"""
|
|
if isinstance(hook_name, str):
|
|
try:
|
|
return HookName(hook_name)
|
|
except ValueError as err:
|
|
raise ValueError(f"Invalid hook name: {hook_name}") from err
|
|
return hook_name
|
|
|
|
def emit(self, hook_name: HookName, *args: Any, **kwargs: Any) -> None:
|
|
"""
|
|
Generic method to emit events for any hook type.
|
|
|
|
Args:
|
|
hook_name: The hook to emit
|
|
*args: Positional arguments to pass to handlers
|
|
**kwargs: Keyword arguments to pass to handlers
|
|
"""
|
|
for handler in self._handlers[hook_name]:
|
|
try:
|
|
if kwargs:
|
|
try:
|
|
sig = inspect.signature(handler)
|
|
sig.bind(*args, **kwargs)
|
|
except TypeError:
|
|
handler(*args)
|
|
continue
|
|
handler(*args, **kwargs)
|
|
except Exception:
|
|
error_traceback = traceback.format_exc()
|
|
warnings.warn(
|
|
f"Error in {hook_name.value} handler:\n{error_traceback}",
|
|
stacklevel=2,
|
|
)
|
|
|
|
def emit_completion_arguments(self, *args: Any, **kwargs: Any) -> None:
|
|
"""
|
|
Emit a completion arguments event.
|
|
|
|
Args:
|
|
*args: Positional arguments to pass to handlers
|
|
**kwargs: Keyword arguments to pass to handlers
|
|
"""
|
|
self.emit(HookName.COMPLETION_KWARGS, *args, **kwargs)
|
|
|
|
def emit_completion_response(self, response: Any) -> None:
|
|
"""
|
|
Emit a completion response event.
|
|
|
|
Args:
|
|
response: The completion response to pass to handlers
|
|
"""
|
|
self.emit(HookName.COMPLETION_RESPONSE, response)
|
|
|
|
def emit_completion_error(self, error: Exception, **kwargs: Any) -> None:
|
|
"""
|
|
Emit a completion error event.
|
|
|
|
Args:
|
|
error: The exception to pass to handlers
|
|
**kwargs: Optional metadata (attempt_number, max_attempts, is_last_attempt)
|
|
"""
|
|
self.emit(HookName.COMPLETION_ERROR, error, **kwargs)
|
|
|
|
def emit_completion_last_attempt(self, error: Exception, **kwargs: Any) -> None:
|
|
"""
|
|
Emit a completion last attempt event.
|
|
|
|
Args:
|
|
error: The exception to pass to handlers
|
|
**kwargs: Optional metadata (attempt_number, max_attempts, is_last_attempt)
|
|
"""
|
|
self.emit(HookName.COMPLETION_LAST_ATTEMPT, error, **kwargs)
|
|
|
|
def emit_parse_error(self, error: Exception, **kwargs: Any) -> None:
|
|
"""
|
|
Emit a parse error event.
|
|
|
|
Args:
|
|
error: The exception to pass to handlers
|
|
**kwargs: Optional metadata (attempt_number, max_attempts, is_last_attempt)
|
|
"""
|
|
self.emit(HookName.PARSE_ERROR, error, **kwargs)
|
|
|
|
def off(
|
|
self,
|
|
hook_name: HookNameType,
|
|
handler: HandlerType,
|
|
) -> None:
|
|
"""
|
|
Remove a specific handler from an event.
|
|
|
|
Args:
|
|
hook_name: The name of the hook.
|
|
handler: The handler to remove.
|
|
"""
|
|
hook_name = self.get_hook_name(hook_name)
|
|
if hook_name in self._handlers:
|
|
if handler in self._handlers[hook_name]:
|
|
self._handlers[hook_name].remove(handler)
|
|
if not self._handlers[hook_name]:
|
|
del self._handlers[hook_name]
|
|
|
|
def clear(
|
|
self,
|
|
hook_name: HookNameType | None = None,
|
|
) -> None:
|
|
"""
|
|
Clear handlers for a specific event or all events.
|
|
|
|
Args:
|
|
hook_name: The name of the event to clear handlers for.
|
|
If None, all handlers are cleared.
|
|
"""
|
|
if hook_name is not None:
|
|
hook_name = self.get_hook_name(hook_name)
|
|
self._handlers.pop(hook_name, None)
|
|
else:
|
|
self._handlers.clear()
|
|
|
|
def __add__(self, other: Hooks) -> Hooks:
|
|
"""
|
|
Combine two Hooks instances into a new one.
|
|
|
|
This creates a new Hooks instance that contains all handlers from both
|
|
the current instance and the other instance. Handlers are combined by
|
|
appending the other's handlers after the current instance's handlers.
|
|
|
|
Args:
|
|
other: Another Hooks instance to combine with this one.
|
|
|
|
Returns:
|
|
A new Hooks instance containing all handlers from both instances.
|
|
|
|
Example:
|
|
>>> hooks1 = Hooks()
|
|
>>> hooks2 = Hooks()
|
|
>>> hooks1.on("completion:kwargs", lambda **kw: print("Hook 1"))
|
|
>>> hooks2.on("completion:kwargs", lambda **kw: print("Hook 2"))
|
|
>>> combined = hooks1 + hooks2
|
|
>>> combined.emit_completion_arguments() # Prints both "Hook 1" and "Hook 2"
|
|
"""
|
|
if not isinstance(other, Hooks):
|
|
return NotImplemented
|
|
|
|
combined = Hooks()
|
|
|
|
# Copy handlers from self
|
|
for hook_name, handlers in self._handlers.items():
|
|
combined._handlers[hook_name].extend(handlers.copy())
|
|
|
|
# Add handlers from other
|
|
for hook_name, handlers in other._handlers.items():
|
|
combined._handlers[hook_name].extend(handlers.copy())
|
|
|
|
return combined
|
|
|
|
def __iadd__(self, other: Hooks) -> Hooks:
|
|
"""
|
|
Add another Hooks instance to this one in-place.
|
|
|
|
This modifies the current instance by adding all handlers from the other
|
|
instance. The other instance's handlers are appended after the current
|
|
instance's handlers for each event type.
|
|
|
|
Args:
|
|
other: Another Hooks instance to add to this one.
|
|
|
|
Returns:
|
|
This Hooks instance (for method chaining).
|
|
|
|
Example:
|
|
>>> hooks1 = Hooks()
|
|
>>> hooks2 = Hooks()
|
|
>>> hooks1.on("completion:kwargs", lambda **kw: print("Hook 1"))
|
|
>>> hooks2.on("completion:kwargs", lambda **kw: print("Hook 2"))
|
|
>>> hooks1 += hooks2
|
|
>>> hooks1.emit_completion_arguments() # Prints both "Hook 1" and "Hook 2"
|
|
"""
|
|
if not isinstance(other, Hooks):
|
|
return NotImplemented
|
|
|
|
# Add handlers from other to self
|
|
for hook_name, handlers in other._handlers.items():
|
|
self._handlers[hook_name].extend(handlers.copy())
|
|
|
|
return self
|
|
|
|
@classmethod
|
|
def combine(cls, *hooks_instances: Hooks) -> Hooks:
|
|
"""
|
|
Combine multiple Hooks instances into a new one.
|
|
|
|
This class method creates a new Hooks instance that contains all handlers
|
|
from all provided instances. Handlers are combined in the order of the
|
|
provided instances.
|
|
|
|
Args:
|
|
*hooks_instances: Variable number of Hooks instances to combine.
|
|
|
|
Returns:
|
|
A new Hooks instance containing all handlers from all instances.
|
|
|
|
Example:
|
|
>>> hooks1 = Hooks()
|
|
>>> hooks2 = Hooks()
|
|
>>> hooks3 = Hooks()
|
|
>>> hooks1.on("completion:kwargs", lambda **kw: print("Hook 1"))
|
|
>>> hooks2.on("completion:kwargs", lambda **kw: print("Hook 2"))
|
|
>>> hooks3.on("completion:kwargs", lambda **kw: print("Hook 3"))
|
|
>>> combined = Hooks.combine(hooks1, hooks2, hooks3)
|
|
>>> combined.emit_completion_arguments() # Prints all three hooks
|
|
"""
|
|
combined = cls()
|
|
|
|
for hooks_instance in hooks_instances:
|
|
if not isinstance(hooks_instance, cls):
|
|
raise TypeError(f"Expected Hooks instance, got {type(hooks_instance)}")
|
|
combined += hooks_instance
|
|
|
|
return combined
|
|
|
|
def copy(self) -> Hooks:
|
|
"""
|
|
Create a deep copy of this Hooks instance.
|
|
|
|
Returns:
|
|
A new Hooks instance with all the same handlers.
|
|
|
|
Example:
|
|
>>> original = Hooks()
|
|
>>> original.on("completion:kwargs", lambda **kw: print("Hook"))
|
|
>>> copy = original.copy()
|
|
>>> copy.emit_completion_arguments() # Prints "Hook"
|
|
"""
|
|
new_hooks = Hooks()
|
|
for hook_name, handlers in self._handlers.items():
|
|
new_hooks._handlers[hook_name].extend(handlers.copy())
|
|
return new_hooks
|