agent-lab-sdk 0.1.58.dev3__tar.gz → 0.1.59.dev1__tar.gz

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/PKG-INFO

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

agent_lab_sdk-0.1.59.dev1/agent_lab_sdk/core/__init__.py

@@ -0,0 +1,3 @@
+import agent_lab_sdk.core.context
+
+__all__ = ["context"]

agent_lab_sdk-0.1.59.dev1/agent_lab_sdk/core/context.py

@@ -0,0 +1,5 @@
+from contextvars import ContextVar
+from typing import Optional
+
+trace_id_var: ContextVar[Optional[str]] = ContextVar("trace_id_var", default=None)
+request_id_var: ContextVar[Optional[str]] = ContextVar("request_id_var", default=None)

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/core/settings.py

@@ -41,6 +41,9 @@ class Settings(BaseSettings):
 
     GIGACHAT_USE_MTLS: bool = False
 
+    DISABLE_WORKER_CONTEXT: bool = False
+    ENABLE_HANDLE_GIGACHAT_STATUS: bool = False
+
     class Config:
         env_prefix = ENV_PREFIX
 

agent_lab_sdk-0.1.59.dev1/agent_lab_sdk/core/utils.py

@@ -0,0 +1,28 @@
+import contextvars
+
+from agent_lab_sdk.core.context import trace_id_var, request_id_var
+from agent_lab_sdk.core.settings import settings
+
+
+def _get_worker_context():
+    if settings.DISABLE_WORKER_CONTEXT:
+        return {}
+    ctx = contextvars.copy_context()
+    for var, value in ctx.items():
+        if var.name == "worker_context":
+            return value
+    return {}
+
+def get_trace_id() -> str | None:
+    trace_id = trace_id_var.get()
+    if trace_id is None:
+        ctx = _get_worker_context()
+        trace_id = ctx.get("trace_id", None)
+    return trace_id
+
+def get_request_id() -> str | None:
+    request_id = request_id_var.get()
+    if request_id is None:
+        ctx = _get_worker_context()
+        request_id = ctx.get("request_id", None)
+    return request_id

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/langgraph/checkpoint/agw_saver.py

@@ -30,6 +30,7 @@ from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
 
 from .serde import Serializer
 from agent_lab_sdk.metrics import get_metric
+from agent_lab_sdk.core.utils import get_trace_id
 
 __all__ = ["AsyncAGWCheckpointSaver"]
 
@@ -497,9 +498,13 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
         if not cfg:
             return {}
         user_id = cfg.get("configurable", {}).get("user_id")
+        trace_id = cfg.get("configurable", {}).get("trace_id")
+        headers = {}
         if user_id:
-            return {"x-user-id": str(user_id)}
-        return {}
+            headers["x-user-id"] = str(user_id)
+        if trace_id:
+            headers["x-trace-id"] = str(trace_id)
+        return headers
 
     # ----------------------- config <-> api --------------------------
     def _to_api_config(self, cfg: RunnableConfig | None) -> Dict[str, Any]:
@@ -611,6 +616,9 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
     ) -> httpx.Response:
         if headers:
             kw["headers"] = {**kw.get("headers", {}), **headers}
+        trace_id = get_trace_id()
+        if trace_id:
+            kw["headers"] = {**{"x-trace-id": trace_id}, **kw.get("headers", {})}
 
         if "json" in kw:
             payload = kw.pop("json")

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/llm/agw_token_manager.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import Optional
 
 import requests
@@ -11,6 +12,7 @@ import time
 from datetime import datetime, timedelta, timezone
 
 from agent_lab_sdk.core.settings import settings
+from agent_lab_sdk.core.utils import get_trace_id
 
 logger = logging.getLogger(__name__)
 
@@ -37,12 +39,16 @@ class AgwTokenManager:
         params = {"provider": provider}
         if token_type:
             params["type"] = token_type
+
+        headers = {
+            "x-trace-id": get_trace_id() or str(uuid.uuid4()),
+        }
         
         max_retries = TOKEN_PROVIDER_AGW_DEFAULT_MAX_RETRIES
         for attempt in range(1, max_retries + 1):
             try:
                 logger.info(f"Попытка получения токена из AGW ({attempt}/{max_retries})...")
-                resp = requests.post(req_url, params=params, data={}, verify=False, timeout=TOKEN_PROVIDER_AGW_TIMEOUT_SEC)
+                resp = requests.post(req_url, params=params, data={}, verify=False, headers=headers, timeout=TOKEN_PROVIDER_AGW_TIMEOUT_SEC)
                 resp.raise_for_status()
                 result = resp.json()
 

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/llm/clients/agent_service_client.py

@@ -1,10 +1,12 @@
 import logging
+import uuid
 from urllib.parse import urljoin
 
 import httpx
 from httpx import HTTPTransport
 
 from agent_lab_sdk.core.settings import settings, ENV_PREFIX
+from agent_lab_sdk.core.utils import get_trace_id
 
 logger = logging.getLogger(__name__)
 
@@ -45,7 +47,8 @@ class AgentServiceClient:
                 "credentialId": credential_id,
             }
             headers = {
-                "Authorization": f"Bearer {settings.AGENT_SERVICE_TOKEN}"
+                "Authorization": f"Bearer {settings.AGENT_SERVICE_TOKEN}",
+                "x-trace-id": get_trace_id() or str(uuid.uuid4()),
             }
             try:
                 response = client.post(url, headers=headers, json=body)

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/llm/gigachat_token_manager.py

@@ -16,6 +16,7 @@ import urllib3
 
 from agent_lab_sdk.core.settings import settings
 from agent_lab_sdk.llm.managers.ags_token_manager import AgsTokenManager
+from agent_lab_sdk.core.utils import get_trace_id
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -145,7 +146,8 @@ class GigaChatTokenManager:
             "Accept": "application/json",
             "User-Agent": "agent-toolkit",
             "RqUID": str(uuid.uuid4()),  # Уникальный идентификатор запроса
-            "Authorization": f"Basic {gigachat_credentials}"
+            "Authorization": f"Basic {gigachat_credentials}",
+            "x-trace-id": get_trace_id() or str(uuid.uuid4()),
         }
 
         max_retries = DEFAULT_MAX_RETRIES

agent_lab_sdk-0.1.59.dev1/agent_lab_sdk/llm/gigachat_wrapper.py

@@ -0,0 +1,94 @@
+import inspect
+import logging
+import functools
+import uuid
+from http import HTTPStatus
+
+from gigachat.context import trace_id_cvar
+from langchain_gigachat import GigaChat, GigaChatEmbeddings
+from langchain_core.messages import AIMessage
+from gigachat.exceptions import ResponseError
+
+from agent_lab_sdk.core.settings import settings
+from agent_lab_sdk.core.utils import get_trace_id
+
+logger = logging.getLogger(__name__)
+
+
+def _method_wrapper(method):
+    is_async = inspect.iscoroutinefunction(method)
+
+    @functools.wraps(method)
+    async def async_wrapper(self, *args, **kwargs):
+        trace_id = get_trace_id() or str(uuid.uuid4())
+        trace_id_cvar.set(trace_id)
+
+        try:
+            return await method(self, *args, **kwargs)
+        except ResponseError as e:
+            _, status, _, _ = e.args
+
+            if settings.ENABLE_HANDLE_GIGACHAT_STATUS and status in (HTTPStatus.NOT_FOUND, HTTPStatus.FORBIDDEN):
+                logger.warning(
+                    f"Encountered HTTP {status}. Executing 'on_request_error' handler."
+                )
+                return AIMessage(
+                    content=f"Модель '{self.model}' на контуре '{self.base_url}' временно недоступна."
+                )
+
+            logger.error(f"GigaChat error: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Unknown error: {e!r}")
+            raise
+
+    @functools.wraps(method)
+    def sync_wrapper(self, *args, **kwargs):
+        trace_id = get_trace_id() or str(uuid.uuid4())
+        trace_id_cvar.set(trace_id)
+
+        try:
+            return method(self, *args, **kwargs)
+        except ResponseError as e:
+            _, status, _, _ = e.args
+
+            if settings.ENABLE_HANDLE_GIGACHAT_STATUS and status in (HTTPStatus.NOT_FOUND, HTTPStatus.FORBIDDEN):
+                logger.warning(
+                    f"Encountered HTTP {status}. Executing 'on_request_error' handler."
+                )
+                return AIMessage(
+                    content=f"Модель '{self.model}' на контуре '{self.base_url}' временно недоступна."
+                )
+
+            logger.error(f"GigaChat error: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Unknown error: {e!r}")
+            raise
+
+    return async_wrapper if is_async else sync_wrapper
+
+def wrap_gigachat_specific_methods(method_names):
+    def wrap_gigachat(cls):
+        for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
+            if callable(method) and name in method_names:
+                setattr(cls, name, _method_wrapper(method))
+        return cls
+    return wrap_gigachat
+
+
+giga_chat_embeddings_wrap_method_names = ["embed_documents", "embed_query", "aembed_documents", "aembed_query"]
+giga_chat_wrap_method_names = [
+    "invoke", "ainvoke", "stream", "astream", "astream_events", "batch", "abatch",
+    "upload_file", "aupload_file", "batch_as_completed", "abatch_as_completed"
+]
+
+@wrap_gigachat_specific_methods(giga_chat_wrap_method_names)
+class GigaChatWrapper(GigaChat):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+@wrap_gigachat_specific_methods(giga_chat_embeddings_wrap_method_names)
+class GigaChatEmbeddingsWrapper(GigaChatEmbeddings):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/llm/llm.py

@@ -2,6 +2,7 @@ from langchain_gigachat.chat_models import GigaChat
 from langchain_gigachat.embeddings import GigaChatEmbeddings
 from agent_lab_sdk.core.settings import settings
 from agent_lab_sdk.llm.gigachat_token_manager import GigaChatTokenManager
+from agent_lab_sdk.llm.gigachat_wrapper import GigaChatEmbeddingsWrapper, GigaChatWrapper
 from agent_lab_sdk.llm.token_manager_utils import build_token_manager_kwargs
 from agent_lab_sdk.llm.throttled import ThrottledGigaChat, ThrottledGigaChatEmbeddings
 from typing import Union
@@ -48,9 +49,9 @@ def get_model(
     verify_ssl_certs = os.getenv("GIGACHAT_VERIFY_SSL_CERTS", verify_ssl_certs)
 
     if type == "chat":
-        _class = ThrottledGigaChat if throttled else GigaChat
+        _class = ThrottledGigaChat if throttled else GigaChatWrapper
     elif type == "embeddings":
-        _class = ThrottledGigaChatEmbeddings if throttled else GigaChatEmbeddings
+        _class = ThrottledGigaChatEmbeddings if throttled else GigaChatEmbeddingsWrapper
     else:
         raise ValueError(f"unsupported type {type}. possible values: chat, embeddings")
 

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/llm/throttled.py

@@ -3,11 +3,10 @@ import os
 from typing import Optional
 
 import langchain_gigachat.embeddings.gigachat
-from langchain_gigachat import GigaChatEmbeddings
-from langchain_gigachat.chat_models import GigaChat
 from pydantic import PrivateAttr
 
 from agent_lab_sdk.core.settings import settings
+from agent_lab_sdk.llm.gigachat_wrapper import GigaChatEmbeddingsWrapper, GigaChatWrapper
 from agent_lab_sdk.llm.limiting import get_limiter_with_metrics
 from agent_lab_sdk.llm.managers.ags_token_manager import AgsTokenManager
 
@@ -46,7 +45,7 @@ _semaphores = {
     "embed": get_limiter_with_metrics("embed", MAX_EMBED_CONCURRENCY),
 }
 
-class ThrottledGigaChatEmbeddings(GigaChatEmbeddings):
+class ThrottledGigaChatEmbeddings(GigaChatEmbeddingsWrapper):
     _manage_access_token: bool = PrivateAttr(default=True)
     _retry_enabled: bool = PrivateAttr(default=False)
     _base_kwargs: dict = PrivateAttr(default={})
@@ -95,7 +94,7 @@ class ThrottledGigaChatEmbeddings(GigaChatEmbeddings):
         self._use_token_provider_ags = use_token_provider_ags
 
 
-    def _fresh(self) -> GigaChatEmbeddings:
+    def _fresh(self) -> GigaChatEmbeddingsWrapper:
         if self._manage_access_token:
             new_kwargs = _resolve_token_kwargs(
                 self._model,
@@ -104,7 +103,7 @@ class ThrottledGigaChatEmbeddings(GigaChatEmbeddings):
                 self._base_kwargs,
                 self._token_manager_kwargs
             )
-            return GigaChatEmbeddings(**new_kwargs)
+            return GigaChatEmbeddingsWrapper(**new_kwargs)
         else:    
             # возвращаем proxy объект чтобы не ломать цепочку наследования
             return super(ThrottledGigaChatEmbeddings, self)
@@ -142,7 +141,7 @@ class ThrottledGigaChatEmbeddings(GigaChatEmbeddings):
                 lambda: self._fresh().aembed_query(*args, **kwargs),
             )
 
-class ThrottledGigaChat(GigaChat):
+class ThrottledGigaChat(GigaChatWrapper):
     _manage_access_token: bool = PrivateAttr(default=True)
     _retry_enabled: bool = PrivateAttr(default=False)
     _base_kwargs: dict = PrivateAttr(default={})
@@ -190,7 +189,7 @@ class ThrottledGigaChat(GigaChat):
         self._model = model
         self._use_token_provider_ags = use_token_provider_ags
 
-    def _fresh(self) -> GigaChat:
+    def _fresh(self) -> GigaChatWrapper:
         if self._manage_access_token:
             new_kwargs = _resolve_token_kwargs(
                 self._model,
@@ -199,7 +198,7 @@ class ThrottledGigaChat(GigaChat):
                 self._base_kwargs,
                 self._token_manager_kwargs
             )
-            new_gigachat = GigaChat(**new_kwargs)
+            new_gigachat = GigaChatWrapper(**new_kwargs)
             # это поле которое управляется классом BaseChatModel, по другому его пока никак не прокинуть
             new_gigachat.disable_streaming = self.disable_streaming
             return new_gigachat

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/storage/storage.py

@@ -1,10 +1,14 @@
 import base64
 import os
 import mimetypes
+import uuid
 from typing import Optional
 from io import BytesIO
 import requests
 
+from agent_lab_sdk.core.utils import get_trace_id
+
+
 def store_file_in_sd_asset(filename: str, file_base64: str, folder: str = "giga-agents") -> Optional[str]:
     """
     Загружает файл в формате base64 в SD Asset API и возвращает URL загруженного файла.
@@ -32,7 +36,8 @@ def store_file_in_sd_asset(filename: str, file_base64: str, folder: str = "giga-
     
     # Заголовки запроса
     headers = {
-        "X-Api-Key": api_key
+        "X-Api-Key": api_key,
+        "x-trace-id": get_trace_id() or str(uuid.uuid4()),
     }
 
     mimetype, _ = mimetypes.guess_type(filename)

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/storage/storage_v2.py

@@ -1,4 +1,5 @@
 import os
+import uuid
 from typing import Optional
 import logging
 import requests
@@ -6,6 +7,8 @@ from urllib.parse import urljoin
 from pydantic import BaseModel
 import httpx
 
+from agent_lab_sdk.core.utils import get_trace_id
+
 logger = logging.getLogger(__name__)
 
 
@@ -43,7 +46,8 @@ def upload_file(filename: str, file_bytes: bytes) -> Optional[FileUploadResponse
 
     headers = {
         "x-agent-id": ai_agent,
-        "Content-Type": "application/octet-stream"
+        "Content-Type": "application/octet-stream",
+        "x-trace-id": get_trace_id() or str(uuid.uuid4()),
     }
 
     params = {
@@ -94,7 +98,11 @@ async def async_upload_file(filename: str, file_bytes: bytes, timeout=30.) -> Op
     gateway_url = os.getenv("STORAGE_PROVIDER_AGW_URL", "http://localhost")
     url = _get_upload_path(gateway_url)
 
-    headers = {"x-agent-id": ai_agent, "Content-Type": "application/octet-stream"}
+    headers = {
+        "x-agent-id": ai_agent,
+        "Content-Type": "application/octet-stream",
+        "x-trace-id": get_trace_id() or str(uuid.uuid4()),
+    }
     params = {"name": filename}
 
     logger.debug(

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk/tools/search_plugin/search_plugin.py

@@ -6,7 +6,8 @@ from enum import StrEnum
 
 import httpx
 
-from agent_lab_sdk.core.settings import settings, ENV_PREFIX
+from agent_lab_sdk.core.settings import settings
+from agent_lab_sdk.core.utils import get_trace_id
 from agent_lab_sdk.llm.limiting import get_limiter_with_metrics
 from agent_lab_sdk.metrics import get_metric
 import xml.etree.ElementTree as ET
@@ -128,6 +129,7 @@ class SearchPlugin:
             "X-Client-ID": str(uuid.uuid4()),
             "X-Request-ID": str(uuid.uuid4()),
             "X-Session-ID": str(uuid.uuid4()),
+            "x-trace-id": get_trace_id() or str(uuid.uuid4()),
             **headers
         }
 
@@ -190,6 +192,7 @@ class SearchPlugin:
         headers = {
             "Authorization": f"{self._api_key}",
             "User-Agent": "agent-toolkit",
+            "x-trace-id": get_trace_id() or str(uuid.uuid4()),
             **headers
         }
         url = self._base_url or self._get_setting_or_raise("SEARCH_IMAGES_PLUGIN_URL")

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk.egg-info/PKG-INFO

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

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/agent_lab_sdk.egg-info/SOURCES.txt

@@ -8,14 +8,17 @@ agent_lab_sdk.egg-info/dependency_links.txt
 agent_lab_sdk.egg-info/requires.txt
 agent_lab_sdk.egg-info/top_level.txt
 agent_lab_sdk/core/__init__.py
+agent_lab_sdk/core/context.py
 agent_lab_sdk/core/settings.py
 agent_lab_sdk/core/unified_semaphore.py
+agent_lab_sdk/core/utils.py
 agent_lab_sdk/langgraph/checkpoint/__init__.py
 agent_lab_sdk/langgraph/checkpoint/agw_saver.py
 agent_lab_sdk/langgraph/checkpoint/serde.py
 agent_lab_sdk/llm/__init__.py
 agent_lab_sdk/llm/agw_token_manager.py
 agent_lab_sdk/llm/gigachat_token_manager.py
+agent_lab_sdk/llm/gigachat_wrapper.py
 agent_lab_sdk/llm/limiting.py
 agent_lab_sdk/llm/llm.py
 agent_lab_sdk/llm/retry_utils.py

{agent_lab_sdk-0.1.58.dev3 → agent_lab_sdk-0.1.59.dev1}/pyproject.toml

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