agent-lab-sdk 0.1.26__tar.gz → 0.1.28__tar.gz

{agent_lab_sdk-0.1.26 → agent_lab_sdk-0.1.28}/PKG-INFO

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

{agent_lab_sdk-0.1.26 → agent_lab_sdk-0.1.28}/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-0.1.26 → agent_lab_sdk-0.1.28}/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: bool = PrivateAttr(default=True)
+
+    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, **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: bool = PrivateAttr(default=True)
+
+    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, **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.26 → agent_lab_sdk-0.1.28}/agent_lab_sdk.egg-info/PKG-INFO

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

{agent_lab_sdk-0.1.26 → agent_lab_sdk-0.1.28}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "agent-lab-sdk"
-version = "0.1.26"
+version = "0.1.28"
 description = "SDK для работы с Agent Lab"
 readme = "README.md"
 license = { text = "Proprietary and Confidential — All Rights Reserved" }