airbyte-agent-amazon-ads 0.1.6__py3-none-any.whl → 0.1.12__py3-none-any.whl

airbyte_agent_amazon_ads/connector.py

@@ -4,8 +4,11 @@ Amazon-Ads connector.
 
 from __future__ import annotations
 
+import inspect
+import json
 import logging
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, Mapping, TypeVar, overload
 try:
     from typing import Literal
 except ImportError:
@@ -21,6 +24,9 @@ from .types import (
     SponsoredProductCampaignsGetParams,
     SponsoredProductCampaignsListParams,
     SponsoredProductCampaignsListParamsStatefilter,
+    AirbyteSearchParams,
+    ProfilesSearchFilter,
+    ProfilesSearchQuery,
 )
 if TYPE_CHECKING:
     from .models import AmazonAdsAuthConfig
@@ -34,11 +40,47 @@ from .models import (
     Portfolio,
     Profile,
     SponsoredProductCampaign,
+    AirbyteSearchHit,
+    AirbyteSearchResult,
+    ProfilesSearchData,
+    ProfilesSearchResult,
 )
 
 # TypeVar for decorator type preservation
 _F = TypeVar("_F", bound=Callable[..., Any])
 
+DEFAULT_MAX_OUTPUT_CHARS = 50_000  # ~50KB default, configurable per-tool
+
+
+def _raise_output_too_large(message: str) -> None:
+    try:
+        from pydantic_ai import ModelRetry  # type: ignore[import-not-found]
+    except Exception as exc:
+        raise RuntimeError(message) from exc
+    raise ModelRetry(message)
+
+
+def _check_output_size(result: Any, max_chars: int | None, tool_name: str) -> Any:
+    if max_chars is None or max_chars <= 0:
+        return result
+
+    try:
+        serialized = json.dumps(result, default=str)
+    except (TypeError, ValueError):
+        return result
+
+    if len(serialized) > max_chars:
+        truncated_preview = serialized[:500] + "..." if len(serialized) > 500 else serialized
+        _raise_output_too_large(
+            f"Tool '{tool_name}' output too large ({len(serialized):,} chars, limit {max_chars:,}). "
+            "Please narrow your query by: using the 'fields' parameter to select only needed fields, "
+            "adding filters, or reducing the 'limit'. "
+            f"Preview: {truncated_preview}"
+        )
+
+    return result
+
+
 
 
 class AmazonAdsConnector:
@@ -49,7 +91,7 @@ class AmazonAdsConnector:
     """
 
     connector_name = "amazon-ads"
-    connector_version = "1.0.1"
+    connector_version = "1.0.3"
     vendored_sdk_version = "0.1.0"  # Version of vendored connector-sdk
 
     # Map of (entity, action) -> needs_envelope for envelope wrapping decision
@@ -80,7 +122,7 @@ class AmazonAdsConnector:
         airbyte_client_id: str | None = None,
         airbyte_client_secret: str | None = None,
         on_token_refresh: Any | None = None,
-        region_url: str | None = None    ):
+        region: str | None = None    ):
         """
         Initialize a new amazon-ads connector instance.
 
@@ -95,7 +137,7 @@ class AmazonAdsConnector:
             airbyte_client_secret: Airbyte OAuth client secret (required for hosted mode)
             on_token_refresh: Optional callback for OAuth2 token refresh persistence.
                 Called with new_tokens dict when tokens are refreshed. Can be sync or async.
-                Example: lambda tokens: save_to_database(tokens)            region_url: The Amazon Ads API endpoint URL based on region:
+                Example: lambda tokens: save_to_database(tokens)            region: The Amazon Ads API endpoint URL based on region:
 - NA (North America): https://advertising-api.amazon.com
 - EU (Europe): https://advertising-api-eu.amazon.com
 - FE (Far East): https://advertising-api-fe.amazon.com
@@ -142,8 +184,8 @@ class AmazonAdsConnector:
 
             # Build config_values dict from server variables
             config_values: dict[str, str] = {}
-            if region_url:
-                config_values["region_url"] = region_url
+            if region:
+                config_values["region"] = region
 
             self._executor = LocalExecutor(
                 model=AmazonAdsConnectorModel,
@@ -154,8 +196,8 @@ class AmazonAdsConnector:
 
             # Update base_url with server variables if provided
             base_url = self._executor.http_client.base_url
-            if region_url:
-                base_url = base_url.replace("{region_url}", region_url)
+            if region:
+                base_url = base_url.replace("{region}", region)
             self._executor.http_client.base_url = base_url
 
         # Initialize entity query objects
@@ -218,15 +260,15 @@ class AmazonAdsConnector:
     async def execute(
         self,
         entity: str,
-        action: str,
-        params: dict[str, Any]
+        action: Literal["list", "get", "search"],
+        params: Mapping[str, Any]
     ) -> AmazonAdsExecuteResult[Any] | AmazonAdsExecuteResultWithMeta[Any, Any] | Any: ...
 
     async def execute(
         self,
         entity: str,
-        action: str,
-        params: dict[str, Any] | None = None
+        action: Literal["list", "get", "search"],
+        params: Mapping[str, Any] | None = None
     ) -> Any:
         """
         Execute an entity operation with full type safety.
@@ -254,16 +296,17 @@ class AmazonAdsConnector:
         from ._vendored.connector_sdk.executor import ExecutionConfig
 
         # Remap parameter names from snake_case (TypedDict keys) to API parameter names
-        if params:
+        resolved_params = dict(params) if params is not None else None
+        if resolved_params:
             param_map = self._PARAM_MAP.get((entity, action), {})
             if param_map:
-                params = {param_map.get(k, k): v for k, v in params.items()}
+                resolved_params = {param_map.get(k, k): v for k, v in resolved_params.items()}
 
         # Use ExecutionConfig for both local and hosted executors
         config = ExecutionConfig(
             entity=entity,
             action=action,
-            params=params
+            params=resolved_params
         )
 
         result = await self._executor.execute(config)
@@ -290,41 +333,67 @@ class AmazonAdsConnector:
     # ===== INTROSPECTION METHODS =====
 
     @classmethod
-    def describe(cls, func: _F) -> _F:
+    def tool_utils(
+        cls,
+        func: _F | None = None,
+        *,
+        update_docstring: bool = True,
+        max_output_chars: int | None = DEFAULT_MAX_OUTPUT_CHARS,
+    ) -> _F | Callable[[_F], _F]:
         """
-        Decorator that populates a function's docstring with connector capabilities.
-
-        This class method can be used as a decorator to automatically generate
-        comprehensive documentation for AI tool functions.
+        Decorator that adds tool utilities like docstring augmentation and output limits.
 
         Usage:
             @mcp.tool()
-            @AmazonAdsConnector.describe
+            @AmazonAdsConnector.tool_utils
             async def execute(entity: str, action: str, params: dict):
-                '''Execute operations.'''
                 ...
 
-        The decorated function's __doc__ will be updated with:
-        - Available entities and their actions
-        - Parameter signatures with required (*) and optional (?) markers
-        - Response structure documentation
-        - Example questions (if available in OpenAPI spec)
+            @mcp.tool()
+            @AmazonAdsConnector.tool_utils(update_docstring=False, max_output_chars=None)
+            async def execute(entity: str, action: str, params: dict):
+                ...
 
         Args:
-            func: The function to decorate
-
-        Returns:
-            The same function with updated __doc__
+            update_docstring: When True, append connector capabilities to __doc__.
+            max_output_chars: Max serialized output size before raising. Use None to disable.
         """
-        description = generate_tool_description(AmazonAdsConnectorModel)
 
-        original_doc = func.__doc__ or ""
-        if original_doc.strip():
-            func.__doc__ = f"{original_doc.strip()}\n{description}"
-        else:
-            func.__doc__ = description
+        def decorate(inner: _F) -> _F:
+            if update_docstring:
+                description = generate_tool_description(AmazonAdsConnectorModel)
+                original_doc = inner.__doc__ or ""
+                if original_doc.strip():
+                    full_doc = f"{original_doc.strip()}\n{description}"
+                else:
+                    full_doc = description
+            else:
+                full_doc = ""
 
-        return func
+            if inspect.iscoroutinefunction(inner):
+
+                @wraps(inner)
+                async def aw(*args: Any, **kwargs: Any) -> Any:
+                    result = await inner(*args, **kwargs)
+                    return _check_output_size(result, max_output_chars, inner.__name__)
+
+                wrapped = aw
+            else:
+
+                @wraps(inner)
+                def sw(*args: Any, **kwargs: Any) -> Any:
+                    result = inner(*args, **kwargs)
+                    return _check_output_size(result, max_output_chars, inner.__name__)
+
+                wrapped = sw
+
+            if update_docstring:
+                wrapped.__doc__ = full_doc
+            return wrapped  # type: ignore[return-value]
+
+        if func is not None:
+            return decorate(func)
+        return decorate
 
     def list_entities(self) -> list[dict[str, Any]]:
         """
@@ -441,6 +510,65 @@ information about the advertiser's account in a specific marketplace.
 
 
 
+    async def search(
+        self,
+        query: ProfilesSearchQuery,
+        limit: int | None = None,
+        cursor: str | None = None,
+        fields: list[list[str]] | None = None,
+    ) -> ProfilesSearchResult:
+        """
+        Search profiles records from Airbyte cache.
+
+        This operation searches cached data from Airbyte syncs.
+        Only available in hosted execution mode.
+
+        Available filter fields (ProfilesSearchFilter):
+        - account_info: 
+        - country_code: 
+        - currency_code: 
+        - daily_budget: 
+        - profile_id: 
+        - timezone: 
+
+        Args:
+            query: Filter and sort conditions. Supports operators like eq, neq, gt, gte, lt, lte,
+                   in, like, fuzzy, keyword, not, and, or. Example: {"filter": {"eq": {"status": "active"}}}
+            limit: Maximum results to return (default 1000)
+            cursor: Pagination cursor from previous response's next_cursor
+            fields: Field paths to include in results. Each path is a list of keys for nested access.
+                    Example: [["id"], ["user", "name"]] returns id and user.name fields.
+
+        Returns:
+            ProfilesSearchResult with hits (list of AirbyteSearchHit[ProfilesSearchData]) and pagination info
+
+        Raises:
+            NotImplementedError: If called in local execution mode
+        """
+        params: dict[str, Any] = {"query": query}
+        if limit is not None:
+            params["limit"] = limit
+        if cursor is not None:
+            params["cursor"] = cursor
+        if fields is not None:
+            params["fields"] = fields
+
+        result = await self._connector.execute("profiles", "search", params)
+
+        # Parse response into typed result
+        return ProfilesSearchResult(
+            hits=[
+                AirbyteSearchHit[ProfilesSearchData](
+                    id=hit.get("id"),
+                    score=hit.get("score"),
+                    data=ProfilesSearchData(**hit.get("data", {}))
+                )
+                for hit in result.get("hits", [])
+            ],
+            next_cursor=result.get("next_cursor"),
+            took_ms=result.get("took_ms")
+        )
+
 class PortfoliosQuery:
     """
     Query class for Portfolios entity operations.