krons 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

krons/resource/utilities/token_calculator.py

@@ -0,0 +1,185 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from collections.abc import Callable
+from typing import cast
+
+import tiktoken
+
+from krons.errors import KronsError
+
+logger = logging.getLogger(__name__)
+
+
+class TokenCalculationError(KronsError):
+    """Raised when token calculation fails due to encoding/model errors."""
+
+    default_message = "Token calculation failed"
+    default_retryable = False
+
+
+def get_encoding_name(value: str | None) -> str:
+    """Get encoding name for model, with fallback chain.
+
+    Returns:
+        Encoding name (defaults to o200k_base if model/encoding not found or None)
+    """
+    if value is None:
+        return "o200k_base"
+
+    try:
+        enc = tiktoken.encoding_for_model(value)
+        return enc.name
+    except KeyError:
+        # Not a known model name, try as encoding name
+        try:
+            tiktoken.get_encoding(value)
+            return value
+        except Exception as e:
+            logger.warning(
+                f"Unknown model/encoding '{value}', falling back to o200k_base: {e}"
+            )
+            return "o200k_base"
+
+
+class TokenCalculator:
+    @staticmethod
+    def calculate_message_tokens(messages: list[dict], /, **kwargs) -> int:
+        model = kwargs.get("model", "gpt-4o")
+        image_token_cost = kwargs.get("image_token_cost", 500)
+        tokenizer = tiktoken.get_encoding(get_encoding_name(model)).encode
+
+        num_tokens = 0
+        for msg in messages:
+            num_tokens += 4
+            _c = msg.get("content")
+            num_tokens += TokenCalculator._calculate_chatitem(
+                _c,
+                tokenizer=tokenizer,
+                model_name=model,
+                image_token_cost=image_token_cost,
+            )
+        return num_tokens  # buffer for chat
+
+    @staticmethod
+    def calculate_embed_token(inputs: list[str], /, **kwargs) -> int:
+        if not inputs:
+            raise ValueError("inputs must be a non-empty list of strings")
+
+        try:
+            tokenizer = tiktoken.get_encoding(
+                get_encoding_name(kwargs.get("model", "text-embedding-3-small"))
+            ).encode
+
+            return sum(
+                TokenCalculator._calculate_embed_item(i, tokenizer=tokenizer)
+                for i in inputs
+            )
+        except TokenCalculationError:
+            # Re-raise from nested calls
+            raise
+        except Exception as e:
+            logger.error(f"Failed to calculate embed tokens: {e}", exc_info=True)
+            raise TokenCalculationError(f"Embed token calculation failed: {e}") from e
+
+    @staticmethod
+    def tokenize(
+        s_: str | None = None,
+        /,
+        encoding_name: str | None = None,
+        tokenizer: Callable | None = None,
+        decoder: Callable | None = None,
+        return_tokens: bool = False,
+        return_decoded: bool = False,
+    ) -> int | list[int] | tuple[int, str]:
+        if not s_:
+            return 0
+
+        if not callable(tokenizer):
+            encoding_name = get_encoding_name(encoding_name)
+            tokenizer = tiktoken.get_encoding(encoding_name).encode
+        if not callable(decoder):
+            # Use encoding_name if available, otherwise fallback to default
+            decoder_encoding = encoding_name if encoding_name else "o200k_base"
+            decoder = tiktoken.get_encoding(decoder_encoding).decode
+
+        try:
+            if return_tokens:
+                if return_decoded:
+                    a = tokenizer(s_)
+                    return len(a), decoder(a)
+                return tokenizer(s_)
+            return len(tokenizer(s_))
+        except Exception as e:
+            # Actual encoding failure during tokenization - this is an error
+            logger.error(
+                f"Tokenization failed for input (len={len(s_) if s_ else 0}): {e}",
+                exc_info=True,
+            )
+            raise TokenCalculationError(f"Tokenization failed: {e}") from e
+
+    @staticmethod
+    def _calculate_chatitem(
+        i_, tokenizer: Callable, model_name: str, image_token_cost: int = 500
+    ) -> int:
+        try:
+            if isinstance(i_, str):
+                # tokenize returns int when return_tokens=False (default)
+                return cast(int, TokenCalculator.tokenize(i_, tokenizer=tokenizer))
+
+            if isinstance(i_, dict):
+                if "text" in i_:
+                    return TokenCalculator._calculate_chatitem(
+                        str(i_["text"]), tokenizer, model_name, image_token_cost
+                    )
+                elif "image_url" in i_:
+                    return image_token_cost
+
+            if isinstance(i_, list):
+                return sum(
+                    TokenCalculator._calculate_chatitem(
+                        x, tokenizer, model_name, image_token_cost
+                    )
+                    for x in i_
+                )
+
+            # Unknown type - return 0 is valid (no text content)
+            return 0
+        except TokenCalculationError:
+            # Re-raise tokenization errors from nested calls
+            raise
+        except Exception as e:
+            logger.error(
+                f"Failed to calculate chat item tokens (type={type(i_).__name__}): {e}",
+                exc_info=True,
+            )
+            raise TokenCalculationError(
+                f"Chat item token calculation failed: {e}"
+            ) from e
+
+    @staticmethod
+    def _calculate_embed_item(s_, tokenizer: Callable) -> int:
+        try:
+            if isinstance(s_, str):
+                # tokenize returns int when return_tokens=False (default)
+                return cast(int, TokenCalculator.tokenize(s_, tokenizer=tokenizer))
+
+            if isinstance(s_, list):
+                return sum(
+                    TokenCalculator._calculate_embed_item(x, tokenizer) for x in s_
+                )
+
+            # Unknown type - return 0 is valid (no text content)
+            return 0
+        except TokenCalculationError:
+            # Re-raise tokenization errors from nested calls
+            raise
+        except Exception as e:
+            logger.error(
+                f"Failed to calculate embed item tokens (type={type(s_).__name__}): {e}",
+                exc_info=True,
+            )
+            raise TokenCalculationError(
+                f"Embed item token calculation failed: {e}"
+            ) from e

krons/session/__init__.py

@@ -10,23 +10,19 @@ Core types:
     Exchange: Async message router between entity mailboxes.
 
 Validators (raise on failure):
-    resource_must_exist_in_session
-    resource_must_be_accessible_by_branch
-    capabilities_must_be_subset_of_branch
-    resolve_branch_exists_in_session
+    resource_must_exist
+    resource_must_be_accessible
+    capabilities_must_be_granted
 """
 
+from .constraints import (
+    capabilities_must_be_granted,
+    resource_must_be_accessible,
+    resource_must_exist,
+)
 from .exchange import Exchange
 from .message import Message
-from .session import (
-    Branch,
-    Session,
-    SessionConfig,
-    capabilities_must_be_subset_of_branch,
-    resolve_branch_exists_in_session,
-    resource_must_be_accessible_by_branch,
-    resource_must_exist_in_session,
-)
+from .session import Branch, Session, SessionConfig
 
 __all__ = (
     "Branch",
@@ -34,8 +30,7 @@ __all__ = (
     "Message",
     "Session",
     "SessionConfig",
-    "capabilities_must_be_subset_of_branch",
-    "resolve_branch_exists_in_session",
-    "resource_must_be_accessible_by_branch",
-    "resource_must_exist_in_session",
+    "capabilities_must_be_granted",
+    "resource_must_be_accessible",
+    "resource_must_exist",
 )

krons/session/constraints.py

@@ -0,0 +1,70 @@
+from krons.errors import AccessError, ConfigurationError, ExistsError, NotFoundError
+
+__all__ = (
+    "resource_must_exist",
+    "resource_must_be_accessible",
+    "capabilities_must_be_granted",
+    "branch_name_must_be_unique",
+    "genai_model_must_be_configured",
+)
+
+
+def resource_must_exist(session, name: str):
+    """Validate resource exists in session. Raise NotFoundError if not."""
+    if not session.resources.has(name):
+        raise NotFoundError(
+            f"Service '{name}' not found in session services",
+            details={"available": session.resources.list_names()},
+        )
+
+
+def resource_must_be_accessible(branch, name: str) -> None:
+    """Validate branch has resource access. Raise AccessError if not."""
+    if name not in branch.resources:
+        raise AccessError(
+            f"Branch '{branch.name}' has no access to resource '{name}'",
+            details={
+                "branch": branch.name,
+                "resource": name,
+                "available": list(branch.resources),
+            },
+        )
+
+
+def capabilities_must_be_granted(branch, capabilities: set[str]) -> None:
+    """Validate branch has all capabilities. Raise AccessError listing missing."""
+    if not capabilities.issubset(branch.capabilities):
+        missing = capabilities - branch.capabilities
+        raise AccessError(
+            f"Branch '{branch.name}' missing capabilities: {missing}",
+            details={
+                "requested": sorted(capabilities),
+                "available": sorted(branch.capabilities),
+            },
+        )
+
+
+def branch_name_must_be_unique(session, name: str) -> None:
+    try:
+        session.communications.get_progression(name)
+        # If we get here, the name exists - not unique
+        raise ExistsError(f"Branch with name '{name}' already exists")
+    except KeyError:
+        # KeyError means name not found - it's unique, which is good
+        pass
+
+
+def genai_model_must_be_configured(session) -> None:
+    """Validate session has a default GenAI model configured.
+
+    Args:
+        session: Session to check
+
+    Raises:
+        ConfigurationError: If no default model configured
+    """
+    if session.default_gen_model is None:
+        raise ConfigurationError(
+            "Session has no default_gen_model configured",
+            details={"session_id": str(session.id)},
+        )

krons/session/exchange.py

@@ -134,11 +134,17 @@ class Exchange(Element):
                             except Exception:
                                 message_copy = message.model_copy()
                             deliveries.append((other_id, message_copy))
-                elif message.recipient is not None and message.recipient in self._owner_index:
+                elif (
+                    message.recipient is not None
+                    and message.recipient in self._owner_index
+                ):
                     deliveries.append((message.recipient, message))
         if deliveries:
             await concurrency.gather(
-                *[self._deliver_to(recipient_id, message) for recipient_id, message in deliveries],
+                *[
+                    self._deliver_to(recipient_id, message)
+                    for recipient_id, message in deliveries
+                ],
                 return_exceptions=True,
             )
 
@@ -196,7 +202,9 @@ class Exchange(Element):
         if flow is None:
             raise ValueError(f"Sender {sender} not registered")
 
-        message = Message(sender=sender, recipient=recipient, content=content, channel=channel)
+        message = Message(
+            sender=sender, recipient=recipient, content=content, channel=channel
+        )
         flow.add_item(message, progressions=OUTBOX)
         return message
 

krons/session/message.py

@@ -27,7 +27,9 @@ class Message(Node):
     content: Any
     sender: UUID | None = None
     recipient: UUID | None = None
-    channel: str | None = Field(None, description="Optional namespace for message grouping")
+    channel: str | None = Field(
+        None, description="Optional namespace for message grouping"
+    )
 
     @property
     def is_broadcast(self) -> bool:

krons/session/registry.py

@@ -0,0 +1,35 @@
+from krons.core import Pile
+from krons.core.types import ID
+
+from .session import Session
+
+SESSION_REGISTRY: Pile[Session] = Pile(item_type=Session, strict_type=True)
+
+
+async def get_session(session_id: ID[Session]) -> Session:
+    if session_id not in SESSION_REGISTRY:
+        raise ValueError(f"Session with id {session_id} not found in registry.")
+    async with SESSION_REGISTRY:
+        return SESSION_REGISTRY[session_id]
+
+
+async def create_session():
+    session = Session()
+    async with SESSION_REGISTRY:
+        SESSION_REGISTRY.add(session)
+
+
+async def delete_session(session_id: ID[Session]):
+    if session_id not in SESSION_REGISTRY:
+        raise ValueError(f"Session with id {session_id} not found in registry.")
+    async with SESSION_REGISTRY:
+        SESSION_REGISTRY.remove(session_id)
+
+
+async def list_sessions_ids() -> list[ID[Session]]:
+    return list(SESSION_REGISTRY.keys())
+
+
+async def clear_sessions():
+    async with SESSION_REGISTRY:
+        SESSION_REGISTRY.clear()