agent-lab-sdk 0.1.27__py3-none-any.whl → 0.1.29__py3-none-any.whl

agent_lab_sdk/langgraph/checkpoint/agw_saver.py

@@ -105,10 +105,6 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
 
     def _safe_load(self, obj: Any) -> Any:
         if isinstance(obj, (dict, list)):          # уже распакованный JSON
-            # Используем сериализатор напрямую с object_hook/_reviver если это dict
-            if isinstance(obj, dict) and hasattr(self.serde, '_reviver'):
-                return self.serde._reviver(obj)
-            # Для list или если нет _reviver - используем dumps/loads
             return self.serde.loads(orjson.dumps(obj))
         if isinstance(obj, str):
             # сначала plain JSON-строка

agent_lab_sdk/llm/llm.py

@@ -5,10 +5,16 @@ from agent_lab_sdk.llm.throttled import ThrottledGigaChat, ThrottledGigaChatEmbe
 from typing import Union
 import os
 
-def get_model(type="chat", throttled = False, **kwargs) -> Union[GigaChat, GigaChatEmbeddings, ThrottledGigaChat, ThrottledGigaChatEmbeddings]:
+def get_model(
+    type: str = "chat",
+    throttled: bool = False,
+    manage_access_token: bool = True,
+    **kwargs
+) -> Union[GigaChat, GigaChatEmbeddings, ThrottledGigaChat, ThrottledGigaChatEmbeddings]:
     """
     * type - определяет тип моледи : chat | embeddings
     * throttled - оборачивает модель в класс обертку с дополнительным регулированием через семафор
+    * manage_access_token - включает режим авто-обновления токена в обёртках (только если throttled=True)
     * kwargs - прокидываются в модель
     """
     access_token = kwargs.pop("access_token", None)
@@ -22,17 +28,14 @@ def get_model(type="chat", throttled = False, **kwargs) -> Union[GigaChat, GigaC
     verify_ssl_certs = os.getenv("GIGACHAT_VERIFY_SSL_CERTS", verify_ssl_certs)
 
     if type == "chat":
-        _class = GigaChat
-        if throttled:
-            _class = ThrottledGigaChat
+        _class = ThrottledGigaChat if throttled else GigaChat
     elif type == "embeddings":
-        _class = GigaChatEmbeddings
-        if throttled:
-            _class = ThrottledGigaChatEmbeddings
+        _class = ThrottledGigaChatEmbeddings if throttled else GigaChatEmbeddings
     else:
         raise ValueError(f"unsupported type {type}. possible values: chat, embeddings")
-    return _class(
-        access_token=access_token,
-        verify_ssl_certs=verify_ssl_certs,
-        **kwargs
-        )
+
+    if throttled and manage_access_token:
+        # Включаем режим авто-обновления токена в обёртках
+        kwargs["manage_access_token"] = True
+
+    return _class(verify_ssl_certs=verify_ssl_certs, **kwargs)

agent_lab_sdk/llm/throttled.py

@@ -2,6 +2,7 @@ import os
 import asyncio
 import threading
 import time
+from pydantic import PrivateAttr
 from langchain_gigachat.chat_models import GigaChat
 from langchain_gigachat import GigaChatEmbeddings
 import langchain_gigachat.embeddings.gigachat
@@ -13,6 +14,7 @@ MAX_EMBED_CONCURRENCY = int(os.getenv("MAX_EMBED_CONCURRENCY", "100000"))
 
 
 from agent_lab_sdk.metrics import get_metric
+from agent_lab_sdk.llm.gigachat_token_manager import GigaChatTokenManager
 
 def create_metrics(prefix: str):
     in_use = get_metric(
@@ -107,71 +109,101 @@ _semaphores = {
 }
 
 class ThrottledGigaChatEmbeddings(GigaChatEmbeddings):
+    _manage_access_token: bool = PrivateAttr(default=True)
+    _base_kwargs: dict = PrivateAttr(default={})
+
+    def __init__(self, *args, manage_access_token=True, **kwargs):
+        print(self)
+        self._manage_access_token = manage_access_token
+        self._base_kwargs = dict(kwargs)
+
+        if self._manage_access_token and "access_token" not in kwargs:
+            token = GigaChatTokenManager.get_token()
+            super().__init__(access_token=token, **kwargs)
+        else:
+            super().__init__(**kwargs)
+
+    def _fresh(self) -> GigaChatEmbeddings:
+        if self._manage_access_token:
+            return GigaChatEmbeddings(access_token=GigaChatTokenManager.get_token(), **self._base_kwargs)
+        else:    
+            return self
+
+
     def embed_documents(self, *args, **kwargs):
         with _semaphores["embed"]:
-            return super().embed_documents(*args, **kwargs)
+            return self._fresh().embed_documents(*args, **kwargs)
 
     def embed_query(self, *args, **kwargs):
-        # здесь семафор не нужен, под капотом вызвается embed_documents, семафор уже там
-        return super().embed_query(*args, **kwargs)
+        with _semaphores["embed"]:
+            return self._fresh().embed_query(*args, **kwargs)
 
     async def aembed_documents(self, *args, **kwargs):
         async with _semaphores["embed"]:
-            return await super().aembed_documents(*args, **kwargs)
+            return await self._fresh().aembed_documents(*args, **kwargs)
 
     async def aembed_query(self, *args, **kwargs):
-        # здесь семафор не нужен, под капотом вызвается aembed_documents, семафор уже там
-        return await super().aembed_query(*args, **kwargs)
+        async with _semaphores["embed"]:
+            return await self._fresh().aembed_query(*args, **kwargs)
 
-# по хорошему бы переопределять клиент гигачата или манкипатчить его, но это не так просто
 class ThrottledGigaChat(GigaChat):
+    _manage_access_token: bool = PrivateAttr(default=True)
+    _base_kwargs: dict = PrivateAttr(default={})
+
+    def __init__(self, *args, manage_access_token=True, **kwargs):
+        print(self)
+        self._manage_access_token = manage_access_token
+        self._base_kwargs = dict(kwargs)
+
+        if self._manage_access_token and "access_token" not in kwargs:
+            token = GigaChatTokenManager.get_token()
+            super().__init__(access_token=token, **kwargs)
+        else:
+            super().__init__(**kwargs)
+
+    def _fresh(self) -> GigaChat:
+        if self._manage_access_token:
+            return GigaChat(access_token=GigaChatTokenManager.get_token(), **self._base_kwargs)
+        else:    
+            return self
+
     def invoke(self, *args, **kwargs):
         with _semaphores["chat"]:
-            return super().invoke(*args, **kwargs)
+            return self._fresh().invoke(*args, **kwargs)
 
     async def ainvoke(self, *args, **kwargs):
         async with _semaphores["chat"]:
-            return await super().ainvoke(*args, **kwargs)
+            return await self._fresh().ainvoke(*args, **kwargs)
 
     def stream(self, *args, **kwargs):
-        if super()._should_stream(async_api=False, **{**kwargs, "stream": True}):
-            with _semaphores["chat"]:
-                for chunk in super().stream(*args, **kwargs):
-                    yield chunk
-        else:
-            # здесь есть проблема когда внутри stream вызывается invoke, поэтому без семафора
-            for chunk in super().stream(*args, **kwargs):
-                    yield chunk
+        with _semaphores["chat"]:
+            for chunk in self._fresh().stream(*args, **kwargs):
+                yield chunk
 
     async def astream(self, *args, **kwargs):
-        if super()._should_stream(async_api=True, **{**kwargs, "stream": True}):
-            async with _semaphores["chat"]:
-                async for chunk in super().astream(*args, **kwargs):
-                    yield chunk
-        else:
-            # здесь есть проблема когда внутри stream вызывается ainvoke, поэтому без семафора
-            async for chunk in super().astream(*args, **kwargs):
+        async with _semaphores["chat"]:
+            async for chunk in self._fresh().astream(*args, **kwargs):
                 yield chunk
 
     async def astream_events(self, *args, **kwargs):
         async with _semaphores["chat"]:
-            async for ev in super().astream_events(*args, **kwargs):
+            async for ev in self._fresh().astream_events(*args, **kwargs):
                 yield ev
 
     def batch(self, *args, **kwargs):
-        # здесь семафор не нужен, под капотом вызывается invoke, семафор уже там
-        return super().batch(*args, **kwargs)
+        with _semaphores["chat"]:
+            return self._fresh().batch(*args, **kwargs)
 
     async def abatch(self, *args, **kwargs):
-        # здесь семафор не нужен, под капотом вызывается ainvoke, семафор уже там
-        return await super().abatch(*args, **kwargs)
+        async with _semaphores["chat"]:
+            return await self._fresh().abatch(*args, **kwargs)
 
     def batch_as_completed(self, *args, **kwargs):
-        # здесь семафор не нужен, под капотом вызывается invoke, семафор уже там
-        for item in super().batch_as_completed(*args, **kwargs):
-            yield item
+        with _semaphores["chat"]:
+            for item in self._fresh().batch_as_completed(*args, **kwargs):
+                yield item
 
     async def abatch_as_completed(self, *args, **kwargs):
-        # здесь семафор не нужен, под капотом вызывается ainvoke, семафор уже там
-        async for item in super().abatch_as_completed(*args, **kwargs):
-            yield item
+        async with _semaphores["chat"]:
+            async for item in self._fresh().abatch_as_completed(*args, **kwargs):
+                yield item

{agent_lab_sdk-0.1.27.dist-info → agent_lab_sdk-0.1.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agent-lab-sdk
-Version: 0.1.27
+Version: 0.1.29
 Summary: SDK для работы с Agent Lab
 Author-email: Andrew Ohurtsov <andermirik@yandex.com>
 License: Proprietary and Confidential — All Rights Reserved

{agent_lab_sdk-0.1.27.dist-info → agent_lab_sdk-0.1.29.dist-info}/RECORD

@@ -1,11 +1,11 @@
 agent_lab_sdk/__init__.py,sha256=1Dlmv-wuz1QuciymKtYtX7jXzr_fkeGTe7aENfEDl3E,108
 agent_lab_sdk/langgraph/checkpoint/__init__.py,sha256=DnKwR1LwbaQ3qhb124lE-tnojrUIVcCdNzHEHwgpL5M,86
-agent_lab_sdk/langgraph/checkpoint/agw_saver.py,sha256=c-qCC--_1SGAHYfFA9o-zCMY_hydVqawnVZzgOQwWoo,14100
+agent_lab_sdk/langgraph/checkpoint/agw_saver.py,sha256=K4ni4x6bMibkC9PphEll7RLtbKbaHCGTz0uM6ouQaV0,13763
 agent_lab_sdk/llm/__init__.py,sha256=Yo9MbYdHS1iX05A9XiJGwWN1Hm4IARGav9mNFPrtDeA,376
 agent_lab_sdk/llm/agw_token_manager.py,sha256=_bPPI8muaEa6H01P8hHQOJHiiivaLd8N_d3OT9UT_80,4787
 agent_lab_sdk/llm/gigachat_token_manager.py,sha256=nlOxHcwJovsmM4cpI4fwMrYjoSeMjelDaHTipXsrUuA,8282
-agent_lab_sdk/llm/llm.py,sha256=cPZll_NoovPXztyGnXthwBxK0YBZ5bRfHe5VI_Z-2jM,1648
-agent_lab_sdk/llm/throttled.py,sha256=9_nm1i3Uuep0VEWsY1KNCllZA-vM202XVdlgXhgC8BA,7005
+agent_lab_sdk/llm/llm.py,sha256=HIp97-rs8Xcxokgptgh6bCfkIx7UJUFOYdGbpi7SsSg,1949
+agent_lab_sdk/llm/throttled.py,sha256=XyY0i_xb71Mt-zsZ6VSXltRK11Bi7CeWVR3mxbIRr4I,7272
 agent_lab_sdk/metrics/__init__.py,sha256=G4VSlzKwupPMM4c6vZaF1rnd0KusKarezDMjli9pVFw,57
 agent_lab_sdk/metrics/metrics.py,sha256=_XTT9vMG7T0u_D2pL371wm8GoBU5fodJ45D2RACnBJw,3439
 agent_lab_sdk/schema/__init__.py,sha256=bHSyXQYkcB9fWBlziWodXR_IzC5nKrdKzrCpyVWNY9o,521
@@ -13,8 +13,8 @@ agent_lab_sdk/schema/input_types.py,sha256=A8yRSFdCo-BSZM50sj-oXge8HdcVz1AkqbBWg
 agent_lab_sdk/schema/log_message.py,sha256=nadi6lZGRuDSPmfbYs9QPpRJUT9Pfy8Y7pGCvyFF5Mw,638
 agent_lab_sdk/storage/__init__.py,sha256=ik1_v1DMTwehvcAEXIYxuvLuCjJCa3y5qAuJqoQpuSA,81
 agent_lab_sdk/storage/storage.py,sha256=ELpt7GRwFD-aWa6ctinfA_QwcvzWLvKS0Wz8FlxVqAs,2075
-agent_lab_sdk-0.1.27.dist-info/licenses/LICENSE,sha256=_TRXHkF3S9ilWBPdZcHLI_S-PRjK0L_SeOb2pcPAdV4,417
-agent_lab_sdk-0.1.27.dist-info/METADATA,sha256=sSXmzRMrpdM2nPgWYTyhuofybxFFT1ImP4LtHJQC9KQ,17884
-agent_lab_sdk-0.1.27.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agent_lab_sdk-0.1.27.dist-info/top_level.txt,sha256=E1efqkJ89KNmPBWdLzdMHeVtH0dYyCo4fhnSb81_15I,14
-agent_lab_sdk-0.1.27.dist-info/RECORD,,
+agent_lab_sdk-0.1.29.dist-info/licenses/LICENSE,sha256=_TRXHkF3S9ilWBPdZcHLI_S-PRjK0L_SeOb2pcPAdV4,417
+agent_lab_sdk-0.1.29.dist-info/METADATA,sha256=uyc3EoyiOojFugnXayZuNI2TX6qppx1Te93pUGWEtrk,17884
+agent_lab_sdk-0.1.29.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+agent_lab_sdk-0.1.29.dist-info/top_level.txt,sha256=E1efqkJ89KNmPBWdLzdMHeVtH0dYyCo4fhnSb81_15I,14
+agent_lab_sdk-0.1.29.dist-info/RECORD,,