lionagi 0.17.4__py3-none-any.whl → 0.17.6__py3-none-any.whl

lionagi/service/connections/mcp/wrapper.py

@@ -7,15 +7,7 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Any, Dict
-
-try:
-    from fastmcp import Client as FastMCPClient
-
-    FASTMCP_AVAILABLE = True
-except ImportError:
-    FASTMCP_AVAILABLE = False
-    FastMCPClient = None
+from typing import Any
 
 # Suppress MCP server logging by default
 logging.getLogger("mcp").setLevel(logging.WARNING)
@@ -88,11 +80,6 @@ class MCPConnectionPool:
     @classmethod
     async def get_client(cls, server_config: dict[str, Any]) -> Any:
         """Get or create a pooled MCP client."""
-        if not FASTMCP_AVAILABLE:
-            raise ImportError(
-                "FastMCP not installed. Run: pip install fastmcp"
-            )
-
         # Generate unique key for this config
         if "server" in server_config:
             # Server reference from .mcp.json
@@ -149,6 +136,13 @@ class MCPConnectionPool:
         if not any(k in config for k in ["url", "command"]):
             raise ValueError("Config must have either 'url' or 'command' key")
 
+        try:
+            from fastmcp import Client as FastMCPClient
+        except ImportError:
+            raise ImportError(
+                "FastMCP not installed. Run: pip install fastmcp"
+            )
+
         # Handle different config formats
         if "url" in config:
             # Direct URL connection
@@ -172,7 +166,6 @@ class MCPConnectionPool:
                 # Common environment variables to suppress logging
                 env.setdefault("LOG_LEVEL", "ERROR")
                 env.setdefault("PYTHONWARNINGS", "ignore")
-                env.setdefault("KHIVEMCP_LOG_LEVEL", "ERROR")
                 # Suppress FastMCP server logs
                 env.setdefault("FASTMCP_QUIET", "true")
                 env.setdefault("MCP_QUIET", "true")

lionagi/service/hooks/__init__.py

@@ -1,19 +1,10 @@
 # Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
 # SPDX-License-Identifier: Apache-2.0
 
-from lionagi.protocols.types import DataLogger
-
 from ._types import AssosiatedEventInfo, HookDict, HookEventTypes
 from .hook_event import HookEvent
 from .hook_registry import HookRegistry
-
-global_hook_logger = DataLogger(
-    persist_dir="./data/logs",
-    subfolder="hooks",
-    file_prefix="hook",
-    capacity=1000,
-)
-
+from .hooked_event import HookedEvent, global_hook_logger
 
 __all__ = (
     "HookEventTypes",
@@ -22,4 +13,5 @@ __all__ = (
     "HookEvent",
     "HookRegistry",
     "global_hook_logger",
+    "HookedEvent",
 )

lionagi/service/hooks/_types.py

@@ -37,6 +37,7 @@ class HookDict(TypedDict):
 
 
 StreamHandlers = dict[str, Callable[[SC], Awaitable[None]]]
+"""Mapping of chunk type names to their respective asynchronous handler functions."""
 
 
 class AssosiatedEventInfo(TypedDict, total=False):

lionagi/service/hooks/hooked_event.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import anyio
+from pydantic import PrivateAttr
+
+from lionagi.ln import get_cancelled_exc_class
+from lionagi.protocols.types import DataLogger, Event, EventStatus, Log
+from lionagi.service.hooks import HookEvent, HookEventTypes
+
+global_hook_logger = DataLogger(
+    persist_dir="./data/logs",
+    subfolder="hooks",
+    file_prefix="hook",
+    capacity=100,
+)
+
+
+class HookedEvent(Event):
+    """Handles asynchronous API calls with automatic token usage tracking.
+
+    This class manages API calls through endpoints, handling both regular
+    and streaming responses with optional token usage tracking.
+    """
+
+    _pre_invoke_hook_event: HookEvent = PrivateAttr(None)
+    _post_invoke_hook_event: HookEvent = PrivateAttr(None)
+
+    async def _stream(self):
+        raise NotImplementedError
+
+    async def _invoke(self):
+        raise NotImplementedError
+
+    async def invoke(self) -> None:
+        """Execute the API call through the endpoint.
+
+        Updates execution status and stores the response or error.
+        """
+        start = anyio.current_time()
+
+        try:
+            self.execution.status = EventStatus.PROCESSING
+            if h_ev := self._pre_invoke_hook_event:
+                await h_ev.invoke()
+                if h_ev._should_exit:
+                    raise h_ev._exit_cause or RuntimeError(
+                        "Pre-invocation hook requested exit without a cause"
+                    )
+                await global_hook_logger.alog(Log.create(h_ev))
+
+            response = await self._invoke()
+
+            if h_ev := self._post_invoke_hook_event:
+                await h_ev.invoke()
+                if h_ev._should_exit:
+                    raise h_ev._exit_cause or RuntimeError(
+                        "Post-invocation hook requested exit without a cause"
+                    )
+                await global_hook_logger.alog(Log.create(h_ev))
+
+            self.execution.response = response
+            self.execution.status = EventStatus.COMPLETED
+
+        except get_cancelled_exc_class():
+            self.execution.error = "Invocation cancelled"
+            self.execution.status = EventStatus.CANCELLED
+            raise
+
+        except Exception as e:
+            self.execution.error = str(e)
+            self.execution.status = EventStatus.FAILED
+
+        finally:
+            self.execution.duration = anyio.current_time() - start
+
+    async def stream(self):
+        """Stream the API response through the endpoint.
+
+        Yields:
+            Streaming chunks from the API.
+        """
+        start = anyio.current_time()
+
+        response = []
+
+        try:
+            self.execution.status = EventStatus.PROCESSING
+
+            async for chunk in self._stream():
+                response.append(chunk)
+                yield chunk
+
+            self.execution.response = response
+            self.execution.status = EventStatus.COMPLETED
+
+        except get_cancelled_exc_class():
+            self.execution.error = "Streaming cancelled"
+            self.execution.status = EventStatus.CANCELLED
+            raise
+
+        except Exception as e:
+            self.execution.error = str(e)
+            self.execution.status = EventStatus.FAILED
+
+        finally:
+            self.execution.duration = anyio.current_time() - start
+
+    def create_pre_invoke_hook(
+        self,
+        hook_registry,
+        exit_hook: bool = None,
+        hook_timeout: float = 30.0,
+        hook_params: dict = None,
+    ):
+        h_ev = HookEvent(
+            hook_type=HookEventTypes.PreInvokation,
+            event_like=self,
+            registry=hook_registry,
+            exit=exit_hook,
+            timeout=hook_timeout,
+            params=hook_params or {},
+        )
+        self._pre_invoke_hook_event = h_ev
+
+    def create_post_invoke_hook(
+        self,
+        hook_registry,
+        exit_hook: bool = None,
+        hook_timeout: float = 30.0,
+        hook_params: dict = None,
+    ):
+        h_ev = HookEvent(
+            hook_type=HookEventTypes.PostInvokation,
+            event_like=self,
+            registry=hook_registry,
+            exit=exit_hook,
+            timeout=hook_timeout,
+            params=hook_params or {},
+        )
+        self._post_invoke_hook_event = h_ev

lionagi/service/imodel.py

@@ -7,10 +7,10 @@ from collections.abc import AsyncGenerator, Callable
 
 from pydantic import BaseModel
 
+from lionagi.ln import is_coro_func, now_utc
 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
@@ -106,7 +106,7 @@ class iModel:
                 raise ValueError("created_at must be a float timestamp.")
             self.created_at = created_at
         else:
-            self.created_at = time()
+            self.created_at = now_utc().timestamp()
 
         # 2. Configure Endpoint ---------------------------------------------
         model = kwargs.get("model", None)

lionagi/session/branch.py

@@ -3,15 +3,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import AsyncGenerator, Callable
-from enum import Enum
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
 from pydantic import BaseModel, Field, JsonValue, PrivateAttr, field_serializer
 
 from lionagi.config import settings
-from lionagi.fields import Instruct
-from lionagi.libs.schema.as_readable import as_readable
+from lionagi.ln.types import Unset
 from lionagi.models.field_model import FieldModel
+from lionagi.operations.flow import AlcallParams
 from lionagi.operations.manager import OperationManager
 from lionagi.protocols.action.manager import ActionManager
 from lionagi.protocols.action.tool import FuncTool, Tool, ToolRef
@@ -42,21 +41,22 @@ from lionagi.protocols.types import (
 )
 from lionagi.service.connections.endpoint import Endpoint
 from lionagi.service.types import iModel, iModelManager
-from lionagi.settings import Settings
 from lionagi.tools.base import LionTool
-from lionagi.utils import UNDEFINED
-from lionagi.utils import alcall as alcall_legacy
 from lionagi.utils import copy
 
 from .prompts import LION_SYSTEM_MESSAGE
 
 if TYPE_CHECKING:
+    from lionagi.fields import Instruct
     from lionagi.protocols.operatives.operative import Operative
 
 
 __all__ = ("Branch",)
 
 
+_DEFAULT_ALCALL_PARAMS = None
+
+
 class Branch(Element, Communicatable, Relational):
     """
     Manages a conversation 'branch' with messages, tools, and iModels.
@@ -150,10 +150,10 @@ class Branch(Element, Communicatable, Relational):
                 Sender to attribute to the system message if it is added.
             chat_model (iModel, optional):
                 The primary "chat" iModel for conversation. If not provided,
-                a default from `Settings.iModel.CHAT` is used.
+                uses default provider and model from settings.
             parse_model (iModel, optional):
                 The "parse" iModel for structured data parsing.
-                Defaults to `Settings.iModel.PARSE`.
+                Defaults to chat_model if not provided.
             imodel (iModel, optional):
                 Deprecated. Alias for `chat_model`.
             tools (FuncTool | list[FuncTool], optional):
@@ -233,7 +233,7 @@ class Branch(Element, Communicatable, Relational):
                 log_config = LogManagerConfig(**log_config)
             self._log_manager = LogManager.from_config(log_config, logs=logs)
         else:
-            self._log_manager = LogManager(**Settings.Config.LOG, logs=logs)
+            self._log_manager = LogManager(**settings.LOG_CONFIG, logs=logs)
 
         self._operation_manager = OperationManager()
 
@@ -703,12 +703,12 @@ class Branch(Element, Communicatable, Relational):
             Branch: A new `Branch` instance based on the deserialized data.
         """
         dict_ = {
-            "messages": data.pop("messages", UNDEFINED),
-            "logs": data.pop("logs", UNDEFINED),
-            "chat_model": data.pop("chat_model", UNDEFINED),
-            "parse_model": data.pop("parse_model", UNDEFINED),
-            "system": data.pop("system", UNDEFINED),
-            "log_config": data.pop("log_config", UNDEFINED),
+            "messages": data.pop("messages", Unset),
+            "logs": data.pop("logs", Unset),
+            "chat_model": data.pop("chat_model", Unset),
+            "parse_model": data.pop("parse_model", Unset),
+            "system": data.pop("system", Unset),
+            "log_config": data.pop("log_config", Unset),
         }
         params = {}
 
@@ -721,8 +721,8 @@ class Branch(Element, Communicatable, Relational):
                 params[k] = v
 
         params.update(dict_)
-        # Remove placeholders (UNDEFINED) so we don't incorrectly assign them
-        return cls(**{k: v for k, v in params.items() if v is not UNDEFINED})
+        # Remove placeholders (Unset) so we don't incorrectly assign them
+        return cls(**{k: v for k, v in params.items() if v is not Unset})
 
     def dump_logs(self, clear: bool = True, persist_path=None):
         """
@@ -913,7 +913,7 @@ class Branch(Element, Communicatable, Relational):
     async def operate(
         self,
         *,
-        instruct: Instruct = None,
+        instruct: "Instruct" = None,
         instruction: Instruction | JsonValue = None,
         guidance: JsonValue = None,
         context: JsonValue = None,
@@ -934,7 +934,7 @@ class Branch(Element, Communicatable, Relational):
         ] = None,  # alias of operative.request_type
         actions: bool = False,
         reason: bool = False,
-        action_kwargs: dict = None,
+        call_params: AlcallParams = None,
         action_strategy: Literal["sequential", "concurrent"] = "concurrent",
         verbose_action: bool = False,
         field_models: list[FieldModel] = None,
@@ -1001,8 +1001,6 @@ class Branch(Element, Communicatable, Relational):
                 If `True`, signals that function-calling or "action" usage is expected.
             reason (bool, optional):
                 If `True`, signals that the LLM should provide chain-of-thought or reasoning (where applicable).
-            action_kwargs (dict | None, optional):
-                Additional parameters for the `branch.act()` call if tools are invoked.
             action_strategy (Literal["sequential","concurrent"], optional):
                 The strategy for invoking tools (default: "concurrent").
             verbose_action (bool, optional):
@@ -1052,7 +1050,7 @@ class Branch(Element, Communicatable, Relational):
             response_format=response_format,
             actions=actions,
             reason=reason,
-            action_kwargs=action_kwargs,
+            call_params=call_params,
             action_strategy=action_strategy,
             verbose_action=verbose_action,
             field_models=field_models,
@@ -1162,21 +1160,9 @@ class Branch(Element, Communicatable, Relational):
     async def _act(
         self,
         action_request: ActionRequest | BaseModel | dict,
-        suppress_errors: bool = False,
-        verbose_action: bool = False,
+        suppress_errors: bool,
+        verbose_action: bool,
     ) -> ActionResponse:
-        """
-        Internal method to invoke a tool (action) asynchronously.
-
-        Args:
-            action_request (ActionRequest|BaseModel|dict):
-                Must contain `function` and `arguments`.
-            suppress_errors (bool, optional):
-                If True, errors are logged instead of raised.
-
-        Returns:
-            ActionResponse: Result of the tool invocation or `None` if suppressed.
-        """
         from lionagi.operations._act.act import _act
 
         return await _act(
@@ -1193,103 +1179,38 @@ class Branch(Element, Communicatable, Relational):
         strategy: Literal["concurrent", "sequential"] = "concurrent",
         verbose_action: bool = False,
         suppress_errors: bool = True,
-        sanitize_input: bool = False,
-        unique_input: bool = False,
-        num_retries: int = 0,
-        initial_delay: float = 0,
-        retry_delay: float = 0,
-        backoff_factor: float = 1,
-        retry_default: Any = UNDEFINED,
-        retry_timeout: float | None = None,
-        max_concurrent: int | None = None,
-        throttle_period: float | None = None,
-        flatten: bool = True,
-        dropna: bool = True,
-        unique_output: bool = False,
-        flatten_tuple_set: bool = False,
+        call_params: AlcallParams = None,
     ) -> list[ActionResponse]:
-        """
-        Public, potentially batched, asynchronous interface to run one or multiple action requests.
-
-        Args:
-            action_request (list|ActionRequest|BaseModel|dict):
-                A single or list of action requests, each requiring
-                `function` and `arguments`.
-            strategy (Literal["concurrent","sequential","batch"]):
-                The execution strategy to use.
-            verbose_action (bool):
-                If True, log detailed information about the action.
-            suppress_errors (bool):
-                If True, log errors instead of raising exceptions.
-            sanitize_input (bool):
-                Reserved. Potentially sanitize the action arguments.
-            unique_input (bool):
-                Reserved. Filter out duplicate requests.
-            num_retries (int):
-                Number of times to retry on failure (default 0).
-            initial_delay (float):
-                Delay before first attempt (seconds).
-            retry_delay (float):
-                Base delay between retries.
-            backoff_factor (float):
-                Multiplier for the `retry_delay` after each attempt.
-            retry_default (Any):
-                Fallback value if all retries fail (if suppressing errors).
-            retry_timeout (float|None):
-                Overall timeout for all attempts (None = no limit).
-            max_concurrent (int|None):
-                Maximum concurrent tasks.
-            throttle_period (float|None):
-                Minimum spacing (in seconds) between requests.
-            flatten (bool):
-                If a list of results is returned, flatten them if possible.
-            dropna (bool):
-                Remove `None` or invalid results from final output if True.
-            unique_output (bool):
-                Only return unique results if True.
-            flatten_tuple_set (bool):
-                Flatten nested tuples in results if True.
+        global _DEFAULT_ALCALL_PARAMS
+        if call_params is None:
+            if _DEFAULT_ALCALL_PARAMS is None:
+                _DEFAULT_ALCALL_PARAMS = AlcallParams(output_dropna=True)
+            call_params = _DEFAULT_ALCALL_PARAMS
+
+        kw = {
+            "suppress_errors": suppress_errors,
+            "verbose_action": verbose_action,
+        }
 
-        Returns:
-            Any:
-                The result or results from the invoked tool(s).
-        """
         match strategy:
             case "concurrent":
                 return await self._concurrent_act(
-                    action_request,
-                    verbose_action=verbose_action,
-                    suppress_errors=suppress_errors,
-                    sanitize_input=sanitize_input,
-                    unique_input=unique_input,
-                    num_retries=num_retries,
-                    initial_delay=initial_delay,
-                    retry_delay=retry_delay,
-                    backoff_factor=backoff_factor,
-                    retry_default=retry_default,
-                    retry_timeout=retry_timeout,
-                    max_concurrent=max_concurrent,
-                    throttle_period=throttle_period,
-                    flatten=flatten,
-                    dropna=dropna,
-                    unique_output=unique_output,
-                    flatten_tuple_set=flatten_tuple_set,
+                    action_request, call_params, **kw
                 )
             case "sequential":
-                return await self._sequential_act(
-                    action_request,
-                    verbose_action=verbose_action,
-                    suppress_errors=suppress_errors,
-                )
+                return await self._sequential_act(action_request, **kw)
             case _:
-                raise
+                raise ValueError(
+                    "Invalid strategy. Choose 'concurrent' or 'sequential'."
+                )
 
     async def _concurrent_act(
         self,
         action_request: ActionRequest | BaseModel | dict,
+        call_params: AlcallParams = None,
         **kwargs,
     ) -> list:
-        return await alcall_legacy(action_request, self._act, **kwargs)
+        return await call_params(action_request, self._act, **kwargs)
 
     async def _sequential_act(
         self,
@@ -1305,56 +1226,10 @@ class Branch(Element, Communicatable, Relational):
         results = []
         for req in action_request:
             results.append(
-                await self._act(
-                    req,
-                    verbose_action=verbose_action,
-                    suppress_errors=suppress_errors,
-                )
+                await self._act(req, verbose_action, suppress_errors)
             )
         return results
 
-    async def select(
-        self,
-        instruct: Instruct | dict[str, Any],
-        choices: list[str] | type[Enum] | dict[str, Any],
-        max_num_selections: int = 1,
-        branch_kwargs: dict[str, Any] | None = None,
-        verbose: bool = False,
-        **kwargs: Any,
-    ):
-        """
-        Performs a selection operation from given choices using an LLM-driven approach.
-
-        Args:
-            instruct (Instruct|dict[str, Any]):
-                The instruction model or dictionary for the LLM call.
-            choices (list[str]|type[Enum]|dict[str,Any]):
-                The set of options to choose from.
-            max_num_selections (int):
-                Maximum allowed selections (default = 1).
-            branch_kwargs (dict[str, Any]|None):
-                Extra arguments to create or configure a new branch if needed.
-            verbose (bool):
-                If True, prints debug info.
-            **kwargs:
-                Additional arguments for the underlying `operate()` call.
-
-        Returns:
-            Any:
-                A `SelectionModel` or similar that indicates the user's choice(s).
-        """
-        from lionagi.operations.select.select import select
-
-        return await select(
-            branch=self,
-            instruct=instruct,
-            choices=choices,
-            max_num_selections=max_num_selections,
-            branch_kwargs=branch_kwargs,
-            verbose=verbose,
-            **kwargs,
-        )
-
     async def interpret(
         self,
         text: str,
@@ -1414,7 +1289,7 @@ class Branch(Element, Communicatable, Relational):
 
     async def instruct(
         self,
-        instruct: Instruct,
+        instruct: "Instruct",
         /,
         **kwargs,
     ):
@@ -1441,7 +1316,7 @@ class Branch(Element, Communicatable, Relational):
 
     async def ReAct(
         self,
-        instruct: Instruct | dict[str, Any],
+        instruct: "Instruct | dict[str, Any]",
         interpret: bool = False,
         interpret_domain: str | None = None,
         interpret_style: str | None = None,
@@ -1569,7 +1444,7 @@ class Branch(Element, Communicatable, Relational):
 
     async def ReActStream(
         self,
-        instruct: Instruct | dict[str, Any],
+        instruct: "Instruct | dict[str, Any]",
         interpret: bool = False,
         interpret_domain: str | None = None,
         interpret_style: str | None = None,
@@ -1621,6 +1496,8 @@ class Branch(Element, Communicatable, Relational):
         ):
             analysis, str_ = result
             if verbose:
+                from lionagi.libs.schema.as_readable import as_readable
+
                 str_ += "\n---------\n"
                 as_readable(str_, md=True, display_str=True)
             yield analysis

lionagi/session/session.py

@@ -416,48 +416,5 @@ class Session(Node, Communicatable, Relational):
             alcall_params=alcall_params,
         )
 
-    def cleanup_memory(
-        self, clear_branches: bool = True, clear_mail: bool = True
-    ):
-        """
-        Clean up session memory to prevent memory accumulation.
-
-        Args:
-            clear_branches: Whether to clear branch logs and memory
-            clear_mail: Whether to clear mail transfer history
-        """
-        if clear_branches and self.branches:
-            for branch in self.branches:
-                if hasattr(branch, "dump_logs"):
-                    branch.dump_logs(clear=True)
-
-        if clear_mail and self.mail_transfer:
-            # Clear mail transfer history if available
-            if hasattr(self.mail_transfer, "clear"):
-                self.mail_transfer.clear()
-
-    async def acleanup_memory(
-        self, clear_branches: bool = True, clear_mail: bool = True
-    ):
-        """
-        Asynchronously clean up session memory to prevent memory accumulation.
-
-        Args:
-            clear_branches: Whether to clear branch logs and memory
-            clear_mail: Whether to clear mail transfer history
-        """
-        if clear_branches and self.branches:
-            for branch in self.branches:
-                if hasattr(branch, "adump_logs"):
-                    await branch.adump_logs(clear=True)
-
-        if clear_mail and self.mail_transfer:
-            # Clear mail transfer history if available
-            if hasattr(self.mail_transfer, "aclear"):
-                await self.mail_transfer.aclear()
-            elif hasattr(self.mail_transfer, "clear"):
-                self.mail_transfer.clear()
-
 
-__all__ = ["Session"]
-# File: autoos/session/session.py
+__all__ = ("Session",)