lionagi 0.14.8__py3-none-any.whl → 0.14.10__py3-none-any.whl

lionagi/service/hooks/hook_registry.py

@@ -0,0 +1,221 @@
+# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, TypeVar
+
+from lionagi.libs.concurrency import get_cancelled_exc_class
+from lionagi.protocols.types import Event, EventStatus
+from lionagi.utils import UNDEFINED
+
+from ._types import HookDict, HookEventTypes, StreamHandlers
+from ._utils import validate_hooks, validate_stream_handlers
+
+E = TypeVar("E", bound=Event)
+
+
+class HookRegistry:
+    def __init__(
+        self,
+        hooks: HookDict = None,
+        stream_handlers: StreamHandlers = None,
+    ):
+        _hooks = {}
+        _stream_handlers = {}
+
+        if hooks is not None:
+            validate_hooks(hooks)
+            _hooks.update(hooks)
+
+        if stream_handlers is not None:
+            validate_stream_handlers(stream_handlers)
+            _stream_handlers.update(stream_handlers)
+
+        self._hooks = _hooks
+        self._stream_handlers = _stream_handlers
+
+    async def _call(
+        self,
+        ht_: HookEventTypes,
+        ct_: str | type,
+        ch_: Any,
+        ev_: E | type[E],
+        /,
+        **kw,
+    ) -> tuple[Any | Exception, bool]:
+        if ht_ is None and ct_ is None:
+            raise RuntimeError(
+                "Either hook_type or chunk_type must be provided"
+            )
+        if ht_ and (h := self._hooks.get(ht_)):
+            validate_hooks({ht_: h})
+            return await h(ev_, **kw)
+        elif not ct_:
+            raise RuntimeError(
+                "Hook type is required when chunk_type is not provided"
+            )
+        else:
+            validate_stream_handlers(
+                {ct_: (h := self._stream_handlers.get(ct_))}
+            )
+            return await h(ev_, ct_, ch_, **kw)
+
+    async def _call_stream_handler(
+        self,
+        ct_: str | type,
+        ch_: Any,
+        ev_,
+        /,
+        **kw,
+    ):
+        handler = self._stream_handlers.get(ct_)
+        validate_stream_handlers({ct_: handler})
+        return await handler(ev_, ct_, ch_, **kw)
+
+    async def pre_event_create(
+        self, event_type: type[E], /, exit: bool = False, **kw
+    ) -> tuple[E | Exception | None, bool, EventStatus]:
+        """Hook to be called before an event is created.
+
+        Typically used to modify or validate the event creation parameters.
+
+        The hook function takes an event type and any additional keyword arguments.
+        It can:
+            - return an instance of the event type
+            - return None if no event should be created during handling, event will be
+                created in corresponding default manner
+            - raise an exception if this event should be cancelled
+                (status: cancelled, reason: f"pre-event-create hook aborted this event: {e}")
+        """
+        try:
+            res = await self._call(
+                HookEventTypes.PreEventCreate,
+                None,
+                None,
+                event_type,
+                **kw,
+            )
+            return (res, False, EventStatus.COMPLETED)
+        except get_cancelled_exc_class() as e:
+            return ((UNDEFINED, e), True, EventStatus.CANCELLED)
+        except Exception as e:
+            return (e, exit, EventStatus.CANCELLED)
+
+    async def pre_invokation(
+        self, event: E, /, exit: bool = False, **kw
+    ) -> tuple[Any, bool, EventStatus]:
+        """Hook to be called when an event is dequeued and right before it is invoked.
+
+        Typically used to check permissions.
+
+        The hook function takes the content of the event as a dictionary.
+        It can either raise an exception to abort the event invokation or pass to continue (status: cancelled).
+        It cannot modify the event itself, and won't be able to access the event instance.
+        """
+        try:
+            res = await self._call(
+                HookEventTypes.PreInvokation,
+                None,
+                None,
+                event,
+                **kw,
+            )
+            return (res, False, EventStatus.COMPLETED)
+        except get_cancelled_exc_class() as e:
+            return ((UNDEFINED, e), True, EventStatus.CANCELLED)
+        except Exception as e:
+            return (e, exit, EventStatus.CANCELLED)
+
+    async def post_invokation(
+        self, event: E, /, exit: bool = False, **kw
+    ) -> tuple[None | Exception, bool, EventStatus, EventStatus]:
+        """Hook to be called right after event finished its execution.
+        It can either raise an exception to abort the event invokation or pass to continue (status: aborted).
+        It cannot modify the event itself, and won't be able to access the event instance.
+        """
+        try:
+            res = await self._call(
+                HookEventTypes.PostInvokation,
+                None,
+                None,
+                event,
+                **kw,
+            )
+            return (res, False, EventStatus.COMPLETED)
+        except get_cancelled_exc_class() as e:
+            return ((UNDEFINED, e), True, EventStatus.CANCELLED)
+        except Exception as e:
+            return (e, exit, EventStatus.ABORTED)
+
+    async def handle_streaming_chunk(
+        self, chunk_type: str | type, chunk: Any, /, exit: bool = False, **kw
+    ) -> tuple[Any, bool, EventStatus | None]:
+        """Hook to be called to consume streaming chunks.
+
+        Typically used for logging or stream event abortion.
+
+        The handler function signature should be: `async def handler(chunk: Any) -> None`
+        It can either raise an exception to mark the event invokation as "failed" or pass to continue (status: aborted).
+        """
+        try:
+            res = await self._call_stream_handler(
+                chunk_type,
+                chunk,
+                None,
+                **kw,
+            )
+            return (res, False, None)
+        except get_cancelled_exc_class() as e:
+            return ((UNDEFINED, e), True, EventStatus.CANCELLED)
+        except Exception as e:
+            return (e, exit, EventStatus.ABORTED)
+
+    async def call(
+        self,
+        event_like: Event | type[Event],
+        /,
+        *,
+        hook_type: HookEventTypes = None,
+        chunk_type=None,
+        chunk=None,
+        exit=False,
+        **kw,
+    ):
+        """Call a hook or stream handler.
+
+        If method is provided, it will call the corresponding hook.
+        If chunk_type is provided, it will call the corresponding stream handler.
+        If both are provided, method will be used.
+        """
+        if hook_type is None and chunk_type is None:
+            raise ValueError("Either method or chunk_type must be provided")
+        if hook_type:
+            meta = {}
+            meta["event_type"] = event_like.class_name(full=True)
+            match hook_type:
+                case HookEventTypes.PreEventCreate:
+                    return await self.pre_event_create(event_like, **kw), meta
+                case HookEventTypes.PreInvokation:
+                    meta["event_id"] = str(event_like.id)
+                    meta["event_created_at"] = event_like.created_at
+                    return await self.pre_invokation(event_like, **kw), meta
+                case HookEventTypes.PostInvokation:
+                    meta["event_id"] = str(event_like.id)
+                    meta["event_created_at"] = event_like.created_at
+                    return await self.post_invokation(**kw), meta
+        return await self.handle_streaming_chunk(chunk_type, chunk, exit, **kw)
+
+    def _can_handle(
+        self,
+        /,
+        *,
+        ht_: HookEventTypes = None,
+        ct_=None,
+    ) -> bool:
+        """Check if the registry can handle the given event or chunk type."""
+        if ht_:
+            return ht_ in self._hooks
+        if ct_:
+            return ct_ in self._stream_handlers
+        return False

lionagi/service/imodel.py

@@ -7,12 +7,15 @@ from collections.abc import AsyncGenerator, Callable
 
 from pydantic import BaseModel
 
-from lionagi.protocols.generic.event import EventStatus
-from lionagi.utils import is_coro_func
+from lionagi.protocols.generic.log import Log
+from lionagi.protocols.types import ID, Event, EventStatus, IDType
+from lionagi.service.hooks.hook_event import HookEventTypes
+from lionagi.utils import is_coro_func, time
 
 from .connections.api_calling import APICalling
 from .connections.endpoint import Endpoint
 from .connections.match_endpoint import match_endpoint
+from .hooks import HookEvent, HookRegistry, global_hook_logger
 from .rate_limited_processor import RateLimitedAPIExecutor
 
 
@@ -47,6 +50,10 @@ class iModel:
         concurrency_limit: int | None = None,
         streaming_process_func: Callable = None,
         provider_metadata: dict | None = None,
+        hook_registry: HookRegistry | dict | None = None,
+        exit_hook: bool = False,
+        id: IDType | str = None,
+        created_at: float | None = None,
         **kwargs,
     ) -> None:
         """Initializes the iModel instance.
@@ -86,6 +93,22 @@ class iModel:
                 Additional keyword arguments, such as `model`, or any other
                 provider-specific fields.
         """
+
+        # 1. put in ID and timestamp -----------------------------------------
+        self.id = None
+        self.created_at = None
+        if id is not None:
+            self.id = ID.get_id(id)
+        else:
+            self.id = IDType.create()
+        if created_at is not None:
+            if not isinstance(created_at, float):
+                raise ValueError("created_at must be a float timestamp.")
+            self.created_at = created_at
+        else:
+            self.created_at = time()
+
+        # 2. Configure Endpoint ---------------------------------------------
         model = kwargs.get("model", None)
         if model:
             if not provider:
@@ -96,7 +119,6 @@ class iModel:
                 else:
                     raise ValueError("Provider must be provided")
 
-        # Pass api_key to endpoint if provided
         if api_key is not None:
             kwargs["api_key"] = api_key
         if isinstance(endpoint, Endpoint):
@@ -112,6 +134,7 @@ class iModel:
         if base_url:
             self.endpoint.config.base_url = base_url
 
+        # 3. Configure executor ---------------------------------------------
         self.executor = RateLimitedAPIExecutor(
             queue_capacity=queue_capacity,
             capacity_refresh_time=capacity_refresh_time,
@@ -120,11 +143,91 @@ class iModel:
             limit_tokens=limit_tokens,
             concurrency_limit=concurrency_limit,
         )
-        # Use provided streaming_process_func or default to None
-        self.streaming_process_func = streaming_process_func
 
-        # Provider-specific metadata storage (e.g., session_id for Claude Code)
+        # 4. other configurations --------------------------------------------
+        self.streaming_process_func = streaming_process_func
         self.provider_metadata = provider_metadata or {}
+        self.hook_registry = hook_registry or HookRegistry()
+        if isinstance(self.hook_registry, dict):
+            self.hook_registry = HookRegistry(**self.hook_registry)
+        self.exit_hook: bool = exit_hook
+
+    async def create_event(
+        self,
+        create_event_type: type[Event] = APICalling,
+        create_event_exit_hook: bool = None,
+        create_event_hook_timeout: float = 10.0,
+        create_event_hook_params: dict = None,
+        pre_invoke_event_exit_hook: bool = None,
+        pre_invoke_event_hook_timeout: float = 30.0,
+        pre_invoke_event_hook_params: dict = None,
+        post_invoke_event_exit_hook: bool = None,
+        post_invoke_event_hook_timeout: float = 30.0,
+        post_invoke_event_hook_params: dict = None,
+        **kwargs,
+    ) -> tuple[HookEvent | None, APICalling]:
+        h_ev = None
+        if self.hook_registry._can_handle(ht_=HookEventTypes.PreEventCreate):
+            h_ev = HookEvent(
+                hook_type=HookEventTypes.PreEventCreate,
+                registry=self.hook_registry,
+                event_like=create_event_type,
+                params=create_event_hook_params or {},
+                exit=(
+                    self.exit_hook
+                    if create_event_exit_hook is None
+                    else create_event_exit_hook
+                ),
+                timeout=create_event_hook_timeout,
+            )
+            await h_ev.invoke()
+            if h_ev._should_exit:
+                raise h_ev._exit_cause or RuntimeError(
+                    "PreEventCreate hook requested exit without a cause"
+                )
+
+        if create_event_type is APICalling:
+            api_call = self.create_api_calling(**kwargs)
+            if h_ev:
+                h_ev.assosiated_event_info["event_id"] = str(api_call.id)
+                h_ev.assosiated_event_info["event_created_at"] = (
+                    api_call.created_at
+                )
+                await global_hook_logger.alog(Log(content=h_ev.to_dict()))
+
+            if self.hook_registry._can_handle(
+                ht_=HookEventTypes.PreInvokation
+            ):
+                api_call.create_pre_invoke_hook(
+                    hook_registry=self.hook_registry,
+                    exit_hook=(
+                        self.exit_hook
+                        if pre_invoke_event_exit_hook is None
+                        else pre_invoke_event_exit_hook
+                    ),
+                    hook_timeout=pre_invoke_event_hook_timeout,
+                    hook_params=pre_invoke_event_hook_params or {},
+                )
+
+            if self.hook_registry._can_handle(
+                ht_=HookEventTypes.PostInvokation
+            ):
+                api_call.create_post_invoke_hook(
+                    hook_registry=self.hook_registry,
+                    exit_hook=(
+                        self.exit_hook
+                        if post_invoke_event_exit_hook is None
+                        else post_invoke_event_exit_hook
+                    ),
+                    hook_timeout=post_invoke_event_hook_timeout,
+                    hook_params=post_invoke_event_hook_params or {},
+                )
+
+            return api_call
+
+        raise ValueError(
+            f"Unsupported event type: {create_event_type}. Only APICalling is supported."
+        )
 
     def create_api_calling(
         self, include_token_usage_to_model: bool = False, **kwargs
@@ -178,12 +281,7 @@ class iModel:
                 return await self.streaming_process_func(chunk)
             return self.streaming_process_func(chunk)
 
-    async def stream(
-        self,
-        api_call=None,
-        include_token_usage_to_model: bool = False,
-        **kwargs,
-    ) -> AsyncGenerator:
+    async def stream(self, api_call=None, **kw) -> AsyncGenerator:
         """Performs a streaming API call with the given arguments.
 
         Args:
@@ -196,11 +294,8 @@ class iModel:
                 goes wrong.
         """
         if api_call is None:
-            kwargs["stream"] = True
-            api_call = self.create_api_calling(
-                include_token_usage_to_model=include_token_usage_to_model,
-                **kwargs,
-            )
+            kw["stream"] = True
+            api_call = await self.create_event(**kw)
             await self.executor.append(api_call)
 
         if (
@@ -231,9 +326,7 @@ class iModel:
             finally:
                 yield self.executor.pile.pop(api_call.id)
 
-    async def invoke(
-        self, api_call: APICalling = None, **kwargs
-    ) -> APICalling | None:
+    async def invoke(self, api_call: APICalling = None, **kw) -> APICalling:
         """Invokes a rate-limited API call with the given arguments.
 
         Args:
@@ -251,8 +344,9 @@ class iModel:
         """
         try:
             if api_call is None:
-                kwargs.pop("stream", None)
-                api_call = self.create_api_calling(**kwargs)
+                kw.pop("stream", None)
+                api_call = await self.create_event(**kw)
+
             if (
                 self.executor.processor is None
                 or self.executor.processor.is_stopped()
@@ -291,18 +385,6 @@ class iModel:
         except Exception as e:
             raise ValueError(f"Failed to invoke API call: {e}")
 
-    @property
-    def allowed_roles(self) -> set[str]:
-        """list[str]: Roles that are permissible for this endpoint.
-
-        Returns:
-            If the endpoint has an `allowed_roles` attribute, returns that;
-            otherwise, defaults to `{"system", "user", "assistant"}`.
-        """
-        if hasattr(self.endpoint, "allowed_roles"):
-            return self.endpoint.allowed_roles
-        return {"system", "user", "assistant"}
-
     @property
     def model_name(self) -> str:
         """str: The name of the model used by the endpoint.
@@ -323,6 +405,8 @@ class iModel:
 
     def to_dict(self):
         return {
+            "id": str(self.id) if self.id else None,
+            "created_at": self.created_at,
             "endpoint": self.endpoint.to_dict(),
             "processor_config": self.executor.config,
             "provider_metadata": self.provider_metadata,
@@ -343,5 +427,7 @@ class iModel:
         return cls(
             endpoint=e1,
             provider_metadata=data.get("provider_metadata"),
+            id=data.get("id"),
+            created_at=data.get("created_at"),
             **data.get("processor_config", {}),
         )