agent-lab-sdk 0.1.35__py3-none-any.whl → 0.1.49__py3-none-any.whl

agent_lab_sdk/langgraph/checkpoint/serde.py

@@ -0,0 +1,172 @@
+import asyncio
+import re
+import uuid
+from base64 import b64encode
+from collections import deque
+from collections.abc import Mapping
+from datetime import timedelta, timezone
+from decimal import Decimal
+from ipaddress import (
+    IPv4Address,
+    IPv4Interface,
+    IPv4Network,
+    IPv6Address,
+    IPv6Interface,
+    IPv6Network,
+)
+from pathlib import Path
+from re import Pattern
+from typing import Any, NamedTuple
+from zoneinfo import ZoneInfo
+
+import cloudpickle
+import orjson
+import logging
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+logger = logging.getLogger(__name__)
+
+
+class Fragment(NamedTuple):
+    buf: bytes
+
+
+def decimal_encoder(dec_value: Decimal) -> int | float:
+    """
+    Encodes a Decimal as int of there's no exponent, otherwise float
+
+    This is useful when we use ConstrainedDecimal to represent Numeric(x,0)
+    where a integer (but not int typed) is used. Encoding this as a float
+    results in failed round-tripping between encode and parse.
+    Our Id type is a prime example of this.
+
+    >>> decimal_encoder(Decimal("1.0"))
+    1.0
+
+    >>> decimal_encoder(Decimal("1"))
+    1
+    """
+    if dec_value.as_tuple().exponent >= 0:
+        return int(dec_value)
+    else:
+        return float(dec_value)
+
+
+def default(obj):
+    # Only need to handle types that orjson doesn't serialize by default
+    # https://github.com/ijl/orjson#serialize
+    if isinstance(obj, Fragment):
+        return orjson.Fragment(obj.buf)
+    if (
+        hasattr(obj, "model_dump")
+        and callable(obj.model_dump)
+        and not isinstance(obj, type)
+    ):
+        return obj.model_dump()
+    elif hasattr(obj, "dict") and callable(obj.dict) and not isinstance(obj, type):
+        return obj.dict()
+    elif (
+        hasattr(obj, "_asdict") and callable(obj._asdict) and not isinstance(obj, type)
+    ):
+        return obj._asdict()
+    elif isinstance(obj, BaseException):
+        return {"error": type(obj).__name__, "message": str(obj)}
+    elif isinstance(obj, (set, frozenset, deque)):  # noqa: UP038
+        return list(obj)
+    elif isinstance(obj, (timezone, ZoneInfo)):  # noqa: UP038
+        return obj.tzname(None)
+    elif isinstance(obj, timedelta):
+        return obj.total_seconds()
+    elif isinstance(obj, Decimal):
+        return decimal_encoder(obj)
+    elif isinstance(obj, uuid.UUID):
+        return str(obj)
+    elif isinstance(  # noqa: UP038
+        obj,
+        (
+            IPv4Address,
+            IPv4Interface,
+            IPv4Network,
+            IPv6Address,
+            IPv6Interface,
+            IPv6Network,
+            Path,
+        ),
+    ):
+        return str(obj)
+    elif isinstance(obj, Pattern):
+        return obj.pattern
+    elif isinstance(obj, bytes | bytearray):
+        return b64encode(obj).decode()
+    return None
+
+
+_option = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS
+
+_SURROGATE_RE = re.compile(r"[\ud800-\udfff]")
+
+
+def _strip_surr(s: str) -> str:
+    return s if _SURROGATE_RE.search(s) is None else _SURROGATE_RE.sub("", s)
+
+
+def _sanitise(o: Any) -> Any:
+    if isinstance(o, str):
+        return _strip_surr(o)
+    if isinstance(o, Mapping):
+        return {_sanitise(k): _sanitise(v) for k, v in o.items()}
+    if isinstance(o, list | tuple | set):
+        ctor = list if isinstance(o, list) else type(o)
+        return ctor(_sanitise(x) for x in o)
+    return o
+
+
+def json_dumpb(obj) -> bytes:
+    try:
+        dumped = orjson.dumps(obj, default=default, option=_option)
+    except TypeError as e:
+        if "surrogates not allowed" not in str(e):
+            raise
+        dumped = orjson.dumps(_sanitise(obj), default=default, option=_option)
+    return (
+        # Unfortunately simply doing ``.replace(rb"\\u0000", b"")`` on
+        # the dumped bytes can leave an **orphaned back-slash** (e.g. ``\\q``)
+        # which makes the resulting JSON invalid.  The fix is to delete the *double*
+        # back-slash form **first**, then (optionally) the single-escapes.
+        dumped.replace(rb"\\u0000", b"").replace(rb"\u0000", b"")
+    )
+
+
+def json_loads(content: bytes | Fragment | dict) -> Any:
+    if isinstance(content, Fragment):
+        content = content.buf
+    if isinstance(content, dict):
+        return content
+    return orjson.loads(content)
+
+
+async def ajson_loads(content: bytes | Fragment) -> Any:
+    return await asyncio.to_thread(json_loads, content)
+
+
+class Serializer(JsonPlusSerializer):
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
+        try:
+            return super().dumps_typed(obj)
+        except TypeError:
+            return "pickle", cloudpickle.dumps(obj)
+
+    def dumps(self, obj: Any) -> bytes:
+        # See comment above (in json_dumpb)
+        return super().dumps(obj).replace(rb"\\u0000", b"").replace(rb"\u0000", b"")
+
+    def loads_typed(self, data: tuple[str, bytes]) -> Any:
+        if data[0] == "pickle":
+            try:
+                return cloudpickle.loads(data[1])
+            except Exception as e:
+                logger.warning(
+                    "Failed to unpickle object, replacing w None", exc_info=e
+                )
+                return None
+        return super().loads_typed(data)

agent_lab_sdk/llm/llm.py

@@ -38,4 +38,7 @@ def get_model(
         # Включаем режим авто-обновления токена в обёртках
         kwargs["manage_access_token"] = True
 
-    return _class(verify_ssl_certs=verify_ssl_certs, **kwargs)
+    if access_token:
+        kwargs["access_token"] = access_token
+
+    return _class(verify_ssl_certs=verify_ssl_certs, **kwargs)

agent_lab_sdk/schema/__init__.py

@@ -10,13 +10,14 @@ from .input_types import (
     SwitchInput,
     FileInput,
     FilesInput,
-    SelectOption
+    SelectOption,
+    Visibility
 )
 
 __all__ = [
     "LogMessage",
     "MainInput",
-    "StringInput", 
+    "StringInput",
     "StringArrayInput",
     "StringArrayInputInline",
     "NumberInput",
@@ -25,5 +26,6 @@ __all__ = [
     "SwitchInput",
     "FileInput",
     "FilesInput",
-    "SelectOption"
+    "SelectOption",
+    "Visibility"
 ]

agent_lab_sdk/schema/input_types.py

@@ -1,52 +1,77 @@
-from typing import List, Optional, Any
-from pydantic import BaseModel
-from pydantic.json_schema import WithJsonSchema
+from typing import List, Optional, Any, Literal
+from enum import Enum
+from pydantic import BaseModel, Field
+import os
 
+def _as_bool(s: str | None) -> bool:
+    return (s or "").lower() in {"1", "true", "yes", "on"}
 
-def MainInput(placeholder: str | None = None) -> type:
+UI_TYPE_KEY = "inputType" if _as_bool(os.getenv("USE_UI_TYPE_INSTEAD_OF_TYPE", "False")) else "type"
+
+class Visibility(str, Enum):
+    """
+    Enum representing visibility states for input fields.
+
+    Attributes:
+        ALWAYS: Field is always visible
+        START: Field is visible only at start
+        AFTER_START: Field is visible after start
+    """
+    ALWAYS = "always"
+    START = "start"
+    AFTER_START = "after_start"
+
+
+def MainInput(placeholder: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a main input type with optional placeholder.
-    
+
     Args:
         placeholder: Optional placeholder text for the input field
-        
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for main input field
     """
-    return WithJsonSchema({
-        "type": "main-input",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "main-input",
         "placeholder": placeholder,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
-def StringInput(default: str | None = None, title: str | None = None, description: str | None = None, hidden: bool | None = False, depends: str | None = None) -> type:
+def StringInput(default: str | None = None, title: str | None = None, description: str | None = None,
+                hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a string input type.
-    
+
     Args:
         default: Default value for the string input
         title: Title for the string input
         description: Description text for the string input
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
-        
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for string input field
     """
-    return WithJsonSchema({
-        "type": "input-string",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "input-string",
         "default": default,
         "title": title,
         "description": description,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
-def StringArrayInput(placeholder: str | None = None, title: str | None = None, description: str | None = None, group: str | None = None, hidden: bool | None = False, depends: str | None = None) -> type:
+def StringArrayInput(placeholder: str | None = None, title: str | None = None, description: str | None = None,
+                     group: str | None = None, hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a string array input type.
-    
+
     Args:
         placeholder: Placeholder text for the input field
         title: Title for the string array input
@@ -54,23 +79,25 @@ def StringArrayInput(placeholder: str | None = None, title: str | None = None, d
         group: Group name for organizing inputs in the UI
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
-        
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for string array input field
     """
-    return WithJsonSchema({
-        "type": "string[]",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "string[]",
         "placeholder": placeholder,
         "title": title,
         "description": description,
         "group": group,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
 def StringArrayInputInline(placeholder: str | None = None, title: str | None = None, description: str | None = None,
-                     group: str | None = None, hidden: bool | None = False, depends: str | None = None) -> type:
+                     group: str | None = None, hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a string array input inline type.
 
@@ -81,42 +108,47 @@ def StringArrayInputInline(placeholder: str | None = None, title: str | None = N
         group: Group name for organizing inputs in the UI
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
+        visibility: Visibility state of the field (default: "always")
 
     Returns:
         Type annotation for string array input field
     """
-    return WithJsonSchema({
-        "type": "string[]-inline",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "string[]-inline",
         "placeholder": placeholder,
         "title": title,
         "description": description,
         "group": group,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
-def NumberInput(default: float | None = None, title: str | None = None, description: str | None = None, hidden: bool | None = False, depends: str | None = None) -> type:
+def NumberInput(default: float | None = None, title: str | None = None, description: str | None = None,
+                hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a number input type.
-    
+
     Args:
         default: Default value for the number input
         title: Title for the number input
         description: Description text for the number input
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
-        
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for number input field
     """
-    return WithJsonSchema({
-        "type": "input-number",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "input-number",
         "default": default,
         "title": title,
         "description": description,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
@@ -126,7 +158,7 @@ class SelectOption(BaseModel):
     
     Attributes:
         label: Display label for the option
-        value: Actual value of the option
+        value: Actual value for the option
         description: Optional description for the option
     """
     label: str
@@ -134,10 +166,11 @@ class SelectOption(BaseModel):
     description: Optional[str] = None
 
 
-def SelectInput(items: List[Any] = [], title: str | None = None, group: str | None = None, default: str | None = None, hidden: bool | None = False, depends: str | None = None) -> type:
+def SelectInput(items: List[Any] = [], title: str | None = None, group: str | None = None, default: str | None = None,
+                hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a select input type.
-    
+
     Args:
         items: List of SelectOption objects or dictionaries
         title: Title for the select input
@@ -145,25 +178,28 @@ def SelectInput(items: List[Any] = [], title: str | None = None, group: str | No
         default: Default selected value
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
-        
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for select input field
     """
-    return WithJsonSchema({
-        "type": "select",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "select",
         "title": title,
         "items": items,
         "group": group,
         "default": default,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
-def CheckboxInput(title: str | None = None, group: str | None = None, description: str | None = None, default: bool | None = False, hidden: bool | None = False, depends: str | None = None) -> type:
+def CheckboxInput(title: str | None = None, group: str | None = None, description: str | None = None,
+                  default: bool | None = False, hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a checkbox input type.
-    
+
     Args:
         title: Title for the checkbox
         group: Group name for organizing inputs in the UI
@@ -171,23 +207,25 @@ def CheckboxInput(title: str | None = None, group: str | None = None, descriptio
         default: Default checked state
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
-        
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for checkbox input field
     """
-    return WithJsonSchema({
-        "type": "checkbox",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "checkbox",
         "title": title,
         "group": group,
         "description": description,
         "default": default,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
 def SwitchInput(title: str | None = None, group: str | None = None, description: str | None = None,
-                  default: bool | None = False, hidden: bool | None = False, depends: str | None = None) -> type:
+                  default: bool | None = False, hidden: bool | None = False, depends: str | None = None, visibility: str | Visibility = Visibility.ALWAYS) -> type:
     """
     Factory function for creating a switch input type.
 
@@ -198,50 +236,60 @@ def SwitchInput(title: str | None = None, group: str | None = None, description:
         default: Default checked state
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
+        visibility: Visibility state of the field (default: "always")
 
     Returns:
         Type annotation for switch input field
     """
-    return WithJsonSchema({
-        "type": "switch",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "switch",
         "title": title,
         "group": group,
         "description": description,
         "default": default,
         "hidden": hidden,
         "depends": depends,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
     })
 
 
-def FileInput(title: str | None = None, file_extensions: str | None = None, group: str | None = None, hidden: bool | None = False, depends: str | None = None) -> type:
+def FileInput(title: str | None = None, file_extensions: str | None = None, group: str | None = None,
+              hidden: bool | None = False, depends: str | None = None, view: Literal["button", "dropzone"] | None = None,
+              visibility: str | Visibility = Visibility.ALWAYS, max_size_mb: float | None = 100.0) -> type:
     """
     Factory function for creating a single file input type.
-    
+
     Args:
         title: Title for the file input
         file_extensions: Comma-separated list of allowed file extensions (e.g., ".pdf,.txt")
         group: Group name for organizing inputs in the UI
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
-        
+        view: View mode for the file input ("button" or "dropzone")
+        visibility: Visibility state of the field (default: "always")
+
     Returns:
         Type annotation for file input field
     """
-    return WithJsonSchema({
-        "type": "file",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "file",
         "title": title,
         "fileExtensions": file_extensions,
         "group": group,
         "hidden": hidden,
         "depends": depends,
+        "view": view,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
+        "maxSizeMB": max_size_mb,
     })
 
 
 def FilesInput(title: str | None = None, file_extensions: str | None = None, group: str | None = None,
-               hidden: bool | None = False, depends: str | None = None, limit: int | None = 10) -> type:
+               hidden: bool | None = False, depends: str | None = None, limit: int | None = 10,
+               view: Literal["button", "dropzone"] | None = None, visibility: str | Visibility = Visibility.ALWAYS, max_size_mb: float | None = 100.0) -> type:
     """
     Factory function for creating a multiple files input type.
-    
+
     Args:
         title: Title for the files input
         file_extensions: Comma-separated list of allowed file extensions (e.g., ".pdf,.txt")
@@ -249,16 +297,21 @@ def FilesInput(title: str | None = None, file_extensions: str | None = None, gro
         hidden: Whether the input should be hidden in the UI
         depends: Specifies the parameter that this value depends on or is derived from.
         limit: Limit count files.
+        view: View mode for the files input ("button" or "dropzone")
+        visibility: Visibility state of the field (default: "always")
 
     Returns:
         Type annotation for files input field
     """
-    return WithJsonSchema({
-        "type": "files",
+    return Field(json_schema_extra={
+        UI_TYPE_KEY: "files",
         "title": title,
         "fileExtensions": file_extensions,
         "group": group,
         "hidden": hidden,
         "depends": depends,
         "limit": limit,
-    })
+        "view": view,
+        "visibility": visibility.value if isinstance(visibility, Visibility) else visibility,
+        "maxSizeMB": max_size_mb,
+    })

agent_lab_sdk/storage/__init__.py

@@ -1,3 +1,4 @@
 from .storage import store_file_in_sd_asset
+from .storage_v2 import upload_file, FileUploadResponse
 
-__all__ = ["store_file_in_sd_asset"]
+__all__ = ["store_file_in_sd_asset", "upload_file", "FileUploadResponse"]