-**Build, deploy, and govern autonomous AI agents — in a few lines of Python.**
+**Build, deploy, and run autonomous AI agents — governed by default, in a few lines of Python.**
-_The fastest way to deploy managed, multi-agent systems with runtime governance — on any infrastructure._
+aixplain is the operating system for autonomous AI: multi-agent orchestration with **runtime governance on every action**, across cloud, on-prem, edge, and local. The full lifecycle — build → evaluate → deploy → monitor → evolve — on one runtime, instead of stitching tools together.
-aixplain is the operating system for autonomous AI — **multi-agent orchestration and runtime governance across cloud, on-prem, edge, and fully air-gapped environments**. It spans the full lifecycle — **development → evaluation → deployment → monitoring → evolution** — so you ship faster instead of stitching tools together. Agents plan, call models and tools, run code, and adapt at runtime, with **agents-in-the-loop** governance enforced on every execution.
+**On your terms** — **your data** in your perimeter, **your cost** free on local models and tools, pay as you go in the cloud, **your independence** across any model or infrastructure, no lock-in.
-**Build any kind of agent on one runtime** — knowledge (RAG), data, custom-logic, integration, and team agents — with built-in memory, multimodality, and file handling.
-
-Develop your way — **SDK, API, CLI, and MCP** environments all run on the same runtime layer, backed by an integrated marketplace of **900+ AI models, tools, and integrations from 70+ vendors** (OpenAI, Anthropic, Gemini, Firecrawl, Tavily, SerpAPI, a code-execution sandbox, and more). **One API key, one bill**, pay as you go — no per-vendor keys or contracts.
+Build any agent — knowledge (RAG), data, custom-logic, integration, and team — via **SDK, API, CLI, or MCP**, on a marketplace of **900+ models, tools, and integrations**.
## Why aixplain
Less to build, less to operate:
- **Deploy with one call** — `agent.save()` promotes an agent to a persistent, versioned endpoint; no Dockerfiles, queues, or autoscaling to manage.
-- **No integration glue** — reach [900+ models, tools, and integrations](#mcp-server-marketplace) through one key; skip per-provider SDKs, auth, and rate-limit handling.
-- **Governance you don't have to code** — allow-lists, per-asset permissions, rate and usage limits, and enterprise RBAC enforced at runtime.
+- **No integration glue** — reach [900+ models, tools, and integrations](#marketplace) through one key; skip per-provider SDKs, auth, and rate-limit handling.
+- **Guardrails you don't have to build** — allow-lists, per-asset permissions, rate and usage limits, and access control enforced at runtime.
- **Self-debugging** — step-level traces of every plan, tool call, and outcome.
-- **Run it anywhere** — the same definition runs serverless, on-prem, at the edge, or fully air-gapped.
+- **Run it anywhere** — the same definition runs in the cloud, on-prem, at the edge, or locally.
+- **Works with your coding agent** — native [MCP support](#marketplace) for MCP-compatible IDEs and coding agents.
-| | aixplain SDK | Other agent frameworks |
-|---|---|---|
-| Governance | Asset/action allow-lists, per-asset permissions, rate and usage limits, and enterprise RBAC | Usually custom code or external guardrails |
-| Models and tools | 900+ models and tools with one API key | Provider-by-provider setup |
-| Deployment | Cloud (instant) or on-prem | Usually self-assembled runtime and infra |
-| Observability | Built-in traces and dashboards | Varies by framework |
-| [Coding-agent workflows](https://github.com/aixplain/aiXplain/tree/main/skills/aixplain-agent-builder) | Works natively with MCP-compatible coding agents and IDEs | Usually not a first-class workflow target |
+## How it works
-## Agentic OS
+The portable runtime behind aixplain agents: orchestration, governed asset serving, and observability across cloud, on-prem, edge, and local. See the [documentation](https://docs.aixplain.com) for the full architecture.
-Agentic OS is the portable runtime platform behind aixplain agents — orchestration, governed asset serving, and observability across Cloud (instant), on-prem, edge, and air-gapped deployments. See the [documentation](https://docs.aixplain.com) for the full architecture.
+
+
+
---
## Quick start
-> **This README documents SDK v2, the default API.** SDK v1 (the legacy factory API) keeps working until **end of July 2026**, after which v2 is the only supported surface.
+> **This README documents SDK v2, the default API.** SDK v1 (the legacy factory API) keeps working until **August 1, 2026**, after which v2 is the only supported surface.
```bash
pip install aixplain
```
-Get your API key from your [aixplain account](https://console.aixplain.com/settings/keys), then expose it to the SDK:
+Get your API key from your aixplain account, then expose it to the SDK:
```bash
export AIXPLAIN_API_KEY=
@@ -66,7 +60,6 @@ export AIXPLAIN_API_KEY=
### Create and run your first agent
```python
-from uuid import uuid4
from aixplain import Aixplain
aix = Aixplain() # reads AIXPLAIN_API_KEY from the environment
@@ -75,7 +68,7 @@ search_tool = aix.Tool.get("tavily/tavily-web-search/tavily")
search_tool.allowed_actions = ["search"]
agent = aix.Agent(
- name=f"Research agent {uuid4().hex[:8]}",
+ name="Research agent",
description="Answers questions with concise web-grounded findings.",
instructions="Use the search tool when needed and cite key findings.",
tools=[search_tool],
@@ -93,9 +86,8 @@ print(result.data.output)
### Build a multi-agent team
```python
-from uuid import uuid4
from aixplain import Aixplain
-from aixplain.v2 import EditorConfig, EvaluatorConfig, EvaluatorType, Inspector, InspectorAction, InspectorActionConfig, InspectorSeverity, InspectorTarget
+from aixplain.v2 import Inspector
aix = Aixplain() # reads AIXPLAIN_API_KEY from the environment
search_tool = aix.Tool.get("tavily/tavily-web-search/tavily")
@@ -107,29 +99,25 @@ def never_edit(text: str) -> bool:
def passthrough(text: str) -> str:
return text
+# Config is plain data — strings for action/targets/severity, and a Metric,
+# asset-id string, or callable for the `metric` (the universal judge).
noop_inspector = Inspector(
- name=f"noop-output-inspector-{uuid4().hex[:8]}",
- severity=InspectorSeverity.LOW,
- targets=[InspectorTarget.OUTPUT],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function=never_edit,
- ),
- editor=EditorConfig(
- type=EvaluatorType.FUNCTION,
- function=passthrough,
- ),
+ name="noop-output-inspector",
+ severity="low",
+ targets=["output"],
+ action="edit",
+ metric=never_edit,
+ editor=passthrough,
)
researcher = aix.Agent(
- name=f"Researcher {uuid4().hex[:8]}",
+ name="Researcher",
instructions="Find and summarize reliable sources.",
tools=[search_tool],
)
team_agent = aix.Agent(
- name=f"Research team {uuid4().hex[:8]}",
+ name="Research team",
instructions="Research the topic and return exactly 5 concise bullet points.",
subagents=[researcher],
inspectors=[noop_inspector],
@@ -142,10 +130,6 @@ response = team_agent.run(
print(response.data.output)
```
-
-
-
-
Execution order:
```text
@@ -156,20 +140,19 @@ Team agent
├── Orchestrator: routes work to the right subagent
├── Researcher subagent
│ └── Tavily search tool: finds and summarizes reliable sources
-├── Inspector: checks the final output through a simple runtime policy
-├── Orchestrator: decides whether another pass is needed
-└── Responder: returns one final answer
+├── Inspector: validates the output against a runtime policy
+└── Orchestrator: composes and returns the final answer
```
-> **SDK 1 (legacy):** available until end of July 2026 — see the [SDK 1 docs](https://docs.aixplain.com/1.0/).
+> **SDK v1 (legacy):** available until August 1, 2026 — see the [SDK v1 docs](https://docs.aixplain.com/1.0/).
---
-## MCP Server Marketplace
+## Marketplace
-[aixplain Marketplace](https://studio.aixplain.com/browse) now also exposes MCP servers for **900+ models and tools**, allowing external clients to access selected **tool, integration, and model assets**, for example **Opus 4.6, Kimi, Qwen, Airtable, and Slack**, through **aixplain-hosted MCP endpoints** with a single API key 🔑.
+The [aixplain Marketplace](https://studio.aixplain.com/browse) is a catalog of **900+ models, tools, and integrations**. Every asset is reachable through the same three outlets — **SDK, API, and MCP** — with a single API key 🔑.
-Read the full MCP setup guide in the [MCP servers docs](https://docs.aixplain.com/api-reference/mcp-servers).
+For MCP-compatible clients and IDEs, assets (for example Opus 4.6, Kimi, Qwen, Airtable, Slack) are served through aixplain-hosted MCP endpoints. See the [MCP servers docs](https://docs.aixplain.com/api-reference/mcp-servers).
```json
{
@@ -185,32 +168,11 @@ Read the full MCP setup guide in the [MCP servers docs](https://docs.aixplain.co
---
-## Core concepts
-
-| Concept | What it is |
-|---|---|
-| **Agent** | An autonomous entity that reasons, plans, and uses tools to complete a task. |
-| **Team agent** | Multiple specialized agents coordinated by a planner and orchestrator at runtime. |
-| **Tool** | A capability an agent can invoke — a model, an integration, or code (Python/SQL). |
-| **Model** | An LLM, utility, or index asset from the [marketplace](https://studio.aixplain.com/browse). |
-| **Integration** | A connector to an external service (Slack, Airtable, Gmail, and more) that an agent can act through. |
-| **Micro-agents** | Built-in governance components that run inline on every execution: Mentalist (planning), Orchestrator (routing), Inspector (validation — e.g. block pip installs or network access), Bodyguard (security), Responder (formatting). |
-| **Meta-agent** | An agent that improves other agents — the Evolver monitors KPIs and refines behavior over time. |
-
-See the [documentation](https://docs.aixplain.com) for the full API reference.
-
----
-
## Data handling and deployment
-aixplain applies runtime governance and enterprise controls by default:
-
-- **We do not train on your data** — your data is not used to train foundation models.
-- **No data retained by default** — agent memory is opt-in (short-term and long-term).
-- **SOC 2 Type II certified** — enterprise security and compliance posture.
-- **Runtime policy enforcement** — Inspector and Bodyguard govern every agent execution, with allow-lists for assets and actions, per-asset permissions, rate and usage limits, and enterprise RBAC.
-- **Portable deployment options** — Cloud (instant) or on-prem (including VPC and air-gapped environments).
-- **Encryption** — TLS 1.2+ in transit and encrypted storage at rest.
+- **Your data stays yours** — never used to train foundation models; agent memory is opt-in. SOC 2 Type II; TLS 1.2+ in transit, encrypted at rest.
+- **Governed at runtime** — Inspector and Bodyguard enforce allow-lists, per-asset permissions, rate and usage limits, and access control on every execution.
+- **Deploy anywhere** — cloud, on-prem, edge, or local; air-gapped and VPC available on-prem or local.
Learn more at aixplain [Security](https://aixplain.com/security/) and aixplain [pricing](https://aixplain.com/pricing/).
diff --git a/aixplain/utils/user_info_utils.py b/aixplain/utils/user_info_utils.py
index 51eca79e..88fe0e0d 100644
--- a/aixplain/utils/user_info_utils.py
+++ b/aixplain/utils/user_info_utils.py
@@ -33,7 +33,12 @@
@lru_cache(maxsize=1)
def _fetch_ipinfo() -> Dict[str, Any]:
- """Fetch and cache the ipinfo.io payload for the current public IP."""
+ """Fetch and cache the ipinfo.io payload for the current public IP.
+
+ All errors degrade silently to ``{}`` — this is auxiliary metadata
+ and an outage / network-isolation / test-mock should never break a
+ caller's run.
+ """
try:
response = requests.get(_IPINFO_URL, timeout=_IPINFO_TIMEOUT)
response.raise_for_status()
diff --git a/aixplain/v1/modules/pipeline/asset.py b/aixplain/v1/modules/pipeline/asset.py
index 0d6ba2e2..e59add23 100644
--- a/aixplain/v1/modules/pipeline/asset.py
+++ b/aixplain/v1/modules/pipeline/asset.py
@@ -260,12 +260,16 @@ def run(
if response_version == "v1":
polling_response["elapsed_time"] = end - start
return polling_response
+ # `status`/`completed` were already read above via subscript access, which
+ # works for both a PipelineResponse and the plain dict __polling may return
+ # when polling fails. Use subscript/`in` access here too so a failed run
+ # doesn't raise AttributeError on a dict.
return PipelineResponse(
status=status,
completed=completed,
- error=polling_response.get("error"),
+ error=polling_response["error"] if "error" in polling_response else None,
elapsed_time=end - start,
- data=polling_response.get("data", {}),
+ data=polling_response["data"] if "data" in polling_response else {},
**kwargs,
)
diff --git a/aixplain/v2/__init__.py b/aixplain/v2/__init__.py
index 4aab70dc..27e2094e 100644
--- a/aixplain/v2/__init__.py
+++ b/aixplain/v2/__init__.py
@@ -3,22 +3,20 @@
from .core import Aixplain
from .rlm import RLM, RLMResult
from .utility import Utility
-from .agent import Agent, ContextOverflowStrategy
+from .agent import Agent, Budget, ContextOverflowStrategy
from .tool import Tool
+from .skill import Skill
from .actions import Input, Inputs, Action, Actions
+from .integration import TriggerTypeSpec, TriggerEventOption, TriggerTypes
+from .trigger import Trigger, TriggerConfiguration, TriggerRepeatRule
from .file import Resource
from .upload_utils import FileUploader, upload_file, validate_file_for_upload
-from .inspector import (
- Inspector,
- InspectorTarget,
- InspectorAction,
- InspectorOnExhaust,
- InspectorSeverity,
- InspectorActionConfig,
- EvaluatorType,
- EvaluatorConfig,
- EditorConfig,
- PrebuiltInspector,
+from .inspector import Inspector
+from .session import (
+ ExecutionConfig,
+ Session,
+ SessionMessage,
+ SessionMessageAttachment,
)
from .meta_agents import Debugger, DebugResult
from .agent_progress import AgentProgressTracker, ProgressFormat
@@ -84,6 +82,11 @@
EvolveType,
CodeInterpreterModel,
SplittingOptions,
+ SessionStatus,
+ RunStatus,
+ MessageRole,
+ Reaction,
+ AttachmentType,
)
__all__ = [
@@ -92,24 +95,21 @@
"RLMResult",
"Utility",
"Agent",
+ "Budget",
"ContextOverflowStrategy",
"Tool",
+ "Skill",
"Resource",
"FileUploader",
"upload_file",
"validate_file_for_upload",
- # Inspector classes
+ # Session classes
+ "Session",
+ "SessionMessage",
+ "SessionMessageAttachment",
+ "ExecutionConfig",
+ # Inspector
"Inspector",
- "InspectorTarget",
- "InspectorAction",
- "InspectorOnExhaust",
- "InspectorSeverity",
- "InspectorActionConfig",
- "EvaluatorType",
- "EvaluatorConfig",
- "EditorConfig",
- "PrebuiltInspector",
- "ModelResponse",
# Meta-agents
"Debugger",
"DebugResult",
@@ -180,9 +180,21 @@
"EvolveType",
"CodeInterpreterModel",
"SplittingOptions",
+ "SessionStatus",
+ "RunStatus",
+ "MessageRole",
+ "Reaction",
+ "AttachmentType",
# Actions / Inputs hierarchy
"Input",
"Inputs",
"Action",
"Actions",
+ # Triggers
+ "Trigger",
+ "TriggerConfiguration",
+ "TriggerRepeatRule",
+ "TriggerTypeSpec",
+ "TriggerEventOption",
+ "TriggerTypes",
]
diff --git a/aixplain/v2/actions.py b/aixplain/v2/actions.py
index d38a7cc2..cbd21bcf 100644
--- a/aixplain/v2/actions.py
+++ b/aixplain/v2/actions.py
@@ -529,6 +529,20 @@ def __getitem__(self, key: str) -> Action:
raise KeyError(f"Action '{key}' not found")
+ def __getattr__(self, name: str) -> Action:
+ """Dot-notation read: ``tool.actions.search`` (case-insensitive).
+
+ Only consulted when normal attribute lookup fails, so it never shadows
+ the real instance attributes (``_actions`` etc.). Private names are
+ rejected to avoid recursion during initialization.
+ """
+ if name.startswith("_"):
+ raise AttributeError(name)
+ try:
+ return self[name]
+ except KeyError as exc:
+ raise AttributeError(f"Action '{name}' not found") from exc
+
def __contains__(self, key: object) -> bool:
"""Return whether *key* matches an available action name."""
if not isinstance(key, str):
diff --git a/aixplain/v2/agent.py b/aixplain/v2/agent.py
index 5455a8be..fbbd8526 100644
--- a/aixplain/v2/agent.py
+++ b/aixplain/v2/agent.py
@@ -7,7 +7,7 @@
from datetime import datetime
from enum import Enum
from dataclasses import dataclass, field
-from typing import ClassVar, List, Optional, Any, Dict, Tuple, Union, Text
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Any, Dict, Tuple, Union, Text
from typing_extensions import Unpack, NotRequired, TypedDict, Literal
from dataclasses_json import dataclass_json, config
@@ -15,6 +15,7 @@
from .enums import AssetStatus, ResponseStatus
from .model import Model
+from .skill import Skill
from .mixins import ToolableMixin
from ..utils.user_info_utils import build_run_metadata
@@ -33,6 +34,12 @@
with_hooks,
)
+if TYPE_CHECKING:
+ from .session import ExecutionConfig, Session
+
+
+logger = logging.getLogger(__name__)
+
# Type definitions for conversation history
class ConversationMessage(TypedDict):
@@ -41,10 +48,15 @@ class ConversationMessage(TypedDict):
Attributes:
role: The role of the message sender, either 'user' or 'assistant'
content: The text content of the message
+ attachments: Optional attachments — hosted-URL/local-path strings or dicts
+ with ``url`` or ``path`` (plus optional type/name/mimeType).
+ files: Deprecated. Local file paths to upload — pass through ``attachments``.
"""
role: Literal["user", "assistant"]
content: str
+ attachments: NotRequired[Optional[List[Union[str, Dict[str, Any]]]]]
+ files: NotRequired[Optional[List[Any]]]
def validate_history(history: List[Dict[str, Any]]) -> bool:
@@ -219,42 +231,89 @@ class AgentRunParams(BaseRunParams):
"""Parameters for running an agent.
Attributes:
- session_id: Session ID for conversation continuity
+ session: Conversation thread to run within. A
+ :class:`~aixplain.v2.session.Session` instance or a session id
+ string. Omit for a one-shot, stateless run. Replaces the removed
+ ``via_session`` flag and id-only ``session_id``.
query: The query to run
variables: Variables to replace {{variable}} placeholders in instructions and description.
The backend performs the actual substitution.
- allow_history_and_session_id: Allow both history and session ID
tasks: List of tasks for the agent
prompt: Custom prompt override
history: Conversation history
- execution_params: Execution parameters (maxTokens, etc.)
+ execution_params: Execution parameters (maxTokens, etc.). Passing
+ ``max_iterations`` here is deprecated; set ``agent.budget.max_iterations``
+ instead. A deprecated value is folded into ``budget.max_iterations``
+ (the agent's budget wins on conflict) and the standalone key is not
+ emitted.
criteria: Criteria for evaluation
evolve: Evolution parameters
inspectors: Inspector configurations
run_response_generation: Whether to run response generation. Defaults to False.
+ attachments: Multimodal attachments for the turn.
+ Each entry is a hosted-URL/local-path string or a dict with ``url`` or
+ ``path`` (plus optional ``type``/``name``/``mimeType``). Local paths are
+ uploaded to aiXplain storage automatically.
+ files: Deprecated. Local file paths to upload — pass through ``attachments`` instead.
progress_format: Display format - "status" (single line) or "logs" (timeline).
If None (default), progress tracking is disabled.
progress_verbosity: Detail level - 1 (minimal), 2 (thoughts), 3 (full I/O)
progress_truncate: Whether to truncate long text in progress display
"""
- session_id: NotRequired[Optional[Text]]
+ session: NotRequired[Optional[Union["Session", Text]]]
query: NotRequired[Optional[Union[Dict, Text]]]
variables: NotRequired[Optional[Dict[str, Any]]]
- allow_history_and_session_id: NotRequired[Optional[bool]]
tasks: NotRequired[Optional[List[Any]]]
prompt: NotRequired[Optional[Text]]
history: NotRequired[Optional[List[ConversationMessage]]]
execution_params: NotRequired[Optional[Dict[str, Any]]]
criteria: NotRequired[Optional[Text]]
evolve: NotRequired[Optional[Text]]
+ identifier: NotRequired[Optional[Text]]
inspectors: NotRequired[Optional[List[Dict]]]
run_response_generation: NotRequired[Optional[bool]]
+ attachments: NotRequired[Optional[List[Union[str, Dict[str, Any]]]]]
+ files: NotRequired[Optional[List[Any]]]
progress_format: NotRequired[Optional[Text]]
progress_verbosity: NotRequired[Optional[int]]
progress_truncate: NotRequired[Optional[bool]]
+@dataclass_json
+@dataclass
+class Budget:
+ """Budget caps governing an agent run (cost / duration / iterations).
+
+ Every :class:`Agent` owns a ``budget`` (defaulting to an empty ``Budget()``),
+ mutated in place via attribute access — mirroring ``model.inputs``::
+
+ agent.budget.max_cost = 0.5
+ agent.budget.max_iterations = 10
+
+ The same object serves two roles: ``agent.save()`` persists it as the agent's
+ default budget, and ``agent.run(...)`` sends its current state as the run-time
+ budget (the backend merges the run-time budget field-by-field over the
+ persisted default). The Python API is snake_case; serialization produces the
+ agreed camelCase wire keys (``maxCost`` / ``maxDurationSeconds`` /
+ ``maxIterations``). All fields are optional and ``None`` fields are dropped
+ from ``to_dict()``.
+ """
+
+ max_cost: Optional[float] = field(
+ default=None,
+ metadata=config(field_name="maxCost", exclude=lambda v: v is None),
+ )
+ max_duration_seconds: Optional[float] = field(
+ default=None,
+ metadata=config(field_name="maxDurationSeconds", exclude=lambda v: v is None),
+ )
+ max_iterations: Optional[int] = field(
+ default=None,
+ metadata=config(field_name="maxIterations", exclude=lambda v: v is None),
+ )
+
+
@dataclass_json
@dataclass
class AgentResponseData:
@@ -265,7 +324,27 @@ class AgentResponseData:
steps: Optional[List[Dict[str, Any]]] = field(default_factory=list)
session_id: Optional[str] = None
execution_stats: Optional[Dict[str, Any]] = field(default=None, metadata=config(field_name="executionStats"))
+ diagnostic_error_codes: List[str] = field(default_factory=list, metadata=config(field_name="diagnosticErrorCodes"))
critiques: Optional[str] = ""
+ governance: Optional[Dict[str, Any]] = None
+ _governance_status: Optional[str] = field(
+ default=None, repr=False, metadata=config(field_name="governanceStatus", exclude=lambda x: True)
+ )
+ _governance_source: Optional[str] = field(
+ default=None, repr=False, metadata=config(field_name="governanceSource", exclude=lambda x: True)
+ )
+ _governance_reason: Optional[str] = field(
+ default=None, repr=False, metadata=config(field_name="governanceReason", exclude=lambda x: True)
+ )
+
+ def __post_init__(self) -> None:
+ """Assemble the nested ``governance`` dict from the flat wire fields."""
+ if self.governance is None:
+ self.governance = {
+ "status": self._governance_status,
+ "source": self._governance_source,
+ "reason": self._governance_reason,
+ }
@dataclass_json
@@ -280,6 +359,35 @@ class AgentRunResult(Result):
run_time: float = field(default=0.0, metadata=config(field_name="runTime"))
diagnostic_error_codes: List[str] = field(default_factory=list, metadata=config(field_name="diagnosticErrorCodes"))
+ def __post_init__(self) -> None:
+ """Promote diagnostic codes the backend nests under ``data``.
+
+ The poll body carries them at ``data.diagnosticErrorCodes`` (or only
+ inside ``executionStats`` on older builds), never top-level.
+ """
+ if not self.diagnostic_error_codes:
+ self.diagnostic_error_codes = self._codes_from_data()
+
+ def _codes_from_data(self) -> List[str]:
+ """Extract diagnostic codes from ``data`` or its execution stats."""
+ data = self.data
+ if isinstance(data, AgentResponseData):
+ if data.diagnostic_error_codes:
+ return list(data.diagnostic_error_codes)
+ stats = data.execution_stats
+ elif isinstance(data, dict):
+ codes = data.get("diagnosticErrorCodes") or data.get("diagnostic_error_codes")
+ if codes:
+ return list(codes)
+ stats = data.get("executionStats") or data.get("execution_stats")
+ else:
+ return []
+ if isinstance(stats, dict):
+ codes = stats.get("diagnostic_error_codes") or stats.get("diagnosticErrorCodes")
+ if codes:
+ return list(codes)
+ return []
+
# Internal reference to client context for debug() method
_context: Optional[Any] = field(
default=None,
@@ -433,6 +541,10 @@ class Agent(
metadata=config(exclude=lambda x: True),
)
+ # Skills (knowledge bundles) attached to the agent — Skill objects or ids,
+ # the same way `tools` and `agents` are passed.
+ skills: Optional[List[Union[str, "Skill"]]] = field(default_factory=list, metadata=config(field_name="skills"))
+
# Output and execution fields
output_format: Optional[Union[str, OutputFormat]] = field(
default=OutputFormat.TEXT.value, metadata=config(field_name="outputFormat")
@@ -448,7 +560,24 @@ class Agent(
max_inspectors: Optional[int] = field(default=None, metadata=config(field_name="maxInspectors"))
inspectors: Optional[List[Any]] = field(default_factory=list)
resource_info: Optional[Dict[str, Any]] = field(default_factory=dict, metadata=config(field_name="resourceInfo"))
- max_iterations: Optional[int] = field(default=5, metadata=config(field_name="maxIterations"))
+ # Deprecated: persisted iteration cap. Use ``budget=Budget(max_iterations=...)``
+ # instead. Defaults to ``None`` (was previously ``5``) so a plain ``Agent(...)``
+ # does not warn; any non-None value (explicit or via ``from_dict``) is folded
+ # into ``budget`` and is never serialized as a standalone ``maxIterations``
+ # (see ``build_save_payload``).
+ max_iterations: Optional[int] = field(default=None, metadata=config(field_name="maxIterations"))
+ # Budget (cost / duration / iterations) for this agent. Always a ``Budget``
+ # instance — defaults to an empty ``Budget()``, never ``None`` — so callers
+ # can mutate ``agent.budget.max_cost`` etc. without a None check (mirrors
+ # ``model.inputs``). Assigning a dict/Budget/None is coerced back to a
+ # ``Budget`` by ``__setattr__``. Excluded from ``to_dict()`` (so it never
+ # counts toward ``is_modified``); serialized manually in ``build_save_payload``
+ # (persisted default) and ``build_run_payload`` (run-time budget) via
+ # ``_normalize_budget``.
+ budget: "Budget" = field(
+ default_factory=lambda: Budget(),
+ metadata=config(field_name="budget", exclude=lambda v: True),
+ )
max_tokens: Optional[int] = field(default=2048, metadata=config(field_name="maxTokens"))
context_overflow_strategy: Optional[str] = field(
default=None,
@@ -469,15 +598,13 @@ def __post_init__(self) -> None:
self.tasks = [Task.from_dict(task) for task in self.tasks]
# Deserialize inspectors to Inspector objects so mutate-and-save round-trips.
- # Prebuilt inspectors travel as a lightweight {presetId, ...} reference
- # (no "name"/"evaluator"), so dispatch on shape before deserializing.
+ # Prebuilt guards and custom inspectors are the same Inspector type, so a
+ # single deserialization path covers both.
if self.inspectors:
- from .inspector import Inspector, PrebuiltInspector
+ from .inspector import Inspector
self.inspectors = [
- (PrebuiltInspector.from_dict(inspector) if "presetId" in inspector else Inspector.from_dict(inspector))
- if isinstance(inspector, dict)
- else inspector
+ Inspector.from_dict(inspector) if isinstance(inspector, dict) else inspector
for inspector in self.inspectors
]
@@ -492,34 +619,59 @@ def __post_init__(self) -> None:
self.agents = self.subagents
self.subagents = None
+ # ``self.budget`` is already a ``Budget`` instance here: the generated
+ # ``__init__`` assignment routed through ``__setattr__``, which coerces
+ # any dict/Budget/None into a (never-None) ``Budget``.
+
+ # Deprecated persisted ``max_iterations``: fold into ``budget.max_iterations``.
+ # ``None`` (the default) means "not provided" so a plain ``Agent(...)`` never
+ # warns; any non-None value (explicit or via ``from_dict``) is folded.
+ if self.max_iterations is not None:
+ warnings.warn(
+ "Agent 'max_iterations' is deprecated; set agent.budget.max_iterations instead. "
+ "It will be removed in a future release.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ if self.budget.max_iterations is None:
+ self.budget.max_iterations = self.max_iterations
+ else:
+ warnings.warn(self._BUDGET_ITER_CONFLICT_MSG, UserWarning, stacklevel=2)
+
# Store original agent objects to resolve IDs at save time
self._original_agents = list(self.agents)
# Convert to IDs for serialization (to_dict), using None as placeholder for unsaved agents
self.agents = [a if isinstance(a, str) else a.get("id") if isinstance(a, dict) else a.id for a in self.agents]
+ # Skills behave exactly like agents: keep the originals to resolve ids
+ # at save time, and serialize as a list of ids.
+ self._original_skills = list(self.skills or [])
+ self.skills = [
+ s if isinstance(s, str) else s.get("id") if isinstance(s, dict) else s.id for s in (self.skills or [])
+ ]
+
if isinstance(self.output_format, OutputFormat):
self.output_format = self.output_format.value
if isinstance(self.context_overflow_strategy, ContextOverflowStrategy):
self.context_overflow_strategy = self.context_overflow_strategy.value
- # Normalize inspector_targets to support both strings and InspectorTarget enums
+ # Inspector targets are plain strings (e.g. "input" | "steps" | "output"
+ # or a sub-agent name); normalize the known stage values to lowercase.
if self.inspector_targets:
- from .inspector import InspectorTarget
+ self.inspector_targets = [
+ target.lower() if isinstance(target, str) and target.lower() in {"input", "steps", "output"} else target
+ for target in self.inspector_targets
+ ]
- normalized_targets = []
- for target in self.inspector_targets:
- if isinstance(target, str):
- # Convert string to InspectorTarget enum
- try:
- normalized_targets.append(InspectorTarget(target.lower()))
- except ValueError:
- # If it's not a valid InspectorTarget value, keep as is
- normalized_targets.append(target)
- else:
- # Already an InspectorTarget or other type
- normalized_targets.append(target)
- self.inspector_targets = normalized_targets
+ # Hydrate plain tool dicts (e.g. from a get()/create response) into
+ # mutable Tool/Model objects so callers can override per-tool parameters
+ # via ``agent.tools[i].actions[...].inputs[...] = value``. Best-effort and
+ # offline — see :meth:`_hydrate_tools`.
+ self._hydrate_tools()
+ # Snapshot the (post-hydration) tool objects so save can restore them
+ # after the create response would otherwise overwrite them with dicts.
+ self._original_tools = list(self.tools) if self.tools else []
# TODO: Re-enable this validation after backend data consistency is fixed
# if self.agents and (self.tasks or self.tools):
@@ -527,6 +679,47 @@ def __post_init__(self) -> None:
# "Team agents cannot have tasks or tools. Please remove the tasks or tools and try again."
# )
+ def __setattr__(self, name: str, value: Any) -> None:
+ """Keep ``self.budget`` a (never-None) ``Budget`` instance.
+
+ Assigning ``agent.budget`` a dict / ``Budget`` / ``None`` is coerced into
+ a ``Budget`` so attribute access (``agent.budget.max_cost = ...``) always
+ works and the invariant "budget is never None" holds (mirrors how
+ ``Model.__setattr__`` coerces bulk ``inputs`` assignment). This runs for
+ the generated ``__init__`` assignment too, so the field is a ``Budget``
+ by the time ``__post_init__`` executes.
+ """
+ if name == "budget":
+ coerced = self._coerce_budget(value)
+ value = coerced if coerced is not None else Budget()
+ super().__setattr__(name, value)
+
+ @classmethod
+ def _fold_legacy_max_iterations(cls, kvs: Any) -> Any:
+ """Fold a legacy top-level ``maxIterations`` into the ``budget`` slot.
+
+ Backend agents may still carry a top-level legacy ``maxIterations``. The
+ public constructor (``Agent(max_iterations=...)``) emits a
+ ``DeprecationWarning`` in ``__post_init__``, but deserialization must NOT
+ warn — loading an agent the caller never configured should be silent.
+ So we fold the legacy ``maxIterations`` straight into the ``budget`` slot
+ (Budget wins silently if it already carries one) and drop the standalone
+ key, so ``__post_init__`` sees ``max_iterations=None`` and the deprecation
+ path never fires on load. Returns a copy; the caller's dict is untouched.
+ """
+ if isinstance(kvs, dict) and kvs.get("maxIterations") is not None:
+ kvs = dict(kvs) # shallow copy; never mutate the caller's dict
+ legacy_iterations = kvs.pop("maxIterations")
+ budget = kvs.get("budget")
+ # Normalize to the camelCase wire shape so the nested fold is uniform
+ # regardless of whether ``budget`` arrived snake_case, camelCase, or absent.
+ normalized = cls._normalize_budget(budget) if budget is not None else {}
+ # Budget wins silently on conflict; otherwise fill the empty slot.
+ if normalized.get("maxIterations") is None:
+ normalized["maxIterations"] = legacy_iterations
+ kvs["budget"] = normalized
+ return kvs
+
def mark_as_deleted(self) -> None:
"""Mark the agent as deleted by setting status to DELETED and calling parent method."""
from .enums import AssetStatus
@@ -534,6 +727,64 @@ def mark_as_deleted(self) -> None:
self.status = AssetStatus.DELETED
super().mark_as_deleted()
+ def _get_serializable_state(self) -> dict:
+ """Serializable state used for change detection (``is_modified``).
+
+ Tools are reduced to identity-only signatures (id + type, excluding
+ per-input values) so that (a) ``to_dict()`` does not recurse into
+ ``Tool`` / ``Model`` objects and crash, and (b) mutating a tool's action
+ inputs does not mark the agent as modified — run-time tool-parameter
+ overrides are ephemeral and must not trigger an auto-save or block an
+ onboarded agent's run. Adding or removing a whole tool is still
+ detected; ``save()`` persists changed values unconditionally.
+ """
+ original_tools = self.tools
+ try:
+ self.tools = [self._tool_identity(tool) for tool in original_tools or []]
+ return super()._get_serializable_state()
+ finally:
+ self.tools = original_tools
+
+ @staticmethod
+ def _tool_identity(tool: Any) -> dict:
+ """Return a stable id/type signature for a tool entry (no param values)."""
+ if isinstance(tool, str):
+ return {"id": tool}
+ if isinstance(tool, dict):
+ return {"id": tool.get("id"), "type": tool.get("type")}
+ return {"id": getattr(tool, "id", None), "type": getattr(tool, "type", None)}
+
+ def _start_progress_tracker(self, kwargs: Dict[str, Any]) -> None:
+ """Initialize ``self._progress_tracker`` from progress kwargs (no-op if disabled)."""
+ progress_format = kwargs.get("progress_format")
+ if progress_format is None:
+ self._progress_tracker = None
+ return
+
+ from .agent_progress import AgentProgressTracker, ProgressFormat
+
+ progress_verbosity = kwargs.get("progress_verbosity", 1)
+ progress_truncate = kwargs.get("progress_truncate", True)
+ fmt = ProgressFormat(progress_format)
+
+ self._progress_tracker = AgentProgressTracker(
+ poll_func=lambda url: self.poll(url),
+ poll_interval=0.05,
+ max_polls=None,
+ )
+ self._progress_tracker.start(
+ format=fmt,
+ verbosity=progress_verbosity,
+ truncate=progress_truncate,
+ )
+
+ def _finish_progress_tracker(self, result: Union[AgentRunResult, Exception]) -> None:
+ """Finalize the progress tracker; safe to call even if it was never started."""
+ if self._progress_tracker is not None:
+ if not isinstance(result, Exception):
+ self._progress_tracker.finish(result)
+ self._progress_tracker = None
+
def before_run(self, *args: Any, **kwargs: Unpack[AgentRunParams]) -> Optional[AgentRunResult]:
"""Hook called before running the agent to validate and prepare state."""
# First, validate that all dependencies are saved before allowing run
@@ -549,30 +800,7 @@ def before_run(self, *args: Any, **kwargs: Unpack[AgentRunParams]) -> Optional[A
if self.is_modified:
raise ValueError("Agent is onboarded and cannot be modified unless you explicitly save it.")
- # Initialize progress tracker if progress_format is provided
- # progress_format being None (default) means no progress tracking
- progress_format = kwargs.get("progress_format")
- if progress_format is not None:
- from .agent_progress import AgentProgressTracker, ProgressFormat
-
- progress_verbosity = kwargs.get("progress_verbosity", 1)
- progress_truncate = kwargs.get("progress_truncate", True)
-
- fmt = ProgressFormat(progress_format)
-
- self._progress_tracker = AgentProgressTracker(
- poll_func=lambda url: self.poll(url),
- poll_interval=0.05,
- max_polls=None,
- )
- self._progress_tracker.start(
- format=fmt,
- verbosity=progress_verbosity,
- truncate=progress_truncate,
- )
- else:
- self._progress_tracker = None
-
+ self._start_progress_tracker(kwargs)
return None
def on_poll(self, response: AgentRunResult, **kwargs: Unpack[AgentRunParams]) -> None:
@@ -595,10 +823,7 @@ def after_run(
) -> Optional[AgentRunResult]:
"""Hook called after running the agent for result transformation."""
# Finish progress tracking if enabled
- if self._progress_tracker is not None:
- if not isinstance(result, Exception):
- self._progress_tracker.finish(result)
- self._progress_tracker = None
+ self._finish_progress_tracker(result)
# Set the context on the result for debug() method support
if not isinstance(result, Exception):
@@ -607,27 +832,108 @@ def after_run(
return None # Return original result
_SNAKE_TO_CAMEL: ClassVar[Dict[str, str]] = {
- "session_id": "sessionId",
- "allow_history_and_session_id": "allowHistoryAndSessionId",
"execution_params": "executionParams",
"run_response_generation": "runResponseGeneration",
}
+ # ``max_iterations`` is intentionally absent: it is deprecated and folded into
+ # ``budget.maxIterations`` (see ``_fold_iter_into_budget`` and the run path),
+ # never emitted as a standalone ``executionParams.maxIterations``. The session
+ # path (``ExecutionConfig`` in session.py) keeps the same invariant.
_EXEC_PARAMS_MAP: ClassVar[Dict[str, str]] = {
"output_format": "outputFormat",
"max_tokens": "maxTokens",
- "max_iterations": "maxIterations",
"max_time": "maxTime",
"expected_output": "expectedOutput",
"context_overflow_strategy": "contextOverflowStrategy",
}
+ # snake_case → camelCase wire keys for the nested executionParams.budget object.
+ _BUDGET_PARAMS_MAP: ClassVar[Dict[str, str]] = {
+ "max_cost": "maxCost",
+ "max_duration_seconds": "maxDurationSeconds",
+ "max_iterations": "maxIterations",
+ }
+
+ # Emitted (as a UserWarning) when a deprecated ``max_iterations`` and an
+ # explicit ``budget.max_iterations`` are both set; Budget wins. Shared by the
+ # constructor, the run path, and the session path so the text never drifts.
+ _BUDGET_ITER_CONFLICT_MSG: ClassVar[str] = (
+ "Both 'max_iterations' and budget.max_iterations are set; budget.max_iterations takes precedence."
+ )
+
+ @classmethod
+ def _normalize_budget(cls, budget: Union[Dict, "Budget"]) -> dict:
+ """Normalize a Budget instance or dict to the camelCase wire shape.
+
+ Accepts a ``Budget`` instance, a snake_case dict, or a camelCase dict and
+ returns a camelCase dict with ``None`` fields dropped. Unknown keys are
+ passed through unchanged (forward compatibility).
+ """
+ if isinstance(budget, Budget):
+ raw = budget.to_dict() # already camelCase, None dropped
+ else:
+ raw = dict(budget)
+ # Map snake_case keys to camelCase; pass camelCase (and unknown) keys through.
+ normalized = {cls._BUDGET_PARAMS_MAP.get(k, k): v for k, v in raw.items()}
+ # Never emit null fields.
+ return {k: v for k, v in normalized.items() if v is not None}
+
+ @classmethod
+ def _coerce_budget(cls, budget: Optional[Union[Dict, "Budget"]]) -> Optional["Budget"]:
+ """Coerce ``None`` / dict / ``Budget`` into a ``Budget`` instance.
+
+ A dict may use snake_case or camelCase keys; it is normalized to the
+ camelCase wire shape first so a single decoder handles both styles.
+ ``None`` passes through as ``None`` (session's ExecutionConfig relies on
+ this); ``Agent.__setattr__`` is what upgrades a ``None`` agent budget to
+ an empty ``Budget()`` to keep ``agent.budget`` never-None.
+ """
+ if budget is None or isinstance(budget, Budget):
+ return budget
+ if isinstance(budget, dict):
+ normalized = cls._normalize_budget(budget)
+ return Budget(
+ max_cost=normalized.get("maxCost"),
+ max_duration_seconds=normalized.get("maxDurationSeconds"),
+ max_iterations=normalized.get("maxIterations"),
+ )
+ raise TypeError(f"budget must be a Budget, dict, or None, got {type(budget)}")
+
+ @classmethod
+ def _fold_iter_into_budget(
+ cls, budget: Optional[Union[Dict, "Budget"]], max_iterations: int
+ ) -> "tuple[dict, bool]":
+ """Fold a deprecated ``max_iterations`` into a budget (Budget wins on conflict).
+
+ Pure: emits no warnings. Returns ``(normalized, conflicted)`` where
+ ``normalized`` is the camelCase wire dict and ``conflicted`` is ``True``
+ when the budget already carried ``maxIterations`` (so the deprecated value
+ was dropped). The caller owns the conflict warning — it knows the right
+ ``stacklevel`` for its own call depth (see the run and session paths).
+ """
+ normalized = cls._normalize_budget(budget) if budget is not None else {}
+ if normalized.get("maxIterations") is not None:
+ return normalized, True
+ normalized["maxIterations"] = max_iterations
+ return normalized, False
+
def run(self, *args: Any, **kwargs: Unpack[AgentRunParams]) -> AgentRunResult:
"""Run the agent with optional progress display.
Args:
*args: Positional arguments (first arg is treated as query)
query: The query to run
+ session: Run within a conversation thread. Accepts a
+ :class:`~aixplain.v2.session.Session` instance or a session id
+ string. When supplied, the run routes through the session path:
+ the user message is posted to
+ ``POST /v1/sessions/{id}/messages`` (carrying the session's
+ ``executionConfig`` plus any per-run execution overrides) and the
+ triggered agent run is awaited. Omit ``session`` for a one-shot,
+ stateless run over ``POST /v2/agents/{id}/run``. There is no
+ ``via_session`` flag and no id-only ``session_id`` — manage
+ threads through ``aix.Session`` and pass them here.
progress_format: Display format - "status" or "logs". If None (default),
progress tracking is disabled.
progress_verbosity: Detail level 1-3 (default: 1)
@@ -641,6 +947,10 @@ def run(self, *args: Any, **kwargs: Unpack[AgentRunParams]) -> AgentRunResult:
kwargs["query"] = args[0]
args = args[1:]
+ session = kwargs.pop("session", None)
+ if session is not None:
+ return self._run_with_session(session, **kwargs)
+
return super().run(*args, **kwargs)
def run_async(self, *args: Any, **kwargs: Unpack[AgentRunParams]) -> AgentRunResult:
@@ -662,6 +972,12 @@ def run_async(self, *args: Any, **kwargs: Unpack[AgentRunParams]) -> AgentRunRes
kwargs["query"] = args[0]
args = args[1:]
+ if kwargs.pop("session", None) is not None:
+ raise NotImplementedError(
+ "session=… runs are sync-only for now; use agent.run(...) or "
+ "session.add_message() + session.messages() directly."
+ )
+
return super().run_async(**kwargs)
def _resolve_poll_url(self, poll_url: str) -> str:
@@ -776,8 +1092,64 @@ def save(self, *args: Any, **kwargs: Any) -> "Agent":
# Validate that all dependencies are saved before proceeding
self._validate_dependencies()
+ # Capture names before save because the backend response can rebuild self.tools without Integration objects.
+ unconnected_integration_names = self._get_unconnected_integration_names()
+
+ # Preserve the in-memory tool objects (carrying any per-tool parameter
+ # overrides the caller set): a create response would otherwise replace
+ # ``self.tools`` with backend dicts and drop those mutations.
+ pre_save_tools = list(self.tools) if self.tools else []
+
# Call the parent save method
- return super().save(*args, **kwargs)
+ saved_agent = super().save(*args, **kwargs)
+
+ # Restore the caller's tool objects, then re-hydrate any dict entries so
+ # ``agent.tools[i]`` stays a mutable Tool/Model object after save.
+ self.tools = pre_save_tools
+ self._hydrate_tools()
+ self._original_tools = list(self.tools) if self.tools else []
+
+ # Re-baseline the saved state against the restored tool objects. The
+ # parent save() captured it mid-flow from the response dicts; without
+ # this, the restored objects would read as "modified" and a subsequent
+ # run() on an onboarded agent would wrongly raise.
+ self._update_saved_state()
+
+ self._warn_for_unconnected_integrations(unconnected_integration_names)
+ return saved_agent
+
+ def _warn_for_unconnected_integrations(self, integration_names: Optional[List[str]] = None) -> None:
+ """Warn when an agent is saved with integration definitions that need connection."""
+ if integration_names is None:
+ integration_names = self._get_unconnected_integration_names()
+ if not integration_names or not self.id:
+ return
+
+ schema_url = f"https://studio.aixplain.com/build/{self.id}/schema"
+ for integration_name in integration_names:
+ warnings.warn(
+ f"Warning: Integration '{integration_name}' is not connected. "
+ f"Connect your unconnected integrations here: {schema_url}",
+ UserWarning,
+ stacklevel=2,
+ )
+
+ def _get_unconnected_integration_names(self) -> List[str]:
+ """Return unique names for Integration objects used directly as tools."""
+ from .integration import Integration
+
+ if not self.tools:
+ return []
+
+ names = []
+ seen = set()
+ for tool in self.tools:
+ if isinstance(tool, Integration):
+ integration_name = tool.name or tool.id
+ if integration_name and integration_name not in seen:
+ names.append(integration_name)
+ seen.add(integration_name)
+ return names
def _save_subcomponents(self) -> None:
"""Recursively save all unsaved child components."""
@@ -805,6 +1177,18 @@ def _save_subcomponents(self) -> None:
agent_name = getattr(agent, "name", f"agent_{i}")
failed_components.append(("agent", agent_name, str(e)))
+ # Save skills
+ if getattr(self, "_original_skills", None):
+ for i, skill in enumerate(self._original_skills):
+ if isinstance(skill, (str, dict)): # Already an ID
+ continue
+ if hasattr(skill, "save") and hasattr(skill, "id") and not skill.id:
+ try:
+ skill.save()
+ except Exception as e:
+ skill_name = getattr(skill, "name", f"skill_{i}")
+ failed_components.append(("skill", skill_name, str(e)))
+
if failed_components:
error_details = "; ".join(
[f"{comp_type} '{name}': {error}" for comp_type, name, error in failed_components]
@@ -830,6 +1214,15 @@ def _validate_run_dependencies(self) -> None:
agent_name = getattr(agent, "name", "unnamed")
unsaved_components.append(f"agent '{agent_name}'")
+ # Check skills
+ if getattr(self, "_original_skills", None):
+ for skill in self._original_skills:
+ if isinstance(skill, (str, dict)): # Already an ID
+ continue
+ if hasattr(skill, "id") and not skill.id:
+ skill_name = getattr(skill, "name", "unnamed")
+ unsaved_components.append(f"skill '{skill_name}'")
+
if unsaved_components:
components_list = ", ".join(unsaved_components)
raise ValueError(
@@ -858,6 +1251,15 @@ def _validate_dependencies(self) -> None:
agent_name = getattr(agent, "name", "unnamed")
unsaved_components.append(f"agent '{agent_name}'")
+ # Check skills
+ if getattr(self, "_original_skills", None):
+ for skill in self._original_skills:
+ if isinstance(skill, (str, dict)): # Already an ID
+ continue
+ if hasattr(skill, "id") and not skill.id:
+ skill_name = getattr(skill, "name", "unnamed")
+ unsaved_components.append(f"skill '{skill_name}'")
+
if unsaved_components:
components_list = ", ".join(unsaved_components)
raise ValueError(
@@ -949,6 +1351,195 @@ def search(
return super().search(**kwargs)
+ @classmethod
+ def _normalize_tool_for_api(cls, tool: Any) -> dict:
+ """Normalize one ``tools`` entry into the API dict shape.
+
+ Per-tool parameter overrides are expressed by mutating the tool object's
+ typed action inputs (``tool.actions[...].inputs[...] = value``); the
+ current values are read back through the object's ``as_tool()`` snapshot
+ (see :meth:`Tool.get_parameters` / :meth:`Model.get_parameters`).
+
+ Accepts:
+
+ - a :class:`ToolableMixin` (``Tool`` / ``Model`` / ``Integration``)
+ whose ``as_tool()`` snapshot carries the correct ``type`` and current
+ parameter values;
+ - a plain string id — resolved to its ``as_tool()`` snapshot so the
+ create payload gets the required ``type``;
+ - an ``as_tool()`` snapshot dict (already carries ``type``) — passed
+ through; a bare ``{"id": ...}`` attach dict is resolved for its type.
+ """
+ if isinstance(tool, ToolableMixin):
+ return cls._normalize_tool_dict_for_api(tool.as_tool())
+ if isinstance(tool, str):
+ return cls._normalize_tool_dict_for_api(cls._resolve_tool_entry(tool))
+ if isinstance(tool, dict):
+ if tool.get("type") or not tool.get("id"):
+ return cls._normalize_tool_dict_for_api(tool)
+ return cls._normalize_tool_dict_for_api(cls._resolve_tool_entry(tool["id"], tool))
+ raise ValueError(
+ f"A tool must be a Tool, Model, ToolableMixin instance, a string id, or a dictionary, got {type(tool)}."
+ )
+
+ @classmethod
+ def _resolve_tool_entry(cls, tool_id: str, attach: Optional[dict] = None) -> dict:
+ """Build a typed tool dict for ``tool_id`` from its ``as_tool()`` snapshot.
+
+ Resolves the asset by id (Tool, then Model) so the entry carries the
+ backend-required ``type`` and other snapshot fields, then overlays any
+ extra keys from the caller-provided attach dict. Falls back to a bare
+ ``{"id": tool_id}`` when the id can't be resolved — preserving the prior
+ behavior offline.
+ """
+ snapshot = cls._resolve_tool_snapshot(tool_id)
+ entry: dict = dict(snapshot) if snapshot else {}
+ entry["id"] = tool_id
+ if isinstance(attach, dict):
+ for key, value in attach.items():
+ if key != "id":
+ entry[key] = value
+ return entry
+
+ @classmethod
+ def _resolve_tool_snapshot(cls, tool_id: str) -> Optional[dict]:
+ """Return the ``as_tool()`` snapshot for ``tool_id``, or ``None``.
+
+ Tries the ``Tool`` resource first, then ``Model`` (mirroring how the
+ platform resolves an asset id). Any failure (no client context,
+ unknown id, network error) yields ``None`` so normalization degrades
+ gracefully instead of raising.
+ """
+ context = getattr(cls, "context", None)
+ if context is None:
+ return None
+ for resource_name in ("Tool", "Model"):
+ resource = getattr(context, resource_name, None)
+ if resource is None:
+ continue
+ try:
+ return dict(resource.get(tool_id).as_tool())
+ except Exception:
+ continue
+ return None
+
+ def _hydrate_tools(self) -> None:
+ """Convert plain tool dicts in ``self.tools`` into mutable objects.
+
+ Runs from :meth:`__post_init__` (so it fires on ``get()``/``create``
+ responses). Entries that are already ``Tool`` / ``Model`` /
+ ``Integration`` objects, plain string ids, or unresolvable dicts are
+ left untouched. Hydration is **offline**: the object's action inputs are
+ reconstructed from the ``parameters`` snapshot embedded in the dict so
+ reads and ``inputs[...] = value`` mutations do not require a network
+ call. When the snapshot carries no parameters, the object is left to
+ load its input specs lazily on first ``.actions`` access (matching a
+ normal ``Tool.get``). Requires a client context; without one (e.g. an
+ unbound ``Agent`` in unit tests) the raw entries are kept as-is.
+ """
+ context = getattr(self, "context", None)
+ if context is None or not self.tools:
+ return
+ self.tools = [self._hydrate_tool_entry(entry, context) for entry in self.tools]
+
+ def _hydrate_tool_entry(self, entry: Any, context: Any) -> Any:
+ """Hydrate one ``tools`` entry into a Tool/Model object (best-effort)."""
+ if not isinstance(entry, dict):
+ return entry # already a Tool/Model/Integration object or a string id
+ tool_id = entry.get("id")
+ if not tool_id:
+ return entry
+
+ parameters = entry.get("parameters")
+ is_model = entry.get("type") == "model"
+ try:
+ if is_model:
+ return self._build_model_tool(entry, context, parameters)
+ return self._build_tool_tool(entry, context, parameters)
+ except Exception:
+ # Never let hydration break construction — fall back to the raw dict.
+ return entry
+
+ @staticmethod
+ def _build_model_tool(entry: dict, context: Any, parameters: Any) -> Any:
+ """Build a bound ``Model`` from a tool dict, seeding inputs from ``parameters``.
+
+ ``parameters`` is the flat ``[{name, value, datatype, required, ...}]``
+ list produced by ``Model.as_tool()``. It is turned into ``params`` so the
+ rebuilt ``Model`` exposes working ``inputs`` (and ``as_tool()`` round-trips
+ the current values offline).
+ """
+ from .model import Parameter
+
+ params = []
+ for param in parameters or []:
+ if not isinstance(param, dict) or not param.get("name"):
+ continue
+ value = param.get("value")
+ params.append(
+ Parameter(
+ name=param["name"],
+ required=bool(param.get("required", False)),
+ data_type=param.get("datatype") or param.get("dataType"),
+ default_values=[value] if value is not None else [],
+ )
+ )
+ model = context.Model(
+ id=entry.get("id"),
+ name=entry.get("name"),
+ description=entry.get("description"),
+ params=params or None,
+ )
+ model.context = context
+ return model
+
+ @staticmethod
+ def _build_tool_tool(entry: dict, context: Any, parameters: Any) -> Any:
+ """Build a bound ``Tool`` from a tool dict, seeding action inputs offline.
+
+ ``parameters`` is the nested ``[{code, name, inputs: {code: {value, ...}}}]``
+ list produced by ``Tool.as_tool()`` / ``Tool.get_parameters()``. When
+ present, the tool's ``actions`` cache is pre-populated so mutations are
+ offline; otherwise the tool loads its input specs lazily on first access.
+ """
+ from .actions import Action, Actions, Input, Inputs
+
+ tool = context.Tool(
+ id=entry.get("id"),
+ name=entry.get("name"),
+ description=entry.get("description"),
+ )
+ tool.context = context
+
+ actions_map: Dict[str, Action] = {}
+ for action_def in parameters or []:
+ if not isinstance(action_def, dict):
+ continue
+ action_name = action_def.get("name") or action_def.get("code")
+ if not action_name:
+ continue
+ input_objs: Dict[str, Input] = {}
+ for code, spec in (action_def.get("inputs") or {}).items():
+ if not isinstance(spec, dict):
+ continue
+ input_objs[code] = Input(
+ name=code,
+ required=bool(spec.get("required", False)),
+ type=spec.get("datatype") or spec.get("dataType"),
+ value=spec.get("value"),
+ description=spec.get("description", "") or "",
+ )
+ actions_map[action_name] = Action(
+ name=action_name, description=action_def.get("description"), inputs=Inputs(input_objs)
+ )
+
+ if actions_map:
+ # Pre-populate the ``actions`` cached_property so reads/mutations are
+ # offline. Also scope allowed_actions so as_tool() serializes them.
+ tool.__dict__["actions"] = Actions(actions_map)
+ tool.allowed_actions = list(actions_map.keys())
+ return tool
+
@staticmethod
def _normalize_tool_dict_for_api(tool_dict: dict) -> dict:
"""Convert snake_case keys in a tool dict to the camelCase the API expects."""
@@ -961,6 +1552,9 @@ def _normalize_tool_dict_for_api(tool_dict: dict) -> dict:
for k, v in tool_dict.items():
api_key = _KEY_MAP.get(k, k)
if api_key == "parameters" and isinstance(v, list):
+ # Snapshot from ``as_tool()`` -> list of full parameter
+ # definitions (current input values included); normalize each
+ # definition's keys to camelCase.
result[api_key] = [Agent._normalize_parameter_for_api(p) for p in v]
else:
result[api_key] = v
@@ -1136,35 +1730,33 @@ def llm_id(self) -> str:
def build_save_payload(self, **kwargs: Any) -> dict:
"""Build the payload for the save action."""
# Import Inspector from v2 module
- from .inspector import Inspector, PrebuiltInspector
+ from .inspector import Inspector
# Pre-serialize inspectors before to_dict() to avoid dataclass_json issues
original_inspectors = self.inspectors
if self.inspectors:
serialized_inspectors = []
for inspector in self.inspectors:
- if isinstance(inspector, (Inspector, PrebuiltInspector)):
+ if isinstance(inspector, Inspector):
serialized_inspectors.append(inspector.to_dict())
elif isinstance(inspector, dict):
serialized_inspectors.append(inspector)
else:
- raise ValueError(f"Inspector must be Inspector, PrebuiltInspector, or dict, got {type(inspector)}")
+ raise ValueError(f"Inspector must be Inspector or dict, got {type(inspector)}")
self.inspectors = serialized_inspectors
- # Pre-serialize inspector_targets to strings (enum values)
- from .inspector import InspectorTarget
-
+ # Pre-serialize inspector_targets to strings
original_inspector_targets = self.inspector_targets
if self.inspector_targets:
- serialized_targets = []
- for target in self.inspector_targets:
- if isinstance(target, InspectorTarget):
- serialized_targets.append(target.value)
- elif isinstance(target, str):
- serialized_targets.append(target)
- else:
- serialized_targets.append(str(target))
- self.inspector_targets = serialized_targets
+ self.inspector_targets = [
+ target if isinstance(target, str) else str(target) for target in self.inspector_targets
+ ]
+
+ # Null out tools before to_dict(): dataclass_json would otherwise recurse
+ # into Tool/Model objects (which raises on their ``context`` descriptor).
+ # The real tools payload is rebuilt from ``self.tools`` below.
+ original_tools = self.tools
+ self.tools = []
# Now call to_dict() with inspectors and inspector_targets already serialized
payload = self.to_dict()
@@ -1172,6 +1764,19 @@ def build_save_payload(self, **kwargs: Any) -> dict:
# Restore original values
self.inspectors = original_inspectors
self.inspector_targets = original_inspector_targets
+ self.tools = original_tools
+
+ # Budget is the single source of truth for the persisted iteration cap.
+ # Drop the deprecated standalone ``maxIterations`` and emit the persisted
+ # ``budget`` (camelCase) only when it carries at least one cap. (The
+ # ``budget`` field is ``exclude=lambda v: True`` so ``to_dict()`` never
+ # emits it — we build the wire shape explicitly below.)
+ payload.pop("maxIterations", None)
+ budget = getattr(self, "budget", None)
+ if budget is not None:
+ normalized_budget = self._normalize_budget(budget)
+ if normalized_budget:
+ payload["budget"] = normalized_budget
# Convert {{var}} to {var} in instructions and description for backend compatibility (v1 format)
# User writes: {{language}} → Backend receives: {language}
@@ -1184,14 +1789,7 @@ def build_save_payload(self, **kwargs: Any) -> dict:
converted_assets = []
if self.tools:
for tool in self.tools:
- if isinstance(tool, ToolableMixin):
- converted_assets.append(self._normalize_tool_dict_for_api(tool.as_tool()))
- elif isinstance(tool, dict):
- converted_assets.append(self._normalize_tool_dict_for_api(tool))
- else:
- raise ValueError(
- "A tool in the agent must be a Tool, Model, ToolableMixin instance, or a dictionary."
- )
+ converted_assets.append(self._normalize_tool_for_api(tool))
# Update the payload with converted assets
payload["tools"] = converted_assets
@@ -1213,13 +1811,33 @@ def build_save_payload(self, **kwargs: Any) -> dict:
converted_agents.append({"id": agent_id, "inspectors": []})
payload["agents"] = converted_agents
- # Handle BaseModel expected_output for save operation
- # We don't send expected_output in the save payload - it's runtime-only
+ # Convert skills to API format. Skills follow the same wire design as
+ # tools: each is sent as an object (via as_tool()), never a bare id.
+ if getattr(self, "_original_skills", None):
+ converted_skills = []
+ for skill in self._original_skills:
+ if isinstance(skill, ToolableMixin):
+ skill_dict = skill.as_tool()
+ elif isinstance(skill, dict):
+ skill_dict = skill
+ elif isinstance(skill, str):
+ skill_dict = {"id": skill, "type": "skill", "asset_id": skill}
+ else:
+ raise ValueError("A skill must be a Skill instance, a dict, or a skill id string.")
+ if not skill_dict.get("id"):
+ raise ValueError("All skills must be saved before saving the agent.")
+ converted_skills.append(self._normalize_tool_dict_for_api(skill_dict))
+ payload["skills"] = converted_skills
+ else:
+ payload.pop("skills", None)
+
+ # Persist expected_output server-side so fetched agents and runs that
+ # don't pass executionParams.expectedOutput (the backend falls back to
+ # the stored value) keep the JSON contract.
if "expectedOutput" in payload:
expected_output = payload["expectedOutput"]
if isinstance(expected_output, type) and issubclass(expected_output, BaseModel):
- # Remove BaseModel classes from save payload - they're not stored server-side
- payload.pop("expectedOutput")
+ payload["expectedOutput"] = json.dumps(expected_output.model_json_schema())
elif isinstance(expected_output, BaseModel):
# Convert BaseModel instance to dict for save
payload["expectedOutput"] = expected_output.model_dump()
@@ -1234,11 +1852,12 @@ def build_run_payload(self, **kwargs: Unpack[AgentRunParams]) -> dict:
# Normalize snake_case keys to camelCase for the API
execution_params = {self._EXEC_PARAMS_MAP.get(k, k): v for k, v in execution_params.items()}
- # Set default values for execution_params if not provided
+ # Set default values for execution_params if not provided.
+ # No ``maxIterations`` default: the iteration cap now travels inside
+ # ``executionParams.budget`` and the backend supplies its own default.
defaults = {
"outputFormat": self.output_format,
"maxTokens": getattr(self, "max_tokens", 2048),
- "maxIterations": getattr(self, "max_iterations", 5),
"maxTime": 300,
"contextOverflowStrategy": getattr(self, "context_overflow_strategy", None),
}
@@ -1266,6 +1885,45 @@ def build_run_payload(self, **kwargs: Unpack[AgentRunParams]) -> dict:
# Backend expects executionParams.expectedOutput as a string.
execution_params["expectedOutput"] = json.dumps(expected_output)
+ # Run-time budget: the agent's current ``budget`` state travels inside
+ # ``executionParams.budget`` (the backend merges it field-by-field over the
+ # persisted default). Set the key only when the budget carries at least one
+ # cap; an empty budget must leave the payload unchanged. ``budget`` is no
+ # longer a run kwarg — drop any stray one so it can't leak to the payload.
+ kwargs.pop("budget", None)
+ # ``getattr`` (not ``self.budget``) mirrors the defensive reads above so a
+ # test-constructed agent (``Agent.__new__`` bypassing ``__init__``) still
+ # works; a missing/None budget is treated as an empty one.
+ budget = getattr(self, "budget", None) or Budget()
+
+ # Deprecated run-time ``max_iterations`` exec param: fold into the run-time
+ # budget (the agent's budget wins on conflict) and stop emitting a standalone
+ # ``executionParams.maxIterations``. ``max_iterations`` is not in
+ # ``_EXEC_PARAMS_MAP``, so a snake_case key passes through unmapped — accept
+ # both spellings here (mirrors ExecutionConfig.to_api_dict in session.py).
+ deprecated_iterations = execution_params.pop("maxIterations", None)
+ if deprecated_iterations is None:
+ deprecated_iterations = execution_params.pop("max_iterations", None)
+ if deprecated_iterations is not None:
+ # Point past the SDK run plumbing (build_run_payload ->
+ # _post_and_handle_run -> RunnableResourceMixin.run -> Agent.run) to
+ # the user's agent.run(...) call site. The conflict warning (below) is
+ # emitted from this same frame, so it shares the stacklevel.
+ warnings.warn(
+ "Execution param 'max_iterations' is deprecated; set agent.budget.max_iterations instead. "
+ "It will be removed in a future release.",
+ DeprecationWarning,
+ stacklevel=5,
+ )
+ normalized_budget, conflicted = self._fold_iter_into_budget(budget, deprecated_iterations)
+ if conflicted:
+ warnings.warn(self._BUDGET_ITER_CONFLICT_MSG, UserWarning, stacklevel=5)
+ else:
+ normalized_budget = self._normalize_budget(budget)
+
+ if normalized_budget:
+ execution_params["budget"] = normalized_budget
+
# Handle run_response_generation with default value of False
run_response_generation = kwargs.pop("run_response_generation", False)
@@ -1273,6 +1931,14 @@ def build_run_payload(self, **kwargs: Unpack[AgentRunParams]) -> dict:
variables = kwargs.pop("variables", None) or {}
query = kwargs.pop("query", None)
+ # Multimodal attachments on the non-session run path: resolve them to
+ # ``{url, name, type, mimeType}`` descriptors (uploading any local paths)
+ # and send them as a structured ``attachments`` field — the same shape
+ # the agent worker consumes. Popped here so they aren't forwarded raw by
+ # the generic kwargs loop below.
+ attachments = kwargs.pop("attachments", None)
+ files = kwargs.pop("files", None)
+
# Build input_data dict with query and variables
if query is not None:
if isinstance(query, dict):
@@ -1299,83 +1965,267 @@ def build_run_payload(self, **kwargs: Unpack[AgentRunParams]) -> dict:
# Add query back if present
if query is not None:
payload["query"] = query
+ if attachments or files:
+ from .session import resolve_attachments
+
+ payload["attachments"] = resolve_attachments(
+ self.context, attachments, files, error_label=f"agent '{self.id}'"
+ )
# Translate remaining snake_case kwargs to camelCase for the API
for key, value in kwargs.items():
if value is not None:
api_key = self._SNAKE_TO_CAMEL.get(key, key)
payload[api_key] = value
+ # Send the agent's current per-tool parameter state as an ephemeral
+ # run-time override. Mutating ``agent.tools[i].actions[...].inputs[...]``
+ # and calling ``run()`` forwards those values without persisting them
+ # (the backend merges by tool id). Only tools that carry parameter
+ # values are included, so the common no-override case stays lightweight.
+ tool_overrides = self._build_tool_overrides()
+ if tool_overrides:
+ payload["tools"] = tool_overrides
+
self._apply_llm_fields_to_run_payload(payload)
return payload
- def generate_session_id(self, history: Optional[List[ConversationMessage]] = None) -> str:
- """Generate a unique session ID for agent conversations.
+ def _build_tool_overrides(self) -> List[dict]:
+ """Serialize ``self.tools`` into run-time per-tool parameter overrides.
- Creates a unique session identifier based on the agent ID and current timestamp.
- If conversation history is provided, it attempts to initialize the session on the
- server to enable context-aware conversations.
+ Reads each tool object's *current* action-input values **offline** and
+ emits the ``[{id, parameters: [{name, value}]}]`` shape the backend
+ merges by tool id. Only tools that carry at least one set value are
+ emitted, so runs without overrides stay lightweight; bare string ids and
+ attach-only dicts are skipped.
+ """
+ overrides: List[dict] = []
+ for tool in getattr(self, "tools", None) or []:
+ if not isinstance(tool, ToolableMixin):
+ continue
+ tool_id = getattr(tool, "id", None)
+ if not tool_id:
+ continue
+ values = self._current_tool_parameters(tool)
+ if values:
+ overrides.append(
+ {
+ "id": tool_id,
+ "parameters": self._params_dict_to_namevalue_list(values),
+ }
+ )
+ return overrides
- Args:
- history: Previous conversation history. Each message should contain
- 'role' (either 'user' or 'assistant') and 'content' keys.
- Defaults to None.
+ @staticmethod
+ def _current_tool_parameters(tool: Any) -> Dict[str, Any]:
+ """Read a tool object's current, non-null input values as ``{name: value}``.
- Returns:
- str: A unique session identifier in the format "{agent_id}_{timestamp}".
+ Fully offline — only inspects input collections that are already
+ materialized (a Model's single ``inputs``, or the actions a caller has
+ touched on a Tool), so it never triggers a lazy spec fetch.
+ """
+ values: Dict[str, Any] = {}
- Raises:
- ValueError: If the history format is invalid.
+ # Model-style: a single ``inputs`` collection (Tools raise on ``.inputs``).
+ try:
+ inputs = tool.inputs
+ except Exception:
+ inputs = None
+ if inputs is not None and hasattr(inputs, "items"):
+ for name, value in inputs.items():
+ if value is not None:
+ values[name] = value
+ return values
+
+ # Tool-style: gather values from already-materialized actions.
+ actions = getattr(tool, "actions", None)
+ materialized = getattr(actions, "_actions", None)
+ if isinstance(materialized, dict):
+ for action in materialized.values():
+ action_inputs = getattr(action, "_inputs", None)
+ if action_inputs is None or not hasattr(action_inputs, "items"):
+ continue
+ for name, value in action_inputs.items():
+ if value is not None:
+ values[name] = value
+ return values
- Example:
- >>> agent = Agent.get("my_agent_id")
- >>> session_id = agent.generate_session_id()
- >>> # Or with history
- >>> history = [
- ... {"role": "user", "content": "Hello"},
- ... {"role": "assistant", "content": "Hi there!"}
- ... ]
- >>> session_id = agent.generate_session_id(history=history)
+ @staticmethod
+ def _apply_run_overrides_to_session(session: "Session", kwargs: Dict[str, Any]) -> None:
+ """Apply per-run execution overrides onto a session.
+
+ When a caller runs within a ``session`` but also passes
+ per-run execution kwargs (``execution_params`` / ``criteria`` /
+ ``evolve`` / ``identifier`` / ``run_response_generation``), those
+ overrides would otherwise be silently dropped — the run would
+ execute with whatever ``executionConfig`` the session was created
+ with. Here we merge the supplied overrides onto the session's
+ stored config (fields not overridden are preserved) and, when the
+ result differs from what's stored, persist it so the overrides
+ take effect.
+
+ We warn because this mutates the session's ``executionConfig`` for
+ every subsequent message in the session, not just this run.
"""
- if not self.id:
- self.save(as_draft=True)
+ from .session import ExecutionConfig
+
+ overrides = {
+ "execution_params": kwargs.get("execution_params"),
+ "criteria": kwargs.get("criteria"),
+ "evolve": kwargs.get("evolve"),
+ "identifier": kwargs.get("identifier"),
+ "run_response_generation": kwargs.get("run_response_generation"),
+ }
+ provided = {key: value for key, value in overrides.items() if value is not None}
+ if not provided:
+ return
+
+ current = session.execution_config
+ base = {
+ "execution_params": getattr(current, "execution_params", None),
+ "criteria": getattr(current, "criteria", None),
+ "evolve": getattr(current, "evolve", None),
+ "identifier": getattr(current, "identifier", None),
+ "run_response_generation": getattr(current, "run_response_generation", None),
+ }
+ merged = ExecutionConfig(**{**base, **provided})
+
+ if current is not None and merged.to_api_dict() == current.to_api_dict():
+ return
+
+ warnings.warn(
+ f"Per-run execution overrides ({', '.join(sorted(provided))}) were "
+ f"passed alongside session '{session.id}'. Updating the session's "
+ f"stored executionConfig so the overrides take effect; this also "
+ f"applies to every subsequent message in this session.",
+ UserWarning,
+ stacklevel=3,
+ )
+ session.execution_config = merged
+ session.save()
+
+ _LEGACY_ONLY_RUN_KWARGS: ClassVar[tuple] = (
+ "tasks",
+ "prompt",
+ "inspectors",
+ "history",
+ "variables",
+ )
+
+ def _resolve_session(self, session: Any) -> "Session":
+ """Resolve the ``session=`` run argument to a bound :class:`Session`.
- if history:
- validate_history(history)
+ Accepts a :class:`~aixplain.v2.session.Session` instance (bound to this
+ agent's context if it isn't already) or a session id string (fetched via
+ ``Session.get``). Any other type is a ``TypeError``.
+ """
+ from .session import Session
+
+ if isinstance(session, Session):
+ if getattr(session, "context", None) is None:
+ session.context = self.context
+ return session
+ if isinstance(session, str):
+ return self.context.Session.get(session)
+ raise TypeError(f"session must be a Session instance or a session id string, got {type(session).__name__}")
+
+ def _run_with_session(self, session: Any, **kwargs: Any) -> AgentRunResult:
+ """Run the agent within a session, awaiting the triggered run.
+
+ Flow:
+ 1. Resolve ``session`` (a Session object or id) to a bound Session and
+ merge any per-run execution overrides onto its ``executionConfig``.
+ 2. POST a ``role="user"`` message via ``session.add_message`` — this
+ triggers the agent run on the backend with the session's
+ ``executionConfig``. Any ``attachments`` / ``files`` passed to
+ ``run`` are forwarded onto the user message so the agent receives
+ them (uploaded and attached by ``add_message``).
+ 3. Pull the agent run's ``requestId`` off the user message and hand it to
+ ``self.sync_poll`` (the ``/sdk/agents/{request_id}/result`` endpoint),
+ which returns a fully populated ``AgentRunResult`` including
+ ``data.steps``, ``execution_stats``, ``used_credits``, and ``run_time``.
+
+ We don't poll session messages for the assistant reply — the run result
+ endpoint is fully populated for the run the user message triggered.
+ """
+ self._validate_run_dependencies()
- timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
- session_id = f"{self.id}_{timestamp}"
+ offending = [k for k in self._LEGACY_ONLY_RUN_KWARGS if kwargs.get(k) is not None]
+ if offending:
+ raise ValueError(
+ f"session=… runs do not support legacy run kwargs: {offending}. Drop them or run without a session."
+ )
- if not history:
- return session_id
+ query = kwargs.get("query")
+ attachments = kwargs.get("attachments")
+ files = kwargs.get("files")
+ # The query is optional when attachments carry the turn's input (e.g. an
+ # audio clip that is itself the prompt); otherwise a text query is required.
+ if query is None and not (attachments or files):
+ raise ValueError("A session run requires a query or attachments.")
+ if query is not None and not isinstance(query, str):
+ raise ValueError("session=… only supports string queries.")
+ query = query or ""
+
+ session = self._resolve_session(session)
+ self._apply_run_overrides_to_session(session, kwargs)
+
+ # The agent's current per-tool parameter state (set by mutating
+ # ``agent.tools[i].actions[...].inputs[...]``) becomes the per-message
+ # override for this run, matching the single-shot run path.
+ message_tools = self._build_tool_overrides() or None
+
+ user_msg = session.add_message(
+ role="user",
+ content=query,
+ attachments=attachments,
+ files=files,
+ tools=message_tools,
+ )
+ if not user_msg.request_id:
+ raise ValueError(
+ f"Backend did not return a requestId on the user message for "
+ f"session '{session.id}'; cannot poll the agent run result."
+ )
+ # Same progress-tracker plumbing as the direct path: sync_poll calls
+ # self.on_poll(...) on every iteration, which forwards to the tracker.
+ self._start_progress_tracker(kwargs)
try:
- # Use the existing run infrastructure to initialize the session
- result = self.run_async(
- query="/",
- session_id=session_id,
- history=history,
- execution_params={
- "max_tokens": 2048,
- "max_iterations": 10,
- "output_format": OutputFormat.TEXT.value,
- "expected_output": None,
- },
- allow_history_and_session_id=True,
+ result = self.sync_poll(
+ user_msg.request_id,
+ timeout=kwargs.get("timeout", 300),
+ wait_time=kwargs.get("wait_time", 0.5),
)
+ except Exception as e:
+ self._finish_progress_tracker(e)
+ raise
+ self._finish_progress_tracker(result)
+
+ # The /sdk/agents/{id}/result response doesn't always echo back
+ # identifiers at the top level — back-fill from what we know locally so
+ # result.session_id / result.request_id are not None for session callers.
+ if not result.session_id:
+ result.session_id = session.id
+ if result.data is not None and getattr(result.data, "session_id", None) in (None, ""):
+ result.data.session_id = session.id
+ if not result.request_id:
+ result.request_id = user_msg.request_id
+ result._context = self.context
+ return result
- # If we got a polling URL, poll for completion
- if result.url and not result.completed:
- final_result = self.sync_poll(result.url, timeout=300, wait_time=0.5)
- if final_result.status == ResponseStatus.SUCCESS:
- return session_id
- else:
- logging.error(f"Session {session_id} initialization failed: {final_result}")
- return session_id
- else:
- # Direct completion or no polling needed
- return session_id
+# ``@dataclass_json`` injects its own ``from_dict`` onto the class, which would
+# clobber any ``from_dict`` defined in the class body. So we wrap the injected
+# decoder here (after decoration) to silently fold a legacy top-level
+# ``maxIterations`` into ``budget`` before decoding — keeping deserialization
+# warning-free while the explicit ``Agent(max_iterations=...)`` constructor path
+# still warns via ``__post_init__``.
+_dataclass_json_agent_from_dict = Agent.from_dict.__func__
- except Exception as e:
- logging.error(f"Failed to initialize session {session_id}: {e}")
- return session_id
+
+def _agent_from_dict(cls, kvs: Any, *, infer_missing: bool = False) -> "Agent":
+ kvs = cls._fold_legacy_max_iterations(kvs)
+ return _dataclass_json_agent_from_dict(cls, kvs, infer_missing=infer_missing)
+
+
+Agent.from_dict = classmethod(_agent_from_dict)
diff --git a/aixplain/v2/core.py b/aixplain/v2/core.py
index 881e9b30..7ed799b1 100644
--- a/aixplain/v2/core.py
+++ b/aixplain/v2/core.py
@@ -9,12 +9,15 @@
from .agent import Agent
from .utility import Utility
from .tool import Tool
+from .skill import Skill
from .agent_evaluator import Eval as EvalClass, Metric as MetricBase
from .integration import Integration
+from .trigger import Trigger
from .file import Resource
from .inspector import Inspector
from .meta_agents import Debugger
from .api_key import APIKey
+from .session import Session
from .rlm import RLM, RLMResult
from .issue import IssueReporter
from . import enums
@@ -24,12 +27,15 @@
AgentType = TypeVar("AgentType", bound=Agent)
UtilityType = TypeVar("UtilityType", bound=Utility)
ToolType = TypeVar("ToolType", bound=Tool)
+SkillType = TypeVar("SkillType", bound=Skill)
MetricType = TypeVar("MetricType", bound=MetricBase)
IntegrationType = TypeVar("IntegrationType", bound=Integration)
+TriggerType = TypeVar("TriggerType", bound=Trigger)
ResourceType = TypeVar("ResourceType", bound=Resource)
InspectorType = TypeVar("InspectorType", bound=Inspector)
DebuggerType = TypeVar("DebuggerType", bound=Debugger)
APIKeyType = TypeVar("APIKeyType", bound=APIKey)
+SessionType = TypeVar("SessionType", bound=Session)
RLMType = TypeVar("RLMType", bound=RLM)
IssueReporterType = TypeVar("IssueReporterType", bound=IssueReporter)
@@ -51,13 +57,16 @@ class Aixplain:
Agent: AgentType = None
Utility: UtilityType = None
Tool: ToolType = None
+ Skill: SkillType = None
Metric: MetricType = None
Eval: type = None
Integration: IntegrationType = None
+ Trigger: TriggerType = None
Resource: ResourceType = None
Inspector: InspectorType = None
Debugger: DebuggerType = None
APIKey: APIKeyType = None
+ Session: SessionType = None
RLM: RLMType = None
issue: IssueReporterType = None
@@ -77,6 +86,12 @@ class Aixplain:
SortOrder = enums.SortOrder
StorageType = enums.StorageType
+ SessionStatus = enums.SessionStatus
+ RunStatus = enums.RunStatus
+ MessageRole = enums.MessageRole
+ Reaction = enums.Reaction
+ AttachmentType = enums.AttachmentType
+
BACKEND_URL = "https://platform-api.aixplain.com"
BENCHMARKS_BACKEND_URL = "https://platform-api.aixplain.com"
MODELS_RUN_URL = "https://models.aixplain.com/api/v2/execute"
@@ -134,12 +149,15 @@ def init_resources(self) -> None:
self.Agent = type("Agent", (Agent,), {"context": self})
self.Utility = type("Utility", (Utility,), {"context": self})
self.Tool = type("Tool", (Tool,), {"context": self})
+ self.Skill = type("Skill", (Skill,), {"context": self})
self.Metric = type("Metric", (MetricBase,), {"context": self})
self.Eval = EvalClass
self.Integration = type("Integration", (Integration,), {"context": self})
+ self.Trigger = type("Trigger", (Trigger,), {"context": self})
self.Resource = type("Resource", (Resource,), {"context": self})
self.Inspector = type("Inspector", (Inspector,), {"context": self})
self.Debugger = type("Debugger", (Debugger,), {"context": self})
self.APIKey = type("APIKey", (APIKey,), {"context": self})
+ self.Session = type("Session", (Session,), {"context": self})
self.RLM = type("RLM", (RLM,), {"context": self})
self.issue = IssueReporter(context=self)
diff --git a/aixplain/v2/enums.py b/aixplain/v2/enums.py
index 7b9f2c81..3c863e42 100644
--- a/aixplain/v2/enums.py
+++ b/aixplain/v2/enums.py
@@ -43,6 +43,7 @@ class Function(str, Enum):
IMAGE_CLASSIFICATION = "IMAGE_CLASSIFICATION"
OBJECT_DETECTION = "OBJECT_DETECTION"
UTILITIES = "utilities" # Add the missing utilities function
+ GUARDRAILS = "guardrails" # Guardrail / inspector guard models
class Language(str, Enum):
@@ -247,6 +248,49 @@ class SplittingOptions(str, Enum):
LINE = "line"
+class SessionStatus(str, Enum):
+ """Session status values."""
+
+ ACTIVE = "active"
+ COMPLETED = "completed"
+ FAILED = "failed"
+ ARCHIVED = "archived"
+
+
+class RunStatus(str, Enum):
+ """Run status values for sessions."""
+
+ IDLE = "idle"
+ RUNNING = "running"
+ COMPLETED = "completed"
+
+
+class MessageRole(str, Enum):
+ """Message role in a session conversation."""
+
+ USER = "user"
+ ASSISTANT = "assistant"
+
+
+class Reaction(str, Enum):
+ """Reaction types for session messages."""
+
+ LIKE = "LIKE"
+ DISLIKE = "DISLIKE"
+
+
+class AttachmentType(str, Enum):
+ """Attachment type for session message attachments."""
+
+ TEXT = "text"
+ IMAGE = "image"
+ VIDEO = "video"
+ AUDIO = "audio"
+ DOCUMENT = "document"
+ CODE = "code"
+ UNKNOWN = "unknown"
+
+
__all__ = [
"AuthenticationScheme",
"FileType",
@@ -268,4 +312,9 @@ class SplittingOptions(str, Enum):
"CodeInterpreterModel",
"DataType",
"SplittingOptions",
+ "SessionStatus",
+ "RunStatus",
+ "MessageRole",
+ "Reaction",
+ "AttachmentType",
]
diff --git a/aixplain/v2/inspector.py b/aixplain/v2/inspector.py
index fae644bd..65671707 100644
--- a/aixplain/v2/inspector.py
+++ b/aixplain/v2/inspector.py
@@ -3,16 +3,43 @@
This module provides inspector functionality for validating team agent operations
at different stages (input, steps, output) with custom policies.
+The public surface is intentionally tiny: construct an :class:`Inspector` with
+plain strings for ``action`` / ``targets`` / ``severity`` and an ``aix.Metric``
+(the universal judge) for ``metric``. No enums or config classes to import.
+
"""
import inspect
+import logging
import textwrap
-from enum import Enum
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, field
+from .enums import Function
+from .exceptions import ResourceError
+from .resource import (
+ BaseResource,
+ BaseGetParams,
+ BaseSearchParams,
+ GetResourceMixin,
+ SearchResourceMixin,
+ Page,
+ encode_resource_id,
+ _flatten_asset_info,
+)
+
+logger = logging.getLogger(__name__)
+
AUTO_DEFAULT_MODEL_ID = "67fd9e2bef0365783d06e2f0"
+# Allowed string values (replacing the former user-facing enums). Kept as plain
+# sets so callers pass strings like action="abort" / targets=["output"] and never
+# import an enum.
+_ACTION_TYPES = frozenset({"continue", "rerun", "abort", "edit"})
+_ON_EXHAUST_VALUES = frozenset({"continue", "abort"})
+_TARGET_VALUES = frozenset({"input", "steps", "output"})
+_SEVERITY_VALUES = frozenset({"low", "medium", "high", "critical"})
+
def _callable_to_string(fn) -> str:
"""Convert a callable to its source string representation."""
@@ -22,61 +49,31 @@ def _callable_to_string(fn) -> str:
return f"{fn.__module__}.{fn.__qualname__}"
-class InspectorTarget(str, Enum):
- """Target stages for inspector validation in the team agent pipeline."""
-
- INPUT = "input"
- STEPS = "steps"
- OUTPUT = "output"
-
- def __str__(self) -> str:
- """Return the string value of the enum."""
- return self.value
-
-
-class InspectorAction(str, Enum):
- """Actions an inspector can take when evaluating content."""
-
- CONTINUE = "continue"
- RERUN = "rerun"
- ABORT = "abort"
- EDIT = "edit"
-
-
-class InspectorOnExhaust(str, Enum):
- """Action to take when max retries are exhausted."""
-
- CONTINUE = "continue"
- ABORT = "abort"
-
-
-class InspectorSeverity(str, Enum):
- """Severity level for inspector findings."""
-
- LOW = "low"
- MEDIUM = "medium"
- HIGH = "high"
- CRITICAL = "critical"
-
-
-class EvaluatorType(str, Enum):
- """Type of evaluator or editor."""
-
- ASSET = "asset"
- FUNCTION = "function"
-
-
@dataclass
-class InspectorActionConfig:
- """Inspector action configuration."""
+class _ActionConfig:
+ """Internal, normalized action policy.
+
+ Users never import or construct this — they pass ``action="abort"`` (a string)
+ or ``action={"type": "rerun", "max_retries": 2, "on_exhaust": "abort"}`` (a
+ dict) and :meth:`coerce` builds this.
+ """
- type: InspectorAction
+ type: str
max_retries: Optional[int] = None
- on_exhaust: Optional[InspectorOnExhaust] = None
+ on_exhaust: Optional[str] = None
def __post_init__(self) -> None:
- """Validate that max_retries and on_exhaust are only used with RERUN."""
- if self.type != InspectorAction.RERUN:
+ """Normalize/validate the action policy."""
+ self.type = str(self.type).strip().lower()
+ if self.type not in _ACTION_TYPES:
+ raise ValueError(f"invalid action {self.type!r}; expected one of {sorted(_ACTION_TYPES)}")
+ if self.on_exhaust is not None:
+ self.on_exhaust = str(self.on_exhaust).strip().lower()
+ if self.on_exhaust not in _ON_EXHAUST_VALUES:
+ raise ValueError(
+ f"invalid on_exhaust {self.on_exhaust!r}; expected one of {sorted(_ON_EXHAUST_VALUES)}"
+ )
+ if self.type != "rerun":
if self.max_retries is not None:
raise ValueError("max_retries is only valid when action type is 'rerun'")
if self.on_exhaust is not None:
@@ -84,46 +81,97 @@ def __post_init__(self) -> None:
if self.max_retries is not None and self.max_retries < 0:
raise ValueError("max_retries must be >= 0")
+ @classmethod
+ def coerce(cls, value: Any) -> "_ActionConfig":
+ """Build an action config from a string, dict, or existing config."""
+ if isinstance(value, cls):
+ return value
+ if isinstance(value, str):
+ return cls(type=value)
+ if isinstance(value, dict):
+ return cls(
+ type=value.get("type"),
+ max_retries=value.get("max_retries", value.get("maxRetries")),
+ on_exhaust=value.get("on_exhaust", value.get("onExhaust")),
+ )
+ raise ValueError(f"action must be a string or dict, got {type(value)!r}")
+
def to_dict(self) -> Dict[str, Any]:
"""Convert the action config to a dictionary for API serialization."""
- d: Dict[str, Any] = {"type": self.type.value}
+ d: Dict[str, Any] = {"type": self.type}
if self.max_retries is not None:
d["maxRetries"] = self.max_retries
if self.on_exhaust is not None:
- d["onExhaust"] = self.on_exhaust.value
+ d["onExhaust"] = self.on_exhaust
return d
- @classmethod
- def from_dict(cls, data: Dict[str, Any]) -> "InspectorActionConfig":
- """Create an InspectorActionConfig from a dictionary."""
- return cls(
- type=InspectorAction(data["type"]),
- max_retries=data.get("maxRetries"),
- on_exhaust=InspectorOnExhaust(data["onExhaust"]) if data.get("onExhaust") else None,
- )
-
@dataclass
-class EvaluatorConfig:
- """Evaluator configuration for an inspector."""
+class _Judge:
+ """Internal, normalized judge (evaluator or editor).
+
+ Users never import or construct this — they pass a Metric, an asset-id string,
+ or a Python callable and :meth:`coerce` builds this. It serializes to the same
+ backend shape the platform has always accepted (``type`` = ``asset`` |
+ ``function``).
+ """
- type: EvaluatorType
asset_id: Optional[str] = None
prompt: Optional[str] = None
function: Optional[str] = None
def __post_init__(self) -> None:
- """Validate and convert callable functions to source strings."""
+ """Convert callable functions to source and validate the judge has a target."""
if callable(self.function):
self.function = _callable_to_string(self.function)
- if self.type == EvaluatorType.ASSET and not self.asset_id:
- raise ValueError("asset_id is required when evaluator type is 'asset'")
- if self.type == EvaluatorType.FUNCTION and not self.function:
- raise ValueError("function is required when evaluator type is 'function'")
+ if not self.asset_id and not self.function:
+ raise ValueError("a judge needs an asset id (or Metric) or a function")
+
+ @property
+ def type(self) -> str:
+ """``"function"`` for callable judges, otherwise ``"asset"``."""
+ return "function" if self.function else "asset"
+
+ @classmethod
+ def coerce(cls, value: Any, *, prompt: Optional[str] = None) -> Optional["_Judge"]:
+ """Build a judge from a Metric, asset-id string, callable, dict, or judge.
+
+ Accepts:
+ - an ``aix.Metric`` (or any asset-backed object exposing ``id``) — the
+ universal judge; its id and prompt flow into the evaluator payload;
+ - a plain asset-id ``str``;
+ - a Python ``callable`` (or its source) — a custom function judge;
+ - a ``dict`` in either snake_case or the backend's camelCase.
+ """
+ if value is None:
+ return None
+ if isinstance(value, cls):
+ return value
+ if isinstance(value, dict):
+ return cls.from_dict(value)
+ if isinstance(value, str):
+ return cls(asset_id=value, prompt=prompt)
+ # A Metric (or any asset-backed judge) carries an ``id``.
+ asset_id = getattr(value, "id", None)
+ if asset_id:
+ resolved_prompt = prompt
+ if resolved_prompt is None:
+ resolved_prompt = getattr(value, "prompt_template", None)
+ if resolved_prompt is None:
+ cfg = getattr(value, "config", None)
+ if isinstance(cfg, dict):
+ resolved_prompt = cfg.get("prompt")
+ return cls(asset_id=asset_id, prompt=resolved_prompt)
+ if callable(value):
+ return cls(function=_callable_to_string(value))
+ if hasattr(value, "prompt_template"):
+ # A Metric-like object that hasn't been onboarded yet has no id.
+ raise ValueError("metric must be created before use as a judge — call Metric.create(...) so it has an id")
+ raise ValueError(f"metric must be a Metric, asset-id string, or callable; got {type(value)!r}")
def to_dict(self) -> Dict[str, Any]:
"""Convert to a dictionary for API serialization."""
- d: Dict[str, Any] = {"type": self.type.value}
+ d: Dict[str, Any] = {"type": self.type}
if self.asset_id is not None:
d["assetId"] = self.asset_id
if self.prompt is not None:
@@ -133,71 +181,188 @@ def to_dict(self) -> Dict[str, Any]:
return d
@classmethod
- def from_dict(cls, data: Dict[str, Any]) -> "EvaluatorConfig":
- """Create an EvaluatorConfig from a dictionary."""
+ def from_dict(cls, data: Dict[str, Any]) -> "_Judge":
+ """Create a judge from a backend (or user) dict."""
return cls(
- type=EvaluatorType(data["type"]),
- asset_id=data.get("assetId"),
+ asset_id=data.get("assetId", data.get("asset_id")),
prompt=data.get("prompt"),
function=data.get("function"),
)
-@dataclass
-class EditorConfig:
- """Editor configuration for an inspector."""
+# Default action/targets applied when a guard is fetched by its canonical
+# marketplace path. Keyed by the guard's asset name so the SDK never has to
+# hardcode a guard's asset id — a new guard ships purely as a marketplace entry
+# and inherits the safe default below until/unless it gets a tuned entry here.
+_DEFAULT_GUARD_CONFIG: Dict[str, Any] = {
+ "targets": ["input"],
+ "action": {"type": "abort"},
+}
- type: EvaluatorType
- asset_id: Optional[str] = None
- prompt: Optional[str] = None
- function: Optional[str] = None
+# Keys are the marketplace ``assetName`` (the middle segment of a guard's
+# ``host/asset-name/instance`` path), verified against the onboarded AWS guards.
+_GUARD_CONFIG_DEFAULTS: Dict[str, Dict[str, Any]] = {
+ "detect-prompt-attacks-guardrail": {
+ "targets": ["input"],
+ "action": {"type": "abort"},
+ },
+ "sensitive-information-guardrail": {
+ "targets": ["input"],
+ "action": {"type": "edit"},
+ },
+ "contextual-grounding-check-guardrail": {
+ "targets": ["output"],
+ "action": {"type": "rerun", "maxRetries": 2, "onExhaust": "abort"},
+ },
+}
- def __post_init__(self) -> None:
- """Validate and convert callable functions to source strings."""
- if callable(self.function):
- self.function = _callable_to_string(self.function)
- def to_dict(self) -> Dict[str, Any]:
- """Convert to a dictionary for API serialization."""
- d: Dict[str, Any] = {"type": self.type.value}
- if self.asset_id is not None:
- d["assetId"] = self.asset_id
- if self.prompt is not None:
- d["prompt"] = self.prompt
- if self.function is not None:
- d["function"] = self.function
- return d
+def _guard_slugs(value: Optional[str]) -> List[str]:
+ """Return every path segment of a guard path/id, lowercased.
- @classmethod
- def from_dict(cls, data: Dict[str, Any]) -> "EditorConfig":
- """Create an EditorConfig from a dictionary."""
- return cls(
- type=EvaluatorType(data["type"]),
- asset_id=data.get("assetId"),
- prompt=data.get("prompt"),
- function=data.get("function"),
+ Marketplace guard paths are ``host/asset-name/instance`` (the ``instanceId``,
+ e.g. ``"aws/sensitive-information-guardrail/aws"``) or ``host/asset-name``
+ (the ``assetPath``). The guard's identity is the *asset-name* segment, not
+ the trailing host — so we return all segments and let the registry match the
+ asset name wherever it sits. A bare asset id has no ``"/"`` and yields a
+ single non-matching slug, falling back to the safe default.
+ """
+ if not value:
+ return []
+ return [seg.lower() for seg in str(value).strip().strip("/").split("/") if seg]
+
+
+def _resolve_guard_defaults(*candidates: Optional[str]) -> Dict[str, Any]:
+ """Pick the tuned guard config matching the first known asset-name segment.
+
+ Tries each candidate (asset name, path/id) in order and, within each, every
+ path segment, preferring whichever matches a key in
+ :data:`_GUARD_CONFIG_DEFAULTS`. This matters when a guard is retrieved by its
+ bare asset id (which never matches the registry) but the payload also carries
+ its canonical ``path`` — the path's asset-name segment still selects the
+ tuned ``action``/``targets`` instead of the safe fallback. Falls back to the
+ safe default when nothing matches.
+ """
+ for candidate in candidates:
+ for slug in _guard_slugs(candidate):
+ if slug in _GUARD_CONFIG_DEFAULTS:
+ return _GUARD_CONFIG_DEFAULTS[slug]
+ return _DEFAULT_GUARD_CONFIG
+
+
+_GUARD_PRESET_IDS: Dict[str, str] = {
+ "detect-prompt-attacks-guardrail": "prompt_injection_guard",
+ "sensitive-information-guardrail": "pii_redaction",
+ "contextual-grounding-check-guardrail": "hallucination_guard",
+}
+
+
+def _resolve_guard_preset_id(*candidates: Optional[str]) -> Optional[str]:
+ """Return the backend preset id for a guard, matching its asset-name slug.
+
+ Mirrors :func:`_resolve_guard_defaults`: tries each candidate (asset name,
+ path/id) and every path segment, returning the first known preset id.
+ Returns ``None`` when nothing matches (unknown/future guards).
+ """
+ for candidate in candidates:
+ for slug in _guard_slugs(candidate):
+ if slug in _GUARD_PRESET_IDS:
+ return _GUARD_PRESET_IDS[slug]
+ return None
+
+
+@dataclass(repr=False)
+class Inspector(
+ BaseResource,
+ GetResourceMixin[BaseGetParams, "Inspector"],
+ SearchResourceMixin[BaseSearchParams, "Inspector"],
+):
+ """Inspector v2 configuration object.
+
+ An ``Inspector`` is the single type the ``aix.Agent(inspectors=[...])`` slot
+ accepts — whether it is hand-built or retrieved from the marketplace via
+ :meth:`get` / :meth:`search`. Prebuilt guards are ordinary marketplace assets
+ under the ``guardrails`` :class:`~aixplain.v2.enums.Function`; retrieving one
+ returns a fully-configured ``Inspector`` whose ``metric`` points at the guard
+ model, so a fetched guard and a custom inspector are indistinguishable to the
+ agent.
+
+ Configuration is plain data — no enums or config classes to import:
+
+ - ``action``: a string (``"continue" | "rerun" | "abort" | "edit"``) or, when
+ you need retry parameters, a dict
+ ``{"type": "rerun", "max_retries": 2, "on_exhaust": "abort"}``.
+ - ``targets``: a list of strings (``"input" | "steps" | "output"`` or a
+ sub-agent name).
+ - ``severity``: a string (``"low" | "medium" | "high" | "critical"``).
+ - ``metric``: the universal judge — an ``aix.Metric``, an asset-id string, or
+ a Python callable. Required.
+ - ``editor``: required when ``action`` is ``"edit"``; same accepted types as
+ ``metric``.
+
+ Example::
+
+ from aixplain import Aixplain
+
+ aix = Aixplain(api_key="")
+
+ # A custom inspector judged by a Metric (the universal judge)
+ inspector = aix.Inspector(
+ name="grounded_output",
+ severity="high",
+ targets=["output"],
+ action="abort",
+ metric=aix.Metric.create(
+ name="grounded",
+ llm_path="",
+ prompt_template="Abort if the answer is not grounded in the context.",
+ ),
)
+ # Discover and retrieve prebuilt guards like any other asset
+ aix.Inspector.search("guard")
+ guard = aix.Inspector.get("aws/detect-prompt-attacks-guardrail/aws")
+ redactor = aix.Inspector.get("aws/sensitive-information-guardrail/aws")
+ redactor.targets = ["output"] # config as an inspectable attribute
-@dataclass
-class Inspector:
- """Inspector v2 configuration object."""
+ team = aix.Agent(name="team", agents=[...], inspectors=[guard, inspector])
+ """
- name: str
- action: InspectorActionConfig
- evaluator: EvaluatorConfig
+ # Guards are marketplace models under function=guardrails, so retrieval is
+ # backed by the models endpoint and filtered to that function.
+ RESOURCE_PATH = "v2/models"
- description: Optional[str] = None
- severity: Optional[InspectorSeverity] = None
+ # ``name`` / ``description`` / ``id`` / ``path`` are inherited from BaseResource.
+ action: Any = None
+ metric: Any = None
+ severity: Optional[str] = None
targets: List[str] = field(default_factory=list)
- editor: Optional[EditorConfig] = None
+ editor: Any = None
+ # Backend inspector-preset id. Set automatically for fetched marketplace
+ # guards (see :meth:`from_guard_model`); ``None`` for custom inspectors.
+ preset_id: Optional[str] = None
def __post_init__(self) -> None:
- """Validate inspector configuration after initialization."""
+ """Normalize and validate inspector configuration after initialization."""
if not self.name or not str(self.name).strip():
raise ValueError("name cannot be empty")
- self.targets = [t for t in (self.targets or []) if t and str(t).strip()]
- if self.action.type == InspectorAction.EDIT and self.editor is None:
+ if self.action is None:
+ raise ValueError("action is required")
+ if self.metric is None:
+ raise ValueError("metric is required")
+
+ self.action = _ActionConfig.coerce(self.action)
+ self.metric = _Judge.coerce(self.metric)
+ self.editor = _Judge.coerce(self.editor)
+
+ if self.severity is not None:
+ self.severity = str(self.severity).strip().lower()
+ if self.severity not in _SEVERITY_VALUES:
+ raise ValueError(f"invalid severity {self.severity!r}; expected one of {sorted(_SEVERITY_VALUES)}")
+
+ self.targets = [str(t).strip() for t in (self.targets or []) if t and str(t).strip()]
+
+ if self.action.type == "edit" and self.editor is None:
raise ValueError("editor is required when action type is 'edit'")
def to_dict(self) -> Dict[str, Any]:
@@ -206,302 +371,161 @@ def to_dict(self) -> Dict[str, Any]:
"name": self.name,
"targets": list(self.targets),
"action": self.action.to_dict(),
- "evaluator": self.evaluator.to_dict(),
+ # The judge serializes under the backend's long-standing "evaluator" key.
+ "evaluator": self.metric.to_dict(),
}
if self.description is not None:
d["description"] = self.description
if self.severity is not None:
- d["severity"] = self.severity.value
+ d["severity"] = self.severity
if self.editor is not None:
d["editor"] = self.editor.to_dict()
+ if self.preset_id is not None:
+ d["presetId"] = self.preset_id
return d
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "Inspector":
- """Create an Inspector from a dictionary."""
+ """Create an Inspector from an inspector-shaped dictionary."""
if not isinstance(data, dict):
raise ValueError("Inspector data must be a dict")
- severity = data.get("severity")
editor_data = data.get("editor")
return cls(
name=data["name"],
description=data.get("description"),
- severity=InspectorSeverity(severity) if severity else None,
+ severity=data.get("severity"),
targets=data.get("targets") or [],
- action=InspectorActionConfig.from_dict(data.get("action") or {}),
- evaluator=EvaluatorConfig.from_dict(data["evaluator"]),
- editor=EditorConfig.from_dict(editor_data) if editor_data else None,
- )
-
-
-PROMPT_INJECTION_GUARDRAIL_ASSET_ID = "69a9974367d506543103ca18"
-PII_REDACTION_GUARDRAIL_ASSET_ID = "69cbf63cd74e334a6bacfeb1"
-HALLUCINATION_GUARD_ASSET_ID = "69a9a38967d506543103ca1d"
-
-_PREBUILT_REGISTRY: Dict[str, Dict[str, Any]] = {
- "prompt_injection_guard": {
- "name": "Prompt Injection Guard",
- "category": "protection",
- "description": "Detects prompt attacks before they influence planning or execution.",
- "default_targets": [InspectorTarget.INPUT.value],
- "default_action": {"type": InspectorAction.ABORT.value},
- "evaluator_asset_id": PROMPT_INJECTION_GUARDRAIL_ASSET_ID,
- "supported_actions": [
- InspectorAction.CONTINUE.value,
- InspectorAction.RERUN.value,
- InspectorAction.ABORT.value,
- ],
- "vendor": "aws",
- },
- "pii_redaction": {
- "name": "PII Redaction",
- "category": "redaction",
- "description": "Finds sensitive information and returns redacted content from the guardrail evaluator.",
- "default_targets": [InspectorTarget.INPUT.value],
- "default_action": {"type": InspectorAction.EDIT.value},
- "evaluator_asset_id": PII_REDACTION_GUARDRAIL_ASSET_ID,
- "supported_actions": [
- InspectorAction.CONTINUE.value,
- InspectorAction.EDIT.value,
- InspectorAction.ABORT.value,
- ],
- "vendor": "aws",
- },
- "hallucination_guard": {
- "name": "Hallucination Guard",
- "category": "quality",
- "description": "Detects hallucinations in the agent's final response by verifying it against the intermediate step results and the original user query.",
- "default_targets": [InspectorTarget.OUTPUT.value],
- "default_action": {"type": InspectorAction.RERUN.value, "maxRetries": 2, "onExhaust": "abort"},
- "evaluator_asset_id": HALLUCINATION_GUARD_ASSET_ID,
- "supported_actions": [
- InspectorAction.CONTINUE.value,
- InspectorAction.RERUN.value,
- InspectorAction.ABORT.value,
- ],
- "vendor": None,
- },
-}
-
-
-@dataclass
-class PrebuiltInspector:
- """A lightweight preset reference that the backend resolves into a full Inspector.
-
- Instead of manually configuring an evaluator, action, and editor, users can
- reference one of the platform's pre-built inspector presets by ID. The
- backend's ``normalize_prebuilt_inspectors`` validator expands the reference
- before the agent graph is constructed.
-
- Example::
-
- from aixplain.v2 import PrebuiltInspector, InspectorTarget
-
- team = client.Agent(
- name="Safe Agent",
- agents=[agent1, agent2],
- inspectors=[
- PrebuiltInspector.prompt_injection_guard(),
- PrebuiltInspector.pii_redaction(targets=[InspectorTarget.OUTPUT]),
- ],
+ action=data.get("action") or {},
+ metric=_Judge.from_dict(data["evaluator"]),
+ editor=_Judge.from_dict(editor_data) if editor_data else None,
+ preset_id=data.get("presetId") or data.get("preset_id"),
)
- """
-
- preset_id: str
- targets: Optional[List[str]] = None
- action: Optional[Dict[str, Any]] = None
- severity: Optional[InspectorSeverity] = None
- description: Optional[str] = None
- name: Optional[str] = None
- config: Optional[Dict[str, Any]] = None
-
- def __post_init__(self) -> None:
- """Validate the inspector configuration after initialization."""
- if self.preset_id not in _PREBUILT_REGISTRY:
- available = ", ".join(sorted(_PREBUILT_REGISTRY))
- raise ValueError(f"Unknown inspector preset '{self.preset_id}'. Available presets: {available}")
- if self.targets is not None:
- self.targets = [t.value if isinstance(t, InspectorTarget) else str(t).lower() for t in self.targets]
- if self.action is not None:
- action_type = self.action.get("type", "")
- supported = _PREBUILT_REGISTRY[self.preset_id]["supported_actions"]
- if str(action_type).lower() not in supported:
- raise ValueError(
- f"Action '{action_type}' is not supported by preset '{self.preset_id}'. "
- f"Supported actions: {supported}"
- )
-
- def to_dict(self) -> Dict[str, Any]:
- """Serialize to the lightweight reference format expected by the backend."""
- d: Dict[str, Any] = {"presetId": self.preset_id}
- if self.name is not None:
- d["name"] = self.name
- if self.description is not None:
- d["description"] = self.description
- if self.targets is not None:
- d["targets"] = list(self.targets)
- if self.action is not None:
- d["action"] = dict(self.action)
- if self.severity is not None:
- d["severity"] = self.severity.value
- if self.config is not None:
- d["config"] = dict(self.config)
- return d
+ # ------------------------------------------------------------------
+ # Marketplace retrieval
+ # ------------------------------------------------------------------
@classmethod
- def from_dict(cls, data: Dict[str, Any]) -> "PrebuiltInspector":
- """Create a PrebuiltInspector from a dictionary."""
- severity_raw = data.get("severity")
- return cls(
- preset_id=data["presetId"],
- name=data.get("name"),
- description=data.get("description"),
- targets=data.get("targets"),
- action=data.get("action"),
- severity=InspectorSeverity(severity_raw) if severity_raw else None,
- config=data.get("config"),
- )
+ def from_guard_model(cls, payload: Dict[str, Any], requested_path: Optional[str] = None) -> "Inspector":
+ """Adapt a ``guardrails`` marketplace model payload into a configured Inspector.
- @staticmethod
- def prompt_injection_guard(
- *,
- targets: Optional[List[Any]] = None,
- action: Optional[Dict[str, Any]] = None,
- severity: Optional[InspectorSeverity] = None,
- name: Optional[str] = None,
- description: Optional[str] = None,
- ) -> "PrebuiltInspector":
- """Create a Prompt Injection Guard inspector.
-
- Detects prompt injection attacks before they influence planning or
- execution. Defaults to ``ABORT`` on ``INPUT``.
+ The guard model becomes the inspector's ``metric`` (``asset`` judge), and
+ sensible default ``action`` / ``targets`` are applied based on the guard's
+ canonical path slug (see :data:`_GUARD_CONFIG_DEFAULTS`). Unknown guards
+ fall back to a safe ``abort`` on ``input`` default, so future guards need
+ no SDK change.
Args:
- targets: Override default targets (default: ``[InspectorTarget.INPUT]``).
- action: Override default action dict (default: ``{"type": "abort"}``).
- severity: Optional severity level.
- name: Optional custom name for the inspector node.
- description: Optional custom description.
+ payload: The guard-model dict returned by the marketplace.
+ requested_path: The path/id the caller passed to :meth:`get`, used to
+ resolve default config when the payload omits a path.
Returns:
- A PrebuiltInspector configured for prompt injection detection.
+ A fully-configured Inspector ready to attach to an agent.
"""
- return PrebuiltInspector(
- preset_id="prompt_injection_guard",
- targets=targets,
+ if not isinstance(payload, dict):
+ raise ValueError("guard model payload must be a dict")
+
+ payload = _flatten_asset_info(dict(payload))
+ model_id = payload.get("id")
+ asset_name = (payload.get("assetInfo") or {}).get("assetName")
+ defaults = _resolve_guard_defaults(asset_name, requested_path, payload.get("path"))
+ preset_id = _resolve_guard_preset_id(asset_name, requested_path, payload.get("path"))
+
+ action = _ActionConfig.coerce(defaults["action"])
+ editor = None
+ if action.type == "edit":
+ # The guard model performs the edit/redaction itself.
+ editor = _Judge(asset_id=model_id)
+
+ inspector = cls(
+ id=model_id,
+ name=payload.get("name") or "Guardrail",
+ description=payload.get("description"),
+ metric=_Judge(asset_id=model_id),
action=action,
- severity=severity,
- name=name,
- description=description,
+ targets=list(defaults["targets"]),
+ editor=editor,
+ preset_id=preset_id,
)
+ inspector.path = payload.get("path") or requested_path
+ return inspector
- @staticmethod
- def pii_redaction(
- *,
- targets: Optional[List[Any]] = None,
- action: Optional[Dict[str, Any]] = None,
- severity: Optional[InspectorSeverity] = None,
- name: Optional[str] = None,
- description: Optional[str] = None,
- ) -> "PrebuiltInspector":
- """Create a PII Redaction inspector.
-
- Finds sensitive information (PII) and returns redacted content from the
- guardrail evaluator. Defaults to ``EDIT`` on ``INPUT``.
+ @classmethod
+ def get(cls, id: Any, **kwargs: Any) -> "Inspector":
+ """Retrieve a prebuilt guard by human-readable path (IDs also accepted).
Args:
- targets: Override default targets (default: ``[InspectorTarget.INPUT]``).
- action: Override default action dict (default: ``{"type": "edit"}``).
- severity: Optional severity level.
- name: Optional custom name for the inspector node.
- description: Optional custom description.
+ id: The guard's marketplace path (e.g.
+ ``"aws/sensitive-information-guardrail/aws"``) or its asset id.
+ **kwargs: Additional request parameters (e.g. ``resource_path``)
+ forwarded to the underlying client call.
Returns:
- A PrebuiltInspector configured for PII redaction.
+ A fully-configured Inspector backed by the guard model.
"""
- return PrebuiltInspector(
- preset_id="pii_redaction",
- targets=targets,
- action=action,
- severity=severity,
- name=name,
- description=description,
- )
+ context = getattr(cls, "context", None)
+ if context is None:
+ raise ResourceError("Context is required for resource operations")
- @staticmethod
- def hallucination_guard(
- *,
- targets: Optional[List[Any]] = None,
- action: Optional[Dict[str, Any]] = None,
- severity: Optional[InspectorSeverity] = None,
- name: Optional[str] = None,
- description: Optional[str] = None,
- ) -> "PrebuiltInspector":
- """Create a Hallucination Guard inspector.
-
- Detects hallucinations in the agent's final response by verifying it
- against the intermediate step results and the original user query.
- Defaults to ``RERUN`` (max 2 retries, abort on exhaust) on ``OUTPUT``.
+ resource_path = kwargs.pop("resource_path", None) or getattr(cls, "RESOURCE_PATH", "")
+ encoded_id = encode_resource_id(id)
+ payload = context.client.get(f"{resource_path}/{encoded_id}", **kwargs)
- Args:
- targets: Override default targets (default: ``[InspectorTarget.OUTPUT]``).
- action: Override default action dict (default: ``{"type": "rerun", "maxRetries": 2, "onExhaust": "abort"}``).
- severity: Optional severity level.
- name: Optional custom name for the inspector node.
- description: Optional custom description.
+ inspector = cls.from_guard_model(payload, requested_path=str(id))
+ setattr(inspector, "context", context)
+ inspector._update_saved_state()
+ return inspector
- Returns:
- A PrebuiltInspector configured for hallucination detection.
- """
- return PrebuiltInspector(
- preset_id="hallucination_guard",
- targets=targets,
- action=action,
- severity=severity,
- name=name,
- description=description,
- )
+ @classmethod
+ def search(cls, query: Optional[str] = None, **kwargs: Any) -> Page["Inspector"]:
+ """Search available guards, returning the standard paginated shape.
- @staticmethod
- def list_presets() -> Dict[str, Dict[str, Any]]:
- """Return metadata for all available pre-built inspector presets.
+ Args:
+ query: Optional free-text query (e.g. ``"guard"``).
+ **kwargs: Additional pagination/search parameters.
Returns:
- A dict mapping preset IDs to their metadata (name, category,
- description, default targets/action, supported actions, vendor).
+ A ``Page`` of configured Inspectors.
"""
- return {
- pid: {
- "name": meta["name"],
- "category": meta["category"],
- "description": meta["description"],
- "default_targets": list(meta["default_targets"]),
- "default_action": dict(meta["default_action"]),
- "supported_actions": list(meta["supported_actions"]),
- "vendor": meta.get("vendor"),
- }
- for pid, meta in _PREBUILT_REGISTRY.items()
- }
+ if query is not None:
+ kwargs["query"] = query
+ return super().search(**kwargs)
+ @classmethod
+ def _populate_filters(cls, params: Dict[str, Any]) -> dict:
+ """Pin the search to the ``guardrails`` function (guards are models).
+
+ The ``v2/models/paginate`` ``functions`` filter expects a list of
+ function-code strings (the backend matches with ``function: {$in: [...]}``
+ and validates each element with ``@IsString({each: true})``). Sending the
+ object shape ``[{"id": }]`` fails validation with
+ ``each value in functions must be a string``.
+ """
+ filters = super()._populate_filters(params)
+ filters["functions"] = [Function.GUARDRAILS.value]
+ # The v2/models/paginate endpoint requires a sort array.
+ filters.setdefault("sort", [{}])
+ return filters
-def is_prebuilt_inspector(obj: Any) -> bool:
- """Return True if *obj* is a PrebuiltInspector or a dict preset reference."""
- if isinstance(obj, PrebuiltInspector):
- return True
- return isinstance(obj, dict) and isinstance(obj.get("presetId"), str)
+ @classmethod
+ def _build_resources(cls, items: List[dict], context: Any) -> List["Inspector"]:
+ """Adapt each guard-model item into a configured Inspector."""
+ resources: List["Inspector"] = []
+ for item in items:
+ if not isinstance(item, dict):
+ continue
+ try:
+ inspector = cls.from_guard_model(item)
+ except Exception as e: # pragma: no cover - defensive
+ logger.warning("Skipping guard during Inspector deserialization: %s", e)
+ continue
+ setattr(inspector, "context", context)
+ inspector._update_saved_state()
+ resources.append(inspector)
+ return resources
__all__ = [
"Inspector",
- "InspectorTarget",
- "InspectorAction",
- "InspectorOnExhaust",
- "InspectorSeverity",
- "InspectorActionConfig",
- "EvaluatorType",
- "EvaluatorConfig",
- "EditorConfig",
- "PrebuiltInspector",
- "is_prebuilt_inspector",
]
diff --git a/aixplain/v2/integration.py b/aixplain/v2/integration.py
index b233ba45..662f640a 100644
--- a/aixplain/v2/integration.py
+++ b/aixplain/v2/integration.py
@@ -90,6 +90,108 @@ class IntegrationSearchParams(BaseSearchParams):
pass
+@dataclass_json
+@dataclass
+class TriggerTypeSpec:
+ """Backend spec for an available external trigger type (deserialization only)."""
+
+ slug: Optional[str] = None
+ name: Optional[str] = None
+ description: Optional[str] = None
+ config: Optional[Any] = None
+ payload: Optional[Any] = None
+ type: Optional[str] = None
+
+
+class TriggerEventOption:
+ """A selectable external event option (e.g. ``gmail.triggers["NEW_EMAIL"]``).
+
+ Carries the event ``slug`` and its config schema. When sourced from a
+ *connected* tool it also carries ``connection_id`` (the tool id), which is
+ required to activate the trigger. Pass it to ``aix.Trigger(event=...)``.
+ """
+
+ def __init__(
+ self,
+ slug: str,
+ name: Optional[str] = None,
+ description: Optional[str] = None,
+ config: Optional[Any] = None,
+ connection_id: Optional[str] = None,
+ ) -> None:
+ """Initialize an event option."""
+ self.slug = slug
+ self.name = name or slug
+ self.description = description
+ self.config = config
+ self.connection_id = connection_id
+ self.values: Dict[str, Any] = {}
+
+ def configure(self, **values: Any) -> "TriggerEventOption":
+ """Set config values passed to the trigger on activation."""
+ self.values.update(values)
+ return self
+
+ def __repr__(self) -> str:
+ """Return a concise representation."""
+ conn = f", connected={bool(self.connection_id)}" if self.connection_id is not None else ""
+ return f"TriggerEventOption(slug={self.slug!r}{conn})"
+
+
+class TriggerTypes:
+ """Browsable collection of :class:`TriggerEventOption` for an integration/tool.
+
+ Supports ``integration.triggers["NEW_EMAIL"]`` (case-insensitive), iteration,
+ ``in``, and ``len``.
+ """
+
+ def __init__(self, specs: List[TriggerTypeSpec], connection_id: Optional[str] = None) -> None:
+ """Initialize from backend specs and an optional connection id."""
+ self._by_slug: Dict[str, TriggerTypeSpec] = {}
+ for spec in specs:
+ key = (spec.slug or spec.name or "").strip()
+ if key:
+ self._by_slug[key] = spec
+ self._connection_id = connection_id
+
+ def _resolve(self, key: str) -> Optional[str]:
+ normalized = str(key).strip().lower()
+ for slug in self._by_slug:
+ if slug.lower() == normalized:
+ return slug
+ return None
+
+ def __getitem__(self, key: str) -> TriggerEventOption:
+ """Return the :class:`TriggerEventOption` for *key* (case-insensitive)."""
+ slug = self._resolve(key)
+ if slug is None:
+ raise KeyError(f"Trigger event '{key}' not found")
+ spec = self._by_slug[slug]
+ return TriggerEventOption(
+ slug=spec.slug or slug,
+ name=spec.name,
+ description=spec.description,
+ config=spec.config,
+ connection_id=self._connection_id,
+ )
+
+ def __contains__(self, key: object) -> bool:
+ """Return whether *key* matches an available trigger event."""
+ return isinstance(key, str) and self._resolve(key) is not None
+
+ def __iter__(self):
+ """Iterate over available trigger event slugs."""
+ return iter(self._by_slug.keys())
+
+ def __len__(self) -> int:
+ """Return the number of available trigger events."""
+ return len(self._by_slug)
+
+ def __repr__(self) -> str:
+ """Return ``TriggerTypes(['SLUG', ...])``."""
+ return f"TriggerTypes({list(self._by_slug.keys())})"
+
+
@dataclass
class ActionMixin:
"""Mixin class providing action-related functionality for integrations and tools."""
@@ -220,6 +322,47 @@ def _load_inputs() -> Inputs:
_actions_lister=_list_action_names,
)
+ def list_trigger_types(self) -> List[TriggerTypeSpec]:
+ """List available external event trigger types for the integration/tool.
+
+ Uses the same model-execute mechanism as :meth:`list_actions`.
+
+ Returns:
+ List of :class:`TriggerTypeSpec` objects from the backend.
+ """
+ run_url = self.build_run_url()
+ response = self.context.client.request("post", run_url, json={"action": "list_trigger_types", "data": ""})
+
+ data = self._poll_for_data(response)
+ if not data or not isinstance(data, list):
+ return []
+
+ specs: List[TriggerTypeSpec] = []
+ for item in data:
+ if isinstance(item, dict):
+ try:
+ specs.append(TriggerTypeSpec.from_dict(item))
+ except Exception:
+ continue
+ return specs
+
+ @cached_property
+ def triggers(self) -> TriggerTypes:
+ """Browsable collection of external event options (like :attr:`actions`).
+
+ Pick an option and pass it to ``aix.Trigger(event=...)``. When accessed on
+ a *connected* tool, each option carries the connection id needed to
+ activate the trigger; on an unconnected integration it is discovery-only.
+
+ Returns:
+ :class:`TriggerTypes` collection (e.g. ``gmail.triggers["NEW_EMAIL"]``).
+ """
+ specs = self.list_trigger_types()
+ # A connected tool can activate triggers; an integration is discovery-only.
+ is_connection = str(getattr(self, "RESOURCE_PATH", "")).endswith("tools")
+ connection_id = self.id if is_connection else None
+ return TriggerTypes(specs, connection_id=connection_id)
+
def set_inputs(self, inputs_dict: Dict[str, Dict[str, Any]]) -> None:
"""Set multiple action inputs in bulk using a dictionary tree structure.
diff --git a/aixplain/v2/model.py b/aixplain/v2/model.py
index 0197bed1..da7224a9 100644
--- a/aixplain/v2/model.py
+++ b/aixplain/v2/model.py
@@ -4,6 +4,7 @@
import json
import logging
+import re
import time
from typing import Dict, Union, List, Optional, Any, TYPE_CHECKING, Iterator
from typing_extensions import NotRequired, Unpack
@@ -31,6 +32,13 @@
logger = logging.getLogger(__name__)
+_MODEL_POLL_URL_RE = re.compile(
+ r"^(?:https?://[^/?#]+)?/?(?:"
+ r"api/v1/data/[A-Za-z0-9_-]+|"
+ r"sdk/(?:models|runs)/[A-Za-z0-9_-]+(?:/result)?"
+ r")(?:[?#].*)?$"
+)
+
@dataclass_json
@dataclass
@@ -39,6 +47,7 @@ class Message:
role: str
content: Optional[str] = None
+ reasoning_content: Optional[str] = None
tool_calls: Optional[List[dict[str, Any]]] = None
refusal: Optional[str] = None
annotations: List[Any] = field(default_factory=list)
@@ -148,6 +157,7 @@ class StreamChunk:
Attributes:
status: The current status of the streaming operation (IN_PROGRESS or SUCCESS)
data: The content/token of this chunk
+ reasoning_content: Reasoning-model chain-of-thought text delta, when provided
tool_calls: Tool call deltas when stream uses OpenAI-style chunk format
usage: Usage payload when provided in a stream chunk
finish_reason: Completion reason for the current choice, when provided
@@ -155,6 +165,7 @@ class StreamChunk:
status: ResponseStatus
data: str
+ reasoning_content: Optional[str] = None
tool_calls: Optional[List[dict[str, Any]]] = None
usage: Optional[dict[str, Any]] = None
finish_reason: Optional[str] = None
@@ -291,6 +302,9 @@ def __next__(self) -> StreamChunk:
content = delta.get("content")
content = content if isinstance(content, str) else ""
+ reasoning_content = delta.get("reasoning_content")
+ reasoning_content = reasoning_content if isinstance(reasoning_content, str) else None
+
tool_calls = delta.get("tool_calls")
if tool_calls is not None and not isinstance(tool_calls, list):
tool_calls = [tool_calls]
@@ -304,6 +318,7 @@ def __next__(self) -> StreamChunk:
return StreamChunk(
status=self.status,
data=content,
+ reasoning_content=reasoning_content,
tool_calls=tool_calls,
usage=usage,
finish_reason=finish_reason,
@@ -678,13 +693,18 @@ def run(self, **kwargs: Unpack[ModelRunParams]) -> ModelResult:
if self.is_sync_only:
result = self._run_sync_v2(**effective_params)
- if result.url and not result.completed:
+ if result.url and result.status == "IN_PROGRESS" and not result.completed and self._is_poll_url(result.url):
result = self.sync_poll(result.url, **effective_params)
return result
else:
# Async-capable models: Use base run() which calls run_async() and polls
return super().run(**effective_params)
+ @staticmethod
+ def _is_poll_url(url: str) -> bool:
+ """Return True when a model URL matches a known polling endpoint."""
+ return bool(_MODEL_POLL_URL_RE.match(url))
+
def _run_sync_v2(self, **kwargs: Unpack[ModelRunParams]) -> ModelResult:
"""Run the model synchronously using V2 endpoint directly.
@@ -886,6 +906,13 @@ def _validate_params(self, **kwargs) -> List[str]:
if not self.params:
return []
+ # A multimodal ``data`` payload (OpenAI-style messages carrying text and
+ # media) is the complete input on its own. When it is present the flat
+ # ``text``/``prompt`` inputs are superseded, so missing-required checks
+ # for them must not fire — otherwise a vision call (``run(data=...)``)
+ # is rejected before it can reach the backend.
+ has_data = kwargs.get("data") is not None
+
errors = []
# Validate all parameters (required and optional)
@@ -899,7 +926,7 @@ def _validate_params(self, **kwargs) -> List[str]:
f"Expected {param.data_type}, "
f"got {type(value).__name__}"
)
- elif param.required:
+ elif param.required and not has_data:
errors.append(f"Required parameter '{param.name}' is missing")
return errors
@@ -1045,15 +1072,15 @@ def _populate_filters(cls, params: dict) -> dict:
if params.get("path") is not None:
filters["path"] = params["path"]
- # functions - accept list of strings and convert to backend shape
- # Use v2 format: array of objects with "id" key (required by v2/models/paginate endpoint)
+ # functions - backend validates "each value in functions must be a string"
if params.get("functions") is not None:
functions_param = params["functions"]
if isinstance(functions_param, list):
- filters["functions"] = [{"id": (f.value if hasattr(f, "value") else str(f))} for f in functions_param]
+ filters["functions"] = [(f.value if hasattr(f, "value") else str(f)) for f in functions_param]
else:
- value = functions_param.value if hasattr(functions_param, "value") else str(functions_param)
- filters["functions"] = [{"id": value}]
+ filters["functions"] = [
+ functions_param.value if hasattr(functions_param, "value") else str(functions_param)
+ ]
# suppliers - should be array of strings
if params.get("vendors") is not None:
diff --git a/aixplain/v2/resource.py b/aixplain/v2/resource.py
index 62a43310..9f0ee9af 100644
--- a/aixplain/v2/resource.py
+++ b/aixplain/v2/resource.py
@@ -224,7 +224,11 @@ class BaseResource:
path: Full path identifier (e.g., "openai/whisper-large/groq").
"""
- context: Any = field(repr=False, compare=False, metadata=config(exclude=lambda x: True), init=False)
+ # default=None so the attribute always exists: instances created without a
+ # bound context (e.g. Inspector.from_dict nested inside an Agent) must still
+ # be serializable — dataclasses_json's _asdict calls getattr(obj, "context")
+ # before applying `exclude`, so a missing attribute raises AttributeError.
+ context: Any = field(default=None, repr=False, compare=False, metadata=config(exclude=lambda x: True), init=False)
RESOURCE_PATH: str = field(
default="",
repr=False,
@@ -1253,13 +1257,24 @@ def _payload_kwargs_for_run(self, kwargs: dict) -> dict:
"""Kwargs for payload/URL builders excluding SDK orchestration keys."""
return {k: v for k, v in kwargs.items() if k not in self._RUN_CONTROL_KEYS}
+ def _headers_for_run(self, kwargs: dict) -> Optional[dict]:
+ """Build per-run headers from optional runtime metadata."""
+ identifier = kwargs.get("identifier")
+ if identifier is None:
+ return None
+ return {"x-user-id": str(identifier)}
+
def _post_and_handle_run(self, **kwargs: Unpack[RunParamsT]) -> ResultT:
"""Single POST + handle_run_response (no retries, no before_run)."""
self._ensure_valid_state()
payload_input = self._payload_kwargs_for_run(kwargs)
payload = self.build_run_payload(**payload_input)
run_url = self.build_run_url(**payload_input)
- response = self.context.client.request("post", run_url, json=payload)
+ request_kwargs = {"json": payload}
+ headers = self._headers_for_run(kwargs)
+ if headers:
+ request_kwargs["headers"] = headers
+ response = self.context.client.request("post", run_url, **request_kwargs)
return self.handle_run_response(response, **kwargs)
def _apply_after_run(self, result: ResultT, **kwargs: Unpack[RunParamsT]) -> ResultT:
@@ -1315,31 +1330,36 @@ def handle_run_response(self, response: dict, **kwargs: Unpack[RunParamsT]) -> R
Returns:
Response instance from the configured response class
"""
- # Check for polling URL in data field (legacy format)
- if response.get("data") and isinstance(response["data"], str) and response["data"].startswith("http"):
+ status = response.get("status", "IN_PROGRESS")
+ data = response.get("data")
+
+ # Check for polling URL in data field (legacy format). A successful
+ # response may also contain a URL as final data, such as generated media.
+ if status == "IN_PROGRESS" and data and isinstance(data, str) and data.startswith("http"):
# This is a polling URL case
response_class = getattr(self, "RESPONSE_CLASS", Result)
return response_class.from_dict(
{
- "status": response.get("status", "IN_PROGRESS"),
- "url": response["data"],
+ "status": status,
+ "url": data,
"completed": False,
+ "requestId": response.get("requestId"),
}
)
- elif response.get("status") == "IN_PROGRESS" and response.get("data"):
+ elif status == "IN_PROGRESS" and data:
# This is a polling URL case
response_class = getattr(self, "RESPONSE_CLASS", Result)
return response_class.from_dict(
{
- "status": response["status"],
- "url": response["data"],
+ "status": status,
+ "url": data,
"completed": False,
+ "requestId": response.get("requestId"),
}
)
else:
# Direct response case - pass the entire response to let dataclass_json handle field mapping
# Check for failed status and raise appropriate error
- status = response.get("status", "IN_PROGRESS")
if status == "FAILED":
raise create_operation_failed_error(response)
@@ -1417,7 +1437,10 @@ def run(self, *args: Any, **kwargs: Unpack[RunParamsT]) -> ResultT:
try:
result = self._post_and_handle_run(**kwargs)
if result.url and not result.completed:
+ request_id = getattr(result, "request_id", None)
result = self.sync_poll(result.url, **kwargs)
+ if request_id is not None and hasattr(result, "request_id") and not result.request_id:
+ result.request_id = request_id
return self._apply_after_run(result, **kwargs)
except APIError as e:
if not self._is_retryable_run_error(e) or attempt >= run_retries:
@@ -1479,14 +1502,17 @@ def poll(self, poll_url: str) -> ResultT:
# Handle polling response - use camelCase keys (what backend sends)
# dataclass_json with config(field_name=...) handles mapping to snake_case
run_time, used_credits = _extract_run_time_and_used_credits(response)
+ data = response.get("data") or {}
+ data_error = data.get("error") if isinstance(data, dict) else None
+ error_message = response.get("errorMessage") or data_error
filtered_response = {
"status": response.get("status", "IN_PROGRESS"),
"completed": response.get("completed", False),
- "errorMessage": response.get("errorMessage"),
+ "errorMessage": error_message,
"url": response.get("url"),
"result": response.get("result"),
"supplierError": response.get("supplierError"),
- "data": response.get("data") or {},
+ "data": data,
"sessionId": response.get("sessionId"),
"usedCredits": used_credits,
"runTime": run_time,
diff --git a/aixplain/v2/session.py b/aixplain/v2/session.py
new file mode 100644
index 00000000..72e87565
--- /dev/null
+++ b/aixplain/v2/session.py
@@ -0,0 +1,846 @@
+"""Session module for aiXplain v2 SDK."""
+
+import os
+import logging
+import mimetypes
+import warnings
+from dataclasses import dataclass, field, InitVar
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
+from pathlib import Path
+
+from dataclasses_json import dataclass_json, config
+
+from .enums import AttachmentType
+from .exceptions import APIError, ResourceError
+from .resource import (
+ BaseResource,
+ GetResourceMixin,
+ DeleteResourceMixin,
+ SearchResourceMixin,
+ BaseGetParams,
+ BaseDeleteParams,
+ BaseSearchParams,
+ Page,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# Maps snake_case execution param keys to the camelCase the backend accepts.
+# Kept in this module (not on Agent) so Session create/update — which carries
+# the same params as legacy agent.run — uses the same normalization.
+# NOTE: ``max_iterations`` is intentionally absent — it is deprecated and folded
+# into ``executionParams.budget.maxIterations`` (see ``ExecutionConfig``), never
+# emitted as a standalone ``executionParams.maxIterations``.
+EXECUTION_PARAMS_MAP: Dict[str, str] = {
+ "output_format": "outputFormat",
+ "max_tokens": "maxTokens",
+ "max_time": "maxTime",
+ "expected_output": "expectedOutput",
+}
+
+
+def _normalize_execution_params(params: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+ """Translate snake_case execution param keys to camelCase for the API."""
+ if not params:
+ return params
+ return {EXECUTION_PARAMS_MAP.get(k, k): v for k, v in params.items()}
+
+
+def _is_hosted_url(value: str) -> bool:
+ """True if ``value`` is an already-hosted location (not a local path).
+
+ Used by the unified ``attachments`` parameter to decide whether a string
+ entry is a URL to attach as-is or a local file path to upload.
+ """
+ return value.startswith(("http://", "https://", "s3://"))
+
+
+def _mime_to_attachment_type(mime_type: str) -> str:
+ """Map a MIME type string to an AttachmentType value."""
+ if not mime_type:
+ return AttachmentType.UNKNOWN.value
+
+ main_type = mime_type.split("/")[0] if "/" in mime_type else ""
+ sub_type = mime_type.split("/")[1] if "/" in mime_type else ""
+
+ # Code types
+ code_subtypes = {
+ "x-python",
+ "javascript",
+ "x-javascript",
+ "typescript",
+ "x-java-source",
+ "x-c",
+ "x-c++",
+ "x-ruby",
+ "x-go",
+ "x-rust",
+ "x-shellscript",
+ }
+ if sub_type in code_subtypes:
+ return AttachmentType.CODE.value
+
+ # Document types
+ document_subtypes = {
+ "pdf",
+ "msword",
+ "vnd.openxmlformats-officedocument.wordprocessingml.document",
+ "vnd.ms-excel",
+ "vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+ "vnd.ms-powerpoint",
+ "vnd.openxmlformats-officedocument.presentationml.presentation",
+ }
+ if sub_type in document_subtypes:
+ return AttachmentType.DOCUMENT.value
+
+ # Main type mapping
+ type_map = {
+ "image": AttachmentType.IMAGE.value,
+ "video": AttachmentType.VIDEO.value,
+ "audio": AttachmentType.AUDIO.value,
+ "text": AttachmentType.TEXT.value,
+ }
+ if main_type in type_map:
+ return type_map[main_type]
+
+ return AttachmentType.UNKNOWN.value
+
+
+def _infer_type_and_mime(ref: str) -> tuple[Optional[str], Optional[str]]:
+ """Infer ``(attachment_type, mimeType)`` from a filename or URL by extension.
+
+ Used to auto-detect the ``type`` of a hosted URL (or a path) when the caller
+ doesn't supply one — the agent worker keys off ``type`` (``image`` / ``audio``
+ / ``file``), so a URL with no type would otherwise be treated as a generic
+ file. Query/fragment are stripped so ``clip.wav?sig=…`` still resolves.
+ Returns ``(None, None)`` when nothing can be inferred (caller leaves the
+ fields unset → worker defaults to ``file``).
+ """
+ from .upload_utils import MimeTypeDetector
+
+ path = ref.split("?", 1)[0].split("#", 1)[0]
+ ext = Path(path).suffix.lower()
+ mime = MimeTypeDetector.EXTENSION_MAPPING.get(ext) or mimetypes.guess_type(path)[0]
+ if not mime:
+ return None, None
+ return _mime_to_attachment_type(mime), mime
+
+
+def _augment_hosted_attachment(att: Dict[str, Any]) -> Dict[str, Any]:
+ """Fill a hosted-URL attachment's missing ``type``/``mimeType`` by inference.
+
+ Explicit caller-supplied values are preserved; only absent fields are filled,
+ inferred from the attachment's ``name`` (preferred) or ``url`` extension.
+ """
+ if att.get("type") and att.get("mimeType"):
+ return att
+ att_type, mime = _infer_type_and_mime(att.get("name") or att.get("url") or "")
+ if att_type and not att.get("type"):
+ att["type"] = att_type
+ if mime and not att.get("mimeType"):
+ att["mimeType"] = mime
+ return att
+
+
+def resolve_attachments(
+ context: Any,
+ attachments: Optional[List[Union[str, Path, Dict[str, Any]]]],
+ files: Optional[List[Union[str, Path]]],
+ *,
+ error_label: str = "",
+) -> List[Dict[str, Any]]:
+ """Normalize the unified ``attachments`` (plus deprecated ``files``) for the API.
+
+ Each entry becomes a ``{url, name, type, mimeType}`` dict. URL entries (``http(s)://``
+ / ``s3://`` strings, or dicts carrying a ``url``) pass through unchanged; local paths
+ (plain strings, or dicts carrying a ``path``) are uploaded to aiXplain storage and the
+ resulting download link is attached. The ``FileUploader`` is created lazily, only when
+ an upload is actually needed. Shared by ``Session.add_message`` and ``Agent`` runs.
+
+ Args:
+ context: An object exposing ``backend_url`` and ``api_key`` (the SDK context).
+ attachments: The unified attachments list.
+ files: Deprecated local-path list (merged in, with a warning).
+ error_label: Optional context for upload-error messages (e.g. ``"session 's1'"``).
+ """
+ resolved: List[Dict[str, Any]] = []
+ uploader = None # lazily created on first upload
+
+ def _upload(path_str: str) -> Dict[str, Any]:
+ nonlocal uploader
+ from .upload_utils import FileUploader, MimeTypeDetector
+
+ if uploader is None:
+ uploader = FileUploader(backend_url=context.backend_url, api_key=context.api_key)
+ try:
+ download_url = uploader.upload(path_str, is_temp=True, return_download_link=True)
+ except Exception as e:
+ where = f" for {error_label}" if error_label else ""
+ raise ResourceError(f"Failed to upload file '{path_str}'{where}: {e}")
+ mime_type = MimeTypeDetector.detect_mime_type(path_str)
+ return {
+ "url": download_url,
+ "name": os.path.basename(path_str),
+ "type": _mime_to_attachment_type(mime_type),
+ "mimeType": mime_type or None,
+ }
+
+ for entry in attachments or []:
+ if isinstance(entry, dict):
+ if entry.get("url"):
+ # Auto-detect a missing type/mimeType from the url/name extension
+ # so a bare ``{"url": "...wav"}`` is recognized as audio (not a file).
+ resolved.append(_augment_hosted_attachment(dict(entry)))
+ elif entry.get("path"):
+ att = _upload(str(entry["path"]))
+ # Let caller-supplied keys (type/name/mimeType) override detection.
+ att.update({k: v for k, v in entry.items() if k != "path" and v is not None})
+ resolved.append(att)
+ else:
+ raise ResourceError("attachment dict must have a 'url' or a 'path' key")
+ elif isinstance(entry, (str, Path)):
+ value = str(entry)
+ if _is_hosted_url(value):
+ # A bare URL string: attach as-is, auto-detecting type/mimeType
+ # (and name) from the extension.
+ att: Dict[str, Any] = {"url": value, "name": os.path.basename(value.split("?", 1)[0])}
+ resolved.append(_augment_hosted_attachment(att))
+ else:
+ resolved.append(_upload(value))
+ else:
+ raise ResourceError(f"unsupported attachment entry type: {type(entry).__name__}")
+
+ if files:
+ warnings.warn(
+ "`files` is deprecated; pass local paths (or URLs) through `attachments` instead.",
+ DeprecationWarning,
+ stacklevel=3,
+ )
+ for file_path in files:
+ resolved.append(_upload(str(file_path)))
+
+ return resolved
+
+
+def _parse_list_response(response: Any, item_type: str) -> list:
+ """Parse a response that should be a bare JSON array.
+
+ Args:
+ response: The API response (should be a list).
+ item_type: Description of items for error messages.
+
+ Returns:
+ The response as a list.
+
+ Raises:
+ ResourceError: If response is not a list.
+ """
+ if isinstance(response, list):
+ return response
+ raise ResourceError(
+ f"Expected a list of {item_type} from the API, got {type(response).__name__}: {str(response)[:200]}"
+ )
+
+
+def _deserialize(cls, data: dict, description: str):
+ """Safely deserialize a dict into a dataclass.
+
+ Args:
+ cls: The dataclass type with from_dict.
+ data: The dict to deserialize.
+ description: Description for error messages.
+
+ Returns:
+ The deserialized instance.
+
+ Raises:
+ ResourceError: If deserialization fails.
+ """
+ try:
+ return cls.from_dict(data)
+ except Exception as e:
+ raise ResourceError(f"Failed to parse {description}: {e}. Response data: {str(data)[:200]}")
+
+
+def _resolve_agent_id(agent: Any) -> Optional[str]:
+ """Resolve an ``agent`` argument (Agent instance or id string) to its id.
+
+ Accepts ``None`` (returns ``None``), a plain id string, or any object
+ exposing an ``id`` attribute (an :class:`~aixplain.v2.agent.Agent`). Used by
+ both ``Session(agent=…)`` and ``Session.search(agent=…)`` so callers can pass
+ an agent object interchangeably with its id.
+ """
+ if agent is None:
+ return None
+ if isinstance(agent, str):
+ return agent
+ agent_id = getattr(agent, "id", None)
+ if not agent_id:
+ raise ResourceError("agent must be a saved Agent (with an id) or an agent id string")
+ return agent_id
+
+
+def _to_iso(value: Any) -> str:
+ """Serialize a date filter value (``datetime`` or string) to an ISO string."""
+ if isinstance(value, datetime):
+ return value.isoformat()
+ return str(value)
+
+
+@dataclass_json
+@dataclass
+class SessionMessageAttachment:
+ """Attachment on a session message."""
+
+ url: str
+ name: Optional[str] = None
+ type: Optional[str] = None
+ # Precise MIME type (e.g. ``audio/wav``). Carried so the worker can derive
+ # an audio codec instead of guessing from the filename extension. ``type``
+ # stays the coarse class (image/audio/...) for backward compatibility.
+ mimeType: Optional[str] = field(default=None, metadata=config(field_name="mimeType"))
+
+
+@dataclass_json
+@dataclass
+class SessionMessage:
+ """A message within a session (not a resource — all ops go through Session)."""
+
+ id: str = ""
+ session_id: str = field(default="", metadata=config(field_name="sessionId"))
+ user_id: str = field(default="", metadata=config(field_name="userId"))
+ agent_id: str = field(default="", metadata=config(field_name="agentId"))
+ role: str = ""
+ content: str = ""
+ sequence: int = 0
+ request_id: Optional[str] = field(default=None, metadata=config(field_name="requestId"))
+ reaction: Optional[str] = None
+ attachments: Optional[List[SessionMessageAttachment]] = None
+ created_at: str = field(default="", metadata=config(field_name="createdAt"))
+
+
+@dataclass_json
+@dataclass
+class ExecutionConfig:
+ """Per-session execution configuration.
+
+ Mirrors the run-time agent parameters historically passed to
+ ``agent.run`` so that messages posted to a session execute the agent
+ with the same configuration. The backend reads ``executionConfig``
+ on session create/update and applies it when subsequent user
+ messages trigger agent runs.
+
+ Attributes:
+ execution_params: Backend execution params (output format, max
+ tokens, etc.). Both snake_case and camelCase keys are
+ accepted; snake_case is normalized to camelCase on send.
+ criteria: Free-form evaluation criteria sent to the agent.
+ evolve: Evolution config as a JSON string (kept as a string to
+ match the backend contract).
+ identifier: Free-form identifier the backend can echo back on
+ messages (e.g. for client-side correlation).
+ run_response_generation: Whether the agent should run its final
+ response-generation step.
+ budget: Per-session run budget (cost / duration / iterations). Accepts a
+ ``Budget`` instance or a snake_case/camelCase dict. Serialized into
+ ``executionParams.budget`` so messages posted to the session run the
+ agent with this budget — the session-scoped equivalent of the agent's
+ own ``agent.budget``.
+ """
+
+ execution_params: Optional[Dict[str, Any]] = field(default=None, metadata=config(field_name="executionParams"))
+ criteria: Optional[str] = None
+ evolve: Optional[str] = None
+ identifier: Optional[str] = None
+ run_response_generation: Optional[bool] = field(default=None, metadata=config(field_name="runResponseGeneration"))
+ # Per-session run budget. Serialized manually into ``executionParams.budget``
+ # by ``to_api_dict`` (never as a top-level field), so it is excluded from the
+ # auto-generated dataclass_json serialization.
+ budget: Optional[Any] = field(default=None, metadata=config(field_name="budget", exclude=lambda v: True))
+
+ def __post_init__(self) -> None:
+ """Coerce a dict/Budget ``budget`` into a ``Budget`` instance."""
+ from .agent import Agent
+
+ self.budget = Agent._coerce_budget(self.budget)
+
+ def to_api_dict(self) -> Dict[str, Any]:
+ """Build the camelCase API payload, normalizing nested params.
+
+ Only fields the caller set are included so the backend keeps
+ existing values on partial updates. The deprecated
+ ``execution_params['max_iterations']`` is folded into
+ ``executionParams.budget.maxIterations`` (Budget wins on conflict) and a
+ standalone ``executionParams.maxIterations`` is never emitted — mirroring
+ the agent run path so sessions and direct runs behave identically.
+ """
+ from .agent import Agent
+
+ out: Dict[str, Any] = {}
+
+ normalized = _normalize_execution_params(self.execution_params) or {}
+ # Resolve the run-time budget first so the deprecated fold can defer to it.
+ budget = self.budget
+
+ # Deprecated run-time ``max_iterations`` exec param: fold into the budget
+ # (Budget wins on conflict) and stop emitting a standalone maxIterations.
+ # ``max_iterations`` is not in EXECUTION_PARAMS_MAP, so a snake_case key
+ # passes through unmapped — accept both spellings here (mirrors the run
+ # path in Agent.build_run_payload). ``_fold_iter_into_budget`` is pure; we
+ # own the conflict warning so it resolves to this to_api_dict() call site.
+ deprecated_iterations = normalized.pop("maxIterations", None)
+ if deprecated_iterations is None:
+ deprecated_iterations = normalized.pop("max_iterations", None)
+ if deprecated_iterations is not None:
+ warnings.warn(
+ "Execution param 'max_iterations' is deprecated; use budget=Budget(max_iterations=...). "
+ "It will be removed in a future release.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ budget, conflicted = Agent._fold_iter_into_budget(budget, deprecated_iterations)
+ if conflicted:
+ warnings.warn(Agent._BUDGET_ITER_CONFLICT_MSG, UserWarning, stacklevel=2)
+
+ # Serialize the budget into executionParams.budget (drops None fields).
+ if budget is not None:
+ normalized_budget = Agent._normalize_budget(budget)
+ if normalized_budget:
+ normalized["budget"] = normalized_budget
+
+ if normalized:
+ out["executionParams"] = normalized
+ if self.criteria is not None:
+ out["criteria"] = self.criteria
+ if self.evolve is not None:
+ out["evolve"] = self.evolve
+ if self.identifier is not None:
+ out["identifier"] = self.identifier
+ if self.run_response_generation is not None:
+ out["runResponseGeneration"] = self.run_response_generation
+ return out
+
+ @classmethod
+ def _fold_legacy_max_iterations(cls, kvs: Any) -> Any:
+ """Fold a legacy ``executionParams.maxIterations`` into ``budget`` silently.
+
+ Backend session payloads may still carry a legacy
+ ``executionParams.maxIterations``. Deserialization must NOT warn — only
+ the explicit ``ExecutionConfig(execution_params={'max_iterations': ...})``
+ send path (via ``to_api_dict``) does. So we fold any legacy
+ ``maxIterations`` straight into the ``budget`` slot here (Budget wins
+ silently if already set) and drop the standalone key. Returns a copy;
+ the caller's dict is untouched.
+ """
+ if isinstance(kvs, dict):
+ ep = kvs.get("executionParams") or kvs.get("execution_params")
+ if isinstance(ep, dict) and ep.get("maxIterations") is not None:
+ from .agent import Agent
+
+ kvs = dict(kvs) # shallow copy; never mutate the caller's dict
+ ep = dict(ep)
+ legacy_iterations = ep.pop("maxIterations")
+ budget = kvs.get("budget")
+ normalized = Agent._normalize_budget(budget) if budget is not None else {}
+ if normalized.get("maxIterations") is None:
+ normalized["maxIterations"] = legacy_iterations
+ kvs["budget"] = normalized
+ # Write the cleaned exec params back under whichever key was present.
+ if "executionParams" in kvs:
+ kvs["executionParams"] = ep
+ else:
+ kvs["execution_params"] = ep
+ return kvs
+
+ @classmethod
+ def coerce(cls, value: Any) -> Optional["ExecutionConfig"]:
+ """Accept an ExecutionConfig, dict, or None and return a config or None."""
+ if value is None:
+ return None
+ if isinstance(value, cls):
+ return value
+ if isinstance(value, dict):
+ return cls.from_dict(value)
+ raise TypeError(f"execution_config must be ExecutionConfig, dict, or None; got {type(value).__name__}")
+
+
+# ``@dataclass_json`` injects its own ``from_dict`` onto ExecutionConfig, which
+# would clobber any ``from_dict`` defined in the class body. So we wrap the
+# injected decoder here (after decoration) to silently fold a legacy
+# ``executionParams.maxIterations`` into ``budget`` before decoding — keeping
+# deserialization warning-free while the explicit ``to_api_dict`` send path
+# still warns. Mirrors the same pattern in ``aixplain/v2/agent.py``.
+_dataclass_json_execution_config_from_dict = ExecutionConfig.from_dict.__func__
+
+
+def _execution_config_from_dict(cls, kvs: Any, *, infer_missing: bool = False) -> "ExecutionConfig":
+ kvs = cls._fold_legacy_max_iterations(kvs)
+ return _dataclass_json_execution_config_from_dict(cls, kvs, infer_missing=infer_missing)
+
+
+ExecutionConfig.from_dict = classmethod(_execution_config_from_dict)
+
+
+@dataclass_json
+@dataclass(repr=False)
+class Session(
+ BaseResource,
+ GetResourceMixin[BaseGetParams, "Session"],
+ DeleteResourceMixin[BaseDeleteParams, "Session"],
+ SearchResourceMixin[BaseSearchParams, "Session"],
+):
+ """Session resource for managing agent conversation sessions.
+
+ Sessions are the single entry point for conversation threads: create with
+ ``aix.Session(agent=…)`` + ``save()``, find with ``aix.Session.search(agent=…)``,
+ and drive with ``agent.run(query, session=…)``. Each session is bound to one
+ agent (``agent_id``).
+ """
+
+ RESOURCE_PATH = "v1/sessions"
+
+ # The session list endpoint is a plain ``GET /v1/sessions`` with query-param
+ # filters — not the standard ``POST {path}/paginate``. ``search`` is fully
+ # overridden below to honor that shape while still returning a ``Page``.
+ PAGINATE_PATH = ""
+ PAGINATE_METHOD = "get"
+
+ # ``agent`` is an init-only convenience: ``Session(agent=agent_or_id)`` sets
+ # ``agent_id``. Not a serialized field (InitVar), so it never appears on the
+ # wire and deserialization (which never passes it) keeps the stored id.
+ agent: InitVar[Optional[Any]] = None
+
+ user_id: str = field(default="", metadata=config(field_name="userId"))
+ agent_id: str = field(default="", metadata=config(field_name="agentId"))
+ name: Optional[str] = None
+ status: str = "active"
+ run_status: str = field(default="", metadata=config(field_name="runStatus"))
+ message_count: int = field(default=0, metadata=config(field_name="messageCount"))
+ last_message_preview: Optional[str] = field(default=None, metadata=config(field_name="lastMessagePreview"))
+ last_message_at: Optional[str] = field(default=None, metadata=config(field_name="lastMessageAt"))
+ created_at: str = field(default="", metadata=config(field_name="createdAt"))
+ updated_at: str = field(default="", metadata=config(field_name="updatedAt"))
+ execution_config: Optional[ExecutionConfig] = field(default=None, metadata=config(field_name="executionConfig"))
+
+ def __post_init__(self, agent: Optional[Any] = None) -> None:
+ """Resolve the ``agent`` convenience arg and coerce ``execution_config``."""
+ if agent is not None:
+ resolved = _resolve_agent_id(agent)
+ if resolved:
+ self.agent_id = resolved
+ if self.execution_config is not None and not isinstance(self.execution_config, ExecutionConfig):
+ self.execution_config = ExecutionConfig.coerce(self.execution_config)
+
+ @classmethod
+ def _fold_legacy_execution_config(cls, kvs: Any) -> Any:
+ """Pre-fold a legacy ``executionConfig`` payload before decoding.
+
+ dataclass_json decodes the nested ``executionConfig`` field with its own
+ decoder, bypassing the overridden ``ExecutionConfig.from_dict``. So a
+ backend session carrying a legacy ``executionConfig.executionParams
+ .maxIterations`` would otherwise survive unflattened and trigger a
+ spurious ``DeprecationWarning`` later on serialization. Folding it here
+ (silently, via ``ExecutionConfig._fold_legacy_max_iterations``) keeps the
+ load path warning-free. Returns a copy; the caller's dict is untouched.
+ """
+ if isinstance(kvs, dict):
+ ec = kvs.get("executionConfig") or kvs.get("execution_config")
+ if isinstance(ec, dict):
+ folded = ExecutionConfig._fold_legacy_max_iterations(ec)
+ if folded is not ec: # only copy when a fold actually happened
+ kvs = dict(kvs)
+ if "executionConfig" in kvs:
+ kvs["executionConfig"] = folded
+ else:
+ kvs["execution_config"] = folded
+ return kvs
+
+ def build_save_payload(self, **kwargs: Any) -> dict:
+ """Build payload with only mutable fields."""
+ payload: Dict[str, Any] = {}
+ if self.agent_id:
+ payload["agentId"] = self.agent_id
+ if self.name is not None:
+ payload["name"] = self.name
+ if self.status:
+ payload["status"] = self.status
+ if self.execution_config is not None:
+ payload["executionConfig"] = self.execution_config.to_api_dict()
+ return payload
+
+ @classmethod
+ def search(
+ cls,
+ agent: Optional[Any] = None,
+ status: Optional[str] = None,
+ user_id: Optional[str] = None,
+ created_after: Optional[Union[str, datetime]] = None,
+ created_before: Optional[Union[str, datetime]] = None,
+ memory_enabled: Optional[bool] = None,
+ page_number: int = 0,
+ page_size: int = 20,
+ **kwargs: Any,
+ ) -> Page["Session"]:
+ """Search sessions with optional filters, returning a paginated ``Page``.
+
+ The single, standard way to list sessions (there is no ``agent.list_sessions()``
+ and no bespoke ``Session.list()``). Mirrors the ``search`` on every other
+ asset, but hits the session list endpoint (a plain ``GET /v1/sessions``
+ with query-param filters) and wraps the result in a ``Page``.
+
+ Args:
+ agent: Filter by agent — an :class:`~aixplain.v2.agent.Agent` instance
+ or an agent id string.
+ status: Filter by session status (e.g. ``"active"``).
+ user_id: Filter by owning user id.
+ created_after: Lower bound on the session's creation time
+ (``datetime`` or ISO string).
+ created_before: Upper bound on the session's creation time.
+ memory_enabled: When ``True`` return memory-on threads (sessions with
+ persisted traces); when ``False`` return memory-off runs. ``None``
+ (default) applies no memory filter. *(Backend filter delivery is a
+ follow-up; the SDK forwards the parameter today.)*
+ page_number: Zero-indexed page number (default 0).
+ page_size: Page size (default 20).
+ **kwargs: Accepted for forward compatibility with the standard
+ search signature; ignored by the session list endpoint.
+
+ Returns:
+ Page[Session]: A page of Session instances.
+
+ Raises:
+ ResourceError: If the API response cannot be parsed or
+ deserialization fails.
+ APIError: If the API request fails.
+ """
+ context = getattr(cls, "context", None)
+ if context is None:
+ raise ResourceError("Context is required for resource operations")
+
+ params: Dict[str, Any] = {}
+ agent_id = _resolve_agent_id(agent)
+ if agent_id is not None:
+ params["agentId"] = agent_id
+ if status is not None:
+ params["status"] = status
+ if user_id is not None:
+ params["userId"] = user_id
+ if created_after is not None:
+ params["createdAfter"] = _to_iso(created_after)
+ if created_before is not None:
+ params["createdBefore"] = _to_iso(created_before)
+ if memory_enabled is not None:
+ params["memoryEnabled"] = memory_enabled
+ params["pageNumber"] = page_number
+ params["pageSize"] = page_size
+
+ try:
+ response = context.client.request("get", cls.RESOURCE_PATH, params=params)
+ except APIError:
+ raise
+ except Exception as e:
+ raise ResourceError(f"Failed to search sessions: {e}")
+
+ # The backend returns a bare JSON array today; tolerate a paginated dict
+ # too so a future backend change (envelope with total/pageTotal) still works.
+ if isinstance(response, dict):
+ items = response.get(cls.PAGINATE_ITEMS_KEY) or response.get("items") or []
+ total = response.get(cls.PAGINATE_TOTAL_KEY, len(items))
+ page_total = response.get(cls.PAGINATE_PAGE_TOTAL_KEY, 1)
+ else:
+ items = _parse_list_response(response, "sessions")
+ total = len(items)
+ page_total = 1
+
+ results: List["Session"] = []
+ for item in items:
+ session = _deserialize(cls, item, "session")
+ session.context = context
+ session._update_saved_state()
+ results.append(session)
+
+ return Page(results=results, page_number=page_number, page_total=page_total, total=total)
+
+ def messages(self) -> List[SessionMessage]:
+ """Get all messages in this session.
+
+ Returns:
+ List of SessionMessage instances.
+
+ Raises:
+ ResourceError: If the API response is not a list or
+ deserialization fails.
+ APIError: If the API request fails.
+ """
+ self._ensure_valid_state()
+ path = f"{self.RESOURCE_PATH}/{self.encoded_id}/messages"
+
+ try:
+ response = self.context.client.request("get", path)
+ except APIError:
+ raise
+ except Exception as e:
+ raise ResourceError(f"Failed to list messages for session '{self.id}': {e}")
+
+ items = _parse_list_response(response, "messages")
+ return [_deserialize(SessionMessage, item, "session message") for item in items]
+
+ def add_message(
+ self,
+ role: str,
+ content: str,
+ request_id: Optional[str] = None,
+ attachments: Optional[List[Union[str, Path, Dict[str, Any]]]] = None,
+ files: Optional[List[Union[str, Path]]] = None,
+ tools: Optional[List[Dict[str, Any]]] = None,
+ ) -> SessionMessage:
+ """Add a message to this session.
+
+ Args:
+ role: Message role ("user" or "assistant").
+ content: Message content. May be empty when ``attachments`` carry the
+ turn's input (e.g. an audio clip that is itself the prompt).
+ request_id: Optional request ID to associate with the message.
+ attachments: The message's attachments. Each entry may be:
+
+ * a hosted-URL dict ``{"url", "type"?, "name"?, "mimeType"?}`` — used as-is;
+ * a local-path dict ``{"path": "/...", "type"?, ...}`` — uploaded;
+ * a string URL (``http(s)://`` / ``s3://``) — attached as-is;
+ * a string local path — uploaded to aiXplain storage.
+
+ files: Deprecated. Local file paths to upload and attach — pass these
+ through ``attachments`` instead.
+ tools: Per-message per-tool parameter overrides in the platform
+ ``[{id, parameters: [{name, value}]}]`` shape, applied to the
+ run this message triggers. Normally populated automatically from
+ the agent's tool objects by ``agent.run(query, session=…)``.
+
+ Returns:
+ The created SessionMessage.
+
+ Raises:
+ ResourceError: If the operation fails.
+ APIError: If the API request fails.
+ FileUploadError: If a file upload fails.
+ """
+ self._ensure_valid_state()
+
+ all_attachments = self._resolve_attachments(attachments, files)
+
+ payload: Dict[str, Any] = {"role": role, "content": content}
+ if request_id is not None:
+ payload["requestId"] = request_id
+ if all_attachments:
+ payload["attachments"] = all_attachments
+ if tools:
+ payload["tools"] = tools
+
+ path = f"{self.RESOURCE_PATH}/{self.encoded_id}/messages"
+ try:
+ response = self.context.client.request("post", path, json=payload)
+ except APIError:
+ raise
+ except Exception as e:
+ raise ResourceError(f"Failed to add message to session '{self.id}': {e}")
+
+ return _deserialize(SessionMessage, response, "session message")
+
+ def _resolve_attachments(
+ self,
+ attachments: Optional[List[Union[str, Path, Dict[str, Any]]]],
+ files: Optional[List[Union[str, Path]]],
+ ) -> List[Dict[str, Any]]:
+ """Resolve the unified ``attachments`` (+ deprecated ``files``) for this session.
+
+ Thin wrapper over :func:`resolve_attachments` that supplies this session's
+ context and an error label.
+ """
+ return resolve_attachments(self.context, attachments, files, error_label=f"session '{self.id}'")
+
+ def get_message(self, message_id: str) -> SessionMessage:
+ """Get a specific message by ID.
+
+ Args:
+ message_id: The message ID.
+
+ Returns:
+ The SessionMessage.
+
+ Raises:
+ ResourceError: If deserialization fails.
+ APIError: If the API request fails (e.g., message not found).
+ """
+ self._ensure_valid_state()
+ path = f"{self.RESOURCE_PATH}/{self.encoded_id}/messages/{message_id}"
+ try:
+ response = self.context.client.request("get", path)
+ except APIError:
+ raise
+ except Exception as e:
+ raise ResourceError(f"Failed to get message '{message_id}' from session '{self.id}': {e}")
+ return _deserialize(SessionMessage, response, "session message")
+
+ def delete_message(self, message_id: str) -> None:
+ """Delete a message from this session.
+
+ Args:
+ message_id: The message ID to delete.
+
+ Raises:
+ APIError: If the API request fails (e.g., message not found).
+ ResourceError: If the session is in an invalid state.
+ """
+ self._ensure_valid_state()
+ path = f"{self.RESOURCE_PATH}/{self.encoded_id}/messages/{message_id}"
+ try:
+ self.context.client.request_raw("delete", path)
+ except APIError:
+ raise
+ except Exception as e:
+ raise ResourceError(f"Failed to delete message '{message_id}' from session '{self.id}': {e}")
+
+ def react(self, message_id: str, reaction: Optional[str]) -> SessionMessage:
+ """React to a message or clear a reaction.
+
+ Only assistant messages can be reacted to.
+
+ Args:
+ message_id: The message ID to react to.
+ reaction: "LIKE", "DISLIKE", or None to clear.
+
+ Returns:
+ The updated SessionMessage.
+
+ Raises:
+ APIError: If the API request fails (e.g., reacting to a
+ non-assistant message).
+ ResourceError: If deserialization fails.
+ """
+ self._ensure_valid_state()
+ path = f"{self.RESOURCE_PATH}/{self.encoded_id}/messages/{message_id}/reaction"
+ payload: Dict[str, Any] = {"reaction": reaction}
+ try:
+ response = self.context.client.request("post", path, json=payload)
+ except APIError:
+ raise
+ except Exception as e:
+ raise ResourceError(f"Failed to react to message '{message_id}' in session '{self.id}': {e}")
+ return _deserialize(SessionMessage, response, "session message")
+
+
+# ``@dataclass_json`` injects its own ``from_dict`` onto Session, which would
+# clobber any ``from_dict`` defined in the class body. We wrap the injected
+# decoder here to silently pre-fold a legacy nested ``executionConfig`` before
+# decoding — so loading a backend session never emits a spurious deprecation
+# warning. Mirrors the pattern used for ``ExecutionConfig`` and ``Agent``.
+_dataclass_json_session_from_dict = Session.from_dict.__func__
+
+
+def _session_from_dict(cls, kvs: Any, *, infer_missing: bool = False) -> "Session":
+ kvs = cls._fold_legacy_execution_config(kvs)
+ return _dataclass_json_session_from_dict(cls, kvs, infer_missing=infer_missing)
+
+
+Session.from_dict = classmethod(_session_from_dict)
diff --git a/aixplain/v2/skill.py b/aixplain/v2/skill.py
new file mode 100644
index 00000000..db73e565
--- /dev/null
+++ b/aixplain/v2/skill.py
@@ -0,0 +1,312 @@
+"""Skill resource module.
+
+A ``Skill`` is a Claude-style skill — a ``SKILL.md`` (YAML frontmatter + markdown
+instructions), optionally alongside ``scripts/`` and ``resources/`` — registered as
+an aiXplain asset and attachable to agents. It is authored from a local path, either
+a folder containing ``SKILL.md`` or a single ``.md`` file; the file tree is uploaded
+and managed internally.
+
+The frontmatter ``description`` is the routing signal an agent sees; the body and
+resources are loaded just-in-time at runtime (progressive disclosure). Skills are
+attached to agents the same way tools are::
+
+ skill = aix.Skill(file_path="./skills/pdf-filler") # folder
+ skill = aix.Skill(file_path="calculator.md") # single file
+ skill.save() # upload bundle + register asset
+
+ agent = aix.Agent(name="analyst", skills=[skill])
+ agent.save()
+
+ aix.Skill.get("my-workspace/pdf-filler") # retrieve (path or id)
+ aix.Skill.search("pdf form") # search
+ skill.download() # download the bundle to ./{name}.zip
+ skill.download(file_path="./pdf-filler.zip") # ...or an explicit path
+"""
+
+import os
+from typing import Any, List, Optional, Tuple
+from dataclasses import dataclass, field
+from typing_extensions import NotRequired, Unpack
+from dataclasses_json import dataclass_json, config
+
+import yaml
+
+from .resource import (
+ BaseResource,
+ BaseGetParams,
+ BaseSearchParams,
+ BaseDeleteParams,
+ SearchResourceMixin,
+ GetResourceMixin,
+ DeleteResourceMixin,
+ Page,
+)
+from .enums import Privacy
+from .mixins import ToolableMixin
+from .upload_utils import FileUploader
+
+
+def _exclude(_: Any) -> bool:
+ """Marker to drop a field from the serialized payload."""
+ return True
+
+
+def _parse_skill_md(text: str) -> Tuple[Optional[str], Optional[str], List[str], str]:
+ """Parse a ``SKILL.md`` into (name, description, required_tools, body).
+
+ Frontmatter is an optional ``---``-delimited YAML block at the top of the
+ file carrying ``name``, ``description``, and ``requires`` (tool paths).
+ """
+ name = description = None
+ requires: List[str] = []
+ body = text
+ if text.lstrip().startswith("---"):
+ stripped = text.lstrip()
+ end = stripped.find("\n---", 3)
+ if end != -1:
+ front = stripped[3:end].strip()
+ body = stripped[end + 4 :].lstrip("\n")
+ meta = yaml.safe_load(front) or {}
+ name = meta.get("name")
+ description = meta.get("description")
+ requires = meta.get("requires") or meta.get("required_tools") or []
+ if isinstance(requires, str):
+ requires = [requires]
+ return name, description, list(requires), body
+
+
+class SkillSearchParams(BaseSearchParams):
+ """Search parameters for skills.
+
+ Attributes:
+ tags: Filter by tags.
+ suppliers: Filter by suppliers.
+ saved: Only return skills the caller has saved.
+ """
+
+ tags: NotRequired[List[str]]
+ suppliers: NotRequired[List[str]]
+ saved: NotRequired[bool]
+
+
+@dataclass_json
+@dataclass(repr=False)
+class Skill(
+ BaseResource,
+ SearchResourceMixin[SkillSearchParams, "Skill"],
+ GetResourceMixin[BaseGetParams, "Skill"],
+ DeleteResourceMixin[BaseDeleteParams, "Skill"],
+ ToolableMixin,
+):
+ """A Claude-style skill registered as an aiXplain asset.
+
+ Authored from a local path via ``aix.Skill(file_path=...)`` — either a folder
+ containing ``SKILL.md`` or a single ``.md`` file; the bundle's file tree is
+ uploaded internally on ``save()``. Attach to agents with
+ ``aix.Agent(skills=[skill_or_id])``.
+ """
+
+ RESOURCE_PATH = "sdk/skill"
+
+ tags: List[str] = field(default_factory=list)
+ privacy: Privacy = Privacy.PRIVATE
+ whitelist: List[str] = field(default_factory=list)
+
+ # Authoring input: a local Claude-style skill folder or a single ``.md`` file.
+ # Parsed on construction; never sent to the backend.
+ file_path: Optional[str] = field(default=None, repr=False, metadata=config(exclude=_exclude))
+
+ # Parsed from SKILL.md frontmatter/body when authored from a folder. Read-only
+ # to a developer and excluded from the create/update payload.
+ required_tools: List[str] = field(default_factory=list, metadata=config(exclude=_exclude))
+ instructions: Optional[str] = field(default=None, repr=False, metadata=config(exclude=_exclude))
+
+ # backend-populated, read-only metadata
+ team: Optional[int] = None
+ user: Optional[int] = None
+ status: Optional[str] = None
+ asset_type: Optional[str] = field(default=None, metadata=config(field_name="assetType"))
+ file_type: Optional[str] = field(default=None, metadata=config(field_name="fileType"))
+ created_at: Optional[str] = field(default=None, metadata=config(field_name="createdAt"))
+ updated_at: Optional[str] = field(default=None, metadata=config(field_name="updatedAt"))
+
+ def __post_init__(self) -> None:
+ """Load skill metadata from the local path when authoring a new skill."""
+ self._local_path = None
+ self._local_is_file = False
+ if self.file_path and not self.id:
+ self._load_from_path(self.file_path)
+
+ def _load_from_path(self, path: str) -> None:
+ """Parse the skill markdown and stage the source for upload on save.
+
+ Accepts either a Claude-style skill folder (containing ``SKILL.md``) or a
+ single ``.md`` file that serves as the skill's ``SKILL.md``.
+ """
+ path = os.path.abspath(path)
+ if os.path.isdir(path):
+ skill_md = os.path.join(path, "SKILL.md")
+ if not os.path.isfile(skill_md):
+ raise ValueError(f"A skill folder must contain a SKILL.md file: {path}")
+ fallback_name = os.path.basename(path)
+ self._local_is_file = False
+ elif os.path.isfile(path):
+ skill_md = path
+ fallback_name = os.path.splitext(os.path.basename(path))[0]
+ self._local_is_file = True
+ else:
+ raise ValueError(f"Skill path not found: {path}")
+ with open(skill_md, "r", encoding="utf-8") as handle:
+ name, description, requires, body = _parse_skill_md(handle.read())
+ # Precedence: explicit name= > SKILL.md frontmatter > folder/file name.
+ self.name = self.name or name or fallback_name
+ if not self.name:
+ raise ValueError("Could not determine a skill name (pass name= or set it in SKILL.md).")
+ self.description = self.description or description or ""
+ self.required_tools = requires
+ self.instructions = body
+ self._local_path = path
+
+ # ------------------------------------------------------------------ #
+ # Retrieval / search
+ # ------------------------------------------------------------------ #
+ @classmethod
+ def get(cls: type["Skill"], id: str, **kwargs: Unpack[BaseGetParams]) -> "Skill":
+ """Get a skill by path or id."""
+ return super().get(id, **kwargs)
+
+ @classmethod
+ def search(
+ cls: type["Skill"],
+ query: Optional[str] = None,
+ **kwargs: Unpack[SkillSearchParams],
+ ) -> Page["Skill"]:
+ """Search skills with an optional free-text query and filters."""
+ if query is not None:
+ kwargs["query"] = query
+ return super().search(**kwargs)
+
+ @classmethod
+ def _populate_filters(cls, params: dict) -> dict:
+ """Add skill-specific filters on top of the standard pagination ones."""
+ filters = super()._populate_filters(params)
+ if params.get("tags") is not None:
+ filters["tags"] = params["tags"]
+ if params.get("suppliers") is not None:
+ filters["suppliers"] = params["suppliers"]
+ if params.get("saved") is not None:
+ filters["saved"] = params["saved"]
+ return filters
+
+ # ------------------------------------------------------------------ #
+ # Lifecycle
+ # ------------------------------------------------------------------ #
+ def save(self, *args: Any, **kwargs: Any) -> "Skill":
+ """Save the skill, uploading the bundle when authored from a local path.
+
+ Args:
+ *args: Positional arguments passed to the base save method.
+ **kwargs: Attributes to set before saving (passed to base save).
+ """
+ super().save(*args, **kwargs)
+ if getattr(self, "_local_path", None):
+ if self._local_is_file:
+ self._upload_file_as_skill(self._local_path)
+ else:
+ self._upload_folder(self._local_path)
+ self._local_path = None
+ return self
+
+ def refresh(self) -> "Skill":
+ """Reload the skill's metadata from the backend."""
+ fresh = type(self).get(self.id)
+ self.status = fresh.status
+ self.updated_at = fresh.updated_at
+ return self
+
+ def download(self, file_path: Optional[str] = None) -> str:
+ """Download the skill bundle to a local path. Returns the written path.
+
+ Args:
+ file_path: Where to write the bundle. Defaults to ``./{name}.zip``.
+ """
+ self._ensure_valid_state()
+ file_path = file_path or f"./{self.name}.zip"
+ url = f"{self.RESOURCE_PATH}/{self.encoded_id}/download"
+ response = self.context.client.request_raw("get", url)
+ with open(file_path, "wb") as handle:
+ handle.write(response.content)
+ return file_path
+
+ def as_tool(self) -> dict:
+ """Serialize this skill as a tool object for agent attachment.
+
+ Skills follow the same wire design as tools: attached as objects (not bare
+ ids), with ``type="skill"``.
+ """
+ return {
+ "id": self.id,
+ "name": self.name,
+ "description": self.description or "",
+ "supplier": "aixplain",
+ "type": "skill",
+ "version": None,
+ "asset_id": self.id,
+ }
+
+ # ------------------------------------------------------------------ #
+ # Internal: upload the local source as the skill's file tree
+ # ------------------------------------------------------------------ #
+ def _upload_file_as_skill(self, path: str) -> None:
+ """Upload a single ``.md`` file as the skill's ``SKILL.md`` at the root."""
+ self._ensure_valid_state()
+ base = f"{self.RESOURCE_PATH}/{self.encoded_id}"
+ url = self._upload(path)
+ self.context.client.request(
+ "post",
+ f"{base}/file",
+ json={"name": "SKILL.md", "url": url, "description": "", "parentId": None},
+ )
+
+ def _upload_folder(self, root: str) -> None:
+ """Walk the local folder and create the backend file/folder tree.
+
+ Folder structure is preserved: each subdirectory becomes a folder node and
+ each file is uploaded and registered under its parent. Node management is
+ entirely internal — it is not part of the developer-facing surface.
+ """
+ self._ensure_valid_state()
+ base = f"{self.RESOURCE_PATH}/{self.encoded_id}"
+ folder_ids = {"": None} # relative dir -> backend folder id (root -> None)
+
+ for dirpath, _dirnames, filenames in os.walk(root):
+ rel = os.path.relpath(dirpath, root)
+ rel = "" if rel == "." else rel
+
+ if rel: # create a folder node for this subdirectory
+ parent_id = folder_ids.get(os.path.dirname(rel))
+ result = self.context.client.request(
+ "post",
+ f"{base}/folder",
+ json={"name": os.path.basename(rel), "description": "", "parentId": parent_id},
+ )
+ folder_ids[rel] = result.get("id")
+
+ parent_id = folder_ids.get(rel)
+ for filename in sorted(filenames):
+ url = self._upload(os.path.join(dirpath, filename))
+ self.context.client.request(
+ "post",
+ f"{base}/file",
+ json={"name": filename, "url": url, "description": "", "parentId": parent_id},
+ )
+
+ def _upload(self, file_path: str) -> str:
+ """Upload a local file to S3 and return its download URL."""
+ if not os.path.exists(file_path):
+ raise ValueError(f"File not found: {file_path}")
+ uploader = FileUploader(
+ api_key=self.context.client.team_api_key,
+ backend_url=self.context.backend_url,
+ )
+ return uploader.upload(file_path, is_temp=True, return_download_link=True)
diff --git a/aixplain/v2/trigger.py b/aixplain/v2/trigger.py
new file mode 100644
index 00000000..5d5cc0af
--- /dev/null
+++ b/aixplain/v2/trigger.py
@@ -0,0 +1,508 @@
+"""Trigger management module for the aiXplain v2 API.
+
+A :class:`Trigger` fires an agent with an ``input`` (the query) either on a time
+schedule (once / daily / weekly / monthly / interval) or on an external event
+(e.g. a Composio integration event such as a new Gmail email).
+
+Time triggers map straight onto ``POST/GET/PUT/DELETE /v1/triggers``. Event
+triggers are orchestrated over the existing endpoints: the SDK activates the
+Composio trigger on a connected tool via the model-execute endpoint (the same
+mechanism used by ``integration.actions``), then persists the returned trigger id
+through ``/v1/triggers``.
+"""
+
+from __future__ import annotations
+
+import warnings
+from dataclasses import dataclass, field
+from dataclasses_json import dataclass_json, config as dj_config
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+from .resource import (
+ BaseResource,
+ SearchResourceMixin,
+ GetResourceMixin,
+ DeleteResourceMixin,
+ DeleteResult,
+ Page,
+ BaseSearchParams,
+ BaseGetParams,
+ BaseDeleteParams,
+)
+from .exceptions import ResourceError
+
+if TYPE_CHECKING:
+ from .core import Aixplain
+
+
+# Backend-accepted string values (see platform-backend src/trigger/enums).
+_SCHEDULE_TYPES = {"once", "daily", "weekly", "monthly", "recurring"}
+_TIME_UNITS = {"minute", "hour", "day", "week"} # NOTE: no "month" unit on the backend
+_WEEKDAYS = {"sun", "mon", "tue", "wed", "thu", "fri", "sat"}
+_WEEKDAY_ALIASES = {
+ "sunday": "sun",
+ "monday": "mon",
+ "tuesday": "tue",
+ "wednesday": "wed",
+ "thursday": "thu",
+ "friday": "fri",
+ "saturday": "sat",
+}
+
+
+def _strip_none(data: Any) -> Any:
+ """Recursively drop ``None`` values from dicts/lists for clean payloads."""
+ if isinstance(data, dict):
+ return {k: _strip_none(v) for k, v in data.items() if v is not None}
+ if isinstance(data, list):
+ return [_strip_none(v) for v in data]
+ return data
+
+
+def _normalize_weekdays(value: Any) -> List[str]:
+ """Normalize ``on=`` weekdays to backend codes (``mon``..``sun``)."""
+ if value is None:
+ return []
+ if isinstance(value, str):
+ value = [value]
+ days: List[str] = []
+ for raw in value:
+ key = str(raw).strip().lower()
+ key = _WEEKDAY_ALIASES.get(key, key[:3])
+ if key not in _WEEKDAYS:
+ raise ValueError(f"Invalid weekday {raw!r}; use one of {sorted(_WEEKDAYS)} (or full names).")
+ days.append(key)
+ return days
+
+
+def _normalize_monthdays(value: Any) -> List[int]:
+ """Normalize ``on=`` days-of-month to ints in ``1..31``."""
+ if value is None:
+ return []
+ if isinstance(value, (int, str)):
+ value = [value]
+ days: List[int] = []
+ for raw in value:
+ try:
+ day = int(raw)
+ except (TypeError, ValueError):
+ raise ValueError(f"Invalid day-of-month {raw!r}; expected an integer 1-31.")
+ if not 1 <= day <= 31:
+ raise ValueError(f"Invalid day-of-month {day}; must be between 1 and 31.")
+ days.append(day)
+ return days
+
+
+@dataclass_json
+@dataclass
+class TriggerRepeatRule:
+ """Interval rule for a ``recurring`` schedule (e.g. every 2 hours)."""
+
+ every: Optional[int] = None
+ unit: Optional[str] = None
+
+
+@dataclass_json
+@dataclass
+class TriggerConfiguration:
+ """Structured time-schedule configuration (mirrors the backend config)."""
+
+ type: Optional[str] = None
+ time: Optional[str] = None
+ timezone: Optional[str] = None
+ days_of_week: Optional[List[str]] = field(default=None, metadata=dj_config(field_name="daysOfWeek"))
+ days_of_month: Optional[List[int]] = field(default=None, metadata=dj_config(field_name="daysOfMonth"))
+ run_at: Optional[str] = field(default=None, metadata=dj_config(field_name="runAt"))
+ start_at: Optional[str] = field(default=None, metadata=dj_config(field_name="startAt"))
+ repeat: Optional[TriggerRepeatRule] = None
+
+
+class TriggerSearchParams(BaseSearchParams):
+ """Search parameters for triggers (filter by agent)."""
+
+ pass
+
+
+class TriggerGetParams(BaseGetParams):
+ """Get parameters for triggers."""
+
+ pass
+
+
+class TriggerDeleteParams(BaseDeleteParams):
+ """Delete parameters for triggers."""
+
+ pass
+
+
+@dataclass_json
+@dataclass(repr=False)
+class Trigger(
+ BaseResource,
+ SearchResourceMixin[TriggerSearchParams, "Trigger"],
+ GetResourceMixin[TriggerGetParams, "Trigger"],
+ DeleteResourceMixin[TriggerDeleteParams, DeleteResult],
+):
+ """A schedule/event trigger that fires an agent with a fixed input.
+
+ Time trigger examples::
+
+ aix.Trigger(name="Launch reminder", agent=agent, input="Remind the team.",
+ run_at="2026-01-26T12:00:00Z").save() # once
+ aix.Trigger(name="Daily digest", agent=agent, input="Summarise the news.",
+ every="day", at="09:00", timezone="Europe/London").save() # daily
+ aix.Trigger(name="Hourly check", agent=agent, input="Check the queue.",
+ every="hour", interval=2).save() # every 2 hours
+ aix.Trigger(name="Weekly report", agent=agent, input="Compile the report.",
+ every="week", on=["mon", "thu"], at="17:00").save() # weekly
+ aix.Trigger(name="Invoice run", agent=agent, input="Generate invoices.",
+ every="month", on=[1, 15], at="09:00").save() # monthly
+
+ Event trigger example (requires a connected tool)::
+
+ gmail = aix.Integration.get("composio/gmail")
+ tool = gmail.connect(...)
+ aix.Trigger(name="Triage inbox", agent=agent, input="Triage this email.",
+ event=tool.triggers["NEW_EMAIL"]).save()
+
+ Manage::
+
+ t = aix.Trigger.get("")
+ aix.Trigger.search(agent=agent) # -> Page
+ t.enabled = False; t.save() # disable (re-enable with True)
+ t.delete()
+ """
+
+ RESOURCE_PATH = "v1/triggers"
+
+ # Simple list endpoint: GET /v1/triggers?agentId= (no /paginate suffix)
+ PAGINATE_PATH = ""
+ PAGINATE_METHOD = "get"
+ PAGINATE_ITEMS_KEY = None # bare array response
+
+ # --- Backend-shaped, persisted fields (round-tripped via from_dict/to_dict) ---
+ input: Optional[str] = None
+ asset_id: Optional[str] = field(default=None, metadata=dj_config(field_name="assetId"))
+ asset_type: Optional[str] = field(default="agent", metadata=dj_config(field_name="assetType"))
+ trigger_type: Optional[str] = field(default=None, metadata=dj_config(field_name="triggerType"))
+ trigger_id: Optional[str] = field(default=None, metadata=dj_config(field_name="triggerId"))
+ configuration: Optional[TriggerConfiguration] = None
+ enabled: Optional[bool] = None
+ notifications: Optional[bool] = None
+ retry_count: Optional[int] = field(default=None, metadata=dj_config(field_name="retryCount"))
+
+ # Read-only backend fields (never sent; excluded from payload).
+ next_run_at: Optional[str] = field(default=None, metadata=dj_config(field_name="nextRunAt", exclude=lambda x: True))
+ last_run_at: Optional[str] = field(default=None, metadata=dj_config(field_name="lastRunAt", exclude=lambda x: True))
+ enqueued: Optional[bool] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ created_at: Optional[str] = field(default=None, metadata=dj_config(field_name="createdAt", exclude=lambda x: True))
+ updated_at: Optional[str] = field(default=None, metadata=dj_config(field_name="updatedAt", exclude=lambda x: True))
+
+ # Local-only fields for the event lifecycle (not returned by the REST DTO).
+ connection_id: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ trigger_slug: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+
+ # --- Friendly, init-only kwargs (never serialized; translated in __post_init__) ---
+ agent: Optional[Any] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ run_at: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ every: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ interval: int = field(default=1, metadata=dj_config(exclude=lambda x: True))
+ at: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ on: Optional[Any] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ timezone: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ start_at: Optional[str] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ event: Optional[Any] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ event_config: Optional[Dict[str, Any]] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+ connection: Optional[Any] = field(default=None, metadata=dj_config(exclude=lambda x: True))
+
+ def __post_init__(self) -> None:
+ """Translate friendly construction kwargs into backend-shaped fields.
+
+ Skipped when rehydrating from the backend (``id`` already set), so that
+ ``from_dict`` in get/search/create is left untouched.
+ """
+ if self.id is not None:
+ return # rehydration path — nothing to translate
+
+ # Resolve the target agent -> assetId.
+ if self.agent is not None and not self.asset_id:
+ self.asset_id = getattr(self.agent, "id", None) or self.agent
+ self.asset_type = "agent"
+
+ # New triggers default to enabled so they fire on schedule.
+ if self.enabled is None:
+ self.enabled = True
+ if self.notifications is None:
+ self.notifications = False
+
+ if self.event is not None:
+ self._configure_from_event()
+ elif self.run_at is not None or self.every is not None or self.start_at is not None:
+ self.trigger_type = "time"
+ self.configuration = self._build_configuration()
+
+ # ------------------------------------------------------------------
+ # Event configuration
+ # ------------------------------------------------------------------
+
+ def _configure_from_event(self) -> None:
+ """Capture slug / config / connection from an event option."""
+ self.trigger_type = "external"
+ self.trigger_slug = getattr(self.event, "slug", None) or getattr(self.event, "name", None)
+ if not self.trigger_slug:
+ raise ValueError("event= must be a trigger option from integration.triggers/tool.triggers.")
+ if self.event_config is None:
+ self.event_config = getattr(self.event, "values", None) or {}
+ # Resolve a connection id for activation.
+ if self.connection is not None:
+ self.connection_id = getattr(self.connection, "id", None) or self.connection
+ if not self.connection_id:
+ self.connection_id = getattr(self.event, "connection_id", None)
+
+ # ------------------------------------------------------------------
+ # Schedule mapping (friendly kwargs -> backend configuration)
+ # ------------------------------------------------------------------
+
+ def _build_configuration(self) -> TriggerConfiguration:
+ """Map the schedule kwargs to a backend ``TriggerConfiguration``.
+
+ Uses the *semantic* schedule types (daily/weekly/monthly) whenever a
+ time-of-day (``at``) or day selection (``on``) is given, because the
+ backend scheduler only honours those on the semantic types; the generic
+ ``recurring`` type ignores them and has no ``month`` unit.
+ """
+ tz = self.timezone
+ if self.run_at is not None:
+ return TriggerConfiguration(type="once", run_at=self.run_at, timezone=tz)
+
+ every = (self.every or "").strip().lower()
+ interval = int(self.interval or 1)
+ if not every:
+ raise ValueError("Provide run_at=... for a one-off trigger, or every=... for a recurring one.")
+
+ if every in ("minute", "hour"):
+ return TriggerConfiguration(
+ type="recurring",
+ timezone=tz,
+ start_at=self.start_at,
+ repeat=TriggerRepeatRule(every=interval, unit=every),
+ )
+
+ if every == "day":
+ if self.at is not None and interval == 1:
+ return TriggerConfiguration(type="daily", time=self.at, timezone=tz)
+ if self.at is not None and interval > 1:
+ warnings.warn(
+ "`at` (time-of-day) is ignored for every='day' with interval>1; "
+ "the backend fires every N days relative to creation time.",
+ stacklevel=2,
+ )
+ return TriggerConfiguration(
+ type="recurring",
+ timezone=tz,
+ start_at=self.start_at,
+ repeat=TriggerRepeatRule(every=interval, unit="day"),
+ )
+
+ if every == "week":
+ if interval > 1:
+ warnings.warn(
+ "`on`/`at` are ignored for every='week' with interval>1; "
+ "the backend fires every N weeks relative to creation time.",
+ stacklevel=2,
+ )
+ return TriggerConfiguration(
+ type="recurring",
+ timezone=tz,
+ start_at=self.start_at,
+ repeat=TriggerRepeatRule(every=interval, unit="week"),
+ )
+ days = _normalize_weekdays(self.on)
+ if not days:
+ raise ValueError('Weekly triggers require on=[...] weekdays, e.g. on=["mon", "thu"].')
+ return TriggerConfiguration(type="weekly", days_of_week=days, time=self.at, timezone=tz)
+
+ if every == "month":
+ if interval > 1:
+ raise ValueError("Monthly triggers don't support interval>1 (the backend has no month interval unit).")
+ days = _normalize_monthdays(self.on)
+ if not days:
+ raise ValueError("Monthly triggers require on=[...] day numbers, e.g. on=[1, 15].")
+ return TriggerConfiguration(type="monthly", days_of_month=days, time=self.at, timezone=tz)
+
+ raise ValueError(f"Unknown every={self.every!r}; use one of minute/hour/day/week/month.")
+
+ # ------------------------------------------------------------------
+ # Save
+ # ------------------------------------------------------------------
+
+ def before_save(self, *args: Any, **kwargs: Any) -> None:
+ """Activate the Composio trigger before persisting a new event trigger."""
+ if not self.id and self.trigger_type == "external" and not self.trigger_id:
+ self._activate_event()
+ return None
+
+ def build_save_payload(self, **kwargs: Any) -> Dict[str, Any]:
+ """Build the whitelisted payload for ``POST``/``PUT`` /v1/triggers.
+
+ The backend uses ``forbidNonWhitelisted`` validation, so only the fields
+ accepted by ``TriggerInput`` are sent (never read-only fields).
+ """
+ trigger_type = self.trigger_type or "time"
+ payload: Dict[str, Any] = {
+ "name": self.name,
+ "input": self.input,
+ "assetType": self.asset_type or "agent",
+ "triggerType": trigger_type,
+ "enabled": bool(self.enabled),
+ "notifications": bool(self.notifications),
+ }
+ if self.id:
+ payload["id"] = self.id
+ if self.description is not None:
+ payload["description"] = self.description
+ if self.asset_id:
+ payload["assetId"] = self.asset_id
+ if self.retry_count is not None:
+ payload["retryCount"] = self.retry_count
+
+ if trigger_type == "external":
+ payload["triggerId"] = self.trigger_id
+ elif self.configuration is not None:
+ payload["configuration"] = _strip_none(self.configuration.to_dict())
+ else:
+ raise ValueError(
+ "A time trigger needs a schedule. Pass run_at=... (one-off) or "
+ "every=... (recurring) when constructing the Trigger."
+ )
+
+ return payload
+
+ # ------------------------------------------------------------------
+ # Search (GET /v1/triggers?agentId=)
+ # ------------------------------------------------------------------
+
+ @classmethod
+ def search(cls, agent: Optional[Any] = None, agent_id: Optional[str] = None, **kwargs: Any) -> Page["Trigger"]:
+ """List triggers for the team, optionally filtered by agent.
+
+ Args:
+ agent: An Agent instance (or anything with ``.id``) to filter by.
+ agent_id: An agent id string to filter by.
+ **kwargs: Additional options forwarded to page building (e.g. ``page_number``).
+
+ Returns:
+ Page[Trigger]
+ """
+ context = getattr(cls, "context", None)
+ if context is None:
+ raise ResourceError("Context is required for resource listing")
+
+ resolved_agent_id = agent_id
+ if resolved_agent_id is None and agent is not None:
+ resolved_agent_id = getattr(agent, "id", None) or agent
+
+ request_kwargs: Dict[str, Any] = {}
+ if resolved_agent_id:
+ request_kwargs["params"] = {"agentId": resolved_agent_id}
+
+ response = context.client.get(cls.RESOURCE_PATH, **request_kwargs)
+ kwargs.setdefault("page_number", 0)
+ return cls._build_page(response, context, **kwargs)
+
+ @classmethod
+ def _build_page(cls, response: Any, context: "Aixplain", **kwargs: Any) -> Page["Trigger"]:
+ """Build a Page from the bare-array list response (single page)."""
+ page = super()._build_page(response, context, **kwargs)
+ page.page_total = 1
+ return page
+
+ @classmethod
+ def list(cls, **kwargs: Any) -> List["Trigger"]:
+ """Convenience wrapper returning the results list directly."""
+ return cls.search(**kwargs).results
+
+ # ------------------------------------------------------------------
+ # Delete (deactivate Composio trigger first, then remove the record)
+ # ------------------------------------------------------------------
+
+ def delete(self, *args: Any, **kwargs: Any) -> DeleteResult:
+ """Delete the trigger (deactivating the Composio trigger for events)."""
+ if self.trigger_type == "external":
+ if self.connection_id and self.trigger_id:
+ try:
+ self._deactivate_event()
+ except Exception as e: # best-effort; still remove the record
+ warnings.warn(f"Failed to deactivate the Composio trigger: {e}", stacklevel=2)
+ elif not self.connection_id:
+ warnings.warn(
+ "Deleting the trigger record only; the Composio trigger was not deactivated "
+ "because the connection is unknown (re-fetched event triggers don't carry it).",
+ stacklevel=2,
+ )
+ return super().delete(*args, **kwargs)
+
+ # ------------------------------------------------------------------
+ # Composio activation helpers (via the model-execute endpoint)
+ # ------------------------------------------------------------------
+
+ def _execute_action(self, connection_id: str, action: str, data: Any) -> Any:
+ """POST a connector action to the model-execute endpoint and resolve data."""
+ url = f"{self.context.model_url}/{connection_id}"
+ payload: Dict[str, Any] = {"action": action, "data": data}
+ if action == "activate_trigger":
+ payload["enable"] = True
+ payload["toolkit_versions"] = "latest"
+ response = self.context.client.request("post", url, json=payload)
+ return self._poll_for_data(response)
+
+ @staticmethod
+ def _poll_for_data(response: dict) -> Any:
+ """Return the ``data`` field, polling the URL if the call is async."""
+ import time
+
+ data = response.get("data") if isinstance(response, dict) else None
+ if not isinstance(response, dict):
+ return None
+ if response.get("completed", True) or not isinstance(data, str) or not data.startswith("http"):
+ return data
+ # Not expected for trigger actions, but handle async just in case.
+ return data
+
+ def _activate_event(self) -> None:
+ """Register the Composio trigger and capture its real trigger id."""
+ if not self.connection_id:
+ raise ValueError(
+ "Event triggers require a connection. Create one via integration.connect() and use "
+ "tool.triggers[...], or pass connection= to Trigger(...)."
+ )
+ data = self._execute_action(
+ self.connection_id,
+ "activate_trigger",
+ {"slug": self.trigger_slug, "config": self.event_config or {}},
+ )
+ trigger_id = None
+ if isinstance(data, dict):
+ trigger_id = data.get("trigger_id") or data.get("triggerId")
+ if not trigger_id:
+ raise ResourceError(
+ f"Failed to activate event trigger for slug={self.trigger_slug!r} on connection {self.connection_id}."
+ )
+ self.trigger_id = trigger_id
+
+ def _deactivate_event(self) -> None:
+ """Deactivate the Composio trigger for this event trigger."""
+ self._execute_action(self.connection_id, "delete_trigger", {"trigger_id": self.trigger_id})
+
+ # ------------------------------------------------------------------
+ # Convenience
+ # ------------------------------------------------------------------
+
+ @property
+ def schedule_type(self) -> Optional[str]:
+ """The schedule type (once/daily/weekly/monthly/recurring), if a time trigger."""
+ return self.configuration.type if self.configuration else None
+
+ def __repr__(self) -> str:
+ """Return a concise representation."""
+ return f"Trigger(id={self.id}, name={self.name!r}, type={self.trigger_type})"
diff --git a/docs/README.md b/docs/README.md
index 670f345d..7fe144b0 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,102 +1,74 @@
-
-
-
+
+
+
-
aiXplain SDK
+
aixplain SDK
-
-**Build, deploy, and govern autonomous AI agents for your business operations.**
+**Build, deploy, and run autonomous AI agents — governed by default, in a few lines of Python.**
-aiXplain SDK provides Python and REST APIs for agents that plan, use tools, call models and data, run code, and adapt at runtime. It also works natively with MCP-compatible coding agents and IDEs.
+aixplain is the operating system for autonomous AI: multi-agent orchestration with **runtime governance on every action**, across cloud, on-prem, edge, and local. The full lifecycle — build → evaluate → deploy → monitor → evolve — on one runtime, instead of stitching tools together.
-> **Become an agentic-first organization**
->
-> Designed for business operations: autonomous, governed, MCP-compatible, and built for context management. Your interactive AI assistant is [a click away](https://auth.aixplain.com/).
->
-> _We operate our business with aiXplain agents, using them across product, business development, and marketing._
+**On your terms** — **your data** in your perimeter, **your cost** free on local models and tools, pay as you go in the cloud, **your independence** across any model or infrastructure, no lock-in.
-## Why aiXplain
+Build any agent — knowledge (RAG), data, custom-logic, integration, and team — via **SDK, API, CLI, or MCP**, on a marketplace of **900+ models, tools, and integrations**.
-- **Autonomous runtime loop** — plan, call tools and models, reflect, and continue without fixed flowcharts.
-- **Multi-agent execution** — delegate work to specialized subagents at runtime.
-- **Governance by default** — runtime access and policy enforcement on every run.
-- **Production observability** — inspect step-level traces, tool calls, and outcomes for debugging.
-- **Model and tool portability** — swap assets without rewriting application glue code.
-- **MCP-native access** — connect MCP clients to [900+ aiXplain-hosted assets](#mcp-servers) with one PAYG API key.
-- **Flexible deployment** — run the same agent definition serverless or private.
+## Why aixplain
-| | aiXplain SDK | Other agent frameworks |
-|---|---|---|
-| Governance | Runtime access and policy enforcement built in | Usually custom code or external guardrails |
-| Models and tools | 900+ models and tools with one API key | Provider-by-provider setup |
-| Deployment | Cloud (instant) or on-prem | Usually self-assembled runtime and infra |
-| Observability | Built-in traces and dashboards | Varies by framework |
-| Coding-agent workflows | Works natively with MCP-compatible coding agents and IDEs | Usually not a first-class workflow target |
+Less to build, less to operate:
-## AgenticOS
+- **Deploy with one call** — `agent.save()` promotes an agent to a persistent, versioned endpoint; no Dockerfiles, queues, or autoscaling to manage.
+- **No integration glue** — reach [900+ models, tools, and integrations](#marketplace) through one key; skip per-provider SDKs, auth, and rate-limit handling.
+- **Guardrails you don't have to build** — allow-lists, per-asset permissions, rate and usage limits, and access control enforced at runtime.
+- **Self-debugging** — step-level traces of every plan, tool call, and outcome.
+- **Run it anywhere** — the same definition runs in the cloud, on-prem, at the edge, or locally.
+- **Works with your coding agent** — native [MCP support](#marketplace) for MCP-compatible IDEs and coding agents.
-AgenticOS is the portable runtime platform behind aiXplain agents. AgentEngine orchestrates planning, execution, and delegation for autonomous agents. AssetServing connects agents to models, tools, and data through a governed runtime layer. Observability captures traces, metrics, and monitoring for every production run across Cloud (instant) and on-prem deployments.
+## How it works
-
-
-
+The portable runtime behind aixplain agents: orchestration, governed asset serving, and observability across cloud, on-prem, edge, and local. See the [documentation](https://docs.aixplain.com) for the full architecture.
----
-
-## MCP Server Marketplace
-
-[aiXplain Marketplace](https://studio.aixplain.com/browse) now also exposes MCP servers for **900+ models and tools**, allowing external clients to access selected **tool, integration, and model assets**, for example **Opus 4.6, Kimi, Qwen, Airtable, and Slack**, through **aiXplain-hosted MCP endpoints** with a single API key 🔑.
-
-Read the full MCP setup guide in the [MCP servers docs](https://docs.aixplain.com/api-reference/mcp-servers).
-
-```json
-{
- "ms1": {
- "url": "https://models-mcp.aixplain.com/mcp/",
- "headers": {
- "Authorization": "Bearer ",
- "Accept": "application/json, text/event-stream"
- }
- }
-}
-```
+
+
+
---
## Quick start
+> **This README documents SDK v2, the default API.** SDK v1 (the legacy factory API) keeps working until **August 1, 2026**, after which v2 is the only supported surface.
+
```bash
pip install aixplain
```
-Get your API key from your [aiXplain account](https://console.aixplain.com/settings/keys).
+Get your API key from your aixplain account, then expose it to the SDK:
-
- v2 (default)
+```bash
+export AIXPLAIN_API_KEY=
+```
-### Create and run your first agent (v2)
+### Create and run your first agent
```python
-from uuid import uuid4
from aixplain import Aixplain
-aix = Aixplain(api_key="")
+aix = Aixplain() # reads AIXPLAIN_API_KEY from the environment
search_tool = aix.Tool.get("tavily/tavily-web-search/tavily")
search_tool.allowed_actions = ["search"]
agent = aix.Agent(
- name=f"Research agent {uuid4().hex[:8]}",
+ name="Research agent",
description="Answers questions with concise web-grounded findings.",
instructions="Use the search tool when needed and cite key findings.",
tools=[search_tool],
@@ -109,14 +81,15 @@ result = agent.run(
print(result.data.output)
```
-### Build a multi-agent team (v2)
+> Runs return typed objects — read outputs with `result.data.output`, not dict indexing.
+
+### Build a multi-agent team
```python
-from uuid import uuid4
from aixplain import Aixplain
-from aixplain.v2 import EditorConfig, EvaluatorConfig, EvaluatorType, Inspector, InspectorAction, InspectorActionConfig, InspectorSeverity, InspectorTarget
+from aixplain.v2 import Inspector
-aix = Aixplain(api_key="")
+aix = Aixplain() # reads AIXPLAIN_API_KEY from the environment
search_tool = aix.Tool.get("tavily/tavily-web-search/tavily")
search_tool.allowed_actions = ["search"]
@@ -126,29 +99,25 @@ def never_edit(text: str) -> bool:
def passthrough(text: str) -> str:
return text
+# Config is plain data — strings for action/targets/severity, and a Metric,
+# asset-id string, or callable for the `metric` (the universal judge).
noop_inspector = Inspector(
- name=f"noop-output-inspector-{uuid4().hex[:8]}",
- severity=InspectorSeverity.LOW,
- targets=[InspectorTarget.OUTPUT],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function=never_edit,
- ),
- editor=EditorConfig(
- type=EvaluatorType.FUNCTION,
- function=passthrough,
- ),
+ name="noop-output-inspector",
+ severity="low",
+ targets=["output"],
+ action="edit",
+ metric=never_edit,
+ editor=passthrough,
)
researcher = aix.Agent(
- name=f"Researcher {uuid4().hex[:8]}",
+ name="Researcher",
instructions="Find and summarize reliable sources.",
tools=[search_tool],
)
team_agent = aix.Agent(
- name=f"Research team {uuid4().hex[:8]}",
+ name="Research team",
instructions="Research the topic and return exactly 5 concise bullet points.",
subagents=[researcher],
inspectors=[noop_inspector],
@@ -161,10 +130,6 @@ response = team_agent.run(
print(response.data.output)
```
-
-
-
-
Execution order:
```text
@@ -175,51 +140,41 @@ Team agent
├── Orchestrator: routes work to the right subagent
├── Researcher subagent
│ └── Tavily search tool: finds and summarizes reliable sources
-├── Inspector: checks the final output through a simple runtime policy
-├── Orchestrator: decides whether another pass is needed
-└── Responder: returns one final answer
+├── Inspector: validates the output against a runtime policy
+└── Orchestrator: composes and returns the final answer
```
-
-
- v1 (legacy)
+> **SDK v1 (legacy):** available until August 1, 2026 — see the [SDK v1 docs](https://docs.aixplain.com/1.0/).
-### Create and run your first agent (v1)
+---
-```python
-from aixplain.factories import AgentFactory, ModelFactory
+## Marketplace
-weather_tool = ModelFactory.get("66f83c216eb563266175e201")
+The [aixplain Marketplace](https://studio.aixplain.com/browse) is a catalog of **900+ models, tools, and integrations**. Every asset is reachable through the same three outlets — **SDK, API, and MCP** — with a single API key 🔑.
-agent = AgentFactory.create(
- name="Weather Agent",
- description="Answers weather queries.",
- instructions="Use the weather tool to answer user questions.",
- tools=[weather_tool],
-)
+For MCP-compatible clients and IDEs, assets (for example Opus 4.6, Kimi, Qwen, Airtable, Slack) are served through aixplain-hosted MCP endpoints. See the [MCP servers docs](https://docs.aixplain.com/api-reference/mcp-servers).
-result = agent.run("What is the weather in Liverpool, UK?")
-print(result["data"]["output"])
+```json
+{
+ "ms1": {
+ "url": "https://models-mcp.aixplain.com/mcp/",
+ "headers": {
+ "Authorization": "Bearer ",
+ "Accept": "application/json, text/event-stream"
+ }
+ }
+}
```
-You can still access legacy docs at [docs.aixplain.com/1.0](https://docs.aixplain.com/1.0/).
-
-
-
---
## Data handling and deployment
-aiXplain applies runtime governance and enterprise controls by default:
-
-- **We do not train on your data** — your data is not used to train foundation models.
-- **No data retained by default** — agent memory is opt-in (short-term and long-term).
-- **SOC 2 Type II certified** — enterprise security and compliance posture.
-- **Runtime policy enforcement** — Inspector and Bodyguard govern every agent execution.
-- **Portable deployment options** — Cloud (instant) or on-prem (including VPC and air-gapped environments).
-- **Encryption** — TLS 1.2+ in transit and encrypted storage at rest.
+- **Your data stays yours** — never used to train foundation models; agent memory is opt-in. SOC 2 Type II; TLS 1.2+ in transit, encrypted at rest.
+- **Governed at runtime** — Inspector and Bodyguard enforce allow-lists, per-asset permissions, rate and usage limits, and access control on every execution.
+- **Deploy anywhere** — cloud, on-prem, edge, or local; air-gapped and VPC available on-prem or local.
-Learn more at aiXplain [Security](https://aixplain.com/security/) and aiXplain [pricing](https://aixplain.com/pricing/).
+Learn more at aixplain [Security](https://aixplain.com/security/) and aixplain [pricing](https://aixplain.com/pricing/).
---
@@ -231,7 +186,7 @@ Start free, then scale with usage-based pricing.
- **Subscription plans** — reduce effective consumption-based rates.
- **Custom enterprise pricing** — available for advanced scale and deployment needs.
-Learn more at aiXplain [pricing](https://aixplain.com/pricing/).
+Learn more at aixplain [pricing](https://aixplain.com/pricing/).
---
diff --git a/docs/api-reference/python/aixplain/v2/inspector.md b/docs/api-reference/python/aixplain/v2/inspector.md
index fc899850..aef11414 100644
--- a/docs/api-reference/python/aixplain/v2/inspector.md
+++ b/docs/api-reference/python/aixplain/v2/inspector.md
@@ -8,416 +8,114 @@ Inspector module for v2 API - Team agent inspection and validation.
This module provides inspector functionality for validating team agent operations
at different stages (input, steps, output) with custom policies.
-### InspectorTarget Objects
-
-```python
-class InspectorTarget(str, Enum)
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L25)
-
-Target stages for inspector validation in the team agent pipeline.
-
-#### \_\_str\_\_
-
-```python
-def __str__() -> str
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L32)
-
-Return the string value of the enum.
-
-### InspectorAction Objects
-
-```python
-class InspectorAction(str, Enum)
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L37)
-
-Actions an inspector can take when evaluating content.
-
-### InspectorOnExhaust Objects
-
-```python
-class InspectorOnExhaust(str, Enum)
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L46)
-
-Action to take when max retries are exhausted.
-
-### InspectorSeverity Objects
-
-```python
-class InspectorSeverity(str, Enum)
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L53)
-
-Severity level for inspector findings.
-
-### EvaluatorType Objects
-
-```python
-class EvaluatorType(str, Enum)
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L62)
-
-Type of evaluator or editor.
-
-### InspectorActionConfig Objects
-
-```python
-@dataclass
-class InspectorActionConfig()
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L70)
-
-Inspector action configuration.
-
-#### \_\_post\_init\_\_
-
-```python
-def __post_init__() -> None
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L77)
-
-Validate that max_retries and on_exhaust are only used with RERUN.
-
-#### to\_dict
-
-```python
-def to_dict() -> Dict[str, Any]
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L87)
-
-Convert the action config to a dictionary for API serialization.
-
-#### from\_dict
-
-```python
-@classmethod
-def from_dict(cls, data: Dict[str, Any]) -> "InspectorActionConfig"
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L97)
-
-Create an InspectorActionConfig from a dictionary.
-
-### EvaluatorConfig Objects
-
-```python
-@dataclass
-class EvaluatorConfig()
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L107)
-
-Evaluator configuration for an inspector.
-
-#### \_\_post\_init\_\_
-
-```python
-def __post_init__() -> None
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L115)
-
-Validate and convert callable functions to source strings.
-
-#### to\_dict
-
-```python
-def to_dict() -> Dict[str, Any]
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L124)
-
-Convert to a dictionary for API serialization.
-
-#### from\_dict
-
-```python
-@classmethod
-def from_dict(cls, data: Dict[str, Any]) -> "EvaluatorConfig"
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L136)
-
-Create an EvaluatorConfig from a dictionary.
-
-### EditorConfig Objects
-
-```python
-@dataclass
-class EditorConfig()
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L147)
-
-Editor configuration for an inspector.
-
-#### \_\_post\_init\_\_
-
-```python
-def __post_init__() -> None
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L155)
-
-Validate and convert callable functions to source strings.
-
-#### to\_dict
-
-```python
-def to_dict() -> Dict[str, Any]
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L160)
-
-Convert to a dictionary for API serialization.
-
-#### from\_dict
-
-```python
-@classmethod
-def from_dict(cls, data: Dict[str, Any]) -> "EditorConfig"
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L172)
-
-Create an EditorConfig from a dictionary.
+The public surface is intentionally tiny: construct an `Inspector` with plain
+strings for `action` / `targets` / `severity` and an `aix.Metric` (the universal
+judge) for `metric`. No enums or config classes to import.
### Inspector Objects
```python
-@dataclass
-class Inspector()
+@dataclass(repr=False)
+class Inspector(BaseResource, GetResourceMixin, SearchResourceMixin)
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L183)
+[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py)
Inspector v2 configuration object.
-#### \_\_post\_init\_\_
-
-```python
-def __post_init__() -> None
-```
+An `Inspector` is the single type the `aix.Agent(inspectors=[...])` slot accepts —
+whether it is hand-built or retrieved from the marketplace via `get` / `search`.
+Prebuilt guards are ordinary marketplace assets under the `guardrails` function;
+retrieving one returns a fully-configured `Inspector` whose `metric` points at the
+guard model, so a fetched guard and a custom inspector are indistinguishable to the
+agent.
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L195)
+Configuration is plain data — no enums or config classes to import:
-Validate inspector configuration after initialization.
+- `action`: a string (`"continue" | "rerun" | "abort" | "edit"`) or, when you need
+ retry parameters, a dict `{"type": "rerun", "max_retries": 2, "on_exhaust": "abort"}`.
+- `targets`: a list of strings (`"input" | "steps" | "output"` or a sub-agent name).
+- `severity`: a string (`"low" | "medium" | "high" | "critical"`).
+- `metric`: the universal judge — an `aix.Metric`, an asset-id string, or a Python
+ callable. Required.
+- `editor`: required when `action` is `"edit"`; same accepted types as `metric`.
-#### to\_dict
+Example:
```python
-def to_dict() -> Dict[str, Any]
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L203)
-
-Convert the inspector to a dictionary for API serialization.
-
-#### from\_dict
-
-```python
-@classmethod
-def from_dict(cls, data: Dict[str, Any]) -> "Inspector"
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L220)
-
-Create an Inspector from a dictionary.
-
-### PrebuiltInspector Objects
-
-```python
-@dataclass
-class PrebuiltInspector()
-```
+from aixplain import Aixplain
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L289)
+aix = Aixplain(api_key="")
-A lightweight preset reference that the backend resolves into a full Inspector.
-
-Instead of manually configuring an evaluator, action, and editor, users can
-reference one of the platform's pre-built inspector presets by ID. The
-backend's ``normalize_prebuilt_inspectors`` validator expands the reference
-before the agent graph is constructed.
-
-Example::
-
-from aixplain.v2 import PrebuiltInspector, InspectorTarget
-
-team = client.Agent(
-name="Safe Agent",
-agents=[agent1, agent2],
-inspectors=[
-PrebuiltInspector.prompt_injection_guard(),
-PrebuiltInspector.pii_redaction(targets=[InspectorTarget.OUTPUT]),
-],
+inspector = aix.Inspector(
+ name="grounded_output",
+ severity="high",
+ targets=["output"],
+ action="abort",
+ metric=aix.Metric.create(
+ name="grounded",
+ llm_path="",
+ prompt_template="Abort if the answer is not grounded in the context.",
+ ),
)
-#### \_\_post\_init\_\_
+# Discover and retrieve prebuilt guards like any other asset
+aix.Inspector.search("guard")
+guard = aix.Inspector.get("aws/detect-prompt-attacks-guardrail/aws")
-```python
-def __post_init__() -> None
+team = aix.Agent(name="team", agents=[...], inspectors=[guard, inspector])
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L319)
-
-Validate the inspector configuration after initialization.
-
#### to\_dict
```python
def to_dict() -> Dict[str, Any]
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L335)
-
-Serialize to the lightweight reference format expected by the backend.
+Convert the inspector to a dictionary for API serialization. The judge serializes
+under the backend's long-standing `evaluator` key.
#### from\_dict
```python
@classmethod
-def from_dict(cls, data: Dict[str, Any]) -> "PrebuiltInspector"
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L353)
-
-Create a PrebuiltInspector from a dictionary.
-
-#### prompt\_injection\_guard
-
-```python
-@staticmethod
-def prompt_injection_guard(
- *,
- targets: Optional[List[Any]] = None,
- action: Optional[Dict[str, Any]] = None,
- severity: Optional[InspectorSeverity] = None,
- name: Optional[str] = None,
- description: Optional[str] = None) -> "PrebuiltInspector"
-```
-
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L367)
-
-Create a Prompt Injection Guard inspector.
-
-Detects prompt injection attacks before they influence planning or
-execution. Defaults to ``ABORT`` on ``INPUT``.
-
-**Arguments**:
-
-- `targets` - Override default targets (default: ``[InspectorTarget.INPUT]``).
-- `action` - Override default action dict (default: ``\{"type": "abort"}``).
-- ``0 - Optional severity level.
-- ``1 - Optional custom name for the inspector node.
-- ``2 - Optional custom description.
-
-
-**Returns**:
-
- A PrebuiltInspector configured for prompt injection detection.
-
-#### pii\_redaction
-
-```python
-@staticmethod
-def pii_redaction(*,
- targets: Optional[List[Any]] = None,
- action: Optional[Dict[str, Any]] = None,
- severity: Optional[InspectorSeverity] = None,
- name: Optional[str] = None,
- description: Optional[str] = None) -> "PrebuiltInspector"
+def from_dict(cls, data: Dict[str, Any]) -> "Inspector"
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L400)
-
-Create a PII Redaction inspector.
-
-Finds sensitive information (PII) and returns redacted content from the
-guardrail evaluator. Defaults to ``EDIT`` on ``INPUT``.
-
-**Arguments**:
+Create an Inspector from an inspector-shaped dictionary.
-- `targets` - Override default targets (default: ``[InspectorTarget.INPUT]``).
-- `action` - Override default action dict (default: ``\{"type": "edit"}``).
-- ``0 - Optional severity level.
-- ``1 - Optional custom name for the inspector node.
-- ``2 - Optional custom description.
-
-
-**Returns**:
-
- A PrebuiltInspector configured for PII redaction.
-
-#### hallucination\_guard
+#### from\_guard\_model
```python
-@staticmethod
-def hallucination_guard(
- *,
- targets: Optional[List[Any]] = None,
- action: Optional[Dict[str, Any]] = None,
- severity: Optional[InspectorSeverity] = None,
- name: Optional[str] = None,
- description: Optional[str] = None) -> "PrebuiltInspector"
+@classmethod
+def from_guard_model(cls,
+ payload: Dict[str, Any],
+ requested_path: Optional[str] = None) -> "Inspector"
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L433)
-
-Create a Hallucination Guard inspector.
-
-Detects hallucinations in the agent's final response by verifying it
-against the intermediate step results and the original user query.
-Defaults to ``RERUN`` (max 2 retries, abort on exhaust) on ``OUTPUT``.
+Adapt a `guardrails` marketplace model payload into a configured Inspector. The
+guard model becomes the inspector's `metric` (`asset` judge), and sensible default
+`action` / `targets` are applied based on the guard's canonical path slug. Unknown
+guards fall back to a safe `abort` on `input` default, so future guards need no SDK
+change.
-**Arguments**:
-
-- `targets` - Override default targets (default: ``[InspectorTarget.OUTPUT]``).
-- `action` - Override default action dict (default: ``\{"type": "rerun", "maxRetries": 2, "onExhaust": "abort"}``).
-- ``0 - Optional severity level.
-- ``1 - Optional custom name for the inspector node.
-- ``2 - Optional custom description.
-
-
-**Returns**:
-
- A PrebuiltInspector configured for hallucination detection.
-
-#### list\_presets
+#### get
```python
-@staticmethod
-def list_presets() -> Dict[str, Dict[str, Any]]
+@classmethod
+def get(cls, id: Any, **kwargs: Any) -> "Inspector"
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L467)
-
-Return metadata for all available pre-built inspector presets.
+Retrieve a prebuilt guard by human-readable path (IDs also accepted). Returns a
+fully-configured Inspector backed by the guard model.
-**Returns**:
-
- A dict mapping preset IDs to their metadata (name, category,
- description, default targets/action, supported actions, vendor).
-
-#### is\_prebuilt\_inspector
+#### search
```python
-def is_prebuilt_inspector(obj: Any) -> bool
+@classmethod
+def search(cls,
+ query: Optional[str] = None,
+ **kwargs: Any) -> Page["Inspector"]
```
-[[view_source]](https://github.com/aixplain/aiXplain/blob/main/aixplain/v2/inspector.py#L488)
-
-Return True if *obj* is a PrebuiltInspector or a dict preset reference.
-
+Search available guards, returning the standard paginated shape.
diff --git a/docs/assets/aixplain-agentic-os-architecture.svg b/docs/assets/aixplain-agentic-os-architecture.svg
index ae536e35..337f365c 100644
--- a/docs/assets/aixplain-agentic-os-architecture.svg
+++ b/docs/assets/aixplain-agentic-os-architecture.svg
@@ -1,248 +1,83 @@
-
-
+
\ No newline at end of file
diff --git a/docs/assets/aixplain-logo-dark.svg b/docs/assets/aixplain-logo-dark.svg
new file mode 100644
index 00000000..41ab6457
--- /dev/null
+++ b/docs/assets/aixplain-logo-dark.svg
@@ -0,0 +1,17 @@
+
diff --git a/docs/assets/aixplain-logo-light.svg b/docs/assets/aixplain-logo-light.svg
new file mode 100644
index 00000000..53485000
--- /dev/null
+++ b/docs/assets/aixplain-logo-light.svg
@@ -0,0 +1,17 @@
+
diff --git a/pytest.ini b/pytest.ini
index e618d7a5..83bc576b 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,3 +1,4 @@
[pytest]
+addopts = --import-mode=importlib
testpaths =
tests
diff --git a/skills/aixplain-agent-builder/references/agent-patterns.md b/skills/aixplain-agent-builder/references/agent-patterns.md
index 418b8042..601aba9a 100644
--- a/skills/aixplain-agent-builder/references/agent-patterns.md
+++ b/skills/aixplain-agent-builder/references/agent-patterns.md
@@ -47,31 +47,31 @@ kb_tool.save()
Inspectors attach to **team agents**. Use the **v2** inspector API (`aixplain.v2.inspector`) — it matches the v2 `Aixplain` client used everywhere in this skill. Do **not** import from `aixplain.modules.team_agent.inspector`; that path resolves to the v1 module and its `Inspector(auto=..., policy=...)` shape will fail with `ImportError`/`TypeError` under SDK v0.2.44.
+Config is plain data — no enums or config classes to import. Pass strings for `action` / `targets` / `severity`, and the universal judge (`aix.Metric`, an asset-id string, or a Python callable) for `metric`.
+
```python
-from aixplain.v2.inspector import (
- Inspector, InspectorActionConfig, InspectorAction,
- EvaluatorConfig, EvaluatorType, InspectorTarget, AUTO_DEFAULT_MODEL_ID,
-)
+from aixplain.v2 import Inspector
inspector = Inspector(
name="Content Gate",
- action=InspectorActionConfig(type=InspectorAction.ABORT), # CONTINUE | RERUN | ABORT | EDIT
- evaluator=EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id=AUTO_DEFAULT_MODEL_ID, # default LLM judge; or pass your own model ID
- prompt="Validate output meets policy. Fail if non-compliant.",
+ action="abort", # "continue" | "rerun" | "abort" | "edit"
+ targets=["output"], # "input" | "steps" | "output" (or a sub-agent name)
+ severity="high", # "low" | "medium" | "high" | "critical"
+ metric=aix.Metric.create( # the universal judge
+ name="policy-check",
+ llm_path="", # your judge LLM; wraps the prompt below
+ prompt_template="Validate output meets policy. Fail if non-compliant.",
),
- targets=[InspectorTarget.OUTPUT], # INPUT | STEPS | OUTPUT
)
team.inspectors = [inspector]
-team.inspector_targets = [InspectorTarget.OUTPUT]
+team.inspector_targets = ["output"]
team.save()
```
Notes:
-- `action=RERUN` is the only type that accepts `max_retries` / `on_exhaust` on `InspectorActionConfig`; any other type raises if you set them.
-- `action=EDIT` requires an `editor=EditorConfig(...)`.
-- For a function-based judge instead of an LLM, use `EvaluatorConfig(type=EvaluatorType.FUNCTION, function=)`.
+- `action="rerun"` is the only type that accepts retry params. Pass them via a dict: `action={"type": "rerun", "max_retries": 2, "on_exhaust": "abort"}` (`on_exhaust` is `"continue"` | `"abort"`). Setting retry params on any other action raises.
+- `action="edit"` requires an `editor=...` (same accepted types as `metric`).
+- `metric` accepts an `aix.Metric`, a bare asset-id string, or a Python callable (a custom function judge). For an inline asset + prompt without creating a Metric, pass a dict: `metric={"asset_id": "", "prompt": "..."}`.
After adding inspectors, validate with 3 prompts: allowed, denied, ambiguous. See `references/inspectors.md` for the full validation matrix.
diff --git a/skills/aixplain-agent-builder/references/inspectors.md b/skills/aixplain-agent-builder/references/inspectors.md
index c4ef359c..8837b828 100644
--- a/skills/aixplain-agent-builder/references/inspectors.md
+++ b/skills/aixplain-agent-builder/references/inspectors.md
@@ -9,10 +9,11 @@ See `agent-patterns.md § Add Inspector` for the full v2 construction code. Use
## Policy Guidelines
-- Set the policy via `InspectorActionConfig(type=...)`. Valid `InspectorAction` values: `CONTINUE | RERUN | ABORT | EDIT`.
-- Prefer `ABORT` for hard policy violations.
-- Use `RERUN` for recoverable quality issues — it is the only action type that accepts `max_retries` and `on_exhaust` (`CONTINUE | ABORT`); setting them on any other type raises.
-- `EDIT` requires an `editor=EditorConfig(...)`.
+- Set the policy with a plain string: `action="..."`. Valid values: `"continue" | "rerun" | "abort" | "edit"`.
+- Prefer `"abort"` for hard policy violations.
+- Use `"rerun"` for recoverable quality issues — it is the only action that accepts retry params, passed as a dict: `action={"type": "rerun", "max_retries": 2, "on_exhaust": "abort"}` (`on_exhaust` is `"continue" | "abort"`). Setting retry params on any other action raises.
+- `"edit"` requires an `editor=...` (a Metric, asset-id string, or callable — same as `metric`).
+- The judge is `metric=` (the universal judge): an `aix.Metric`, a bare asset-id string, or a Python callable. No `EvaluatorConfig`/enum imports.
- For existing deployments: fetch the team agent, modify `inspectors`/`inspector_targets`, call `.save()`.
---
@@ -26,7 +27,7 @@ Do not overload `FAILED` for policy blocks — a policy block can still return r
## Inspector Action Values
-`InspectorAction` SDK values: `continue | rerun | abort | edit`.
+`action` SDK values (plain strings): `continue | rerun | abort | edit`.
---
diff --git a/tests/functional/agent/agent_functional_test.py b/tests/functional/agent/agent_functional_test.py
index d0b27df6..32f52519 100644
--- a/tests/functional/agent/agent_functional_test.py
+++ b/tests/functional/agent/agent_functional_test.py
@@ -804,7 +804,7 @@ def test_agent_with_action_tool(slack_token, resource_tracker):
connection = ModelFactory.get(connection_id)
connection.action_scope = [
- action for action in connection.actions if action.code == "SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"
+ action for action in connection.actions if action.code == "SLACK_SEND_MESSAGE"
]
agent_name = f"TA {str(uuid4())[:8]}"
@@ -827,7 +827,7 @@ def test_agent_with_action_tool(slack_token, resource_tracker):
assert response is not None
assert response["status"].lower() == "success"
assert "helsinki" in response.data.output.lower()
- assert "SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL" in [
+ assert "SLACK_SEND_MESSAGE" in [
step["tool"] for step in response.data.intermediate_steps[0]["tool_steps"]
]
@@ -837,7 +837,9 @@ def test_agent_with_action_tool(slack_token, resource_tracker):
FIRECRAWL_CONNECTION_ASSET_ID = "69442021f2e6cb73e286ff0f"
TAVILY_CONNECTION_ASSET_ID = "6931bdf462eb386b7158def3"
GOOGLE_SEARCH_UTILITY_MODEL_ID = "65c51c556eb563350f6e1bb1"
-WEB_SEARCH_TOOL_ASSET_ID = "69fb7750f177c224105dabc6"
+# Google Search API (scale-serp/google-search/Google); the Firecrawl-based
+# Web Search Tool (69fb7750f177c224105dabc6) fails on every invocation.
+WEB_SEARCH_TOOL_ASSET_ID = "692f18557b2cc45d29150cb0"
@pytest.mark.flaky(reruns=2, reruns_delay=2)
diff --git a/tests/functional/agent/agent_mcp_deploy_test.py b/tests/functional/agent/agent_mcp_deploy_test.py
index 2a2248ad..17ae9b69 100644
--- a/tests/functional/agent/agent_mcp_deploy_test.py
+++ b/tests/functional/agent/agent_mcp_deploy_test.py
@@ -21,18 +21,20 @@
@pytest.fixture
def mcp_tool():
"""Create an MCP tool for testing."""
+ # DeepWiki is a public streamable-HTTP MCP server (no auth); the previous
+ # endpoint remote.mcpservers.org no longer resolves (NXDOMAIN).
tool = ToolFactory.create(
integration="686eb9cd26480723d0634d3e", # Remote MCP ID
name=f"Test Remote MCP {uuid4()}",
authentication_schema=AuthenticationSchema.API_KEY,
- data={"url": "https://remote.mcpservers.org/fetch/mcp"},
+ data={"url": "https://mcp.deepwiki.com/mcp"},
)
# Set allowed actions (using ... as in original script)
tool.allowed_actions = [...]
- # Filter actions to only include "fetch" action
- tool.action_scope = [action for action in tool.actions if action.code == "fetch"]
+ # Filter actions to only include the "ask_question" action
+ tool.action_scope = [action for action in tool.actions if action.code == "ask_question"]
yield tool
try:
@@ -46,8 +48,8 @@ def test_agent(mcp_tool):
"""Create a test agent with MCP tool."""
agent = AgentFactory.create(
name=f"Test Agent {uuid4()}",
- description="This agent is used to scrape websites",
- instructions="You are a helpful assistant that can scrape any given website",
+ description="This agent answers questions about GitHub repositories",
+ instructions="You are a helpful assistant that answers questions about GitHub repositories using DeepWiki",
tools=[mcp_tool],
llm="69b7e5f1b2fe44704ab0e7d0",
)
@@ -62,7 +64,7 @@ def test_agent_creation_with_mcp_tool(test_agent, mcp_tool):
"""Test that an agent can be created with an MCP tool."""
assert test_agent is not None
assert test_agent.name.startswith("Test Agent")
- assert test_agent.description == "This agent is used to scrape websites"
+ assert test_agent.description == "This agent answers questions about GitHub repositories"
assert len(test_agent.tools) == 1
assert test_agent.tools[0] == mcp_tool
assert test_agent.status == AssetStatus.DRAFT
@@ -70,7 +72,7 @@ def test_agent_creation_with_mcp_tool(test_agent, mcp_tool):
def test_agent_run_before_deployment(test_agent):
"""Test that an agent can run before being deployed."""
- response = test_agent.run("Give me information about the aixplain website")
+ response = test_agent.run("Give me information about the aixplain/aiXplain GitHub repository")
assert response is not None
assert hasattr(response, "data")
@@ -120,7 +122,7 @@ def test_deployed_agent_can_run(test_agent):
test_agent.deploy()
# Run a query on the deployed agent
- response = test_agent.run("Give me information about the aixplain website")
+ response = test_agent.run("Give me information about the aixplain/aiXplain GitHub repository")
assert response is not None
assert hasattr(response, "data")
@@ -133,7 +135,7 @@ def test_agent_lifecycle_end_to_end(mcp_tool):
agent = AgentFactory.create(
name=f"Test Agent Lifecycle {uuid4()}",
description="This agent is used for lifecycle testing",
- instructions="You are a helpful assistant that can scrape any given website",
+ instructions="You are a helpful assistant that answers questions about GitHub repositories using DeepWiki",
tools=[mcp_tool],
llm="69b7e5f1b2fe44704ab0e7d0",
)
@@ -143,7 +145,7 @@ def test_agent_lifecycle_end_to_end(mcp_tool):
assert agent.status == AssetStatus.DRAFT
# Test run before deployment
- response = agent.run("Give me information about the aixplain website")
+ response = agent.run("Give me information about the aixplain/aiXplain GitHub repository")
assert response is not None
# Deploy agent
@@ -156,7 +158,7 @@ def test_agent_lifecycle_end_to_end(mcp_tool):
assert retrieved_agent.status == AssetStatus.ONBOARDED
# Test run after deployment
- response = retrieved_agent.run("Give me information about the aixplain website")
+ response = retrieved_agent.run("Give me information about the aixplain/aiXplain GitHub repository")
assert response is not None
finally:
try:
@@ -173,6 +175,6 @@ def test_mcp_tool_properties(mcp_tool):
assert hasattr(mcp_tool, "action_scope")
assert len(mcp_tool.action_scope) >= 0 # Should have filtered actions
- # Verify that action_scope only contains "fetch" actions
+ # Verify that action_scope only contains "ask_question" actions
for action in mcp_tool.action_scope:
- assert action.code == "fetch"
+ assert action.code == "ask_question"
diff --git a/tests/functional/model/run_connect_model_test.py b/tests/functional/model/run_connect_model_test.py
index 214cc983..136cd02a 100644
--- a/tests/functional/model/run_connect_model_test.py
+++ b/tests/functional/model/run_connect_model_test.py
@@ -1,4 +1,8 @@
import os
+from uuid import uuid4
+
+import pytest
+
from aixplain.enums import ResponseStatus
from aixplain.factories import ModelFactory
from aixplain.modules.model.integration import Integration, AuthenticationSchema
@@ -37,6 +41,11 @@ def test_run_connect_model():
connection.delete()
+@pytest.mark.skip(
+ reason="The hardcoded Zapier MCP share URL (token committed in plaintext) is "
+ "dead — every POST fails with 'Streamable HTTP error'. Re-enable with a "
+ "fresh, secret-managed endpoint."
+)
def test_run_mcp_connect_model():
# get slack connector
connector = ModelFactory.get("686eb9cd26480723d0634d3e")
@@ -104,13 +113,17 @@ def test_run_script_connection_tool():
def test_function():
return "Hello, world!"
+ # Unique name + finally-cleanup: a fixed name poisons every later run
+ # with "Name already exists" if one run dies before delete().
tool = ModelFactory.create_script_connection_tool(
- name="My Test Tool", code=test_function, function_name="test_function"
+ name=f"My Test Tool {uuid4().hex[:8]}", code=test_function, function_name="test_function"
)
- response = tool.run(inputs={}, action=tool.actions[0])
- assert response.status == ResponseStatus.SUCCESS
- assert response.data["data"] == "Hello, world!"
- tool.delete()
+ try:
+ response = tool.run(inputs={}, action=tool.actions[0])
+ assert response.status == ResponseStatus.SUCCESS
+ assert response.data["data"] == "Hello, world!"
+ finally:
+ tool.delete()
def test_run_script_connection_tool_with_complex_inputs():
@@ -118,13 +131,15 @@ def test_all_types(s: str, i: int, f: float, lst: list, d: dict):
return f"String: {s}\nInt: {i}\nFloat: {f}\nList: {lst}\nDict: {d}"
tool = ModelFactory.create_script_connection_tool(
- name="My Test Tool",
+ name=f"My Test Tool {uuid4().hex[:8]}",
code=test_all_types,
)
- response = tool.run(
- inputs={"s": "test", "i": 1, "f": 1.0, "lst": [1, 2, 3], "d": {"a": 1, "b": 2}},
- action=tool.actions[0],
- )
- assert response.status == ResponseStatus.SUCCESS
- assert response.data["data"] == "String: test\nInt: 1\nFloat: 1\nList: [1, 2, 3]\nDict: {'a': 1, 'b': 2}"
- tool.delete()
+ try:
+ response = tool.run(
+ inputs={"s": "test", "i": 1, "f": 1.0, "lst": [1, 2, 3], "d": {"a": 1, "b": 2}},
+ action=tool.actions[0],
+ )
+ assert response.status == ResponseStatus.SUCCESS
+ assert response.data["data"] == "String: test\nInt: 1\nFloat: 1\nList: [1, 2, 3]\nDict: {'a': 1, 'b': 2}"
+ finally:
+ tool.delete()
diff --git a/tests/functional/model/run_model_test.py b/tests/functional/model/run_model_test.py
index 6f99d123..d0ccd2b6 100644
--- a/tests/functional/model/run_model_test.py
+++ b/tests/functional/model/run_model_test.py
@@ -121,7 +121,12 @@ def run_index_model(index_model, retries):
@pytest.mark.parametrize(
"embedding_model,supplier_params",
[
- pytest.param(None, VectaraParams, id="VECTARA"),
+ pytest.param(
+ None,
+ VectaraParams,
+ id="VECTARA",
+ marks=pytest.mark.skip(reason="Flaky: Vectara integration on test env rejects the configured customer id"),
+ ),
pytest.param(None, ZeroEntropyParams, id="ZERO_ENTROPY"),
pytest.param(EmbeddingModel.OPENAI_ADA002, AirParams, id="AIR - OpenAI Ada 002"),
pytest.param(
diff --git a/tests/functional/team_agent/team_agent_functional_test.py b/tests/functional/team_agent/team_agent_functional_test.py
index 324c4113..ec09d8ca 100644
--- a/tests/functional/team_agent/team_agent_functional_test.py
+++ b/tests/functional/team_agent/team_agent_functional_test.py
@@ -255,6 +255,7 @@ def test_add_remove_agents_from_team_agent(run_input_map, resource_tracker, Team
assert len(team_agent.agents) == len(agents)
+@pytest.mark.flaky(reruns=2, reruns_delay=5)
def test_team_agent_tasks(resource_tracker):
agent_name = f"TSA {str(uuid4())[:8]}"
agent = AgentFactory.create(
@@ -558,7 +559,7 @@ def test_team_agent_with_slack_connector(resource_tracker):
connection = ModelFactory.get(connection_id)
connection.action_scope = [
- action for action in connection.actions if action.code == "SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"
+ action for action in connection.actions if action.code == "SLACK_SEND_MESSAGE"
]
agent_name = f"TA {str(uuid4())[:8]}"
diff --git a/tests/functional/team_agent/test_utils.py b/tests/functional/team_agent/test_utils.py
index 96c24328..a3181611 100644
--- a/tests/functional/team_agent/test_utils.py
+++ b/tests/functional/team_agent/test_utils.py
@@ -67,11 +67,3 @@ def create_team_agent(factory, agents, run_input_map, use_mentalist=True):
)
return team_agent
-
-
-def verify_response_generator(steps: Dict) -> None:
- """Helper function to verify response generator step"""
- response_generator_steps = [step for step in steps if "response_generator" in step.get("agent", "").lower()]
- assert len(response_generator_steps) == 1, (
- f"Expected exactly one response_generator step, found {len(response_generator_steps)}"
- )
diff --git a/tests/functional/v2/inspector_functional_test.py b/tests/functional/v2/inspector_functional_test.py
index 6c11c112..0c117858 100644
--- a/tests/functional/v2/inspector_functional_test.py
+++ b/tests/functional/v2/inspector_functional_test.py
@@ -14,21 +14,10 @@
import aixplain as aix
from aixplain.enums.asset_status import AssetStatus
-from aixplain.v2 import (
- Inspector,
- InspectorTarget,
- InspectorAction,
- InspectorOnExhaust,
- InspectorSeverity,
- InspectorActionConfig,
- EvaluatorType,
- EvaluatorConfig,
- EditorConfig,
-)
+from aixplain.v2 import Inspector
from tests.functional.team_agent.test_utils import (
RUN_FILE,
read_data,
- verify_response_generator,
)
_DEFAULT_INPUT_TARGET = "input"
@@ -97,14 +86,19 @@ def _step_agent_id(step: Dict) -> str:
def _is_inspector_step(step: Dict) -> bool:
- """True if step is an inspector (id is 'inspector' or 'inspector|...')."""
- return _step_agent_id(step).startswith("inspector")
+ """True if step is an inspector.
+
+ The backend now reports the inspector's own name as the step agent id
+ (e.g. 'abort_output_inspector', is_system_agent=True) instead of the old
+ 'inspector' / 'inspector|...' form, so match on substring.
+ """
+ return "inspector" in _step_agent_id(step)
def verify_inspector_steps(
steps: List[Dict],
inspector_names: List[str],
- inspector_targets: List[InspectorTarget],
+ inspector_targets: List[str],
) -> None:
def agent_id(step: Dict) -> str:
a = step.get("agent") or {}
@@ -114,22 +108,11 @@ def agent_name(step: Dict) -> str:
a = step.get("agent") or {}
return (a.get("name") or "").lower()
- rg_indices = [i for i, s in enumerate(steps) if agent_id(s) == "response_generator"]
- assert len(rg_indices) == 1, f"Expected exactly one response_generator step, got {len(rg_indices)}"
- rg_idx = rg_indices[0]
-
+ # The backend no longer emits a response_generator step; inspector steps
+ # are asserted directly wherever they appear in the run.
inspector_indices = [i for i, s in enumerate(steps) if _is_inspector_step(s)]
assert inspector_indices, "Expected at least one inspector step"
- if InspectorTarget.OUTPUT in inspector_targets:
- after = [i for i in inspector_indices if i > rg_idx]
- assert after, "Expected inspector steps after response_generator for OUTPUT target"
-
- last_steps = steps[rg_idx + 1 :]
- assert all(_is_inspector_step(s) for s in last_steps), (
- "Not all steps after response_generator are inspector steps"
- )
-
expected_n = len(inspector_names)
actual_n = len(inspector_indices)
@@ -140,7 +123,7 @@ def agent_name(step: Dict) -> str:
def _run_and_get_steps(team_agent, query: str):
- response = team_agent.run(query)
+ response = team_agent.run(query=query)
assert response is not None
@@ -171,14 +154,13 @@ def test_output_inspector_abort(client, run_input_map, resource_tracker):
inspector = Inspector(
name="always_abort_output_inspector",
- severity=InspectorSeverity.HIGH,
+ severity="high",
targets=[_DEFAULT_OUTPUT_TARGET],
- action=InspectorActionConfig(type=InspectorAction.ABORT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id=run_input_map["llm_id"],
- prompt="ALWAYS abort if the output is in English",
- ),
+ action="abort",
+ metric={
+ "asset_id": run_input_map["llm_id"],
+ "prompt": "ALWAYS abort if the output is in English",
+ },
)
team_agent = _make_team_agent(client, timestamp, agents, [inspector])
@@ -187,16 +169,8 @@ def test_output_inspector_abort(client, run_input_map, resource_tracker):
_, steps = _run_and_get_steps(team_agent, "What's the biggest city in the world?")
- response_generator_steps = [
- s for s in steps if (s.get("agent") or {}).get("id", "").lower() == "response_generator"
- ]
- assert len(response_generator_steps) == 1, (
- f"Expected exactly one response_generator step, got {len(response_generator_steps)}"
- )
- response_generator_index = steps.index(response_generator_steps[0])
-
- inspector_steps = [s for s in steps[response_generator_index + 1 :] if _is_inspector_step(s)]
- assert len(inspector_steps) > 0, "Expected inspector step(s) after response_generator"
+ inspector_steps = [s for s in steps if _is_inspector_step(s)]
+ assert len(inspector_steps) > 0, "Expected inspector step(s) in the run"
assert (inspector_steps[-1].get("action") or "").lower() == "abort", (
f"Expected abort, got {inspector_steps[-1].get('action')}"
@@ -212,18 +186,13 @@ def test_output_inspector_rerun_until_fixed(client, run_input_map, resource_trac
inspector = Inspector(
name="rerun_output_inspector",
- severity=InspectorSeverity.LOW,
+ severity="low",
targets=[_DEFAULT_OUTPUT_TARGET],
- action=InspectorActionConfig(
- type=InspectorAction.RERUN,
- max_retries=2,
- on_exhaust=InspectorOnExhaust.ABORT,
- ),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id=run_input_map["llm_id"],
- prompt="If the output does NOT include the name of the customer (John), instruct to add it.",
- ),
+ action={"type": "rerun", "max_retries": 2, "on_exhaust": "abort"},
+ metric={
+ "asset_id": run_input_map["llm_id"],
+ "prompt": "If the output does NOT include the name of the customer (John), instruct to add it.",
+ },
)
team_agent = _make_team_agent(client, timestamp, agents, [inspector])
@@ -234,12 +203,8 @@ def test_output_inspector_rerun_until_fixed(client, run_input_map, resource_trac
assert "John" in (getattr(response.data, "output", "") or "")
- rg_steps = [s for s in steps if (s.get("agent") or {}).get("id", "").lower() == "response_generator"]
- assert len(rg_steps) == 2
- rg_idx = steps.index(rg_steps[0])
-
- inspector_steps = [s for s in steps[rg_idx + 1 :] if _is_inspector_step(s)]
- assert inspector_steps, "Expected inspector steps after response_generator"
+ inspector_steps = [s for s in steps if _is_inspector_step(s)]
+ assert inspector_steps, "Expected inspector steps in the run"
assert any((s.get("action") or "").lower() == "rerun" for s in inspector_steps), (
f"Expected at least one rerun action, got actions: {[s.get('action') for s in inspector_steps]}"
@@ -254,17 +219,11 @@ def test_edit_steps_always_runs(client, run_input_map, resource_tracker):
inspector = Inspector(
name="edit_steps_inspector",
- severity=InspectorSeverity.MEDIUM,
- targets=[InspectorTarget.STEPS],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function="def evaluator_fn(text: str) -> bool:\n return True",
- ),
- editor=EditorConfig(
- type=EvaluatorType.FUNCTION,
- function='def edit_fn(text: str) -> str:\n return "hello, what\'s the weather in paris like today?"',
- ),
+ severity="medium",
+ targets=["steps"],
+ action="edit",
+ metric={"function": "def evaluator_fn(text: str) -> bool:\n return True"},
+ editor={"function": 'def edit_fn(text: str) -> str:\n return "hello, what\'s the weather in paris like today?"'},
)
team_agent = _make_team_agent(client, timestamp, agents, [inspector])
@@ -298,17 +257,11 @@ def test_edit_with_gate_true(client, run_input_map, resource_tracker):
inspector = Inspector(
name="gated_edit_true",
- severity=InspectorSeverity.MEDIUM,
- targets=[InspectorTarget.INPUT],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function=evaluator_fn,
- ),
- editor=EditorConfig(
- type=EvaluatorType.FUNCTION,
- function=edit_fn,
- ),
+ severity="medium",
+ targets=["input"],
+ action="edit",
+ metric=evaluator_fn,
+ editor=edit_fn,
)
team_agent = _make_team_agent(client, timestamp, agents, [inspector])
@@ -337,17 +290,11 @@ def test_edit_with_gate_false(client, run_input_map, resource_tracker):
inspector = Inspector(
name="gated_edit_false",
- severity=InspectorSeverity.MEDIUM,
- targets=[InspectorTarget.INPUT],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function=evaluator_fn,
- ),
- editor=EditorConfig(
- type=EvaluatorType.FUNCTION,
- function=edit_fn,
- ),
+ severity="medium",
+ targets=["input"],
+ action="edit",
+ metric=evaluator_fn,
+ editor=edit_fn,
)
team_agent = _make_team_agent(client, timestamp, agents, [inspector])
@@ -358,3 +305,111 @@ def test_edit_with_gate_false(client, run_input_map, resource_tracker):
out = (getattr(response.data, "output", "") or "").lower()
assert "paris" not in out
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
+def test_inspector_search_returns_page_of_inspectors(client):
+ """aix.Inspector.search discovers guards like any other marketplace asset."""
+ page = client.Inspector.search("guard")
+
+ # Standard paginated shape.
+ assert hasattr(page, "results")
+ assert isinstance(page.page_number, int)
+ assert isinstance(page.page_total, int)
+ assert isinstance(page.total, int)
+
+ if not page.results:
+ pytest.skip("No onboarded guardrail models available in this environment")
+
+ for guard in page.results:
+ assert isinstance(guard, Inspector)
+ # A discovered guard is a ready-to-use, fully-configured inspector.
+ assert guard.metric is not None
+ assert guard.action is not None
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
+def test_inspector_get_returns_configured_inspector(client):
+ """aix.Inspector.get(path_or_id) returns a configured, agent-ready Inspector."""
+ page = client.Inspector.search("guard")
+ if not page.results:
+ pytest.skip("No onboarded guardrail models available in this environment")
+
+ # Retrieve the same guard by its id/path; a fetched guard and a hand-built
+ # Inspector are the same type, so this slots directly into inspectors=[...].
+ first = page.results[0]
+ fetched = client.Inspector.get(first.path or first.id)
+
+ assert isinstance(fetched, Inspector)
+ assert fetched.metric is not None
+ assert fetched.metric.asset_id == first.metric.asset_id
+
+
+# Canonical marketplace paths for the onboarded AWS guards and their tuned config.
+_PREBUILT_GUARDS = [
+ ("aws/detect-prompt-attacks-guardrail/aws", "abort", [_DEFAULT_INPUT_TARGET]),
+ ("aws/sensitive-information-guardrail/aws", "edit", [_DEFAULT_INPUT_TARGET]),
+ ("aws/contextual-grounding-check-guardrail/aws", "rerun", [_DEFAULT_OUTPUT_TARGET]),
+]
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
+@pytest.mark.parametrize("path,expected_action,expected_targets", _PREBUILT_GUARDS)
+def test_get_prebuilt_guard_by_canonical_path(client, path, expected_action, expected_targets):
+ """aix.Inspector.get() returns the guard with its tuned config.
+
+ The asset-name (middle) segment of the path selects action/targets, so the
+ PII guard resolves to edit (not the safe abort/input fallback).
+ """
+ try:
+ guard = client.Inspector.get(path)
+ except Exception as e:
+ pytest.skip(f"Guard '{path}' not onboarded in this environment: {e}")
+
+ assert isinstance(guard, Inspector)
+ assert guard.path == path
+ # The guard model itself is the judge.
+ assert guard.metric is not None
+ assert guard.metric.asset_id == guard.id
+ assert guard.action.type == expected_action
+ assert guard.targets == expected_targets
+ if expected_action == "edit":
+ # EDIT guards redact via the guard model, so an editor is configured.
+ assert guard.editor is not None
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
+def test_prebuilt_guard_attaches_and_runs_in_team_agent(client, resource_tracker):
+ """A fetched prebuilt guard saves and executes as an input-stage inspector.
+
+ Covers the full path: get a guard by canonical path, attach it to a team
+ agent via inspectors=[...], save (the guard persists as an ordinary
+ inspector), and run — verifying the guard runs as an inspector step.
+ """
+ path = "aws/detect-prompt-attacks-guardrail/aws"
+ try:
+ guard = client.Inspector.get(path)
+ except Exception as e:
+ pytest.skip(f"Guard '{path}' not onboarded in this environment: {e}")
+
+ timestamp = f"{int(time.time())}_{uuid.uuid4().hex[:6]}"
+ agents = _make_two_subagents(client, timestamp)
+ for agent in agents:
+ resource_tracker.append(agent)
+
+ # _make_team_agent saves the team; the fetched guard round-trips as an Inspector.
+ team_agent = _make_team_agent(client, timestamp, agents, [guard])
+ resource_tracker.append(team_agent)
+
+ assert len(team_agent.inspectors) == 1
+ assert isinstance(team_agent.inspectors[0], Inspector)
+
+ _, steps = _run_and_get_steps(team_agent, "What is the capital of France?")
+
+ # The guard runs as an inspector step. The backend may label the step with
+ # the guard's name or the generic 'inspector' id, so accept either.
+ guard_steps = [s for s in steps if _is_inspector_step(s) or _step_agent_id(s) == guard.name.lower()]
+ assert guard_steps, (
+ f"Expected the prebuilt guard to run as an inspector step; "
+ f"got step ids {[_step_agent_id(s) for s in steps]}"
+ )
diff --git a/tests/functional/v2/test_agent.py b/tests/functional/v2/test_agent.py
index 3e1b6b4a..9ffee2a1 100644
--- a/tests/functional/v2/test_agent.py
+++ b/tests/functional/v2/test_agent.py
@@ -239,43 +239,6 @@ def _get_steps(response):
return steps
-def _has_response_generator(steps):
- """Return True if any step has agent id 'response_generator'."""
- return any(((s.get("agent") or {}).get("id") or "").lower() == "response_generator" for s in steps)
-
-
-def test_run_response_generation(client, test_agent):
- """Test that runResponseGeneration controls presence of response_generator step."""
- agent = client.Agent.get(test_agent.id)
-
- # Default (False): response_generator step should NOT be present
- response_default = agent.run("Say hello")
- assert response_default.status == "SUCCESS", f"Default run failed: {response_default.status}"
- steps_default = _get_steps(response_default)
- assert not _has_response_generator(steps_default), (
- f"Expected no response_generator step with default (False), got agent ids: "
- f"{[((s.get('agent') or {}).get('id') or '') for s in steps_default]}"
- )
-
- # Explicit True: response_generator step should be present
- response_true = agent.run("Say hello", run_response_generation=True)
- assert response_true.status == "SUCCESS", f"run_response_generation=True run failed: {response_true.status}"
- steps_true = _get_steps(response_true)
- assert _has_response_generator(steps_true), (
- f"Expected response_generator step with run_response_generation=True, got agent ids: "
- f"{[((s.get('agent') or {}).get('id') or '') for s in steps_true]}"
- )
-
- # False: response_generator step should NOT be present
- response_false = agent.run("Say hello", run_response_generation=False)
- assert response_false.status == "SUCCESS", f"run_response_generation=False run failed: {response_false.status}"
- steps_false = _get_steps(response_false)
- assert not _has_response_generator(steps_false), (
- f"Expected no response_generator step with run_response_generation=False, got agent ids: "
- f"{[((s.get('agent') or {}).get('id') or '') for s in steps_false]}"
- )
-
-
def test_agent_creation_and_deletion(client):
"""Test agent creation and deletion workflow."""
# Create a test agent
@@ -364,13 +327,13 @@ def test_slack_tool_integration_with_agent(client, slack_token):
name=f"test-slack-tool-{int(time.time())}",
integration=integration,
config={"token": slack_token},
- allowed_actions=["SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"],
+ allowed_actions=["SLACK_SEND_MESSAGE"],
)
# Validate tool creation
assert slack_tool.name.startswith("test-slack-tool-")
assert slack_tool.integration.id == "686432941223092cb4294d3f"
- assert "SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL" in slack_tool.allowed_actions
+ assert "SLACK_SEND_MESSAGE" in slack_tool.allowed_actions
# Save tool before running
slack_tool.save()
@@ -378,7 +341,7 @@ def test_slack_tool_integration_with_agent(client, slack_token):
# Test tool execution
test_message = "Hello from aixplain functional test!"
tool_response = slack_tool.run(
- action="SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL",
+ action="SLACK_SEND_MESSAGE",
data={"channel": "#integrations-test", "text": test_message},
)
@@ -470,7 +433,7 @@ def test_slack_tool_integration_with_agent(client, slack_token):
FIRECRAWL_CONNECTION_ASSET_ID = "69442021f2e6cb73e286ff0f"
TAVILY_CONNECTION_ASSET_ID = "6931bdf462eb386b7158def3"
GOOGLE_SEARCH_UTILITY_MODEL_ID = "65c51c556eb563350f6e1bb1"
-WEB_SEARCH_TOOL_ASSET_ID = "69fb7750f177c224105dabc6"
+WEB_SEARCH_TOOL_PATH = "scale-serp/google-search/Google"
def _get_output_text(response) -> str:
@@ -622,7 +585,7 @@ def test_agent_web_search_tool(client):
2. Agent execution succeeds.
3. The response contains information retrieved from the search.
"""
- tool = client.Tool.get(WEB_SEARCH_TOOL_ASSET_ID)
+ tool = client.Tool.get(WEB_SEARCH_TOOL_PATH)
response = _run_platform_tool_agent(
client=client,
diff --git a/tests/functional/v2/test_agent_llm_persistence.py b/tests/functional/v2/test_agent_llm_persistence.py
new file mode 100644
index 00000000..4dd540a1
--- /dev/null
+++ b/tests/functional/v2/test_agent_llm_persistence.py
@@ -0,0 +1,96 @@
+"""Regression tests for the LLM-overwrite bug (0.2.43/0.2.44).
+
+On those releases ``Agent.get()`` never decoded the ``model`` field from the
+API response, so a fetched agent's ``llm`` silently fell back to
+``Agent.DEFAULT_LLM`` and the next ``save()`` — even one that only touched
+``instructions`` — overwrote the agent's configured LLM with the default.
+These tests fail loudly if either the fetch or the fetch→save round-trip
+loses a non-default LLM again.
+"""
+
+import time
+
+import pytest
+
+# GPT-4.1 Nano — a stable, non-default LLM available on the test backend
+# (also used by test_arabic_agent.py). Must differ from Agent.DEFAULT_LLM.
+NON_DEFAULT_LLM_ID = "67fd9e2bef0365783d06e2f0"
+
+
+def llm_ref_id(llm):
+ """Normalize an ``agent.llm`` value (str id, role-ref dict, or Model) to its id."""
+ if llm is None:
+ return None
+ if isinstance(llm, str):
+ return llm
+ if isinstance(llm, dict):
+ return llm.get("id")
+ return getattr(llm, "id", None)
+
+
+@pytest.fixture(scope="module")
+def non_default_llm_agent(client):
+ """Create an agent pinned to a non-default LLM, cleaned up after tests."""
+ assert NON_DEFAULT_LLM_ID != client.Agent.DEFAULT_LLM, "test LLM must not be the default"
+ agent = client.Agent(
+ name=f"LLM Persistence Test Agent {int(time.time())}",
+ description="Temporary agent verifying the LLM survives fetch/save round-trips",
+ instructions="You are a helpful test agent.",
+ llm=NON_DEFAULT_LLM_ID,
+ )
+ agent.save()
+
+ yield agent
+
+ try:
+ agent.delete()
+ except Exception:
+ pass
+
+
+class TestAgentLlmPersistence:
+ def test_create_persists_non_default_llm(self, client, non_default_llm_agent):
+ """The LLM passed at creation must be what the backend stores."""
+ fetched = client.Agent.get(non_default_llm_agent.id)
+ assert llm_ref_id(fetched.llm) == NON_DEFAULT_LLM_ID
+
+ def test_fetched_agent_does_not_fall_back_to_default_llm(self, client, non_default_llm_agent):
+ """``Agent.get()`` must decode the stored LLM, not fall back to DEFAULT_LLM.
+
+ This is the read half of the regression: on 0.2.43/0.2.44 the fetched
+ agent's ``llm`` was always ``DEFAULT_LLM`` regardless of the record.
+ """
+ fetched = client.Agent.get(non_default_llm_agent.id)
+ assert llm_ref_id(fetched.llm) != client.Agent.DEFAULT_LLM
+ # The save payload built from a fetched agent must echo the stored LLM.
+ payload_model = fetched.build_save_payload().get("model") or {}
+ assert payload_model.get("id") == NON_DEFAULT_LLM_ID
+
+ def test_instructions_only_save_preserves_llm(self, client, non_default_llm_agent):
+ """fetch → edit instructions → save must not clobber the agent's LLM.
+
+ This is the write half of the regression: the exact user flow that
+ silently reset agents to GPT-5.4 on 0.2.43/0.2.44.
+ """
+ fetched = client.Agent.get(non_default_llm_agent.id)
+ fetched.instructions = "You are a helpful test agent. (edited)"
+ fetched.save()
+
+ refetched = client.Agent.get(non_default_llm_agent.id)
+ assert refetched.instructions == "You are a helpful test agent. (edited)"
+ assert llm_ref_id(refetched.llm) == NON_DEFAULT_LLM_ID, (
+ "instructions-only save() overwrote the agent's LLM — "
+ "DEFAULT_LLM fallback regression (see 0.2.43/0.2.44)"
+ )
+
+ def test_explicit_llm_update_persists(self, client, non_default_llm_agent):
+ """Explicitly changing ``agent.llm`` on a fetched agent must persist."""
+ fetched = client.Agent.get(non_default_llm_agent.id)
+ fetched.llm = client.Agent.DEFAULT_LLM
+ fetched.save()
+ assert llm_ref_id(client.Agent.get(non_default_llm_agent.id).llm) == client.Agent.DEFAULT_LLM
+
+ # restore the non-default LLM and confirm the flip back also persists
+ fetched.llm = NON_DEFAULT_LLM_ID
+ fetched.save()
+ assert llm_ref_id(client.Agent.get(non_default_llm_agent.id).llm) == NON_DEFAULT_LLM_ID
diff --git a/tests/functional/v2/test_arabic_agent.py b/tests/functional/v2/test_arabic_agent.py
index 2c387d1a..cc764185 100644
--- a/tests/functional/v2/test_arabic_agent.py
+++ b/tests/functional/v2/test_arabic_agent.py
@@ -22,15 +22,7 @@
import pytest
-from aixplain.v2 import (
- EvaluatorConfig,
- EvaluatorType,
- Inspector,
- InspectorAction,
- InspectorActionConfig,
- InspectorSeverity,
- InspectorTarget,
-)
+from aixplain.v2 import Inspector
ARABIC_CHAR_RE = re.compile(r"[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDFF\uFE70-\uFEFF]")
ARABIC_DIACRITICS_RE = re.compile(r"[\u064B-\u065F\u0670]")
@@ -179,9 +171,12 @@ def _build_name(prefix: str, model_name: str) -> str:
return f"{prefix}-{model_name}-{int(time.time())}-{uuid.uuid4().hex[:6]}"
-def _is_inspector_step(step: dict) -> bool:
+def _is_inspector_step(step: dict, inspector_name: str = "") -> bool:
+ # The backend reports the inspector's own name as the step agent id
+ # (e.g. 'ArabicContentValidator-gpt-4o'), not an 'inspector|...' prefix.
agent_info = step.get("agent") or {}
- return (agent_info.get("id") or "").lower().startswith("inspector")
+ step_id = (agent_info.get("id") or "").lower()
+ return "inspector" in step_id or (bool(inspector_name) and step_id == inspector_name.lower())
def _is_inspector_abort_message(output: str) -> bool:
@@ -261,20 +256,20 @@ def _make_team_agent(client, llm_id: str, model_name: str, inspectors=None):
def _make_output_inspector(llm_id: str, model_name: str):
return Inspector(
name=f"ArabicContentValidator-{model_name}",
- severity=InspectorSeverity.HIGH,
- targets=[InspectorTarget.OUTPUT],
- action=InspectorActionConfig(type=InspectorAction.ABORT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id=llm_id,
- prompt=(
+ severity="high",
+ targets=["output"],
+ action="abort",
+ metric={
+ "asset_id": llm_id,
+ "prompt": (
"تحقق من أن الرد مكتوب بالعربية الفصحى وأنه يتعلق بالقانون التجاري السعودي فقط. "
"إذا كان الرد بلغة أخرى أو خارج النطاق، ارفضه."
),
- ),
+ },
)
+@pytest.mark.flaky(reruns=1, reruns_delay=5)
@pytest.mark.parametrize(("model_name", "llm_id"), MODELS)
def test_arabic_single_agent_variants_across_llms(client, resource_tracker, model_name, llm_id):
"""Keep broad LLM coverage, but run only one representative query per agent variant."""
@@ -294,6 +289,7 @@ def test_arabic_single_agent_variants_across_llms(client, resource_tracker, mode
_assert_query_expectations(query_key, output)
+@pytest.mark.flaky(reruns=1, reruns_delay=5)
@pytest.mark.parametrize(("model_name", "llm_id"), MODELS)
def test_arabic_team_agent_across_llms(client, resource_tracker, model_name, llm_id):
"""Use one contract-heavy query to cover the team-agent serialization path."""
@@ -307,6 +303,7 @@ def test_arabic_team_agent_across_llms(client, resource_tracker, model_name, llm
assert steps, "Expected team-agent execution steps for Arabic team flow"
+@pytest.mark.flaky(reruns=1, reruns_delay=5)
@pytest.mark.parametrize(("model_name", "llm_id"), MODELS)
def test_arabic_inspector_agent_across_llms(client, resource_tracker, model_name, llm_id):
"""Verify the inspector actually executes on the Arabic runtime path."""
@@ -317,14 +314,8 @@ def test_arabic_inspector_agent_across_llms(client, resource_tracker, model_name
response = team_agent.run(ARABIC_QUERIES["pure_arabic"])
output, steps = _assert_success_response(response)
- response_generator_steps = [
- step for step in steps if ((step.get("agent") or {}).get("id") or "").lower() == "response_generator"
- ]
- assert len(response_generator_steps) == 1, "Expected exactly one response_generator step"
-
- response_generator_index = steps.index(response_generator_steps[0])
- inspector_steps = [step for step in steps[response_generator_index + 1 :] if _is_inspector_step(step)]
- assert inspector_steps, "Expected inspector step(s) after response_generator"
+ inspector_steps = [step for step in steps if _is_inspector_step(step, inspector.name)]
+ assert inspector_steps, "Expected inspector step(s) in the run"
if _is_inspector_abort_message(output):
return
diff --git a/tests/functional/v2/test_model.py b/tests/functional/v2/test_model.py
index d98629ee..e573c406 100644
--- a/tests/functional/v2/test_model.py
+++ b/tests/functional/v2/test_model.py
@@ -406,7 +406,7 @@ def test_dynamic_validation_slack_integration(client, slack_integration_id, slac
# Test with valid parameters
valid_params = {
- "action": "SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL",
+ "action": "SLACK_SEND_MESSAGE",
"data": {"channel": "#general", "text": "Hello from test!"},
}
@@ -536,10 +536,13 @@ def test_search_models_with_functions_filter(client):
models = client.Model.search(functions=["text-generation"])
assert hasattr(models, "results")
assert isinstance(models.results, list)
- # If results are present, ensure each model has the requested function
+ assert models.results, "Expected the functions filter to return results"
+ # If results are present, ensure each model has the requested function.
+ # The v2 Function enum uses uppercase values (TEXT_GENERATION), while the
+ # backend filter id is the lowercase dashed form.
for model in models.results:
if model.function:
- assert model.function.value == "text-generation"
+ assert model.function.name == "TEXT_GENERATION"
def test_search_models_with_status_filter(client):
@@ -976,6 +979,22 @@ def test_sync_model_run(client, sync_model_id):
assert "Hola" in result.data or "hola" in result.data.lower()
+def test_seedream_run_returns_non_poll_url_without_polling(client):
+ """Sync model run should return Seedream's final output URL without polling it."""
+ model = client.Model.get("69f347e7de823633d9604dfd")
+
+ result = model.run(
+ text="A small red cube on a plain white background.",
+ timeout=120,
+ wait_time=5,
+ )
+
+ assert result.status == "SUCCESS"
+ assert result.completed is True
+ assert result.data
+ assert result.url is None
+
+
def test_sync_model_run_async(client, sync_model_id):
"""Test run_async() on sync model falls back to V1 MS."""
model = client.Model.get(sync_model_id)
diff --git a/tests/functional/v2/test_rlm.py b/tests/functional/v2/test_rlm.py
index 5f511f17..c300d192 100644
--- a/tests/functional/v2/test_rlm.py
+++ b/tests/functional/v2/test_rlm.py
@@ -40,10 +40,14 @@ def test_create_rlm_custom_iterations(self, client):
assert rlm.max_iterations == 3
def test_create_rlm_missing_orchestrator(self, client):
- """Test that RLM raises when orchestrator_id is missing."""
+ """Test that RLM raises when orchestrator_id is missing.
+
+ Only recursive mode requires the orchestrator; auto mode would fall
+ back to parallel/rag (worker-only) and succeed.
+ """
rlm = client.RLM(worker_id=MODEL_ID)
with pytest.raises(Exception):
- rlm.run(data={"context": "test", "query": "test"})
+ rlm.run(data={"context": "test", "query": "test"}, mode="recursive")
def test_create_rlm_missing_worker(self, client):
"""Test that RLM raises when worker_id is missing."""
diff --git a/tests/functional/v2/test_session.py b/tests/functional/v2/test_session.py
new file mode 100644
index 00000000..ebdc53de
--- /dev/null
+++ b/tests/functional/v2/test_session.py
@@ -0,0 +1,361 @@
+"""Functional tests for v2 Session management.
+
+These tests run against a real backend and require valid credentials.
+Set TEAM_API_KEY (or AIXPLAIN_API_KEY) in the environment.
+
+A test agent is created once per module and cleaned up afterwards.
+"""
+
+import os
+import tempfile
+import time
+
+import pytest
+
+from aixplain.v2 import Session, SessionMessage, SessionStatus
+from aixplain.v2.exceptions import APIError
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def test_agent(client):
+ """Create a temporary agent for session tests."""
+ agent = client.Agent(
+ name=f"Session Test Agent {int(time.time())}",
+ description="Temporary agent for session functional tests",
+ instructions="You are a helpful test agent. Keep responses short.",
+ )
+ agent.save()
+ yield agent
+ try:
+ agent.delete()
+ except Exception:
+ pass
+
+
+def _make_session(client, agent, name=None, **kwargs):
+ """Create and persist a session bound to ``agent`` (the single create path)."""
+ session = client.Session(agent=agent, name=name, **kwargs)
+ session.save()
+ return session
+
+
+# ---------------------------------------------------------------------------
+# Session CRUD
+# ---------------------------------------------------------------------------
+
+
+class TestSessionCRUD:
+ """End-to-end session create / get / list / update / delete."""
+
+ def test_create_session_with_agent_object(self, client, test_agent):
+ """Creating a session with ``Session(agent=…)`` returns a saved Session."""
+ session = _make_session(client, test_agent, name="Func Test Session")
+
+ assert session.id is not None
+ assert isinstance(session, Session)
+ assert session.agent_id == test_agent.id
+ assert session.name == "Func Test Session"
+ assert session.status == "active"
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_create_session_with_agent_id(self, client, test_agent):
+ """``Session(agent=…)`` also accepts a bare agent id string."""
+ session = client.Session(agent=test_agent.id, name="Direct Create")
+ session.save()
+
+ assert session.id is not None
+ assert session.agent_id == test_agent.id
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_get_session(self, client, test_agent):
+ """Retrieving a session by ID should return the same session."""
+ session = _make_session(client, test_agent, name="Get Test")
+ fetched = client.Session.get(session.id)
+
+ assert fetched.id == session.id
+ assert fetched.name == session.name
+ assert fetched.agent_id == session.agent_id
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_search_sessions_for_agent(self, client, test_agent):
+ """Session.search(agent=…) returns a Page including the created session."""
+ session = _make_session(client, test_agent, name="List Test")
+
+ page = client.Session.search(agent=test_agent)
+
+ session_ids = [s.id for s in page.results]
+ assert session.id in session_ids
+ assert page.total >= 1
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_search_sessions_with_status_filter(self, client, test_agent):
+ """Filtering by status should only return matching sessions."""
+ session = _make_session(client, test_agent, name="Status Filter Test")
+
+ page = client.Session.search(agent=test_agent, status="active")
+ assert all(s.status == "active" for s in page.results)
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_update_session(self, client, test_agent):
+ """Updating session name via save() should persist the change."""
+ session = _make_session(client, test_agent, name="Before Update")
+ session.name = "After Update"
+ session.save()
+
+ fetched = client.Session.get(session.id)
+ assert fetched.name == "After Update"
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_seed_session_with_messages(self, client, test_agent):
+ """Seeding a session via add_message should persist the transcript."""
+ session = _make_session(client, test_agent, name="History Test")
+ session.add_message(role="user", content="What is 2+2?")
+ session.add_message(role="assistant", content="4")
+
+ assert session.id is not None
+ messages = session.messages()
+ assert len(messages) >= 2
+ contents = [m.content for m in messages]
+ assert "What is 2+2?" in contents
+ assert "4" in contents
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+ def test_delete_session(self, client, test_agent):
+ """Deleting a session should succeed."""
+ session = _make_session(client, test_agent, name="Delete Me")
+ session_id = session.id
+ assert session_id is not None
+
+ result = session.delete()
+ assert result.completed is True
+
+
+# ---------------------------------------------------------------------------
+# Session messages
+# ---------------------------------------------------------------------------
+
+
+class TestSessionMessages:
+ """End-to-end tests for session message operations."""
+
+ @pytest.fixture()
+ def session(self, client, test_agent):
+ """Create a session for message tests and clean up after."""
+ s = _make_session(client, test_agent, name=f"Msg Test {int(time.time())}")
+ yield s
+ try:
+ s.delete()
+ except Exception:
+ pass
+
+ def test_add_and_get_message(self, session):
+ """Adding a message should return a SessionMessage with content."""
+ msg = session.add_message(role="user", content="Hello from test!")
+
+ assert isinstance(msg, SessionMessage)
+ assert msg.id is not None
+ assert msg.content == "Hello from test!"
+ assert msg.role == "user"
+
+ def test_list_messages(self, session):
+ """After adding messages, messages() should return them."""
+ session.add_message(role="user", content="First message")
+ session.add_message(role="assistant", content="Second message")
+
+ messages = session.messages()
+
+ assert isinstance(messages, list)
+ assert len(messages) >= 2
+ contents = [m.content for m in messages]
+ assert "First message" in contents
+ assert "Second message" in contents
+
+ def test_get_single_message(self, session):
+ """get_message() should return the specific message."""
+ created = session.add_message(role="user", content="Specific message")
+
+ fetched = session.get_message(created.id)
+
+ assert fetched.id == created.id
+ assert fetched.content == "Specific message"
+
+ def test_delete_message(self, session):
+ """delete_message() should remove the message."""
+ msg = session.add_message(role="user", content="Delete this message")
+ session.delete_message(msg.id)
+
+ remaining = session.messages()
+ remaining_ids = [m.id for m in remaining]
+ assert msg.id not in remaining_ids
+
+ def test_add_message_with_url_attachments(self, session):
+ """Adding a message with explicit URL attachments should include them."""
+ attachments = [
+ {"url": "https://example.com/test.png", "name": "test.png", "type": "image"},
+ ]
+ msg = session.add_message(
+ role="user",
+ content="See attached",
+ attachments=attachments,
+ )
+
+ assert msg.id is not None
+ if msg.attachments:
+ assert msg.attachments[0].url == "https://example.com/test.png"
+ assert msg.attachments[0].name == "test.png"
+ assert msg.attachments[0].type == "image"
+
+ def test_add_message_with_file_upload(self, session):
+ """Uploading a local file via add_message(files=...) should attach it."""
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", prefix="session_test_", delete=False) as f:
+ f.write("Test content for file upload verification.\n")
+ tmp_path = f.name
+
+ try:
+ msg = session.add_message(
+ role="user",
+ content="Please review this file",
+ files=[tmp_path],
+ )
+
+ assert msg.id is not None
+ assert msg.attachments is not None, "Backend should echo attachments"
+ assert len(msg.attachments) == 1
+
+ att = msg.attachments[0]
+ assert att.url is not None and att.url.startswith("http")
+ assert att.name == os.path.basename(tmp_path)
+ assert att.type == "text"
+
+ # Verify the attachment persists when re-fetching
+ fetched = session.get_message(msg.id)
+ assert fetched.attachments is not None
+ assert len(fetched.attachments) == 1
+ finally:
+ os.remove(tmp_path)
+
+ @pytest.mark.skip(
+ reason="Backend bug: sessions service validates role but stores every "
+ "message as role='user' (verified 2026-07-21), so assistant messages "
+ "can never be reacted to. Unskip when the backend persists the role."
+ )
+ def test_react_like_and_dislike(self, session):
+ """Reacting to an assistant message with LIKE then DISLIKE should work."""
+ session.add_message(role="user", content="Say something")
+ msg = session.add_message(role="assistant", content="Here is my response")
+
+ liked = session.react(msg.id, "LIKE")
+ assert liked.reaction == "LIKE"
+
+ disliked = session.react(msg.id, "DISLIKE")
+ assert disliked.reaction == "DISLIKE"
+
+ @pytest.mark.skip(
+ reason="Backend bug: sessions service stores every message as "
+ "role='user' — see test_react_like_and_dislike."
+ )
+ def test_clear_reaction(self, session):
+ """Passing None to react() should clear the reaction."""
+ session.add_message(role="user", content="Say something")
+ msg = session.add_message(role="assistant", content="Clear reaction response")
+ session.react(msg.id, "DISLIKE")
+
+ cleared = session.react(msg.id, None)
+ assert cleared.reaction is None
+
+
+# ---------------------------------------------------------------------------
+# Agent run + session interaction
+# ---------------------------------------------------------------------------
+
+
+class TestSessionWithAgentRun:
+ """Tests for how agent.run() interacts with sessions."""
+
+ def test_run_with_session_adds_messages(self, client, test_agent):
+ """Running an agent with session=… should add messages to the session."""
+ session = _make_session(client, test_agent, name="Run Test")
+ msgs_before = session.messages()
+
+ result = test_agent.run(
+ "What is the capital of France?",
+ session=session,
+ )
+ assert result.status == "SUCCESS"
+
+ msgs_after = session.messages()
+ assert len(msgs_after) > len(msgs_before), "Backend should auto-add messages to the session during a run"
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
+
+
+# ---------------------------------------------------------------------------
+# Error handling
+# ---------------------------------------------------------------------------
+
+
+class TestSessionErrors:
+ """Tests that backend errors are raised gracefully."""
+
+ def test_get_nonexistent_session(self, client):
+ """Getting a session that doesn't exist should raise APIError."""
+ with pytest.raises(APIError):
+ client.Session.get("nonexistent-session-id-12345")
+
+ def test_react_to_user_message_raises_error(self, client, test_agent):
+ """Reacting to a user message should raise an APIError."""
+ session = _make_session(client, test_agent, name="React Error Test")
+ msg = session.add_message(role="user", content="Can't like this")
+
+ with pytest.raises(APIError, match="assistant"):
+ session.react(msg.id, "LIKE")
+
+ # Cleanup
+ try:
+ session.delete()
+ except Exception:
+ pass
diff --git a/tests/functional/v2/test_snake_case_e2e.py b/tests/functional/v2/test_snake_case_e2e.py
index cd9d6384..eb224f93 100644
--- a/tests/functional/v2/test_snake_case_e2e.py
+++ b/tests/functional/v2/test_snake_case_e2e.py
@@ -33,12 +33,14 @@ def test_asset_id_round_trips_through_backend(self, client):
agent.save()
try:
- fetched = client.Agent.get(agent.id)
- saved_tool = fetched.tools[0] # raw dict from API
-
- # Backend resolves asset_id into the tool's "id" field
- assert saved_tool["id"] == sent_asset_id, (
- f"asset_id we sent ({sent_asset_id}) != id backend returned ({saved_tool.get('id')})"
+ # Read the raw payload: Agent.get hydrates tools into objects, so
+ # the backend's camelCase keys are only visible on the wire shape.
+ raw_agent = client.client.request("get", f"sdk/agents/{agent.id}")
+ saved_tool = raw_agent["assets"][0]
+
+ # Backend stores the tool's asset_id under "assetId"
+ assert saved_tool["assetId"] == sent_asset_id, (
+ f"asset_id we sent ({sent_asset_id}) != assetId backend returned ({saved_tool.get('assetId')})"
)
finally:
try:
@@ -66,13 +68,14 @@ def test_allow_multi_and_supports_variables_round_trip(self, client):
agent.save()
try:
- fetched = client.Agent.get(agent.id)
- saved_tool = fetched.tools[0]
-
- saved_params = saved_tool.get("parameters", [])
+ # Read the raw payload: Agent.get hydrates tools into objects, so
+ # the backend's camelCase keys are only visible on the wire shape.
+ raw_agent = client.client.request("get", f"sdk/agents/{agent.id}")
+ saved_params = raw_agent["assets"][0].get("parameters", [])
assert saved_params, "Backend should return the tool parameters we sent"
- saved_sample = saved_params[0]
+ saved_sample = next((p for p in saved_params if p.get("name") == sample["name"]), None)
+ assert saved_sample is not None, f"Parameter {sample['name']} missing from saved payload"
# Backend returns camelCase keys in the raw dict
assert saved_sample["allowMulti"] == sent_allow_multi, (
f"allow_multi we sent ({sent_allow_multi}) "
@@ -172,16 +175,24 @@ class TestAgentRunParamsKwargs:
call agent.run() with the snake_case kwarg → backend accepts and responds.
"""
- def test_session_id(self, client, test_agent):
- """session_id (renamed from sessionId) is sent to the backend and reflected in the response."""
+ def test_session(self, client, test_agent):
+ """session=… routes the run through the session and reflects the id in the response."""
agent = client.Agent.get(test_agent.id)
- sid = agent.generate_session_id()
+ session = client.Session(agent=agent, name="snake_case_test")
+ session.save()
+ sid = session.id
- response = agent.run("ping", session_id=sid)
+ try:
+ response = agent.run("ping", session=session)
- assert response.status == "SUCCESS"
- assert response.data.session_id is not None
- assert sid in response.data.session_id
+ assert response.status == "SUCCESS"
+ assert response.data.session_id is not None
+ assert sid in response.data.session_id
+ finally:
+ try:
+ session.delete()
+ except Exception:
+ pass
def test_execution_params(self, client, test_agent):
"""execution_params (renamed from executionParams) reaches the backend.
@@ -193,30 +204,25 @@ def test_execution_params(self, client, test_agent):
agent = client.Agent.get(test_agent.id)
- with pytest.raises(APIError, match="(?i)max.?token"):
+ # The backend surfaces a generic failure message for token-limit
+ # errors, so only the failure itself proves the param was honored.
+ with pytest.raises(APIError):
agent.run(
"ping",
execution_params={"max_tokens": 1, "max_iterations": 3, "output_format": "text"},
)
- def test_run_response_generation(self, client, test_agent):
- """run_response_generation (renamed from runResponseGeneration) is accepted by the backend."""
- agent = client.Agent.get(test_agent.id)
-
- response = agent.run("ping", run_response_generation=False)
+ def test_history_on_direct_run(self, client, test_agent):
+ """history (a legacy direct-run kwarg) is accepted on a sessionless run.
- assert response.status == "SUCCESS"
-
- def test_allow_history_and_session_id(self, client, test_agent):
- """allow_history_and_session_id (renamed from allowHistoryAndSessionId) is accepted."""
+ Conversation continuity now flows through session=… rather than
+ combining history with a session in one run.
+ """
agent = client.Agent.get(test_agent.id)
- sid = agent.generate_session_id()
response = agent.run(
"ping",
- session_id=sid,
history=[{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}],
- allow_history_and_session_id=True,
)
assert response.status == "SUCCESS"
diff --git a/tests/functional/v2/test_tool.py b/tests/functional/v2/test_tool.py
index 2fafe607..1bdc524a 100644
--- a/tests/functional/v2/test_tool.py
+++ b/tests/functional/v2/test_tool.py
@@ -4,7 +4,10 @@
from aixplain.v2.integration import Integration
-TAVILY_TOOL_PATH = "tavily/tavily-web-search/tavily"
+# "Tavily Web Search" connector tool (action: search). Fetched by id: the
+# marketplace path tavily/tavily-search-api/Tavily collides with the Legacy
+# Tavily asset (6736411c...), whose single generic 'run' action breaks these tests.
+TAVILY_TOOL_ID = "6931bdf462eb386b7158def3"
@pytest.fixture(scope="module")
@@ -16,7 +19,7 @@ def slack_integration_id():
@pytest.fixture(scope="module")
def single_action_test_agent(client):
"""Create a temporary agent using a single-action tool and clean it up."""
- tool = client.Tool.get(TAVILY_TOOL_PATH)
+ tool = client.Tool.get(TAVILY_TOOL_ID)
tool.allowed_actions = []
agent = client.Agent(
@@ -204,14 +207,14 @@ def test_tool_run(client, slack_integration_id, slack_token):
name=tool_name,
integration=integration,
config={"token": slack_token},
- allowed_actions=["SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"],
+ allowed_actions=["SLACK_SEND_MESSAGE"],
)
# Save tool before running
tool.save()
result = tool.run(
- action="SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL",
+ action="SLACK_SEND_MESSAGE",
data={
"channel": "#integrations-test",
"text": f"Test message from functional test {int(time.time())}",
@@ -349,7 +352,7 @@ def test_tool_run_with_default_params(client):
were sent as raw default dicts instead of extracted primitive values,
causing the backend to reject the request.
"""
- tavily_tool = client.Tool.get(TAVILY_TOOL_PATH)
+ tavily_tool = client.Tool.get(TAVILY_TOOL_ID)
# Verify the action proxy stores extracted primitives, not raw dicts
action_proxy = tavily_tool.actions["search"]
@@ -488,18 +491,18 @@ def test_tool_update_preserves_allowed_actions(client, slack_integration_id, sla
name=tool_name,
integration=slack_integration_id,
config={"token": slack_token},
- allowed_actions=["SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"],
+ allowed_actions=["SLACK_SEND_MESSAGE"],
)
tool.save()
tool_id = tool.id
try:
fetched = client.Tool.get(tool_id)
- fetched.allowed_actions = ["SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"]
+ fetched.allowed_actions = ["SLACK_SEND_MESSAGE"]
fetched.name = f"test-update-actions-renamed-{int(time.time())}"
fetched.save()
- assert fetched.allowed_actions == ["SLACK_SENDS_A_MESSAGE_TO_A_SLACK_CHANNEL"], (
+ assert fetched.allowed_actions == ["SLACK_SEND_MESSAGE"], (
"allowed_actions should be preserved after save()"
)
@@ -514,7 +517,7 @@ def test_tool_update_preserves_allowed_actions(client, slack_integration_id, sla
def test_tool_as_tool_auto_detects_single_action(client):
"""Test that as_tool() auto-includes the action when a tool has exactly one action."""
- tool = client.Tool.get(TAVILY_TOOL_PATH)
+ tool = client.Tool.get(TAVILY_TOOL_ID)
tool.allowed_actions = []
tool_dict = tool.as_tool()
@@ -525,7 +528,7 @@ def test_tool_as_tool_auto_detects_single_action(client):
def test_tool_as_tool_no_mutation(client):
"""Test that as_tool() does NOT mutate self.allowed_actions as a side effect."""
- tool = client.Tool.get(TAVILY_TOOL_PATH)
+ tool = client.Tool.get(TAVILY_TOOL_ID)
tool.allowed_actions = []
tool.as_tool()
@@ -539,7 +542,7 @@ def test_tool_as_tool_caching(client):
"""Test that repeated as_tool() calls reuse cached actions instead of hitting the API again."""
from unittest.mock import patch
- tool = client.Tool.get(TAVILY_TOOL_PATH)
+ tool = client.Tool.get(TAVILY_TOOL_ID)
tool.allowed_actions = []
tool.as_tool()
@@ -563,7 +566,7 @@ def counting_list_actions():
def test_tool_run_auto_detects_single_action(client):
"""Test that run() auto-detects the action for single-action tools without explicit action kwarg."""
- tool = client.Tool.get(TAVILY_TOOL_PATH)
+ tool = client.Tool.get(TAVILY_TOOL_ID)
tool.allowed_actions = []
result = tool.run(data={"query": "friendship paradox", "num_results": 1})
diff --git a/tests/functional/v2/test_trigger.py b/tests/functional/v2/test_trigger.py
new file mode 100644
index 00000000..3daee746
--- /dev/null
+++ b/tests/functional/v2/test_trigger.py
@@ -0,0 +1,249 @@
+"""Functional tests for v2 Triggers (aix.Trigger).
+
+Time-trigger tests run with only TEAM_API_KEY/AIXPLAIN_API_KEY set (a temporary
+agent is created and cleaned up). Event-trigger tests are gated on extra env vars:
+
+- TEST_COMPOSIO_INTEGRATION_ID : an integration id (e.g. composio/gmail resolved id)
+ used for `integration.triggers` discovery.
+- TEST_CONNECTION_ID : a connected tool id used to activate a real event
+ trigger end-to-end.
+"""
+
+import os
+import time
+
+import pytest
+
+from aixplain.v2 import Trigger, TriggerEventOption
+from aixplain.v2.resource import Page
+
+
+# Far-future instant so a "once" trigger is valid/schedulable.
+FUTURE_RUN_AT = "2099-01-26T12:00:00Z"
+
+
+@pytest.fixture(scope="module")
+def test_agent(client):
+ """Create a temporary agent to attach triggers to, and clean it up."""
+ agent = client.Agent(
+ name=f"Trigger Functional Agent {int(time.time())}",
+ description="Temporary agent for trigger functional tests",
+ instructions="You are a helpful test agent. Respond briefly.",
+ )
+ agent.save()
+ yield agent
+ try:
+ agent.delete()
+ except Exception:
+ pass
+
+
+@pytest.fixture
+def cleanup_triggers():
+ """Collect triggers created during a test and delete them afterwards."""
+ created = []
+ yield created
+ for t in created:
+ try:
+ t.delete()
+ except Exception:
+ pass
+
+
+@pytest.fixture(scope="module")
+def composio_integration_id():
+ """Integration id for event-discovery tests (skips if not provided)."""
+ value = os.getenv("TEST_COMPOSIO_INTEGRATION_ID")
+ if not value:
+ pytest.skip("TEST_COMPOSIO_INTEGRATION_ID is required for event-discovery tests")
+ return value
+
+
+@pytest.fixture(scope="module")
+def connection_id():
+ """Connected tool id for end-to-end event-trigger tests (skips if not provided)."""
+ value = os.getenv("TEST_CONNECTION_ID")
+ if not value:
+ pytest.skip("TEST_CONNECTION_ID is required for event-trigger activation tests")
+ return value
+
+
+# =============================================================================
+# Time triggers
+# =============================================================================
+
+
+class TestTimeTriggerLifecycle:
+ def test_create_once_trigger(self, client, test_agent, cleanup_triggers):
+ """A one-off (run_at) trigger is created and enabled by default."""
+ t = client.Trigger(
+ name=f"once-{int(time.time())}",
+ agent=test_agent,
+ input="Remind the team about the launch.",
+ run_at=FUTURE_RUN_AT,
+ )
+ t.save()
+ cleanup_triggers.append(t)
+
+ assert t.id is not None
+ assert t.trigger_type == "time"
+ assert t.schedule_type == "once"
+ assert t.enabled is True
+
+ def test_create_daily_trigger(self, client, test_agent, cleanup_triggers):
+ """A daily trigger maps at/timezone onto a daily schedule."""
+ t = client.Trigger(
+ name=f"daily-{int(time.time())}",
+ agent=test_agent,
+ input="Summarise today's AI news.",
+ every="day",
+ at="09:00",
+ timezone="Europe/London",
+ notifications=True,
+ )
+ t.save()
+ cleanup_triggers.append(t)
+
+ assert t.id is not None
+ assert t.schedule_type == "daily"
+ assert t.notifications is True
+
+ def test_create_interval_and_weekly_and_monthly(self, client, test_agent, cleanup_triggers):
+ """Interval, weekly, and monthly schedules all create successfully."""
+ hourly = client.Trigger(
+ name=f"hourly-{int(time.time())}", agent=test_agent, input="Check the queue.",
+ every="hour", interval=2,
+ )
+ hourly.save()
+ cleanup_triggers.append(hourly)
+ assert hourly.schedule_type == "recurring"
+
+ weekly = client.Trigger(
+ name=f"weekly-{int(time.time())}", agent=test_agent, input="Compile the weekly report.",
+ every="week", on=["mon", "thu"], at="17:00",
+ )
+ weekly.save()
+ cleanup_triggers.append(weekly)
+ assert weekly.schedule_type == "weekly"
+
+ monthly = client.Trigger(
+ name=f"monthly-{int(time.time())}", agent=test_agent, input="Generate invoices.",
+ every="month", on=[1, 15], at="09:00",
+ )
+ monthly.save()
+ cleanup_triggers.append(monthly)
+ assert monthly.schedule_type == "monthly"
+
+ def test_get_trigger(self, client, test_agent, cleanup_triggers):
+ """Trigger.get(id) retrieves a created trigger."""
+ t = client.Trigger(
+ name=f"get-{int(time.time())}", agent=test_agent, input="Ping.", run_at=FUTURE_RUN_AT,
+ )
+ t.save()
+ cleanup_triggers.append(t)
+
+ fetched = client.Trigger.get(t.id)
+ assert fetched.id == t.id
+ assert fetched.name == t.name
+ assert fetched.trigger_type == "time"
+ assert fetched.schedule_type == "once"
+
+ def test_search_by_agent(self, client, test_agent, cleanup_triggers):
+ """Trigger.search(agent=) returns a Page containing the agent's triggers."""
+ t = client.Trigger(
+ name=f"search-{int(time.time())}", agent=test_agent, input="Ping.", run_at=FUTURE_RUN_AT,
+ )
+ t.save()
+ cleanup_triggers.append(t)
+
+ page = client.Trigger.search(agent=test_agent)
+ assert isinstance(page, Page)
+ ids = [item.id for item in page.results]
+ assert t.id in ids
+ for item in page.results:
+ assert isinstance(item, Trigger)
+ assert item.asset_id == test_agent.id
+
+ def test_enable_disable_via_save(self, client, test_agent, cleanup_triggers):
+ """Setting enabled=False and saving disables the trigger."""
+ t = client.Trigger(
+ name=f"toggle-{int(time.time())}", agent=test_agent, input="Ping.", run_at=FUTURE_RUN_AT,
+ )
+ t.save()
+ cleanup_triggers.append(t)
+ assert t.enabled is True
+
+ t.enabled = False
+ t.save()
+ assert client.Trigger.get(t.id).enabled is False
+
+ t.enabled = True
+ t.save()
+ assert client.Trigger.get(t.id).enabled is True
+
+ def test_delete_trigger(self, client, test_agent):
+ """delete() removes the trigger."""
+ t = client.Trigger(
+ name=f"delete-{int(time.time())}", agent=test_agent, input="Ping.", run_at=FUTURE_RUN_AT,
+ )
+ t.save()
+ trigger_id = t.id
+
+ t.delete()
+
+ with pytest.raises(Exception):
+ client.Trigger.get(trigger_id)
+
+
+# =============================================================================
+# Event triggers
+# =============================================================================
+
+
+class TestEventTriggerDiscovery:
+ def test_integration_triggers_lists_options(self, client, composio_integration_id):
+ """integration.triggers lists available event options (like .actions)."""
+ integration = client.Integration.get(composio_integration_id)
+ triggers = integration.triggers
+
+ # Browsable collection: len / iterate / membership / indexing.
+ assert len(triggers) >= 0
+ slugs = list(triggers)
+ for slug in slugs:
+ assert isinstance(slug, str)
+ if slugs:
+ option = triggers[slugs[0]]
+ assert isinstance(option, TriggerEventOption)
+ assert option.slug == slugs[0]
+ # Discovery from an (unconnected) integration carries no connection.
+ assert option.connection_id is None
+
+
+class TestEventTriggerLifecycle:
+ def test_create_and_delete_event_trigger(self, client, test_agent, connection_id):
+ """Create a real Composio event trigger from a connected tool and delete it."""
+ tool = client.Tool.get(connection_id)
+
+ slugs = list(tool.triggers)
+ if not slugs:
+ pytest.skip("Connected tool exposes no trigger types")
+ option = tool.triggers[slugs[0]]
+ assert option.connection_id == tool.id # connected tool carries the connection
+
+ t = client.Trigger(
+ name=f"event-{int(time.time())}",
+ agent=test_agent,
+ input="Handle this event.",
+ event=option,
+ )
+ try:
+ t.save()
+ assert t.id is not None
+ assert t.trigger_type == "external"
+ # Activation filled in the real Composio trigger id.
+ assert t.trigger_id
+ finally:
+ try:
+ t.delete()
+ except Exception:
+ pass
diff --git a/tests/unit/v2/test_agent_budget.py b/tests/unit/v2/test_agent_budget.py
new file mode 100644
index 00000000..a0a6d8e1
--- /dev/null
+++ b/tests/unit/v2/test_agent_budget.py
@@ -0,0 +1,461 @@
+"""Unit tests for agent budget serialization.
+
+Every agent owns a ``budget`` (an always-present, never-None ``Budget``, empty by
+default) mutated via attribute access (``agent.budget.max_cost = ...``). The same
+object serves two roles: ``build_run_payload`` sends its current state inside
+``executionParams`` as a nested camelCase ``budget`` object (run-time budget), and
+``build_save_payload`` persists it at the top level of the create payload. It is
+the single source of truth for the iteration cap, so the deprecated
+``max_iterations`` surfaces fold into it and the standalone ``maxIterations`` key
+is no longer emitted. These tests cover:
+- A ``Budget`` instance, a snake_case dict, and a camelCase dict assigned to
+ ``agent.budget`` all yield the same camelCase payload.
+- Partial budgets emit no null keys; ``warnAtPercent`` is never emitted.
+- An empty budget produces no ``budget`` key (default state is inert).
+- Persisted budget serialization in the create payload.
+- The deprecated persisted/run-time ``max_iterations`` fold into ``budget`` with
+ a ``DeprecationWarning``; the agent's budget wins on conflict.
+- No standalone ``maxIterations`` survives in either emitted payload.
+"""
+
+import warnings
+from unittest.mock import MagicMock
+
+import pytest
+
+from aixplain.v2.agent import Agent, Budget
+
+
+def _create_agent():
+ """Build a V2 Agent through ``from_dict`` with a mocked client context."""
+ agent = Agent.from_dict({"id": "agent-123", "name": "test-agent"})
+ agent.context = MagicMock()
+ return agent
+
+
+class TestBudgetSerialization:
+ """Budget.to_dict() honours the camelCase wire contract."""
+
+ def test_full_budget_to_dict(self):
+ budget = Budget(max_cost=1.0, max_duration_seconds=300, max_iterations=50)
+ assert budget.to_dict() == {
+ "maxCost": 1.0,
+ "maxDurationSeconds": 300,
+ "maxIterations": 50,
+ }
+
+ def test_partial_budget_drops_none_keys(self):
+ assert Budget(max_cost=2.0).to_dict() == {"maxCost": 2.0}
+
+ def test_empty_budget_is_empty(self):
+ assert Budget().to_dict() == {}
+
+ def test_zero_values_are_preserved(self):
+ # Falsy-but-valid values must survive; dropping uses ``is None``, not truthiness.
+ assert Budget(max_cost=0.0, max_iterations=0).to_dict() == {
+ "maxCost": 0.0,
+ "maxIterations": 0,
+ }
+
+
+class TestBudgetInRunPayload:
+ """build_run_payload threads the agent's current ``budget`` into executionParams.budget.
+
+ The per-run budget is no longer a ``run`` kwarg — it is the agent's own
+ ``agent.budget`` state at call time (mutated via attribute access).
+ """
+
+ def test_budget_instance(self):
+ agent = _create_agent()
+ agent.budget = Budget(max_cost=1.0, max_duration_seconds=300, max_iterations=50)
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {
+ "maxCost": 1.0,
+ "maxDurationSeconds": 300,
+ "maxIterations": 50,
+ }
+
+ def test_budget_via_attribute_access(self):
+ # The documented pattern: mutate fields on the always-present budget.
+ agent = _create_agent()
+ agent.budget.max_cost = 1.0
+ agent.budget.max_duration_seconds = 300
+ agent.budget.max_iterations = 50
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {
+ "maxCost": 1.0,
+ "maxDurationSeconds": 300,
+ "maxIterations": 50,
+ }
+
+ def test_budget_dict_snake_case(self):
+ agent = _create_agent()
+ agent.budget = {"max_cost": 1.0, "max_duration_seconds": 300, "max_iterations": 50}
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {
+ "maxCost": 1.0,
+ "maxDurationSeconds": 300,
+ "maxIterations": 50,
+ }
+
+ def test_budget_dict_camel_case(self):
+ agent = _create_agent()
+ agent.budget = {"maxCost": 1.0, "maxDurationSeconds": 300, "maxIterations": 50}
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {
+ "maxCost": 1.0,
+ "maxDurationSeconds": 300,
+ "maxIterations": 50,
+ }
+
+ def test_budget_reflects_latest_mutation(self):
+ # A loop that mutates the budget between runs sends the current value each time.
+ agent = _create_agent()
+ seen = []
+ for cost in [0.1, 0.5, 1.0]:
+ agent.budget.max_cost = cost
+ payload = agent.build_run_payload(query="q")
+ seen.append(payload["executionParams"]["budget"])
+ assert seen == [{"maxCost": 0.1}, {"maxCost": 0.5}, {"maxCost": 1.0}]
+
+ def test_partial_budget_instance(self):
+ agent = _create_agent()
+ agent.budget = Budget(max_cost=2.0)
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {"maxCost": 2.0}
+
+ def test_partial_budget_dict(self):
+ agent = _create_agent()
+ agent.budget = {"max_cost": 2.0}
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {"maxCost": 2.0}
+
+ def test_zero_value_budget_preserved_in_payload(self):
+ # A zero cap is semantically distinct from "no cap" and must be sent.
+ agent = _create_agent()
+ agent.budget.max_cost = 0.0
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {"maxCost": 0.0}
+
+ def test_dict_with_none_value_drops_key(self):
+ agent = _create_agent()
+ agent.budget = {"max_cost": 1.0, "max_iterations": None}
+ payload = agent.build_run_payload(query="q")
+ assert payload["executionParams"]["budget"] == {"maxCost": 1.0}
+
+ def test_budget_does_not_leak_to_payload_top_level(self):
+ agent = _create_agent()
+ agent.budget.max_cost = 1.0
+ payload = agent.build_run_payload(query="q")
+ assert "budget" not in payload
+
+ def test_stray_budget_kwarg_is_ignored(self):
+ # ``budget`` is no longer a run kwarg; a stray one must not leak to the payload.
+ agent = _create_agent()
+ payload = agent.build_run_payload(query="q", budget=Budget(max_cost=9.0))
+ assert "budget" not in payload
+ assert "budget" not in payload["executionParams"]
+
+ def test_no_budget_means_no_budget_key(self):
+ agent = _create_agent()
+ payload = agent.build_run_payload(query="q")
+ assert "budget" not in payload["executionParams"]
+
+ def test_empty_budget_instance_means_no_budget_key(self):
+ agent = _create_agent()
+ agent.budget = Budget()
+ payload = agent.build_run_payload(query="q")
+ assert "budget" not in payload["executionParams"]
+
+ def test_empty_budget_dict_means_no_budget_key(self):
+ agent = _create_agent()
+ agent.budget = {}
+ payload = agent.build_run_payload(query="q")
+ assert "budget" not in payload["executionParams"]
+
+ def test_warn_at_percent_never_emitted(self):
+ agent = _create_agent()
+ agent.budget = Budget(max_cost=1.0, max_duration_seconds=300, max_iterations=50)
+ payload = agent.build_run_payload(query="q")
+ budget = payload["executionParams"]["budget"]
+ assert "warnAtPercent" not in budget
+ assert "warn_at_percent" not in budget
+
+ def test_no_budget_payload_unchanged(self):
+ """A run with an empty budget produces the same executionParams as a set one, minus budget."""
+ with_budget_agent = _create_agent()
+ with_budget_agent.budget.max_cost = 1.0
+ with_budget = with_budget_agent.build_run_payload(query="q")
+ without_budget = _create_agent().build_run_payload(query="q")
+
+ ep_with = dict(with_budget["executionParams"])
+ ep_with.pop("budget", None)
+ assert ep_with == without_budget["executionParams"]
+
+
+def _build_agent(**kwargs):
+ """Construct a V2 Agent (running __post_init__) with a mocked client context."""
+ agent = Agent(name="test-agent", **kwargs)
+ agent.context = MagicMock()
+ return agent
+
+
+class TestPersistedBudgetInCreatePayload:
+ """build_save_payload persists ``budget`` and never emits standalone maxIterations."""
+
+ def test_budget_instance_persisted(self):
+ agent = _build_agent(budget=Budget(max_iterations=10, max_cost=0.50))
+ payload = agent.build_save_payload()
+ assert payload["budget"] == {"maxIterations": 10, "maxCost": 0.50}
+
+ def test_budget_dict_snake_case_persisted(self):
+ agent = _build_agent(budget={"max_iterations": 10})
+ payload = agent.build_save_payload()
+ assert payload["budget"] == {"maxIterations": 10}
+
+ def test_budget_dict_camel_case_persisted(self):
+ agent = _build_agent(budget={"maxIterations": 10})
+ payload = agent.build_save_payload()
+ assert payload["budget"] == {"maxIterations": 10}
+
+ def test_no_budget_means_no_budget_key(self):
+ agent = _build_agent()
+ payload = agent.build_save_payload()
+ assert "budget" not in payload
+
+ def test_empty_budget_means_no_budget_key(self):
+ agent = _build_agent(budget=Budget())
+ payload = agent.build_save_payload()
+ assert "budget" not in payload
+
+ def test_no_standalone_max_iterations_when_unset(self):
+ agent = _build_agent()
+ payload = agent.build_save_payload()
+ assert "maxIterations" not in payload
+
+ def test_zero_iterations_preserved(self):
+ agent = _build_agent(budget=Budget(max_iterations=0))
+ payload = agent.build_save_payload()
+ assert payload["budget"] == {"maxIterations": 0}
+
+
+class TestDeprecatedPersistedMaxIterations:
+ """Agent(max_iterations=N) warns, folds into budget, and drops the standalone key."""
+
+ def test_deprecation_warning_emitted(self):
+ with pytest.warns(DeprecationWarning, match="max_iterations"):
+ _build_agent(max_iterations=7)
+
+ def test_folds_into_budget(self):
+ with pytest.warns(DeprecationWarning):
+ agent = _build_agent(max_iterations=7)
+ assert agent.budget is not None
+ assert agent.budget.max_iterations == 7
+
+ def test_create_payload_has_budget_not_standalone(self):
+ with pytest.warns(DeprecationWarning):
+ agent = _build_agent(max_iterations=7)
+ payload = agent.build_save_payload()
+ assert payload["budget"] == {"maxIterations": 7}
+ assert "maxIterations" not in payload
+
+ def test_no_warning_when_unset(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", DeprecationWarning)
+ _build_agent() # must not raise
+
+ def test_conflict_budget_wins(self):
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ agent = _build_agent(max_iterations=7, budget=Budget(max_iterations=99))
+ categories = {w.category for w in caught}
+ assert DeprecationWarning in categories
+ assert UserWarning in categories
+ assert agent.budget.max_iterations == 99
+ payload = agent.build_save_payload()
+ assert payload["budget"]["maxIterations"] == 99
+ assert "maxIterations" not in payload
+
+
+class TestDeprecatedRunTimeMaxIterations:
+ """run-time execution_params['max_iterations'] folds into executionParams.budget."""
+
+ def test_deprecation_warning_emitted(self):
+ agent = _create_agent()
+ with pytest.warns(DeprecationWarning, match="max_iterations"):
+ agent.build_run_payload(query="q", execution_params={"max_iterations": 7})
+
+ def test_folds_into_execution_budget(self):
+ agent = _create_agent()
+ with pytest.warns(DeprecationWarning):
+ payload = agent.build_run_payload(
+ query="q", execution_params={"max_iterations": 7}
+ )
+ assert payload["executionParams"]["budget"] == {"maxIterations": 7}
+ assert "maxIterations" not in payload["executionParams"]
+
+ def test_camel_case_exec_param_also_folds(self):
+ agent = _create_agent()
+ with pytest.warns(DeprecationWarning):
+ payload = agent.build_run_payload(
+ query="q", execution_params={"maxIterations": 7}
+ )
+ assert payload["executionParams"]["budget"] == {"maxIterations": 7}
+ assert "maxIterations" not in payload["executionParams"]
+
+ def test_no_standalone_max_iterations_by_default(self):
+ agent = _create_agent()
+ payload = agent.build_run_payload(query="q")
+ assert "maxIterations" not in payload["executionParams"]
+ assert "budget" not in payload["executionParams"]
+
+ def test_conflict_budget_wins(self):
+ agent = _create_agent()
+ agent.budget = Budget(max_iterations=99)
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ payload = agent.build_run_payload(
+ query="q",
+ execution_params={"max_iterations": 7},
+ )
+ categories = {w.category for w in caught}
+ assert DeprecationWarning in categories
+ assert UserWarning in categories
+ assert payload["executionParams"]["budget"]["maxIterations"] == 99
+ assert "maxIterations" not in payload["executionParams"]
+
+ def test_fold_combines_with_other_budget_fields(self):
+ agent = _create_agent()
+ agent.budget = Budget(max_cost=1.0)
+ with pytest.warns(DeprecationWarning):
+ payload = agent.build_run_payload(
+ query="q",
+ execution_params={"max_iterations": 7},
+ )
+ assert payload["executionParams"]["budget"] == {
+ "maxCost": 1.0,
+ "maxIterations": 7,
+ }
+
+
+class TestFromDictLegacyMaxIterations:
+ """Loading a backend agent with legacy ``maxIterations`` must NOT warn.
+
+ ``Agent.from_dict`` (used by ``aix.Agent.get(id)``) deserializes backend
+ agents that may still carry a legacy top-level ``maxIterations``. Loading an
+ agent the caller never configured must be silent, but the explicit
+ ``Agent(max_iterations=...)`` constructor must still warn.
+ """
+
+ def test_from_dict_with_max_iterations_emits_no_warning(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error") # any warning becomes an error
+ agent = Agent.from_dict({"id": "a", "name": "t", "maxIterations": 5})
+ assert agent.budget is not None
+ assert agent.budget.max_iterations == 5
+ assert agent.max_iterations is None
+
+ def test_constructor_with_max_iterations_still_warns(self):
+ with pytest.warns(DeprecationWarning, match="max_iterations"):
+ Agent(name="t", max_iterations=5)
+
+ def test_from_dict_folds_into_budget(self):
+ agent = Agent.from_dict({"id": "a", "name": "t", "maxIterations": 5})
+ assert agent.budget.max_iterations == 5
+
+ def test_from_dict_budget_wins_silently_on_conflict(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ agent = Agent.from_dict(
+ {
+ "id": "a",
+ "name": "t",
+ "maxIterations": 5,
+ "budget": {"maxIterations": 99, "maxCost": 1.0},
+ }
+ )
+ assert agent.budget.max_iterations == 99
+ assert agent.budget.max_cost == 1.0
+
+ def test_from_dict_without_max_iterations_empty_budget(self):
+ # ``budget`` always exists as a (never-None) empty Budget, even when the
+ # backend agent carries no budget and no legacy maxIterations.
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ agent = Agent.from_dict({"id": "a", "name": "t"})
+ assert agent.budget == Budget()
+ agent.context = MagicMock()
+ assert agent.build_save_payload().get("budget") is None
+
+ def test_from_dict_does_not_mutate_caller_dict(self):
+ data = {"id": "a", "name": "t", "maxIterations": 5}
+ Agent.from_dict(data)
+ assert data["maxIterations"] == 5 # original dict untouched
+ assert "budget" not in data
+
+ def test_from_dict_serializes_back_to_budget_only(self):
+ agent = Agent.from_dict({"id": "a", "name": "t", "maxIterations": 5})
+ agent.context = MagicMock()
+ payload = agent.build_save_payload()
+ assert payload["budget"] == {"maxIterations": 5}
+ assert "maxIterations" not in payload
+
+ def test_from_dict_preserves_nested_fields(self):
+ agent = Agent.from_dict(
+ {
+ "id": "a",
+ "name": "t",
+ "maxIterations": 3,
+ "tasks": [
+ {"name": "task1", "description": "desc", "expectedOutput": "o"}
+ ],
+ }
+ )
+ assert agent.budget.max_iterations == 3
+ assert len(agent.tasks) == 1
+ assert agent.tasks[0].name == "task1"
+
+
+class TestRunPathWarningStacklevel:
+ """Run-path deprecation/conflict warnings must point at the user's run() call.
+
+ The warnings are emitted deep inside ``build_run_payload`` but the stacklevel
+ is tuned so the displayed location is the caller's ``agent.run(...)`` line,
+ not SDK internals.
+ """
+
+ @staticmethod
+ def _runnable_agent():
+ agent = Agent.from_dict({"id": "a1", "name": "t"})
+ agent.context = MagicMock()
+ agent.status = None
+ fake_result = MagicMock()
+ fake_result.url = None
+ fake_result.completed = True
+ agent.handle_run_response = lambda response, **kw: fake_result
+ agent.before_run = lambda *a, **k: None
+ agent.after_run = lambda result, *a, **k: None
+ return agent
+
+ def test_run_deprecation_warning_points_at_call_site(self):
+ agent = self._runnable_agent()
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ agent.run(query="hi", execution_params={"max_iterations": 7})
+ dep = [w for w in caught if w.category is DeprecationWarning and "max_iterations" in str(w.message)]
+ assert dep, "expected a run-path DeprecationWarning"
+ # The warning must resolve to THIS test file (the user call site), not agent.py.
+ assert dep[0].filename == __file__, dep[0].filename
+
+ def test_run_conflict_warning_points_at_call_site(self):
+ agent = self._runnable_agent()
+ agent.budget = Budget(max_iterations=99)
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ agent.run(
+ query="hi",
+ execution_params={"max_iterations": 7},
+ )
+ conflict = [w for w in caught if w.category is UserWarning and "precedence" in str(w.message)]
+ assert conflict, "expected a run-path conflict UserWarning"
+ assert conflict[0].filename == __file__, conflict[0].filename
diff --git a/tests/unit/v2/test_agent_polling.py b/tests/unit/v2/test_agent_polling.py
index 9a93804e..9cd92e60 100644
--- a/tests/unit/v2/test_agent_polling.py
+++ b/tests/unit/v2/test_agent_polling.py
@@ -58,6 +58,53 @@ def test_diagnostic_error_codes_deserializes_from_backend_camel_case(self):
assert result.diagnostic_error_codes == ["MAX_TOKENS_REACHED", "TOOL_FAILED"]
+ def test_diagnostic_error_codes_promoted_from_data(self):
+ """The backend poll body carries the codes at data.diagnosticErrorCodes,
+ not top-level — they must surface on the result."""
+ result = AgentRunResult.from_dict(
+ {
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {
+ "input": "q",
+ "output": "truncated text",
+ "diagnosticErrorCodes": ["MAX_TOKENS_REACHED"],
+ "executionStats": {"diagnostic_error_codes": ["MAX_TOKENS_REACHED"]},
+ },
+ }
+ )
+
+ assert result.diagnostic_error_codes == ["MAX_TOKENS_REACHED"]
+ assert result.data.diagnostic_error_codes == ["MAX_TOKENS_REACHED"]
+
+ def test_diagnostic_error_codes_fall_back_to_execution_stats(self):
+ """Older backends only carry the codes inside executionStats."""
+ result = AgentRunResult.from_dict(
+ {
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {
+ "input": "q",
+ "output": "truncated text",
+ "executionStats": {"diagnostic_error_codes": ["MAX_TOKENS_REACHED"]},
+ },
+ }
+ )
+
+ assert result.diagnostic_error_codes == ["MAX_TOKENS_REACHED"]
+
+ def test_top_level_diagnostic_error_codes_win_over_data(self):
+ result = AgentRunResult.from_dict(
+ {
+ "status": "SUCCESS",
+ "completed": True,
+ "diagnosticErrorCodes": ["TOOL_FAILED"],
+ "data": {"diagnosticErrorCodes": ["MAX_TOKENS_REACHED"]},
+ }
+ )
+
+ assert result.diagnostic_error_codes == ["TOOL_FAILED"]
+
def test_execution_id_from_request_id(self):
"""Should return request_id directly when available."""
result = AgentRunResult(
@@ -222,6 +269,54 @@ def test_poll_does_not_use_sdk_runs_endpoint(self):
assert "/sdk/agents/" in actual_url
+class TestAgentRunRequestId:
+ """Tests for request ID preservation across sync run polling."""
+
+ def test_run_preserves_request_id_from_initial_poll_url_response(self):
+ """Should keep requestId from initial run response when poll omits it."""
+ agent = _create_agent()
+ agent.before_run = Mock(return_value=None)
+ request_id = "190bab47-30e2b0287e9b"
+ poll_url = f"{BACKEND_URL}/sdk/agents/{request_id}/result"
+ agent.context.client.request = Mock(
+ return_value={
+ "requestId": request_id,
+ "data": poll_url,
+ }
+ )
+ agent.context.client.get = Mock(
+ return_value={
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {"output": "Test received."},
+ }
+ )
+
+ result = agent.run("test")
+
+ assert result.request_id == request_id
+
+ def test_run_sends_identifier_as_x_user_id_header(self):
+ """Should send identifier in payload and as x-user-id header on the run request."""
+ agent = _create_agent()
+ agent.before_run = Mock(return_value=None)
+ identifier = "end-user-123"
+ agent.context.client.request = Mock(
+ return_value={
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {"output": "done"},
+ }
+ )
+
+ agent.run("test", identifier=identifier)
+
+ agent.context.client.request.assert_called_once()
+ _, _, request_kwargs = agent.context.client.request.mock_calls[0]
+ assert request_kwargs["headers"] == {"x-user-id": identifier}
+ assert request_kwargs["json"]["identifier"] == identifier
+
+
class TestAgentSyncPollWithExecutionId:
"""Tests for Agent.sync_poll() accepting execution IDs."""
diff --git a/tests/unit/v2/test_agent_session_tool_parameters.py b/tests/unit/v2/test_agent_session_tool_parameters.py
new file mode 100644
index 00000000..8400cb56
--- /dev/null
+++ b/tests/unit/v2/test_agent_session_tool_parameters.py
@@ -0,0 +1,194 @@
+"""Unit tests for tool-parameter overrides through the session run path (PROD-2481).
+
+Per-tool parameters are set by mutating the agent's tool objects (object API).
+``run(query, session=…)`` serializes the agent's *current* tool parameter state
+into the per-message override, matching the single-shot run path. The
+session-level ``tools`` dict kwargs that PR #967 added were removed in favor of
+this object-based flow.
+"""
+
+from typing import Any, Dict, List
+from unittest.mock import Mock
+
+import pytest
+
+from aixplain.v2.agent import Agent
+from aixplain.v2.model import Model, Parameter
+from aixplain.v2.session import ExecutionConfig, Session, SessionMessage
+from aixplain.v2.tool import Tool
+
+
+def _params_as_dict(parameters: List[dict]) -> Dict[str, Any]:
+ return {item["name"]: item["value"] for item in parameters}
+
+
+def _tool_by_id(tools: List[dict], tool_id: str) -> dict:
+ return next(t for t in tools if t.get("id") == tool_id or t.get("assetId") == tool_id)
+
+
+def _make_mock_context(**overrides):
+ client = Mock()
+ ctx = Mock(client=client, backend_url="https://platform-api.aixplain.com", api_key="test_key")
+ ctx.Model = Model
+ ctx.Tool = Tool
+ for k, v in overrides.items():
+ setattr(ctx, k, v)
+ return ctx
+
+
+def _bound_agent(ctx):
+ class BoundAgent(Agent):
+ context = ctx
+
+ return BoundAgent
+
+
+def _user_message(request_id="req_abc"):
+ return SessionMessage(
+ id="msg_user",
+ session_id="sess_abc",
+ role="user",
+ content="hi",
+ sequence=1,
+ request_id=request_id,
+ )
+
+
+def _model_tool(temperature: Any = None) -> Model:
+ model = Model(id="m2", name="translate", params=[Parameter(name="temperature", data_type="number")])
+ if temperature is not None:
+ model.inputs.temperature = temperature
+ return model
+
+
+# ---------------------------------------------------------------------------
+# ExecutionConfig no longer carries a session-level tools override.
+# ---------------------------------------------------------------------------
+
+
+class TestExecutionConfigNoTools:
+ def test_to_api_dict_has_no_tools(self):
+ assert "tools" not in ExecutionConfig(criteria="x").to_api_dict()
+
+ def test_execution_config_has_no_tools_field(self):
+ assert "tools" not in ExecutionConfig().__dict__
+
+ def test_session_rejects_tools_kwarg(self):
+ # Session carries no session-level tools override; per-tool params flow
+ # through the object API + the per-message tools override instead.
+ with pytest.raises(TypeError):
+ Session(agent_id="agent_99", tools=[{"id": "tool-1"}])
+
+
+# ---------------------------------------------------------------------------
+# add_message still accepts a per-message tools override in the POST payload.
+# ---------------------------------------------------------------------------
+
+
+class TestAddMessageTools:
+ def test_add_message_payload_carries_tools(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {"id": "m1", "role": "user", "content": "hi"}
+
+ session = Session(agent_id="agent_99", name="A")
+ session.context = ctx
+ session.id = "sess_abc"
+ session._update_saved_state()
+
+ tools = [{"id": "tool-1", "parameters": [{"name": "top_k", "value": 9}]}]
+ session.add_message(role="user", content="hi", tools=tools)
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["json"]["tools"] == tools
+
+ def test_add_message_without_tools_omits_key(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {"id": "m1", "role": "user", "content": "hi"}
+
+ session = Session(agent_id="agent_99", name="A")
+ session.context = ctx
+ session.id = "sess_abc"
+ session._update_saved_state()
+
+ session.add_message(role="user", content="hi")
+ _, kwargs = ctx.client.request.call_args
+ assert "tools" not in kwargs["json"]
+
+
+# ---------------------------------------------------------------------------
+# run(query, session=…) forwards the agent's current tool params as the override.
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionTools:
+ def _wire_poll_success(self, ctx, request_id):
+ ctx.client.get.return_value = {
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {"output": "ok", "session_id": "sess_new", "steps": []},
+ "sessionId": "sess_new",
+ "requestId": request_id,
+ "usedCredits": 0.0,
+ "runTime": 0.5,
+ }
+
+ def _bound_session(self, ctx, add_calls):
+ class BoundSession(Session):
+ context = ctx
+
+ def fake_add_message(self, role, content, **kw):
+ add_calls.append(kw)
+ return _user_message(request_id="req_xyz")
+
+ BoundSession.add_message = fake_add_message
+ ctx.Session = BoundSession
+ session = BoundSession(agent_id="agent_99", name="A")
+ session.id = "sess_new"
+ return session
+
+ def test_session_run_routes_agent_tool_params(self):
+ ctx = _make_mock_context()
+ add_calls: List[dict] = []
+ session = self._bound_session(ctx, add_calls)
+ self._wire_poll_success(ctx, "req_xyz")
+
+ agent = _bound_agent(ctx)(id="agent_99", name="A", tools=[_model_tool(temperature=0.7)])
+ agent._update_saved_state()
+
+ agent.run("hi", session=session)
+
+ assert len(add_calls) == 1
+ tool = _tool_by_id(add_calls[0]["tools"], "m2")
+ assert _params_as_dict(tool["parameters"]) == {"temperature": 0.7}
+
+ def test_session_run_without_tool_params_passes_none(self):
+ ctx = _make_mock_context()
+ add_calls: List[dict] = []
+ session = self._bound_session(ctx, add_calls)
+ self._wire_poll_success(ctx, "req_xyz")
+
+ agent = _bound_agent(ctx)(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ agent.run("hi", session=session)
+ assert add_calls[0].get("tools") is None
+
+
+# ---------------------------------------------------------------------------
+# Single-shot (non-session) object-API behavior.
+# ---------------------------------------------------------------------------
+
+
+class TestSingleShot:
+ def test_run_payload_carries_object_tool_params(self):
+ agent = Agent(name="n", description="d", tools=[_model_tool(temperature=0.8)])
+ agent.context = Mock()
+ agent.id = "agent-1"
+ payload = agent.build_run_payload(query="hi")
+ assert _params_as_dict(_tool_by_id(payload["tools"], "m2")["parameters"]) == {"temperature": 0.8}
+
+ def test_save_payload_carries_object_tool_params(self):
+ agent = Agent(name="n", description="d", tools=[_model_tool(temperature=0.5)])
+ agent.context = Mock()
+ payload = agent.build_save_payload()
+ assert _params_as_dict(_tool_by_id(payload["tools"], "m2")["parameters"]) == {"temperature": 0.5}
diff --git a/tests/unit/v2/test_agent_tool_parameters.py b/tests/unit/v2/test_agent_tool_parameters.py
new file mode 100644
index 00000000..135e8822
--- /dev/null
+++ b/tests/unit/v2/test_agent_tool_parameters.py
@@ -0,0 +1,285 @@
+"""Unit tests for object-based per-tool parameter overrides (PROD-2481).
+
+Per-tool parameters are configured by mutating a tool object's typed action
+inputs and attaching it to the agent::
+
+ tool = aix.Tool.get("aixplain/aixplain-web-search")
+ tool.actions.search.inputs.num_results = 2
+ agent = aix.Agent(name="my-agent", tools=[tool])
+ agent.save() # persists num_results=2
+
+ agent.tools[0].actions.search.inputs.num_results = 5
+ agent.run("query") # ephemeral run-time override
+
+These tests cover:
+
+* Hydration — a tool dict from a get()/create response becomes a mutable
+ ``Tool``/``Model`` object, offline (no network).
+* Save payload picks up the current input values.
+* Run payload forwards the current values as an ephemeral override.
+* Mutating a tool input does not mark the agent modified (so an onboarded
+ agent's run is not blocked); adding/removing a whole tool does.
+* Attach-only shapes (string id, ``as_tool()`` snapshot dict) still serialize.
+"""
+
+from typing import Any, Dict, List
+from unittest.mock import Mock, patch
+
+from aixplain.v2.agent import Agent
+from aixplain.v2.model import Model, Parameter
+from aixplain.v2.tool import Tool
+
+
+def _params_as_dict(parameters: List[dict]) -> Dict[str, Any]:
+ """Flatten ``[{name, value}]`` to ``{name: value}`` for ergonomic assertions."""
+ return {item["name"]: item["value"] for item in parameters}
+
+
+def _tool_by_id(tools: List[dict], tool_id: str) -> dict:
+ return next(t for t in tools if t.get("id") == tool_id or t.get("assetId") == tool_id)
+
+
+def _ctx() -> Mock:
+ """A mock client context that exposes the real Tool/Model resource classes."""
+ ctx = Mock(client=Mock(), backend_url="https://platform-api.aixplain.com", api_key="k")
+ ctx.Model = Model
+ ctx.Tool = Tool
+ return ctx
+
+
+def _bound_agent(ctx: Mock):
+ class BoundAgent(Agent):
+ context = ctx
+
+ return BoundAgent
+
+
+def _model_tool(temperature: Any = None) -> Model:
+ model = Model(id="m2", name="translate", params=[Parameter(name="temperature", data_type="number")])
+ if temperature is not None:
+ model.inputs.temperature = temperature
+ return model
+
+
+# ---------------------------------------------------------------------------
+# Hydration: response tool dicts become mutable objects, offline.
+# ---------------------------------------------------------------------------
+
+
+class TestHydration:
+ def test_model_tool_dict_hydrated_offline(self):
+ ctx = _ctx()
+ tool_dict = {
+ "id": "m2",
+ "type": "model",
+ "name": "translate",
+ "parameters": [{"name": "temperature", "value": 0.5, "datatype": "number", "required": False}],
+ }
+ agent = _bound_agent(ctx).from_dict({"name": "n", "description": "d", "tools": [tool_dict]})
+
+ assert isinstance(agent.tools[0], Model)
+ assert agent.tools[0].inputs.temperature.value == 0.5
+ # Hydration must not hit the network.
+ ctx.client.get.assert_not_called()
+ ctx.client.request.assert_not_called()
+
+ def test_action_tool_dict_hydrated_with_dot_access(self):
+ ctx = _ctx()
+ tool_dict = {
+ "id": "t-web",
+ "type": "tool",
+ "name": "web-search",
+ "parameters": [
+ {
+ "code": "SEARCH",
+ "name": "search",
+ "description": "",
+ "inputs": {"num_results": {"name": "num_results", "value": 2, "datatype": "integer"}},
+ }
+ ],
+ }
+ agent = _bound_agent(ctx).from_dict({"name": "n", "description": "d", "tools": [tool_dict]})
+
+ assert isinstance(agent.tools[0], Tool)
+ assert agent.tools[0].actions.search.inputs.num_results.value == 2
+ ctx.client.request.assert_not_called()
+
+ def test_object_tool_passed_through_unchanged(self):
+ ctx = _ctx()
+ model = _model_tool()
+ agent = _bound_agent(ctx)(name="n", description="d", tools=[model])
+ assert agent.tools[0] is model
+
+ def test_string_id_kept_as_is(self):
+ ctx = _ctx()
+ agent = _bound_agent(ctx)(name="n", description="d", tools=["tool-1"])
+ assert agent.tools[0] == "tool-1"
+
+
+# ---------------------------------------------------------------------------
+# Save payload reflects the tool object's current input values.
+# ---------------------------------------------------------------------------
+
+
+class TestSavePayload:
+ def test_save_reflects_mutated_model_input(self):
+ agent = Agent(name="n", description="d", tools=[_model_tool(temperature=0.7)])
+ agent.context = Mock()
+ tool = _tool_by_id(agent.build_save_payload()["tools"], "m2")
+ assert _params_as_dict(tool["parameters"]) == {"temperature": 0.7}
+
+ def test_save_persists_latest_value(self):
+ model = _model_tool(temperature=0.3)
+ agent = Agent(name="n", description="d", tools=[model])
+ agent.context = Mock()
+ first = _tool_by_id(agent.build_save_payload()["tools"], "m2")
+ assert _params_as_dict(first["parameters"]) == {"temperature": 0.3}
+
+ model.inputs.temperature = 0.9
+ second = _tool_by_id(agent.build_save_payload()["tools"], "m2")
+ assert _params_as_dict(second["parameters"]) == {"temperature": 0.9}
+
+ def test_save_does_not_crash_on_object_tools(self):
+ # Regression: to_dict() used to recurse into Model/Tool objects and raise.
+ agent = Agent(name="n", description="d", tools=[_model_tool()])
+ agent.context = Mock()
+ payload = agent.build_save_payload()
+ assert _tool_by_id(payload["tools"], "m2")["id"] == "m2"
+
+
+# ---------------------------------------------------------------------------
+# Run payload forwards the current values as an ephemeral override.
+# ---------------------------------------------------------------------------
+
+
+class TestRunPayload:
+ def test_run_sends_current_model_params(self):
+ agent = Agent(name="n", description="d", tools=[_model_tool(temperature=0.7)])
+ agent.context = Mock()
+ agent.id = "agent-1"
+ payload = agent.build_run_payload(query="hi")
+ assert _params_as_dict(_tool_by_id(payload["tools"], "m2")["parameters"]) == {"temperature": 0.7}
+
+ def test_run_sends_action_tool_params(self):
+ ctx = _ctx()
+ tool_dict = {
+ "id": "t-web",
+ "type": "tool",
+ "name": "web-search",
+ "parameters": [
+ {"code": "SEARCH", "name": "search", "inputs": {"num_results": {"name": "num_results", "value": 2}}}
+ ],
+ }
+ agent = _bound_agent(ctx).from_dict({"name": "n", "description": "d", "tools": [tool_dict]})
+ agent.id = "agent-1"
+ agent.tools[0].actions.search.inputs.num_results = 5
+ payload = agent.build_run_payload(query="hi")
+ assert _params_as_dict(_tool_by_id(payload["tools"], "t-web")["parameters"]) == {"num_results": 5}
+
+ def test_run_without_values_omits_tools_key(self):
+ agent = Agent(name="n", description="d", tools=[_model_tool()]) # no value set
+ agent.context = Mock()
+ agent.id = "agent-1"
+ assert "tools" not in agent.build_run_payload(query="hi")
+
+ def test_run_no_tools_omits_key(self):
+ agent = Agent(name="n", description="d")
+ agent.context = Mock()
+ agent.id = "agent-1"
+ assert "tools" not in agent.build_run_payload(query="hi")
+
+
+# ---------------------------------------------------------------------------
+# is_modified: input mutation is ephemeral; tool add/remove is a change.
+# ---------------------------------------------------------------------------
+
+
+class TestIsModified:
+ def _onboarded_agent_with_model_tool(self, ctx):
+ tool_dict = {
+ "id": "m2",
+ "type": "model",
+ "name": "translate",
+ "parameters": [{"name": "temperature", "value": 0.5, "datatype": "number"}],
+ }
+ agent = _bound_agent(ctx).from_dict(
+ {"name": "n", "description": "d", "tools": [tool_dict], "status": "onboarded"}
+ )
+ agent._update_saved_state()
+ return agent
+
+ def test_input_mutation_not_modified(self):
+ agent = self._onboarded_agent_with_model_tool(_ctx())
+ assert agent.is_modified is False
+ agent.tools[0].inputs.temperature = 0.9
+ assert agent.is_modified is False
+
+ def test_adding_a_tool_is_modified(self):
+ agent = self._onboarded_agent_with_model_tool(_ctx())
+ agent.tools = agent.tools + [_model_tool()]
+ assert agent.is_modified is True
+
+ def test_not_modified_after_save_so_onboarded_run_is_allowed(self):
+ # Regression: save() restored the caller's tool objects but did not
+ # re-baseline the saved state, so is_modified read True and a following
+ # run() on an onboarded agent wrongly raised.
+ ctx = _ctx()
+
+ def fake_request(method, path, **kw):
+ if method == "post":
+ return {
+ "id": "agent-1",
+ "name": "a",
+ "description": "d",
+ "status": "onboarded",
+ "tools": [{"id": "m2", "type": "model", "name": "translate"}],
+ }
+ return {}
+
+ ctx.client.request.side_effect = fake_request
+ ctx.client.get.side_effect = lambda *a, **k: {}
+
+ agent = _bound_agent(ctx)(name="a", description="d", tools=[_model_tool(temperature=2)])
+ agent.save()
+
+ assert agent.is_modified is False
+
+
+# ---------------------------------------------------------------------------
+# Attach-only shapes (no per-tool override) still serialize for create.
+# ---------------------------------------------------------------------------
+
+
+class TestAttachShapes:
+ def test_string_id_resolves_type(self):
+ snap = {"id": "model-1", "type": "model", "name": "Translate"}
+ with patch.object(Agent, "_resolve_tool_snapshot", return_value=snap):
+ entry = Agent._normalize_tool_for_api("model-1")
+ assert entry["type"] == "model"
+ assert entry["id"] == "model-1"
+
+ def test_unresolvable_string_id_falls_back_to_bare_entry(self):
+ with patch.object(Agent, "_resolve_tool_snapshot", return_value=None):
+ entry = Agent._normalize_tool_for_api("tool-x")
+ assert entry == {"id": "tool-x"}
+
+ def test_as_tool_snapshot_dict_passthrough(self):
+ as_tool_dict = {
+ "id": "tool-1",
+ "asset_id": "tool-1",
+ "type": "model",
+ "parameters": [
+ {"name": "temperature", "value": "0.5", "allow_multi": False, "supports_variables": True}
+ ],
+ }
+ # A dict that already carries a type must not trigger asset resolution.
+ with patch.object(Agent, "_resolve_tool_snapshot", side_effect=AssertionError("must not resolve")):
+ agent = Agent(name="n", description="d", tools=[as_tool_dict])
+ agent.context = Mock()
+ entry = _tool_by_id(agent.build_save_payload()["tools"], "tool-1")
+ assert entry["assetId"] == "tool-1"
+ param = entry["parameters"][0]
+ assert param["name"] == "temperature"
+ assert param["value"] == "0.5"
+ assert param["allowMulti"] is False
+ assert param["supportsVariables"] is True
diff --git a/tests/unit/v2/test_agent_via_session.py b/tests/unit/v2/test_agent_via_session.py
new file mode 100644
index 00000000..57adace9
--- /dev/null
+++ b/tests/unit/v2/test_agent_via_session.py
@@ -0,0 +1,445 @@
+"""Unit tests for the `session=` run path on Agent.run().
+
+Passing `session=` (a Session object or id string) posts a user message to that
+Session (carrying its `executionConfig`) and then polls the
+`/sdk/agents/{request_id}/result` endpoint using the request_id the backend
+stamps on the user message. That preserves the full `AgentRunResult` shape
+(steps, execution_stats, used_credits, run_time) while routing the trigger
+through sessions. There is no `via_session` flag and no id-only `session_id=`;
+threads are managed through `aix.Session` and passed here.
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+from aixplain.v2.agent import Agent, AgentRunResult
+from aixplain.v2.session import ExecutionConfig, Session, SessionMessage
+
+
+def _make_mock_context(**overrides):
+ client = Mock()
+ ctx = Mock(
+ client=client,
+ backend_url="https://platform-api.aixplain.com",
+ api_key="test_key",
+ )
+ for k, v in overrides.items():
+ setattr(ctx, k, v)
+ return ctx
+
+
+def _bound_agent(ctx):
+ class BoundAgent(Agent):
+ context = ctx
+
+ return BoundAgent
+
+
+def _user_message(request_id="req_abc", id_="msg_user"):
+ """A user SessionMessage as the backend returns from POST .../messages."""
+ return SessionMessage(
+ id=id_,
+ session_id="sess_abc",
+ user_id="u1",
+ agent_id="agent_99",
+ role="user",
+ content="hi",
+ sequence=1,
+ request_id=request_id,
+ created_at="2025-06-01T10:01:00Z",
+ )
+
+
+def _success_result(session_id="sess_abc", request_id="req_abc"):
+ return {
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {"output": "hi back!", "session_id": session_id, "steps": [{"agent": "agent_99", "thought": "ok"}]},
+ "sessionId": session_id,
+ "requestId": request_id,
+ "usedCredits": 0.42,
+ "runTime": 1.7,
+ }
+
+
+# ---------------------------------------------------------------------------
+# 1. session= a Session object → post to it, then poll the run result
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionObject:
+ def test_session_object_posts_and_polls(self):
+ ctx = _make_mock_context()
+
+ class BoundSession(Session):
+ context = ctx
+
+ add_calls = []
+
+ def fake_add_message(self, role, content, **kw):
+ add_calls.append({"role": role, "content": content, "kwargs": kw})
+ return _user_message(request_id="req_xyz")
+
+ BoundSession.add_message = fake_add_message
+ ctx.Session = BoundSession
+ ctx.client.get.return_value = _success_result(session_id="sess_obj", request_id="req_xyz")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ session = BoundSession(agent=agent, name="thread")
+ session.id = "sess_obj"
+
+ result = agent.run("hi", session=session)
+
+ # The user message was POSTed to the passed-in session.
+ assert len(add_calls) == 1
+ assert add_calls[0]["role"] == "user"
+ assert add_calls[0]["content"] == "hi"
+
+ # The run-result endpoint was hit with the user message's request_id.
+ get_call = ctx.client.get.call_args_list[0]
+ assert "/sdk/agents/req_xyz/result" in get_call[0][0]
+
+ # Result preserves the full result shape.
+ assert isinstance(result, AgentRunResult)
+ assert result.status == "SUCCESS"
+ assert result.completed is True
+ assert result.session_id == "sess_obj"
+ assert result.request_id == "req_xyz"
+ assert result.data.output == "hi back!"
+ assert result.data.steps == [{"agent": "agent_99", "thought": "ok"}]
+ assert result.used_credits == 0.42
+ assert result.run_time == 1.7
+
+ def test_session_object_without_context_is_bound(self):
+ """A Session with no context is bound to the agent's context on run."""
+ ctx = _make_mock_context()
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ session = Session(agent_id="agent_99", name="unbound")
+ session.id = "sess_unbound"
+ session.add_message = Mock(return_value=_user_message(request_id="req_bound"))
+ assert getattr(session, "context", None) is None
+
+ ctx.client.get.return_value = _success_result(session_id="sess_unbound", request_id="req_bound")
+ agent.run("hi", session=session)
+
+ assert session.context is ctx
+ session.add_message.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# 2. session= an id string → fetched via Session.get
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionId:
+ def test_session_id_fetches_via_get(self):
+ ctx = _make_mock_context()
+
+ existing = Mock(spec=Session)
+ existing.id = "sess_abc"
+ existing.context = ctx
+ existing.execution_config = None
+ existing.add_message = Mock(return_value=_user_message(request_id="req_reuse"))
+
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ ctx.client.get.return_value = _success_result(session_id="sess_abc", request_id="req_reuse")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ result = agent.run("hi again", session="sess_abc")
+
+ ctx.Session.get.assert_called_once_with("sess_abc")
+ existing.add_message.assert_called_once()
+ get_call = ctx.client.get.call_args_list[0]
+ assert "/sdk/agents/req_reuse/result" in get_call[0][0]
+ assert result.session_id == "sess_abc"
+
+ def test_bad_session_type_raises(self):
+ ctx = _make_mock_context()
+ ctx.Session = Mock()
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ with pytest.raises(TypeError, match="session must be a Session instance or a session id"):
+ agent.run("hi", session=123)
+
+
+# ---------------------------------------------------------------------------
+# 3. Per-run execution overrides merge onto the session's config
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionOverrides:
+ def test_applies_per_run_overrides_and_warns(self):
+ """Per-run execution kwargs on a session must take effect.
+
+ They are merged onto the session's stored executionConfig and the
+ session is re-saved, with a warning that the config is mutated for
+ all subsequent messages.
+ """
+ ctx = _make_mock_context()
+
+ existing = Mock(spec=Session)
+ existing.id = "sess_abc"
+ existing.context = ctx
+ existing.execution_config = ExecutionConfig(execution_params={"max_tokens": 64})
+ existing.save = Mock()
+ existing.add_message = Mock(return_value=_user_message(request_id="req_override"))
+
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ ctx.client.get.return_value = _success_result(session_id="sess_abc", request_id="req_override")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ with pytest.warns(UserWarning, match="session 'sess_abc'"):
+ agent.run("hi again", session="sess_abc", criteria="be brief")
+
+ existing.save.assert_called_once()
+ assert existing.execution_config.execution_params == {"max_tokens": 64}
+ assert existing.execution_config.criteria == "be brief"
+ existing.add_message.assert_called_once()
+
+ def test_without_overrides_does_not_resave(self):
+ ctx = _make_mock_context()
+
+ existing = Mock(spec=Session)
+ existing.id = "sess_abc"
+ existing.context = ctx
+ existing.execution_config = ExecutionConfig(criteria="be brief")
+ existing.save = Mock()
+ existing.add_message = Mock(return_value=_user_message(request_id="req_noop"))
+
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ ctx.client.get.return_value = _success_result(session_id="sess_abc", request_id="req_noop")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ agent.run("hi again", session="sess_abc")
+ existing.save.assert_not_called()
+
+ def test_matching_overrides_does_not_resave(self):
+ ctx = _make_mock_context()
+
+ existing = Mock(spec=Session)
+ existing.id = "sess_abc"
+ existing.context = ctx
+ existing.execution_config = ExecutionConfig(criteria="be brief")
+ existing.save = Mock()
+ existing.add_message = Mock(return_value=_user_message(request_id="req_same"))
+
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ ctx.client.get.return_value = _success_result(session_id="sess_abc", request_id="req_same")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ import warnings as _warnings
+
+ with _warnings.catch_warnings():
+ _warnings.simplefilter("error")
+ agent.run("hi again", session="sess_abc", criteria="be brief")
+ existing.save.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# 4. Attachments / files / audio-as-prompt
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionAttachments:
+ def test_forwards_attachments_and_files_to_user_message(self):
+ ctx = _make_mock_context()
+
+ existing = Mock(spec=Session)
+ existing.id = "sess_att"
+ existing.context = ctx
+ existing.execution_config = None
+ existing.add_message = Mock(return_value=_user_message(request_id="req_att"))
+
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ ctx.client.get.return_value = _success_result(session_id="sess_att", request_id="req_att")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ attachments = [{"url": "https://s3/a.png", "type": "image"}]
+ files = ["/tmp/report.pdf"]
+ agent.run("describe these", session="sess_att", attachments=attachments, files=files)
+
+ _, kwargs = existing.add_message.call_args
+ assert kwargs["attachments"] == attachments
+ assert kwargs["files"] == files
+
+ def test_audio_as_prompt_allows_missing_query(self):
+ # Audio-as-prompt: no text query, attachments carry the turn's input.
+ ctx = _make_mock_context()
+
+ existing = Mock(spec=Session)
+ existing.id = "sess_aud"
+ existing.context = ctx
+ existing.execution_config = None
+ existing.add_message = Mock(return_value=_user_message(request_id="req_aud"))
+
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ ctx.client.get.return_value = _success_result(session_id="sess_aud", request_id="req_aud")
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ attachments = [{"url": "https://s3/a.wav", "type": "audio"}]
+ agent.run(session="sess_aud", attachments=attachments)
+
+ _, kwargs = existing.add_message.call_args
+ assert kwargs["content"] == "" # empty query coerced to empty content
+ assert kwargs["attachments"] == attachments
+
+ def test_raises_when_no_query_and_no_attachments(self):
+ ctx = _make_mock_context()
+ existing = Mock(spec=Session)
+ existing.id = "sess_x"
+ existing.context = ctx
+ existing.execution_config = None
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ with pytest.raises(ValueError, match="requires a query or attachments"):
+ agent.run(session="sess_x")
+
+
+# ---------------------------------------------------------------------------
+# 5. Missing requestId on the user message → ValueError (defensive)
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionMissingRequestId:
+ def test_raises_when_user_message_has_no_request_id(self):
+ ctx = _make_mock_context()
+
+ class BoundSession(Session):
+ context = ctx
+
+ BoundSession.add_message = lambda self, role, content, **kw: _user_message(request_id=None)
+ ctx.Session = BoundSession
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ session = BoundSession(agent_id="agent_99")
+ session.id = "sess_new"
+
+ with pytest.raises(ValueError, match="requestId"):
+ agent.run("hi", session=session)
+
+
+# ---------------------------------------------------------------------------
+# 6. Rejects legacy-only kwargs when running with a session
+# ---------------------------------------------------------------------------
+
+
+class TestRunWithSessionRejectsLegacyKwargs:
+ @pytest.mark.parametrize(
+ "legacy_kwarg",
+ ["tasks", "prompt", "inspectors", "history", "variables"],
+ )
+ def test_rejects_legacy_kwargs(self, legacy_kwarg):
+ ctx = _make_mock_context()
+ existing = Mock(spec=Session)
+ existing.id = "sess_x"
+ existing.context = ctx
+ existing.execution_config = None
+ ctx.Session = Mock()
+ ctx.Session.get = Mock(return_value=existing)
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ with pytest.raises(ValueError, match=legacy_kwarg):
+ agent.run("hi", session="sess_x", **{legacy_kwarg: ["something"]})
+
+
+# ---------------------------------------------------------------------------
+# 7. run_async with session= is not implemented
+# ---------------------------------------------------------------------------
+
+
+class TestRunAsyncWithSessionNotImplemented:
+ def test_run_async_with_session_not_implemented(self):
+ ctx = _make_mock_context()
+ ctx.Session = Mock()
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ with pytest.raises(NotImplementedError, match="session"):
+ agent.run_async("hi", session="sess_x")
+
+
+# ---------------------------------------------------------------------------
+# 8. Default path (no session) still hits /v2/agents/{id}/run
+# ---------------------------------------------------------------------------
+
+
+class TestDefaultPathUnchanged:
+ def test_run_default_path_hits_direct_endpoint(self):
+ ctx = _make_mock_context()
+ # Direct run: POST returns a polling URL, then GET resolves to a SUCCESS.
+ ctx.client.request.return_value = {
+ "status": "IN_PROGRESS",
+ "data": "https://platform-api.aixplain.com/sdk/agents/exec_42/result",
+ }
+ ctx.client.get.return_value = {
+ "status": "SUCCESS",
+ "completed": True,
+ "data": {"output": "direct reply", "session_id": None, "steps": []},
+ "sessionId": None,
+ "requestId": "req_direct",
+ "usedCredits": 0.5,
+ "runTime": 1.2,
+ }
+
+ BoundAgent = _bound_agent(ctx)
+ agent = BoundAgent(id="agent_99", name="A")
+ agent._update_saved_state()
+
+ with patch("time.sleep"):
+ result = agent.run("hi")
+
+ # Direct POST went to v2/agents/{id}/run.
+ post_call = ctx.client.request.call_args_list[0]
+ assert post_call[0][0] == "post"
+ assert "v2/agents/agent_99/run" in post_call[0][1]
+
+ # Direct poll URL was used.
+ get_call = ctx.client.get.call_args_list[0]
+ assert "/sdk/agents/exec_42/result" in get_call[0][0]
+
+ assert result.status == "SUCCESS"
+ assert result.data.output == "direct reply"
+ assert result.used_credits == 0.5
diff --git a/tests/unit/v2/test_exceptions.py b/tests/unit/v2/test_exceptions.py
index 90bd6b03..6e404777 100644
--- a/tests/unit/v2/test_exceptions.py
+++ b/tests/unit/v2/test_exceptions.py
@@ -332,6 +332,23 @@ def test_extracts_error(self):
assert "Connection refused" in str(error)
+ def test_extracts_nested_data_error(self):
+ """Should extract 'data.error' when documented error fields are empty."""
+ response = {
+ "status": "FAILED",
+ "completed": True,
+ "errorMessage": None,
+ "supplierError": None,
+ "data": {
+ "error": "litellm.APIConnectionError: err.not_enough_balance\nTraceback...",
+ },
+ }
+
+ error = create_operation_failed_error(response)
+
+ assert "err.not_enough_balance" in str(error)
+ assert "err.not_enough_balance" in error.error
+
def test_fallback_message(self):
"""Should use 'Operation failed' as fallback when no error field."""
response = {
diff --git a/tests/unit/v2/test_expected_output_persistence.py b/tests/unit/v2/test_expected_output_persistence.py
new file mode 100644
index 00000000..221061cd
--- /dev/null
+++ b/tests/unit/v2/test_expected_output_persistence.py
@@ -0,0 +1,88 @@
+"""Regression tests: Pydantic ``expected_output`` must persist through ``save()``.
+
+A ``BaseModel``-class ``expected_output`` used to be popped from the save
+payload ("runtime-only"), so the backend stored ``null``. The create-response
+re-hydration in ``BaseResource._create`` then overwrote the local value with
+that ``null``, and every subsequent ``run()`` sent ``expectedOutput: null`` —
+the engine never received the JSON contract, and fetched agents could never
+recover it.
+"""
+
+import json
+from unittest.mock import MagicMock
+from typing import Optional
+
+from pydantic import BaseModel
+
+from aixplain.v2.agent import Agent
+
+
+class ChatReply(BaseModel):
+ content: str
+ artifact: Optional[dict] = None
+
+
+def _json_agent(**overrides):
+ """Unsaved JSON-format agent with a mocked client context."""
+ kwargs = {
+ "name": "Strict JSON Agent",
+ "instructions": "Reply as a JSON array.",
+ "output_format": "json",
+ "expected_output": ChatReply,
+ }
+ kwargs.update(overrides)
+ agent = Agent(**kwargs)
+ agent.context = MagicMock()
+ return agent
+
+
+def _echoing_client(agent):
+ """Make the mocked client echo the create payload back, like the backend.
+
+ The backend persists exactly what it receives and returns the stored
+ document; a field missing from the request comes back as ``None``.
+ """
+
+ def fake_request(method, path, json=None, **kwargs):
+ response = dict(json or {})
+ response.setdefault("id", "agent-123")
+ response.setdefault("expectedOutput", None)
+ return response
+
+ agent.context.client.request.side_effect = fake_request
+
+
+class TestSavePayloadPersistsExpectedOutput:
+ def test_basemodel_class_is_serialized_to_json_schema_string(self):
+ payload = _json_agent().build_save_payload()
+
+ sent = payload.get("expectedOutput")
+ assert isinstance(sent, str), "BaseModel-class expected_output must be persisted as a JSON string"
+ assert json.loads(sent) == ChatReply.model_json_schema()
+
+ def test_basemodel_instance_is_still_dumped_to_dict(self):
+ instance = ChatReply(content="hi", artifact=None)
+ payload = _json_agent(expected_output=instance).build_save_payload()
+
+ assert payload.get("expectedOutput") == instance.model_dump()
+
+
+class TestSaveDoesNotLoseExpectedOutput:
+ def test_expected_output_survives_save_rehydration(self):
+ agent = _json_agent()
+ _echoing_client(agent)
+
+ agent.save()
+
+ assert agent.expected_output is not None, "save() must not wipe expected_output from the local agent"
+
+ def test_run_payload_carries_schema_after_save(self):
+ agent = _json_agent()
+ _echoing_client(agent)
+ agent.save()
+
+ run_payload = agent.build_run_payload(query="hi")
+
+ sent = run_payload["executionParams"]["expectedOutput"]
+ assert sent is not None, "run() after save() must send the JSON contract to the backend"
+ assert json.loads(sent) == ChatReply.model_json_schema()
diff --git a/tests/unit/v2/test_inspector.py b/tests/unit/v2/test_inspector.py
index eb35ce5e..ecc75dd5 100644
--- a/tests/unit/v2/test_inspector.py
+++ b/tests/unit/v2/test_inspector.py
@@ -1,149 +1,121 @@
-"""Unit tests for v2 Inspector payload serialization."""
+"""Unit tests for v2 Inspector payload serialization.
+
+The public surface is intentionally tiny: an :class:`Inspector` built with plain
+strings for ``action`` / ``targets`` / ``severity`` and a Metric / asset-id /
+callable for ``metric`` (the universal judge). No enums or config classes.
+"""
import pytest
from aixplain.v2.inspector import (
Inspector,
- InspectorAction,
- InspectorActionConfig,
- InspectorOnExhaust,
- InspectorSeverity,
- InspectorTarget,
- EvaluatorType,
- EvaluatorConfig,
- EditorConfig,
- PrebuiltInspector,
- is_prebuilt_inspector,
+ _ActionConfig,
+ _Judge,
+ _guard_slugs,
+ _resolve_guard_defaults,
)
+class _FakeMetric:
+ """Duck-typed stand-in for ``aix.Metric`` — an asset-backed judge with an id."""
+
+ def __init__(self, id, prompt_template=None):
+ self.id = id
+ self.prompt_template = prompt_template
+
+
# ---------------------------------------------------------------------------
-# InspectorActionConfig
+# Action policy (string / dict, no InspectorActionConfig import)
# ---------------------------------------------------------------------------
-class TestInspectorActionConfig:
- def test_abort_to_dict(self):
- cfg = InspectorActionConfig(type=InspectorAction.ABORT)
+class TestActionConfig:
+ def test_abort_from_string(self):
+ cfg = _ActionConfig.coerce("abort")
assert cfg.to_dict() == {"type": "abort"}
- def test_rerun_to_dict(self):
- cfg = InspectorActionConfig(
- type=InspectorAction.RERUN,
- max_retries=3,
- on_exhaust=InspectorOnExhaust.ABORT,
- )
+ def test_rerun_from_dict(self):
+ cfg = _ActionConfig.coerce({"type": "rerun", "max_retries": 3, "on_exhaust": "abort"})
assert cfg.to_dict() == {"type": "rerun", "maxRetries": 3, "onExhaust": "abort"}
+ def test_rerun_accepts_camelcase(self):
+ cfg = _ActionConfig.coerce({"type": "rerun", "maxRetries": 1, "onExhaust": "continue"})
+ assert cfg.to_dict() == {"type": "rerun", "maxRetries": 1, "onExhaust": "continue"}
+
def test_rerun_without_optional_fields(self):
- cfg = InspectorActionConfig(type=InspectorAction.RERUN)
- assert cfg.to_dict() == {"type": "rerun"}
+ assert _ActionConfig.coerce("rerun").to_dict() == {"type": "rerun"}
def test_max_retries_invalid_for_non_rerun(self):
with pytest.raises(ValueError, match="max_retries is only valid"):
- InspectorActionConfig(type=InspectorAction.ABORT, max_retries=2)
+ _ActionConfig.coerce({"type": "abort", "max_retries": 2})
def test_on_exhaust_invalid_for_non_rerun(self):
with pytest.raises(ValueError, match="on_exhaust is only valid"):
- InspectorActionConfig(type=InspectorAction.EDIT, on_exhaust=InspectorOnExhaust.ABORT)
+ _ActionConfig.coerce({"type": "edit", "on_exhaust": "abort"})
def test_negative_max_retries(self):
with pytest.raises(ValueError, match="max_retries must be >= 0"):
- InspectorActionConfig(type=InspectorAction.RERUN, max_retries=-1)
+ _ActionConfig.coerce({"type": "rerun", "max_retries": -1})
- def test_from_dict_roundtrip(self):
- original = InspectorActionConfig(
- type=InspectorAction.RERUN,
- max_retries=2,
- on_exhaust=InspectorOnExhaust.CONTINUE,
- )
- restored = InspectorActionConfig.from_dict(original.to_dict())
- assert restored.type == original.type
- assert restored.max_retries == original.max_retries
- assert restored.on_exhaust == original.on_exhaust
+ def test_invalid_action_type(self):
+ with pytest.raises(ValueError, match="invalid action"):
+ _ActionConfig.coerce("bogus")
+
+ def test_invalid_on_exhaust_value(self):
+ with pytest.raises(ValueError, match="invalid on_exhaust"):
+ _ActionConfig.coerce({"type": "rerun", "on_exhaust": "explode"})
# ---------------------------------------------------------------------------
-# EvaluatorConfig
+# Judge (the universal metric: Metric | asset-id str | callable)
# ---------------------------------------------------------------------------
-class TestEvaluatorConfig:
- def test_asset_to_dict(self):
- cfg = EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id="abc123",
- prompt="Check the output",
- )
- assert cfg.to_dict() == {
- "type": "asset",
- "assetId": "abc123",
- "prompt": "Check the output",
- }
-
- def test_function_to_dict(self):
- cfg = EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function="def eval_fn(text: str) -> bool:\n return True",
- )
- assert cfg.to_dict() == {
- "type": "function",
- "function": "def eval_fn(text: str) -> bool:\n return True",
- }
-
- def test_callable_auto_converted(self):
- def my_evaluator(text: str) -> bool:
- return "ok" in text
+class TestJudge:
+ def test_metric_object_becomes_asset(self):
+ judge = _Judge.coerce(_FakeMetric(id="metric-1", prompt_template="Judge it."))
+ assert judge.to_dict() == {"type": "asset", "assetId": "metric-1", "prompt": "Judge it."}
- cfg = EvaluatorConfig(type=EvaluatorType.FUNCTION, function=my_evaluator)
- assert isinstance(cfg.function, str)
- assert "my_evaluator" in cfg.function
+ def test_metric_prompt_falls_back_to_config(self):
+ class _MetricWithConfig:
+ id = "metric-2"
+ prompt_template = None
+ config = {"prompt": "From config."}
- def test_asset_requires_asset_id(self):
- with pytest.raises(ValueError, match="asset_id is required"):
- EvaluatorConfig(type=EvaluatorType.ASSET)
+ judge = _Judge.coerce(_MetricWithConfig())
+ assert judge.to_dict() == {"type": "asset", "assetId": "metric-2", "prompt": "From config."}
- def test_function_requires_function(self):
- with pytest.raises(ValueError, match="function is required"):
- EvaluatorConfig(type=EvaluatorType.FUNCTION)
-
- def test_from_dict_roundtrip(self):
- original = EvaluatorConfig(type=EvaluatorType.ASSET, asset_id="id1", prompt="p")
- restored = EvaluatorConfig.from_dict(original.to_dict())
- assert restored.type == original.type
- assert restored.asset_id == original.asset_id
- assert restored.prompt == original.prompt
+ def test_asset_id_string(self):
+ assert _Judge.coerce("abc123").to_dict() == {"type": "asset", "assetId": "abc123"}
+ def test_dict_snake_and_prompt(self):
+ judge = _Judge.coerce({"asset_id": "id1", "prompt": "p"})
+ assert judge.to_dict() == {"type": "asset", "assetId": "id1", "prompt": "p"}
-# ---------------------------------------------------------------------------
-# EditorConfig
-# ---------------------------------------------------------------------------
+ def test_callable_becomes_function(self):
+ def my_evaluator(text: str) -> bool:
+ return "ok" in text
+ judge = _Judge.coerce(my_evaluator)
+ assert judge.type == "function"
+ assert "my_evaluator" in judge.function
-class TestEditorConfig:
- def test_function_to_dict(self):
- cfg = EditorConfig(
- type=EvaluatorType.FUNCTION,
- function="def edit(text: str) -> str:\n return text",
- )
- assert cfg.to_dict() == {
- "type": "function",
- "function": "def edit(text: str) -> str:\n return text",
- }
+ def test_unonboarded_metric_raises(self):
+ class _Unonboarded:
+ id = None
+ prompt_template = "x"
- def test_callable_auto_converted(self):
- def my_editor(text: str) -> str:
- return text.upper()
+ with pytest.raises(ValueError, match="must be created"):
+ _Judge.coerce(_Unonboarded())
- cfg = EditorConfig(type=EvaluatorType.FUNCTION, function=my_editor)
- assert isinstance(cfg.function, str)
- assert "my_editor" in cfg.function
+ def test_empty_judge_raises(self):
+ with pytest.raises(ValueError, match="needs an asset id"):
+ _Judge()
def test_from_dict_roundtrip(self):
- original = EditorConfig(type=EvaluatorType.FUNCTION, function="def f(): pass")
- restored = EditorConfig.from_dict(original.to_dict())
- assert restored.type == original.type
- assert restored.function == original.function
+ original = _Judge.from_dict({"type": "asset", "assetId": "id1", "prompt": "p"})
+ restored = _Judge.from_dict(original.to_dict())
+ assert restored.to_dict() == original.to_dict()
# ---------------------------------------------------------------------------
@@ -156,17 +128,12 @@ def test_abort_inspector_payload(self):
inspector = Inspector(
name="abort_inspector",
description="Always abort",
- severity=InspectorSeverity.HIGH,
+ severity="high",
targets=["output"],
- action=InspectorActionConfig(type=InspectorAction.ABORT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id="llm-123",
- prompt="Critique the output.",
- ),
+ action="abort",
+ metric=_FakeMetric(id="llm-123", prompt_template="Critique the output."),
)
- payload = inspector.to_dict()
- assert payload == {
+ assert inspector.to_dict() == {
"name": "abort_inspector",
"description": "Always abort",
"severity": "high",
@@ -183,19 +150,10 @@ def test_rerun_inspector_payload(self):
inspector = Inspector(
name="rerun_inspector",
targets=["output"],
- action=InspectorActionConfig(
- type=InspectorAction.RERUN,
- max_retries=2,
- on_exhaust=InspectorOnExhaust.ABORT,
- ),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.ASSET,
- asset_id="llm-456",
- prompt="Check for name.",
- ),
+ action={"type": "rerun", "max_retries": 2, "on_exhaust": "abort"},
+ metric={"asset_id": "llm-456", "prompt": "Check for name."},
)
- payload = inspector.to_dict()
- assert payload == {
+ assert inspector.to_dict() == {
"name": "rerun_inspector",
"targets": ["output"],
"action": {"type": "rerun", "maxRetries": 2, "onExhaust": "abort"},
@@ -206,66 +164,59 @@ def test_rerun_inspector_payload(self):
},
}
- def test_edit_inspector_payload(self):
+ def test_edit_inspector_payload_with_callables(self):
+ def eval_fn(text):
+ return True
+
+ def edit_fn(text):
+ return "edited"
+
inspector = Inspector(
name="edit_inspector",
- severity=InspectorSeverity.MEDIUM,
+ severity="medium",
targets=["steps", "AgentA"],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(
- type=EvaluatorType.FUNCTION,
- function="def eval_fn(text): return True",
- ),
- editor=EditorConfig(
- type=EvaluatorType.FUNCTION,
- function='def edit_fn(text): return "edited"',
- ),
+ action="edit",
+ metric=eval_fn,
+ editor=edit_fn,
)
payload = inspector.to_dict()
- assert payload == {
- "name": "edit_inspector",
- "severity": "medium",
- "targets": ["steps", "AgentA"],
- "action": {"type": "edit"},
- "evaluator": {
- "type": "function",
- "function": "def eval_fn(text): return True",
- },
- "editor": {
- "type": "function",
- "function": 'def edit_fn(text): return "edited"',
- },
- }
+ assert payload["name"] == "edit_inspector"
+ assert payload["severity"] == "medium"
+ assert payload["targets"] == ["steps", "AgentA"]
+ assert payload["action"] == {"type": "edit"}
+ assert payload["evaluator"]["type"] == "function"
+ assert "eval_fn" in payload["evaluator"]["function"]
+ assert payload["editor"]["type"] == "function"
+ assert "edit_fn" in payload["editor"]["function"]
def test_edit_requires_editor(self):
with pytest.raises(ValueError, match="editor is required"):
- Inspector(
- name="bad_edit",
- targets=["steps"],
- action=InspectorActionConfig(type=InspectorAction.EDIT),
- evaluator=EvaluatorConfig(type=EvaluatorType.FUNCTION, function="def f(): pass"),
- )
+ Inspector(name="bad_edit", targets=["steps"], action="edit", metric="x")
def test_empty_name_rejected(self):
with pytest.raises(ValueError, match="name cannot be empty"):
- Inspector(
- name="",
- action=InspectorActionConfig(type=InspectorAction.ABORT),
- evaluator=EvaluatorConfig(type=EvaluatorType.ASSET, asset_id="x"),
- )
+ Inspector(name="", action="abort", metric="x")
+
+ def test_action_required(self):
+ with pytest.raises(ValueError, match="action is required"):
+ Inspector(name="x", metric="y")
+
+ def test_metric_required(self):
+ with pytest.raises(ValueError, match="metric is required"):
+ Inspector(name="x", action="abort")
+
+ def test_invalid_severity_rejected(self):
+ with pytest.raises(ValueError, match="invalid severity"):
+ Inspector(name="x", action="abort", metric="y", severity="urgent")
def test_from_dict_roundtrip(self):
original = Inspector(
name="roundtrip_test",
description="desc",
- severity=InspectorSeverity.LOW,
+ severity="low",
targets=["steps", "AgentB"],
- action=InspectorActionConfig(
- type=InspectorAction.RERUN,
- max_retries=1,
- on_exhaust=InspectorOnExhaust.CONTINUE,
- ),
- evaluator=EvaluatorConfig(type=EvaluatorType.ASSET, asset_id="id1", prompt="p"),
+ action={"type": "rerun", "max_retries": 1, "on_exhaust": "continue"},
+ metric={"asset_id": "id1", "prompt": "p"},
)
restored = Inspector.from_dict(original.to_dict())
assert restored.to_dict() == original.to_dict()
@@ -284,12 +235,12 @@ def test_from_dict_with_editor(self):
assert inspector.to_dict() == data
def test_optional_fields_omitted(self):
- """Verify that None description, severity, and editor are not in the payload."""
+ """None description, severity, and editor are not in the payload."""
inspector = Inspector(
name="minimal",
targets=["output"],
- action=InspectorActionConfig(type=InspectorAction.CONTINUE),
- evaluator=EvaluatorConfig(type=EvaluatorType.ASSET, asset_id="x", prompt="p"),
+ action="continue",
+ metric={"asset_id": "x", "prompt": "p"},
)
payload = inspector.to_dict()
assert "description" not in payload
@@ -298,165 +249,281 @@ def test_optional_fields_omitted(self):
# ---------------------------------------------------------------------------
-# PrebuiltInspector
+# Marketplace retrieval: aix.Inspector.get / .search
# ---------------------------------------------------------------------------
-class TestPrebuiltInspector:
- def test_prompt_injection_guard_defaults(self):
- inspector = PrebuiltInspector.prompt_injection_guard()
- payload = inspector.to_dict()
- assert payload == {"presetId": "prompt_injection_guard"}
+class _FakeClient:
+ """Records calls and returns canned guard payloads."""
- def test_pii_redaction_defaults(self):
- inspector = PrebuiltInspector.pii_redaction()
- payload = inspector.to_dict()
- assert payload == {"presetId": "pii_redaction"}
-
- def test_prompt_injection_guard_with_overrides(self):
- inspector = PrebuiltInspector.prompt_injection_guard(
- targets=[InspectorTarget.OUTPUT],
- severity=InspectorSeverity.CRITICAL,
- name="my_guard",
- description="Custom guard",
- )
- payload = inspector.to_dict()
- assert payload == {
- "presetId": "prompt_injection_guard",
- "name": "my_guard",
- "description": "Custom guard",
- "targets": ["output"],
- "severity": "critical",
- }
+ def __init__(self, get_result=None, paginate_result=None):
+ self._get_result = get_result
+ self._paginate_result = paginate_result
+ self.got_path = None
+ self.request_call = None
- def test_pii_redaction_with_action_override(self):
- inspector = PrebuiltInspector.pii_redaction(
- action={"type": "abort"},
- targets=[InspectorTarget.INPUT, InspectorTarget.OUTPUT],
- )
- payload = inspector.to_dict()
- assert payload == {
- "presetId": "pii_redaction",
- "targets": ["input", "output"],
- "action": {"type": "abort"},
- }
+ def get(self, path, **kwargs):
+ self.got_path = path
+ return self._get_result
- def test_unknown_preset_rejected(self):
- with pytest.raises(ValueError, match="Unknown inspector preset"):
- PrebuiltInspector(preset_id="nonexistent")
+ def request(self, method, path, **kwargs):
+ self.request_call = (method, path, kwargs.get("json"))
+ return self._paginate_result
- def test_unsupported_action_rejected(self):
- with pytest.raises(ValueError, match="not supported by preset"):
- PrebuiltInspector.prompt_injection_guard(action={"type": "edit"})
- def test_pii_redaction_rejects_rerun(self):
- with pytest.raises(ValueError, match="not supported by preset"):
- PrebuiltInspector.pii_redaction(action={"type": "rerun"})
+class _FakeContext:
+ def __init__(self, client):
+ self.client = client
- def test_hallucination_guard_defaults(self):
- inspector = PrebuiltInspector.hallucination_guard()
- payload = inspector.to_dict()
- assert payload == {"presetId": "hallucination_guard"}
-
- def test_hallucination_guard_with_overrides(self):
- inspector = PrebuiltInspector.hallucination_guard(
- targets=[InspectorTarget.OUTPUT],
- severity=InspectorSeverity.HIGH,
- name="my_hallucination_guard",
- description="Custom hallucination guard",
- )
- payload = inspector.to_dict()
- assert payload == {
- "presetId": "hallucination_guard",
- "name": "my_hallucination_guard",
- "description": "Custom hallucination guard",
- "targets": ["output"],
- "severity": "high",
- }
- def test_hallucination_guard_with_action_override(self):
- inspector = PrebuiltInspector.hallucination_guard(
- action={"type": "abort"},
- )
- payload = inspector.to_dict()
- assert payload == {
- "presetId": "hallucination_guard",
- "action": {"type": "abort"},
- }
+def _bind_inspector(client):
+ """Mirror ``aix.Inspector`` — Inspector subclass bound to a context."""
+ ctx = _FakeContext(client)
+ return type("Inspector", (Inspector,), {"context": ctx})
- def test_hallucination_guard_rejects_edit(self):
- with pytest.raises(ValueError, match="not supported by preset"):
- PrebuiltInspector.hallucination_guard(action={"type": "edit"})
- def test_targets_normalized_from_enums(self):
- inspector = PrebuiltInspector.pii_redaction(
- targets=[InspectorTarget.STEPS, InspectorTarget.OUTPUT],
- )
- assert inspector.targets == ["steps", "output"]
+# Canonical marketplace paths for the onboarded AWS guards (host/asset-name/instance),
+# verified against production. The asset-name (middle) segment drives the tuned config.
+_PROMPT_ATTACKS_PATH = "aws/detect-prompt-attacks-guardrail/aws"
+_PII_PATH = "aws/sensitive-information-guardrail/aws"
+_HALLUCINATION_PATH = "aws/contextual-grounding-check-guardrail/aws"
- def test_targets_normalized_from_strings(self):
- inspector = PrebuiltInspector.prompt_injection_guard(
- targets=["INPUT", "Steps"],
- )
- assert inspector.targets == ["input", "steps"]
- def test_from_dict_roundtrip(self):
- original = PrebuiltInspector.prompt_injection_guard(
- targets=[InspectorTarget.INPUT],
- severity=InspectorSeverity.HIGH,
- name="custom_name",
- )
- restored = PrebuiltInspector.from_dict(original.to_dict())
- assert restored.to_dict() == original.to_dict()
+class TestGuardSlugResolution:
+ """Unit coverage for the asset-name-based tuned-config resolution."""
- def test_from_dict_minimal(self):
- data = {"presetId": "pii_redaction"}
- inspector = PrebuiltInspector.from_dict(data)
- assert inspector.preset_id == "pii_redaction"
- assert inspector.targets is None
- assert inspector.action is None
-
- def test_list_presets_returns_all(self):
- presets = PrebuiltInspector.list_presets()
- assert "prompt_injection_guard" in presets
- assert "pii_redaction" in presets
- assert "hallucination_guard" in presets
- assert presets["prompt_injection_guard"]["category"] == "protection"
- assert presets["pii_redaction"]["category"] == "redaction"
- assert presets["hallucination_guard"]["category"] == "quality"
-
- def test_list_presets_contains_expected_fields(self):
- presets = PrebuiltInspector.list_presets()
- for preset_id, meta in presets.items():
- assert "name" in meta
- assert "description" in meta
- assert "default_targets" in meta
- assert "default_action" in meta
- assert "supported_actions" in meta
+ def test_guard_slugs_returns_all_segments_lowercased(self):
+ assert _guard_slugs("aws/Sensitive-Information-Guardrail/aws") == [
+ "aws",
+ "sensitive-information-guardrail",
+ "aws",
+ ]
+ def test_guard_slugs_handles_bare_id_and_empty(self):
+ assert _guard_slugs("69cbf63cd74e334a6bacfeb1") == ["69cbf63cd74e334a6bacfeb1"]
+ assert _guard_slugs(None) == []
+ assert _guard_slugs("") == []
-# ---------------------------------------------------------------------------
-# is_prebuilt_inspector helper
-# ---------------------------------------------------------------------------
+ def test_asset_name_middle_segment_selects_config(self):
+ # The guard identity is the middle segment, not the trailing host.
+ cfg = _resolve_guard_defaults("aws/sensitive-information-guardrail/aws")
+ assert cfg["action"]["type"] == "edit"
+ assert cfg["targets"] == ["input"]
+ def test_trailing_host_segment_does_not_match(self):
+ # An unknown asset name (host segments ignored) falls back to abort/input.
+ cfg = _resolve_guard_defaults("aws/some-future-guard/aws")
+ assert cfg["action"]["type"] == "abort"
+ assert cfg["targets"] == ["input"]
-class TestIsPrebuiltInspector:
- def test_prebuilt_instance(self):
- assert is_prebuilt_inspector(PrebuiltInspector.pii_redaction()) is True
+ def test_bare_id_falls_back_to_safe_default(self):
+ cfg = _resolve_guard_defaults("69cbf63cd74e334a6bacfeb1")
+ assert cfg["action"]["type"] == "abort"
- def test_preset_dict(self):
- assert is_prebuilt_inspector({"presetId": "prompt_injection_guard"}) is True
+ def test_first_matching_candidate_wins(self):
+ # asset_name candidate (checked first) takes precedence over the path.
+ cfg = _resolve_guard_defaults("contextual-grounding-check-guardrail", "aws/ignored/aws")
+ assert cfg["action"]["type"] == "rerun"
+ assert cfg["targets"] == ["output"]
- def test_regular_inspector_instance(self):
- inspector = Inspector(
- name="custom",
- action=InspectorActionConfig(type=InspectorAction.ABORT),
- evaluator=EvaluatorConfig(type=EvaluatorType.ASSET, asset_id="x"),
+
+class TestFromGuardModel:
+ def test_prompt_injection_defaults(self):
+ guard = Inspector.from_guard_model(
+ {"id": "pi-id", "name": "Detect Prompt Attacks Guardrail"},
+ requested_path=_PROMPT_ATTACKS_PATH,
+ )
+ assert guard.to_dict() == {
+ "name": "Detect Prompt Attacks Guardrail",
+ "targets": ["input"],
+ "action": {"type": "abort"},
+ "evaluator": {"type": "asset", "assetId": "pi-id"},
+ "presetId": "prompt_injection_guard",
+ }
+ assert guard.id == "pi-id"
+ assert guard.preset_id == "prompt_injection_guard"
+
+ def test_pii_redaction_defaults_with_editor(self):
+ guard = Inspector.from_guard_model(
+ {"id": "pii-id", "name": "Sensitive Information Guardrail"},
+ requested_path=_PII_PATH,
+ )
+ payload = guard.to_dict()
+ assert payload["action"] == {"type": "edit"}
+ assert payload["targets"] == ["input"]
+ # EDIT requires an editor; the guard model edits itself.
+ assert payload["editor"] == {"type": "asset", "assetId": "pii-id"}
+
+ def test_hallucination_defaults(self):
+ guard = Inspector.from_guard_model(
+ {"id": "h-id", "name": "Contextual Grounding Check Guardrail"},
+ requested_path=_HALLUCINATION_PATH,
+ )
+ payload = guard.to_dict()
+ assert payload["targets"] == ["output"]
+ assert payload["action"] == {"type": "rerun", "maxRetries": 2, "onExhaust": "abort"}
+
+ def test_asset_name_segment_drives_config_not_host(self):
+ # Real paths are host/asset-name/instance, so the trailing segment is the
+ # host ("aws"), not the guard identity. The tuned config must resolve from
+ # the asset-name (middle) segment — a regression guard for the slug bug
+ # where the trailing segment was used and every guard fell back to abort.
+ guard = Inspector.from_guard_model(
+ {"id": "pii-id", "name": "Sensitive Information Guardrail"},
+ requested_path=_PII_PATH,
+ )
+ assert guard.to_dict()["action"] == {"type": "edit"}
+
+ def test_unknown_guard_safe_default(self):
+ guard = Inspector.from_guard_model(
+ {"id": "x-id", "name": "Future Guard"},
+ requested_path="aws/some-future-guard/aws",
+ )
+ payload = guard.to_dict()
+ assert payload["targets"] == ["input"]
+ assert payload["action"] == {"type": "abort"}
+ assert payload["evaluator"] == {"type": "asset", "assetId": "x-id"}
+ # An unknown guard carries no presetId; the backend resolves it itself.
+ assert guard.preset_id is None
+ assert "presetId" not in payload
+
+ def test_pii_guard_emits_preset_id(self):
+ guard = Inspector.from_guard_model(
+ {"id": "pii-id", "name": "Sensitive Information Guardrail"},
+ requested_path=_PII_PATH,
+ )
+ assert guard.preset_id == "pii_redaction"
+ assert guard.to_dict()["presetId"] == "pii_redaction"
+
+ def test_hallucination_guard_emits_preset_id(self):
+ guard = Inspector.from_guard_model(
+ {"id": "h-id", "name": "Contextual Grounding Check Guardrail"},
+ requested_path=_HALLUCINATION_PATH,
+ )
+ assert guard.preset_id == "hallucination_guard"
+ assert guard.to_dict()["presetId"] == "hallucination_guard"
+
+ def test_preset_id_resolved_from_asset_name_field(self):
+ guard = Inspector.from_guard_model(
+ {
+ "id": "pi-id",
+ "name": "Detect Prompt Attacks Guardrail",
+ "assetInfo": {"assetName": "detect-prompt-attacks-guardrail"},
+ }
)
- assert is_prebuilt_inspector(inspector) is False
+ assert guard.preset_id == "prompt_injection_guard"
- def test_regular_dict(self):
- assert is_prebuilt_inspector({"name": "foo", "action": {"type": "abort"}}) is False
+ def test_preset_id_roundtrips_through_from_dict(self):
+ guard = Inspector.from_guard_model(
+ {"id": "pi-id", "name": "Detect Prompt Attacks Guardrail"},
+ requested_path=_PROMPT_ATTACKS_PATH,
+ )
+ restored = Inspector.from_dict(guard.to_dict())
+ assert restored.preset_id == "prompt_injection_guard"
+ assert restored.to_dict() == guard.to_dict()
- def test_none(self):
- assert is_prebuilt_inspector(None) is False
+ def test_slug_resolved_from_payload_path(self):
+ guard = Inspector.from_guard_model(
+ {"id": "h-id", "name": "Contextual Grounding Check Guardrail", "path": _HALLUCINATION_PATH},
+ )
+ assert guard.to_dict()["action"]["type"] == "rerun"
+
+ def test_slug_resolved_from_asset_name_field(self):
+ # When the payload carries assetInfo.assetName directly, it selects the
+ # tuned config regardless of the path shape.
+ guard = Inspector.from_guard_model(
+ {
+ "id": "pii-id",
+ "name": "Sensitive Information Guardrail",
+ "assetInfo": {"assetName": "sensitive-information-guardrail"},
+ }
+ )
+ assert guard.to_dict()["action"] == {"type": "edit"}
+
+ def test_returned_type_is_inspector(self):
+ guard = Inspector.from_guard_model({"id": "id", "name": "G"}, requested_path=_PII_PATH)
+ assert isinstance(guard, Inspector)
+
+ def test_path_extracted_from_asset_info(self):
+ # Paginate/get payloads carry the path under assetInfo.instanceId, like
+ # any other model. It must surface as .path (REPL-browsable) and drive
+ # the tuned defaults even when no requested_path is given (search flow).
+ guard = Inspector.from_guard_model(
+ {
+ "id": "pii-id",
+ "name": "Sensitive Information Guardrail",
+ "assetInfo": {"instanceId": _PII_PATH, "assetName": "sensitive-information-guardrail"},
+ }
+ )
+ assert guard.path == _PII_PATH
+ assert guard.to_dict()["action"] == {"type": "edit"}
+
+ def test_tuned_defaults_when_fetched_by_bare_id(self):
+ # The issue states IDs are also accepted. When retrieved by raw id, the
+ # id never matches the registry, but the payload's path still should
+ # select the tuned config rather than the safe abort/input fallback.
+ guard = Inspector.from_guard_model(
+ {"id": "69cbf63cd74e334a6bacfeb1", "name": "PII", "path": _PII_PATH},
+ requested_path="69cbf63cd74e334a6bacfeb1",
+ )
+ assert guard.to_dict()["action"] == {"type": "edit"}
+
+
+class TestInspectorGet:
+ def test_get_fetches_and_adapts(self):
+ client = _FakeClient(get_result={"id": "pii-id", "name": "PII", "path": _PII_PATH})
+ Bound = _bind_inspector(client)
+
+ guard = Bound.get(_PII_PATH)
+
+ # URL-encodes the (multi-slash) path and hits the models endpoint.
+ assert client.got_path == "v2/models/aws%2Fsensitive-information-guardrail%2Faws"
+ assert isinstance(guard, Inspector)
+ assert guard.metric.asset_id == "pii-id"
+ assert guard.action.type == "edit"
+
+ def test_get_accepts_attribute_override(self):
+ client = _FakeClient(get_result={"id": "pi-id", "name": "PI"})
+ Bound = _bind_inspector(client)
+
+ guard = Bound.get(_PROMPT_ATTACKS_PATH)
+ guard.targets = ["output"]
+ assert guard.to_dict()["targets"] == ["output"]
+
+ def test_get_without_context_raises(self):
+ with pytest.raises(Exception):
+ Inspector.get(_PROMPT_ATTACKS_PATH)
+
+
+class TestInspectorSearch:
+ def test_search_returns_page_shape(self):
+ client = _FakeClient(
+ paginate_result={
+ "results": [
+ {"id": "pi-id", "name": "Detect Prompt Attacks Guardrail", "path": _PROMPT_ATTACKS_PATH},
+ {"id": "pii-id", "name": "Sensitive Information Guardrail", "path": _PII_PATH},
+ ],
+ "total": 6,
+ "pageTotal": 1,
+ }
+ )
+ Bound = _bind_inspector(client)
+
+ page = Bound.search("guard")
+
+ method, path, body = client.request_call
+ assert method == "post"
+ assert path == "v2/models/paginate"
+ # Pinned to the guardrails function (a list of function-id strings, the
+ # shape v2/models/paginate requires) and includes the query.
+ assert body["functions"] == ["guardrails"]
+ assert body["q"] == "guard"
+
+ assert page.total == 6
+ assert page.page_number == 0
+ assert page.page_total == 1
+ assert all(isinstance(r, Inspector) for r in page.results)
+ assert [r.name for r in page.results] == [
+ "Detect Prompt Attacks Guardrail",
+ "Sensitive Information Guardrail",
+ ]
diff --git a/tests/unit/v2/test_model.py b/tests/unit/v2/test_model.py
index 68f38241..acf58678 100644
--- a/tests/unit/v2/test_model.py
+++ b/tests/unit/v2/test_model.py
@@ -212,6 +212,56 @@ def test_supports_structured_output_none_when_function_missing_even_with_markers
assert model.supports_structured_output is None
+# =============================================================================
+# Parameter Validation Tests
+# =============================================================================
+
+
+class TestModelValidateParams:
+ """Tests for _validate_params, including the multimodal ``data`` exemption."""
+
+ @staticmethod
+ def _param(name, required=False, data_type="text", data_sub_type="text"):
+ return SimpleNamespace(
+ name=name,
+ required=required,
+ data_type=data_type,
+ data_sub_type=data_sub_type,
+ default_values=[],
+ )
+
+ def _model(self, params):
+ model = Model.__new__(Model)
+ model.id = "test-model-id"
+ model.name = "Test Model"
+ model.params = params
+ return model
+
+ def test_missing_required_text_is_flagged(self):
+ """A required ``text`` input with no value (and no ``data``) is an error."""
+ model = self._model([self._param("text", required=True)])
+ errors = model._validate_params(max_tokens=10)
+ assert errors == ["Required parameter 'text' is missing"]
+
+ def test_data_payload_exempts_required_text(self):
+ """A multimodal ``data`` payload supersedes required ``text``/``prompt``."""
+ model = self._model([self._param("text", required=True), self._param("prompt", required=True)])
+ data = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "u"}}]}]
+ assert model._validate_params(data=data, max_tokens=10) == []
+
+ def test_data_none_does_not_exempt(self):
+ """An explicit ``data=None`` is not a payload and must not exempt requireds."""
+ model = self._model([self._param("text", required=True)])
+ assert model._validate_params(data=None) == ["Required parameter 'text' is missing"]
+
+ def test_provided_param_type_still_validated_with_data(self):
+ """Type checks on explicitly provided params still run when ``data`` is present."""
+ model = self._model([self._param("text", required=True), self._param("max_tokens", data_type="number")])
+ data = [{"role": "user", "content": []}]
+ errors = model._validate_params(data=data, max_tokens="not-a-number")
+ assert any("max_tokens" in e for e in errors)
+
+
# =============================================================================
# Run Routing Tests
# =============================================================================
@@ -545,6 +595,55 @@ def test_streamer_parses_openai_usage_and_finish_reason(self):
with pytest.raises(StopIteration):
next(streamer)
+ def test_streamer_parses_reasoning_content_deltas(self):
+ """ModelResponseStreamer should preserve delta.reasoning_content from OpenAI chunks.
+
+ Covers three cases: a reasoning-only chunk (empty ``content``) must still be
+ yielded with ``data == ""``, a mixed chunk must carry both, and a plain
+ content chunk must leave ``reasoning_content`` as None.
+ """
+ streamer = self._create_streamer(
+ [
+ 'data: {"id":"chatcmpl-r","choices":[{"index":0,"delta":{"reasoning_content":"Let me think"},"finish_reason":null}],"usage":null}',
+ 'data: {"id":"chatcmpl-r","choices":[{"index":0,"delta":{"content":"Answer","reasoning_content":" more"},"finish_reason":null}],"usage":null}',
+ 'data: {"id":"chatcmpl-r","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}],"usage":null}',
+ "data: [DONE]",
+ ]
+ )
+
+ reasoning_only = next(streamer)
+ mixed = next(streamer)
+ content_only = next(streamer)
+
+ assert reasoning_only.data == ""
+ assert reasoning_only.reasoning_content == "Let me think"
+
+ assert mixed.data == "Answer"
+ assert mixed.reasoning_content == " more"
+
+ assert content_only.data == "!"
+ assert content_only.reasoning_content is None
+
+ with pytest.raises(StopIteration):
+ next(streamer)
+
+ def test_streamer_content_only_chunk_leaves_reasoning_content_none(self):
+ """A content-only OpenAI chunk should stream unchanged with reasoning_content None."""
+ streamer = self._create_streamer(
+ [
+ 'data: {"id":"chatcmpl-c","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}],"usage":null}',
+ "data: [DONE]",
+ ]
+ )
+
+ chunk = next(streamer)
+
+ assert chunk.data == "Hello"
+ assert chunk.reasoning_content is None
+ assert chunk.tool_calls is None
+ assert chunk.usage is None
+ assert chunk.finish_reason is None
+
def test_streamer_normalizes_single_tool_call_object(self):
"""ModelResponseStreamer should normalize a single tool_call object to a list."""
streamer = self._create_streamer(
@@ -701,6 +800,61 @@ def test_message_deserializes_tool_calls_with_null_content(self):
assert message.tool_calls[0]["function"]["name"] == "get_current_time"
assert message.tool_calls[0]["function"]["arguments"] == '{"city":"Tokyo"}'
+ def test_message_deserializes_reasoning_content(self):
+ """ModelResult parsing should preserve reasoning_content from reasoning LLMs.
+
+ For some reasoning models (e.g. Qwen 3.6 35B A3B) the platform routes all
+ completion tokens into ``reasoning_content`` while ``content``/``data`` stay
+ empty; the field must survive deserialization instead of being dropped.
+ """
+ payload = {
+ "details": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "Here's a thinking sequence\n\n1. Deconstruct the prompt: ...",
+ "name": None,
+ "tool_calls": [],
+ },
+ "finish_reason": "length",
+ "logprobs": None,
+ }
+ ],
+ "status": "SUCCESS",
+ "completed": True,
+ "data": "",
+ "usage": {"prompt_tokens": "26", "completion_tokens": "80", "total_tokens": 106},
+ }
+
+ result = ModelResult.from_dict(payload)
+
+ message = result.details[0].message
+ assert isinstance(message, Message)
+ assert message.content == ""
+ assert message.reasoning_content == "Here's a thinking sequence\n\n1. Deconstruct the prompt: ..."
+
+ def test_message_without_reasoning_content_defaults_to_none(self):
+ """A message without reasoning_content should deserialize with reasoning_content is None."""
+ payload = {
+ "status": "SUCCESS",
+ "completed": True,
+ "details": [
+ {
+ "index": 0,
+ "message": {"role": "assistant", "content": "Hello, world!"},
+ "finish_reason": "stop",
+ }
+ ],
+ }
+
+ result = ModelResult.from_dict(payload)
+
+ message = result.details[0].message
+ assert message.content == "Hello, world!"
+ assert message.reasoning_content is None
+
def test_guardrail_response_with_dict_details_deserializes(self):
"""ModelResult parsing should tolerate a non-list ``details`` payload.
diff --git a/tests/unit/v2/test_resource.py b/tests/unit/v2/test_resource.py
index f729a7ba..682e43a3 100644
--- a/tests/unit/v2/test_resource.py
+++ b/tests/unit/v2/test_resource.py
@@ -774,6 +774,27 @@ def test_poll_returns_result(self):
assert isinstance(result, Result)
assert result.status == "SUCCESS"
+ def test_poll_failed_uses_nested_data_error(self):
+ """poll() should surface data.error when top-level error fields are empty."""
+ resource = self._create_runnable_resource()
+ resource.context.client.get = Mock(
+ return_value={
+ "status": "FAILED",
+ "completed": True,
+ "errorMessage": None,
+ "supplierError": None,
+ "data": {
+ "error": "litellm.APIConnectionError: err.not_enough_balance\nTraceback...",
+ },
+ }
+ )
+
+ with pytest.raises(APIError) as exc_info:
+ resource.poll("https://poll.url")
+
+ assert "err.not_enough_balance" in str(exc_info.value)
+ assert "err.not_enough_balance" in exc_info.value.error
+
def test_sync_poll_respects_timeout(self):
"""sync_poll() should raise TimeoutError after timeout."""
resource = self._create_runnable_resource()
@@ -842,6 +863,23 @@ def test_handle_run_response_non_url_in_progress_sets_completed_false(self):
assert result.url == "some-poll-id-not-a-url"
assert result.completed is False
+ def test_handle_run_response_success_url_data_is_final_data(self):
+ """handle_run_response() should not treat successful URL data as a polling URL."""
+ resource = self._create_runnable_resource()
+
+ response = {
+ "status": "SUCCESS",
+ "completed": True,
+ "data": "https://example.com/generated-image.jpeg",
+ }
+ result = resource.handle_run_response(response)
+
+ assert isinstance(result, Result)
+ assert result.status == "SUCCESS"
+ assert result.completed is True
+ assert result.data == "https://example.com/generated-image.jpeg"
+ assert result.url is None
+
def test_run_polls_when_in_progress_with_non_url_data(self):
"""run() should poll when handle_run_response returns IN_PROGRESS with non-URL data."""
resource = self._create_runnable_resource()
diff --git a/tests/unit/v2/test_session.py b/tests/unit/v2/test_session.py
new file mode 100644
index 00000000..006dcf83
--- /dev/null
+++ b/tests/unit/v2/test_session.py
@@ -0,0 +1,1368 @@
+"""Unit tests for the v2 Session module.
+
+This module tests Session, SessionMessage, and SessionMessageAttachment
+dataclasses, serialization, CRUD operations, and the MIME-to-AttachmentType
+mapping — all with mocked HTTP calls.
+"""
+
+import warnings
+from unittest.mock import Mock, patch, call
+
+import pytest
+
+from aixplain.v2.session import (
+ EXECUTION_PARAMS_MAP,
+ ExecutionConfig,
+ Session,
+ SessionMessage,
+ SessionMessageAttachment,
+ _mime_to_attachment_type,
+ _normalize_execution_params,
+)
+from aixplain.v2.enums import (
+ SessionStatus,
+ RunStatus,
+ MessageRole,
+ Reaction,
+ AttachmentType,
+)
+from aixplain.v2.agent import Agent, Budget
+from aixplain.v2.exceptions import ValidationError, APIError, ResourceError
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_context(**overrides):
+ """Create a mock Aixplain context with a mock client."""
+ client = Mock()
+ ctx = Mock(
+ client=client,
+ backend_url="https://platform-api.aixplain.com",
+ api_key="test_key",
+ )
+ for k, v in overrides.items():
+ setattr(ctx, k, v)
+ return ctx
+
+
+def _bound_session_class(ctx=None):
+ """Return a Session subclass bound to a mock context."""
+ ctx = ctx or _make_mock_context()
+
+ class BoundSession(Session):
+ context = ctx
+
+ return BoundSession
+
+
+def _bound_agent_class(ctx=None):
+ """Return an Agent subclass bound to a mock context."""
+ ctx = ctx or _make_mock_context()
+
+ class BoundAgent(Agent):
+ context = ctx
+
+ return BoundAgent
+
+
+# Sample API payloads ---------------------------------------------------
+
+SAMPLE_SESSION_DICT = {
+ "id": "sess_001",
+ "userId": "user_42",
+ "agentId": "agent_99",
+ "name": "My Session",
+ "status": "active",
+ "runStatus": "idle",
+ "messageCount": 3,
+ "lastMessagePreview": "Hello there",
+ "lastMessageAt": "2025-06-01T12:00:00Z",
+ "createdAt": "2025-06-01T10:00:00Z",
+ "updatedAt": "2025-06-01T12:00:00Z",
+}
+
+SAMPLE_MESSAGE_DICT = {
+ "id": "msg_001",
+ "sessionId": "sess_001",
+ "userId": "user_42",
+ "agentId": "agent_99",
+ "role": "user",
+ "content": "Hello!",
+ "sequence": 1,
+ "requestId": "req_001",
+ "reaction": None,
+ "attachments": None,
+ "createdAt": "2025-06-01T10:01:00Z",
+}
+
+
+# =========================================================================
+# Enum tests
+# =========================================================================
+
+
+class TestSessionEnums:
+ """Verify new enums have expected members and values."""
+
+ def test_session_status_values(self):
+ assert SessionStatus.ACTIVE == "active"
+ assert SessionStatus.COMPLETED == "completed"
+ assert SessionStatus.FAILED == "failed"
+ assert SessionStatus.ARCHIVED == "archived"
+
+ def test_run_status_values(self):
+ assert RunStatus.IDLE == "idle"
+ assert RunStatus.RUNNING == "running"
+ assert RunStatus.COMPLETED == "completed"
+
+ def test_message_role_values(self):
+ assert MessageRole.USER == "user"
+ assert MessageRole.ASSISTANT == "assistant"
+
+ def test_reaction_values(self):
+ assert Reaction.LIKE == "LIKE"
+ assert Reaction.DISLIKE == "DISLIKE"
+
+ def test_attachment_type_values(self):
+ assert AttachmentType.TEXT == "text"
+ assert AttachmentType.IMAGE == "image"
+ assert AttachmentType.VIDEO == "video"
+ assert AttachmentType.AUDIO == "audio"
+ assert AttachmentType.DOCUMENT == "document"
+ assert AttachmentType.CODE == "code"
+ assert AttachmentType.UNKNOWN == "unknown"
+
+
+# =========================================================================
+# MIME type mapping
+# =========================================================================
+
+
+class TestMimeToAttachmentType:
+ """Tests for the _mime_to_attachment_type helper."""
+
+ @pytest.mark.parametrize(
+ "mime, expected",
+ [
+ ("image/png", "image"),
+ ("image/jpeg", "image"),
+ ("video/mp4", "video"),
+ ("audio/mpeg", "audio"),
+ ("audio/wav", "audio"),
+ ("text/plain", "text"),
+ ("text/csv", "text"),
+ ],
+ )
+ def test_main_type_mapping(self, mime, expected):
+ assert _mime_to_attachment_type(mime) == expected
+
+ @pytest.mark.parametrize(
+ "mime",
+ [
+ "text/x-python",
+ "application/javascript",
+ "application/x-javascript",
+ "application/typescript",
+ ],
+ )
+ def test_code_types(self, mime):
+ assert _mime_to_attachment_type(mime) == "code"
+
+ @pytest.mark.parametrize(
+ "mime",
+ [
+ "application/pdf",
+ "application/msword",
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+ "application/vnd.ms-excel",
+ "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+ ],
+ )
+ def test_document_types(self, mime):
+ assert _mime_to_attachment_type(mime) == "document"
+
+ def test_unknown_fallback(self):
+ assert _mime_to_attachment_type("application/octet-stream") == "unknown"
+
+ def test_empty_string(self):
+ assert _mime_to_attachment_type("") == "unknown"
+
+ def test_none_like_empty(self):
+ # None is not str, but the function guards against falsy values
+ assert _mime_to_attachment_type("") == "unknown"
+
+
+# =========================================================================
+# SessionMessageAttachment
+# =========================================================================
+
+
+class TestSessionMessageAttachment:
+ """Tests for the SessionMessageAttachment dataclass."""
+
+ def test_create_minimal(self):
+ att = SessionMessageAttachment(url="https://example.com/f.png")
+ assert att.url == "https://example.com/f.png"
+ assert att.name is None
+ assert att.type is None
+
+ def test_create_full(self):
+ att = SessionMessageAttachment(url="https://example.com/f.png", name="f.png", type="image")
+ assert att.name == "f.png"
+ assert att.type == "image"
+
+ def test_serialization_roundtrip(self):
+ att = SessionMessageAttachment(url="https://example.com/f.png", name="f.png", type="image")
+ d = att.to_dict()
+ assert d["url"] == "https://example.com/f.png"
+ restored = SessionMessageAttachment.from_dict(d)
+ assert restored.url == att.url
+ assert restored.name == att.name
+ assert restored.type == att.type
+
+
+# =========================================================================
+# SessionMessage
+# =========================================================================
+
+
+class TestSessionMessage:
+ """Tests for the SessionMessage dataclass."""
+
+ def test_from_dict_maps_camel_case(self):
+ msg = SessionMessage.from_dict(SAMPLE_MESSAGE_DICT)
+ assert msg.id == "msg_001"
+ assert msg.session_id == "sess_001"
+ assert msg.user_id == "user_42"
+ assert msg.agent_id == "agent_99"
+ assert msg.role == "user"
+ assert msg.content == "Hello!"
+ assert msg.sequence == 1
+ assert msg.request_id == "req_001"
+ assert msg.reaction is None
+ assert msg.attachments is None
+ assert msg.created_at == "2025-06-01T10:01:00Z"
+
+ def test_to_dict_uses_camel_case(self):
+ msg = SessionMessage(
+ id="m1",
+ session_id="s1",
+ user_id="u1",
+ agent_id="a1",
+ role="assistant",
+ content="Hi",
+ sequence=2,
+ created_at="2025-01-01",
+ )
+ d = msg.to_dict()
+ assert d["sessionId"] == "s1"
+ assert d["userId"] == "u1"
+ assert d["agentId"] == "a1"
+ assert d["requestId"] is None
+ assert d["createdAt"] == "2025-01-01"
+
+ def test_with_attachments(self):
+ data = {
+ **SAMPLE_MESSAGE_DICT,
+ "attachments": [{"url": "https://example.com/img.png", "name": "img.png", "type": "image"}],
+ }
+ msg = SessionMessage.from_dict(data)
+ assert len(msg.attachments) == 1
+ assert isinstance(msg.attachments[0], SessionMessageAttachment)
+ assert msg.attachments[0].url == "https://example.com/img.png"
+
+
+# =========================================================================
+# Session — dataclass / serialization
+# =========================================================================
+
+
+class TestSessionSerialization:
+ """Tests for Session dataclass creation and serialization."""
+
+ def test_from_dict_maps_camel_case(self):
+ session = Session.from_dict(SAMPLE_SESSION_DICT)
+ assert session.id == "sess_001"
+ assert session.user_id == "user_42"
+ assert session.agent_id == "agent_99"
+ assert session.name == "My Session"
+ assert session.status == "active"
+ assert session.run_status == "idle"
+ assert session.message_count == 3
+ assert session.last_message_preview == "Hello there"
+ assert session.last_message_at == "2025-06-01T12:00:00Z"
+ assert session.created_at == "2025-06-01T10:00:00Z"
+ assert session.updated_at == "2025-06-01T12:00:00Z"
+
+ def test_build_save_payload_create(self):
+ session = Session(agent_id="agent_99", name="New")
+ payload = session.build_save_payload()
+ assert payload == {"agentId": "agent_99", "name": "New", "status": "active"}
+
+ def test_build_save_payload_excludes_readonly(self):
+ session = Session.from_dict(SAMPLE_SESSION_DICT)
+ payload = session.build_save_payload()
+ # Should NOT include read-only fields
+ assert "userId" not in payload
+ assert "runStatus" not in payload
+ assert "messageCount" not in payload
+ assert "createdAt" not in payload
+ assert "updatedAt" not in payload
+ assert "lastMessagePreview" not in payload
+ assert "lastMessageAt" not in payload
+
+ def test_build_save_payload_update(self):
+ session = Session.from_dict(SAMPLE_SESSION_DICT)
+ session.name = "Renamed"
+ session.status = "archived"
+ payload = session.build_save_payload()
+ assert payload["name"] == "Renamed"
+ assert payload["status"] == "archived"
+ assert payload["agentId"] == "agent_99"
+
+ def test_default_status_is_active(self):
+ session = Session(agent_id="a1")
+ assert session.status == "active"
+
+ def test_repr(self):
+ session = Session(id="sess_001", name="My Session", agent_id="a1")
+ r = repr(session)
+ assert "Session" in r
+ assert "sess_001" in r
+
+
+# =========================================================================
+# Session — CRUD operations (mocked)
+# =========================================================================
+
+
+class TestSessionCreate:
+ """Tests for Session.save() (create path)."""
+
+ def test_save_creates_new_session(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {
+ **SAMPLE_SESSION_DICT,
+ "id": "new_sess",
+ }
+ BoundSession = _bound_session_class(ctx)
+ session = BoundSession(agent_id="agent_99", name="Test")
+
+ session.save()
+
+ ctx.client.request.assert_called_once()
+ args, kwargs = ctx.client.request.call_args
+ assert args[0] == "post"
+ assert "v1/sessions" in args[1]
+ assert kwargs["json"]["agentId"] == "agent_99"
+ assert session.id == "new_sess"
+
+ def test_save_updates_existing_session(self):
+ ctx = _make_mock_context()
+ ctx.client.request.side_effect = [
+ # First call: create
+ {**SAMPLE_SESSION_DICT, "id": "sess_001"},
+ # Second call: update
+ None,
+ ]
+ BoundSession = _bound_session_class(ctx)
+ session = BoundSession(agent_id="agent_99", name="Test")
+ session.save()
+ assert session.id == "sess_001"
+
+ # Now update
+ session.name = "Renamed"
+ session.save()
+
+ second_call = ctx.client.request.call_args_list[1]
+ assert second_call[0][0] == "put"
+ assert "sess_001" in second_call[0][1]
+
+
+class TestSessionGet:
+ """Tests for Session.get()."""
+
+ def test_get_by_id(self):
+ ctx = _make_mock_context()
+ ctx.client.get.return_value = SAMPLE_SESSION_DICT
+ BoundSession = _bound_session_class(ctx)
+
+ session = BoundSession.get("sess_001")
+
+ ctx.client.get.assert_called_once()
+ call_path = ctx.client.get.call_args[0][0]
+ assert "v1/sessions/sess_001" in call_path
+ assert session.id == "sess_001"
+ assert session.agent_id == "agent_99"
+
+
+class TestSessionDelete:
+ """Tests for Session.delete()."""
+
+ def test_delete_session(self):
+ ctx = _make_mock_context()
+ ctx.client.request_raw.return_value = Mock(status_code=200, ok=True)
+ BoundSession = _bound_session_class(ctx)
+ session = BoundSession.from_dict(SAMPLE_SESSION_DICT)
+ session.context = ctx
+ session._update_saved_state()
+
+ result = session.delete()
+
+ ctx.client.request_raw.assert_called_once()
+ args = ctx.client.request_raw.call_args[0]
+ assert args[0] == "delete"
+ assert "sess_001" in args[1]
+ assert result.completed is True
+
+
+class TestSessionSearch:
+ """Tests for Session.search() classmethod (returns a Page)."""
+
+ def test_search_returns_page_of_sessions(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = [
+ SAMPLE_SESSION_DICT,
+ {**SAMPLE_SESSION_DICT, "id": "sess_002", "name": "Second"},
+ ]
+ BoundSession = _bound_session_class(ctx)
+
+ page = BoundSession.search()
+
+ ctx.client.request.assert_called_once()
+ args, kwargs = ctx.client.request.call_args
+ assert args[0] == "get"
+ assert "v1/sessions" in args[1]
+ assert page.total == 2
+ assert page.page_total == 1
+ assert page.page_number == 0
+ assert len(page.results) == 2
+ assert page.results[0].id == "sess_001"
+ assert page.results[1].id == "sess_002"
+ # Page is iterable over its results.
+ assert [s.id for s in page] == ["sess_001", "sess_002"]
+
+ def test_search_with_agent_object_filter(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = [SAMPLE_SESSION_DICT]
+ BoundSession = _bound_session_class(ctx)
+
+ class _Agent:
+ id = "agent_99"
+
+ BoundSession.search(agent=_Agent())
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["params"]["agentId"] == "agent_99"
+
+ def test_search_with_agent_id_string_filter(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = [SAMPLE_SESSION_DICT]
+ BoundSession = _bound_session_class(ctx)
+
+ BoundSession.search(agent="agent_99")
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["params"]["agentId"] == "agent_99"
+
+ def test_search_with_status_filter(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = []
+ BoundSession = _bound_session_class(ctx)
+
+ BoundSession.search(status="completed")
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["params"]["status"] == "completed"
+
+ def test_search_with_user_id_filter(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = []
+ BoundSession = _bound_session_class(ctx)
+
+ BoundSession.search(user_id="user_42")
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["params"]["userId"] == "user_42"
+
+ def test_search_with_date_and_memory_filters(self):
+ from datetime import datetime
+
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = []
+ BoundSession = _bound_session_class(ctx)
+
+ BoundSession.search(
+ created_after=datetime(2026, 1, 1),
+ created_before="2026-06-01T00:00:00",
+ memory_enabled=True,
+ )
+
+ _, kwargs = ctx.client.request.call_args
+ params = kwargs["params"]
+ assert params["createdAfter"] == "2026-01-01T00:00:00"
+ assert params["createdBefore"] == "2026-06-01T00:00:00"
+ assert params["memoryEnabled"] is True
+
+ def test_search_forwards_pagination(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = []
+ BoundSession = _bound_session_class(ctx)
+
+ BoundSession.search(page_number=2, page_size=5)
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["params"]["pageNumber"] == 2
+ assert kwargs["params"]["pageSize"] == 5
+
+ def test_search_empty_result(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = []
+ BoundSession = _bound_session_class(ctx)
+
+ page = BoundSession.search()
+ assert page.results == []
+ assert page.total == 0
+
+ def test_search_paginated_dict_envelope(self):
+ """Tolerate a future paginated dict envelope (total/pageTotal)."""
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {
+ "results": [SAMPLE_SESSION_DICT],
+ "total": 17,
+ "pageTotal": 4,
+ }
+ BoundSession = _bound_session_class(ctx)
+
+ page = BoundSession.search(page_number=1)
+ assert page.total == 17
+ assert page.page_total == 4
+ assert len(page.results) == 1
+
+
+# =========================================================================
+# Session — message operations
+# =========================================================================
+
+
+class TestSessionMessages:
+ """Tests for Session message methods."""
+
+ def _make_session(self, ctx=None):
+ ctx = ctx or _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+ session = BoundSession.from_dict(SAMPLE_SESSION_DICT)
+ session.context = ctx
+ session._update_saved_state()
+ return session
+
+ def test_messages_returns_list(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = [SAMPLE_MESSAGE_DICT]
+ session = self._make_session(ctx)
+
+ messages = session.messages()
+
+ ctx.client.request.assert_called_once()
+ args = ctx.client.request.call_args[0]
+ assert args[0] == "get"
+ assert "sess_001/messages" in args[1]
+ assert len(messages) == 1
+ assert isinstance(messages[0], SessionMessage)
+ assert messages[0].id == "msg_001"
+
+ def test_messages_empty(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = []
+ session = self._make_session(ctx)
+
+ messages = session.messages()
+ assert messages == []
+
+ def test_add_message_basic(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ msg = session.add_message(role="user", content="Hello!")
+
+ ctx.client.request.assert_called_once()
+ args, kwargs = ctx.client.request.call_args
+ assert args[0] == "post"
+ assert "sess_001/messages" in args[1]
+ assert kwargs["json"]["role"] == "user"
+ assert kwargs["json"]["content"] == "Hello!"
+ assert "requestId" not in kwargs["json"]
+ assert "attachments" not in kwargs["json"]
+ assert isinstance(msg, SessionMessage)
+
+ def test_add_message_with_request_id(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ session.add_message(role="user", content="Hi", request_id="req_xyz")
+
+ payload = ctx.client.request.call_args[1]["json"]
+ assert payload["requestId"] == "req_xyz"
+
+ def test_add_message_with_attachments(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ attachments = [{"url": "https://example.com/f.png", "name": "f.png", "type": "image"}]
+ session.add_message(role="user", content="See image", attachments=attachments)
+
+ payload = ctx.client.request.call_args[1]["json"]
+ assert len(payload["attachments"]) == 1
+ assert payload["attachments"][0]["url"] == "https://example.com/f.png"
+
+ def test_add_message_with_files(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ with (
+ patch("aixplain.v2.upload_utils.FileUploader") as MockUploader,
+ patch("aixplain.v2.upload_utils.MimeTypeDetector") as MockDetector,
+ ):
+ instance = MockUploader.return_value
+ instance.upload.return_value = "https://cdn.example.com/uploaded.png"
+ MockDetector.detect_mime_type.return_value = "image/png"
+
+ session.add_message(role="user", content="File", files=["/tmp/photo.png"])
+
+ instance.upload.assert_called_once_with(
+ "/tmp/photo.png",
+ is_temp=True,
+ return_download_link=True,
+ )
+ payload = ctx.client.request.call_args[1]["json"]
+ assert len(payload["attachments"]) == 1
+ assert payload["attachments"][0]["url"] == "https://cdn.example.com/uploaded.png"
+ assert payload["attachments"][0]["name"] == "photo.png"
+ assert payload["attachments"][0]["type"] == "image"
+
+ def test_add_message_merges_attachments_and_files(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ with (
+ patch("aixplain.v2.upload_utils.FileUploader") as MockUploader,
+ patch("aixplain.v2.upload_utils.MimeTypeDetector") as MockDetector,
+ ):
+ instance = MockUploader.return_value
+ instance.upload.return_value = "https://cdn.example.com/doc.pdf"
+ MockDetector.detect_mime_type.return_value = "application/pdf"
+
+ existing = [{"url": "https://example.com/a.txt", "name": "a.txt", "type": "text"}]
+ session.add_message(
+ role="user",
+ content="Both",
+ attachments=existing,
+ files=["/tmp/doc.pdf"],
+ )
+
+ payload = ctx.client.request.call_args[1]["json"]
+ assert len(payload["attachments"]) == 2
+ assert payload["attachments"][0]["url"] == "https://example.com/a.txt"
+ assert payload["attachments"][1]["url"] == "https://cdn.example.com/doc.pdf"
+ assert payload["attachments"][1]["type"] == "document"
+
+ def test_add_message_url_string_type_auto_detected(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ with patch("aixplain.v2.upload_utils.FileUploader") as MockUploader:
+ session.add_message(role="user", content="q", attachments=["https://cdn.example.com/a.wav"])
+ MockUploader.assert_not_called() # URLs are not uploaded
+
+ payload = ctx.client.request.call_args[1]["json"]
+ # type/mimeType/name auto-detected from the extension so the worker
+ # treats it as audio, not a generic file.
+ assert payload["attachments"] == [
+ {"url": "https://cdn.example.com/a.wav", "name": "a.wav", "type": "audio", "mimeType": "audio/wav"}
+ ]
+
+ def test_add_message_url_dict_without_type_auto_detected(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ session.add_message(role="user", content="q", attachments=[{"url": "https://cdn.example.com/pic.png"}])
+
+ att = ctx.client.request.call_args[1]["json"]["attachments"][0]
+ assert att["type"] == "image"
+ assert att["mimeType"] == "image/png"
+
+ def test_add_message_explicit_type_preserved(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ # Caller-supplied type must win over inference (e.g. extensionless URL).
+ session.add_message(
+ role="user", content="q", attachments=[{"url": "https://cdn.example.com/clip", "type": "audio"}]
+ )
+ att = ctx.client.request.call_args[1]["json"]["attachments"][0]
+ assert att["type"] == "audio"
+
+ def test_add_message_path_string_uploads_with_mimetype(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ with (
+ patch("aixplain.v2.upload_utils.FileUploader") as MockUploader,
+ patch("aixplain.v2.upload_utils.MimeTypeDetector") as MockDetector,
+ ):
+ MockUploader.return_value.upload.return_value = "https://cdn.example.com/clip.wav"
+ MockDetector.detect_mime_type.return_value = "audio/wav"
+
+ session.add_message(role="user", content="transcribe", attachments=["/tmp/clip.wav"])
+
+ att = ctx.client.request.call_args[1]["json"]["attachments"][0]
+ assert att == {
+ "url": "https://cdn.example.com/clip.wav",
+ "name": "clip.wav",
+ "type": "audio",
+ "mimeType": "audio/wav",
+ }
+
+ def test_add_message_path_dict_uploads_with_type_override(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ with (
+ patch("aixplain.v2.upload_utils.FileUploader") as MockUploader,
+ patch("aixplain.v2.upload_utils.MimeTypeDetector") as MockDetector,
+ ):
+ MockUploader.return_value.upload.return_value = "https://cdn.example.com/clip.bin"
+ MockDetector.detect_mime_type.return_value = "application/octet-stream"
+
+ session.add_message(
+ role="user",
+ content="q",
+ attachments=[{"path": "/tmp/clip.wav", "type": "audio", "name": "voice.wav"}],
+ )
+
+ att = ctx.client.request.call_args[1]["json"]["attachments"][0]
+ assert att["url"] == "https://cdn.example.com/clip.bin"
+ assert att["type"] == "audio" # caller override wins over detection
+ assert att["name"] == "voice.wav"
+
+ def test_add_message_audio_as_prompt_empty_content(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ session.add_message(
+ role="user",
+ content="",
+ attachments=[{"url": "https://cdn.example.com/a.wav", "type": "audio"}],
+ )
+
+ payload = ctx.client.request.call_args[1]["json"]
+ assert payload["content"] == ""
+ assert payload["attachments"][0]["url"] == "https://cdn.example.com/a.wav"
+
+ def test_get_message(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = SAMPLE_MESSAGE_DICT
+ session = self._make_session(ctx)
+
+ msg = session.get_message("msg_001")
+
+ args = ctx.client.request.call_args[0]
+ assert args[0] == "get"
+ assert "sess_001/messages/msg_001" in args[1]
+ assert isinstance(msg, SessionMessage)
+ assert msg.id == "msg_001"
+
+ def test_delete_message(self):
+ ctx = _make_mock_context()
+ ctx.client.request_raw.return_value = Mock(status_code=200, ok=True)
+ session = self._make_session(ctx)
+
+ session.delete_message("msg_001")
+
+ args = ctx.client.request_raw.call_args[0]
+ assert args[0] == "delete"
+ assert "sess_001/messages/msg_001" in args[1]
+
+ def test_react_like(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {**SAMPLE_MESSAGE_DICT, "reaction": "LIKE"}
+ session = self._make_session(ctx)
+
+ msg = session.react("msg_001", "LIKE")
+
+ args, kwargs = ctx.client.request.call_args
+ assert args[0] == "post"
+ assert "sess_001/messages/msg_001/reaction" in args[1]
+ assert kwargs["json"]["reaction"] == "LIKE"
+ assert msg.reaction == "LIKE"
+
+ def test_react_clear(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {**SAMPLE_MESSAGE_DICT, "reaction": None}
+ session = self._make_session(ctx)
+
+ msg = session.react("msg_001", None)
+
+ payload = ctx.client.request.call_args[1]["json"]
+ assert payload["reaction"] is None
+ assert msg.reaction is None
+
+
+class TestSessionValidation:
+ """Tests that session methods enforce valid state."""
+
+ def _make_unsaved_session(self):
+ ctx = _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+ return BoundSession(agent_id="a1")
+
+ def test_messages_requires_saved_session(self):
+ session = self._make_unsaved_session()
+ with pytest.raises(ValidationError):
+ session.messages()
+
+ def test_add_message_requires_saved_session(self):
+ session = self._make_unsaved_session()
+ with pytest.raises(ValidationError):
+ session.add_message(role="user", content="hi")
+
+ def test_get_message_requires_saved_session(self):
+ session = self._make_unsaved_session()
+ with pytest.raises(ValidationError):
+ session.get_message("msg_001")
+
+ def test_delete_message_requires_saved_session(self):
+ session = self._make_unsaved_session()
+ with pytest.raises(ValidationError):
+ session.delete_message("msg_001")
+
+ def test_react_requires_saved_session(self):
+ session = self._make_unsaved_session()
+ with pytest.raises(ValidationError):
+ session.react("msg_001", "LIKE")
+
+
+# =========================================================================
+# Error handling
+# =========================================================================
+
+
+class TestSessionErrorHandling:
+ """Tests that backend errors are handled gracefully."""
+
+ def _make_session(self, ctx=None):
+ ctx = ctx or _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+ session = BoundSession.from_dict(SAMPLE_SESSION_DICT)
+ session.context = ctx
+ session._update_saved_state()
+ return session
+
+ # --- search() errors ---
+
+ def test_search_api_error_propagates(self):
+ ctx = _make_mock_context()
+ ctx.client.request.side_effect = APIError("Unauthorized", status_code=401)
+ BoundSession = _bound_session_class(ctx)
+
+ with pytest.raises(APIError, match="Unauthorized"):
+ BoundSession.search()
+
+ def test_search_dict_without_results_yields_empty_page(self):
+ # A dict response is treated as a (future) paginated envelope; one that
+ # carries no results/items key yields an empty page rather than raising.
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {"error": "something went wrong"}
+ BoundSession = _bound_session_class(ctx)
+
+ page = BoundSession.search()
+ assert page.results == []
+
+ def test_search_malformed_item_raises_resource_error(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = ["not_a_dict"]
+ BoundSession = _bound_session_class(ctx)
+
+ with pytest.raises(ResourceError, match="Failed to parse session"):
+ BoundSession.search()
+
+ # --- messages() errors ---
+
+ def test_messages_api_error_propagates(self):
+ ctx = _make_mock_context()
+ ctx.client.request.side_effect = APIError("Not Found", status_code=404)
+ session = self._make_session(ctx)
+
+ with pytest.raises(APIError, match="Not Found"):
+ session.messages()
+
+ def test_messages_non_list_response_raises_resource_error(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {"error": "bad"}
+ session = self._make_session(ctx)
+
+ with pytest.raises(ResourceError, match="Expected a list of messages"):
+ session.messages()
+
+ # --- add_message() errors ---
+
+ def test_add_message_api_error_propagates(self):
+ ctx = _make_mock_context()
+ ctx.client.request.side_effect = APIError("Bad Request", status_code=400)
+ session = self._make_session(ctx)
+
+ with pytest.raises(APIError, match="Bad Request"):
+ session.add_message(role="user", content="hello")
+
+ def test_add_message_malformed_response_raises_resource_error(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = "not_a_dict"
+ session = self._make_session(ctx)
+
+ with pytest.raises(ResourceError, match="Failed to parse session message"):
+ session.add_message(role="user", content="hello")
+
+ def test_add_message_file_upload_error_wraps_with_context(self):
+ ctx = _make_mock_context()
+ session = self._make_session(ctx)
+
+ with patch("aixplain.v2.upload_utils.FileUploader") as MockUploader:
+ instance = MockUploader.return_value
+ instance.upload.side_effect = Exception("S3 timeout")
+
+ with pytest.raises(ResourceError, match="Failed to upload file.*photo.png.*session.*sess_001"):
+ session.add_message(role="user", content="file", files=["/tmp/photo.png"])
+
+ # --- get_message() errors ---
+
+ def test_get_message_not_found(self):
+ ctx = _make_mock_context()
+ ctx.client.request.side_effect = APIError("Message not found", status_code=404)
+ session = self._make_session(ctx)
+
+ with pytest.raises(APIError, match="Message not found"):
+ session.get_message("nonexistent")
+
+ # --- delete_message() errors ---
+
+ def test_delete_message_not_found(self):
+ ctx = _make_mock_context()
+ ctx.client.request_raw.side_effect = APIError("Message not found", status_code=404)
+ session = self._make_session(ctx)
+
+ with pytest.raises(APIError, match="Message not found"):
+ session.delete_message("nonexistent")
+
+ # --- react() errors ---
+
+ def test_react_to_user_message_error(self):
+ ctx = _make_mock_context()
+ ctx.client.request.side_effect = APIError("Only assistant messages can be liked/disliked", status_code=400)
+ session = self._make_session(ctx)
+
+ with pytest.raises(APIError, match="Only assistant messages"):
+ session.react("msg_001", "LIKE")
+
+ def test_react_malformed_response_raises_resource_error(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = "bad_response"
+ session = self._make_session(ctx)
+
+ with pytest.raises(ResourceError, match="Failed to parse session message"):
+ session.react("msg_001", "LIKE")
+
+
+# =========================================================================
+# Agent — session convenience: create/list live on aix.Session now
+# =========================================================================
+
+
+class TestAgentSessionMethodsRemoved:
+ """The three collapsed create/list/id paths are gone from Agent.
+
+ Sessions are managed through ``aix.Session`` alone: create with
+ ``aix.Session(agent=…)``, find with ``aix.Session.search(agent=…)``.
+ """
+
+ @pytest.mark.parametrize("method_name", ["create_session", "list_sessions", "generate_session_id"])
+ def test_removed_agent_methods(self, method_name):
+ ctx = _make_mock_context()
+ BoundAgent = _bound_agent_class(ctx)
+ agent = BoundAgent(id="agent_99", name="Test Agent")
+ assert not hasattr(agent, method_name)
+
+ def test_session_constructor_accepts_agent(self):
+ ctx = _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+ BoundAgent = _bound_agent_class(ctx)
+ agent = BoundAgent(id="agent_99", name="Test Agent")
+
+ session = BoundSession(agent=agent, name="Chat 1")
+ assert session.agent_id == "agent_99"
+ assert session.name == "Chat 1"
+
+
+# =========================================================================
+# Core registration
+# =========================================================================
+
+
+class TestCoreSessionRegistration:
+ """Tests that Session is properly registered in Aixplain."""
+
+ def test_session_registered_on_init(self):
+ from aixplain.v2.core import Aixplain
+
+ ax = Aixplain(api_key="test_key")
+ assert ax.Session is not None
+ assert issubclass(ax.Session, Session)
+ assert ax.Session.context is ax
+
+ def test_session_unique_per_instance(self):
+ from aixplain.v2.core import Aixplain
+
+ ax1 = Aixplain(api_key="key1")
+ ax2 = Aixplain(api_key="key2")
+ assert ax1.Session is not ax2.Session
+ assert ax1.Session.context is ax1
+ assert ax2.Session.context is ax2
+
+ def test_session_enums_on_aixplain_class(self):
+ from aixplain.v2.core import Aixplain
+
+ assert Aixplain.SessionStatus is SessionStatus
+ assert Aixplain.RunStatus is RunStatus
+ assert Aixplain.MessageRole is MessageRole
+ assert Aixplain.Reaction is Reaction
+ assert Aixplain.AttachmentType is AttachmentType
+
+
+# =========================================================================
+# ExecutionConfig
+# =========================================================================
+
+
+class TestExecutionConfig:
+ """Tests for the ExecutionConfig dataclass."""
+
+ def test_defaults_are_none(self):
+ cfg = ExecutionConfig()
+ assert cfg.execution_params is None
+ assert cfg.criteria is None
+ assert cfg.evolve is None
+ assert cfg.identifier is None
+ assert cfg.run_response_generation is None
+
+ def test_to_api_dict_empty(self):
+ # An empty ExecutionConfig produces {} so callers can omit it.
+ assert ExecutionConfig().to_api_dict() == {}
+
+ def test_to_api_dict_full(self):
+ # ``max_iterations`` is deprecated and folded into budget — see the
+ # dedicated TestExecutionConfigBudget tests; kept out of this mapping test.
+ cfg = ExecutionConfig(
+ execution_params={"max_tokens": 1024, "max_time": 300},
+ criteria="Be concise",
+ evolve='{"toEvolve": true, "type": "adaptive"}',
+ identifier="user-abc",
+ run_response_generation=True,
+ )
+ assert cfg.to_api_dict() == {
+ "executionParams": {
+ "maxTokens": 1024,
+ "maxTime": 300,
+ },
+ "criteria": "Be concise",
+ "evolve": '{"toEvolve": true, "type": "adaptive"}',
+ "identifier": "user-abc",
+ "runResponseGeneration": True,
+ }
+
+ def test_to_api_dict_passes_camel_case_through(self):
+ # If callers already pass camelCase, they're left alone.
+ cfg = ExecutionConfig(execution_params={"outputFormat": "text", "maxTokens": 256})
+ assert cfg.to_api_dict() == {"executionParams": {"outputFormat": "text", "maxTokens": 256}}
+
+ def test_to_api_dict_omits_unset_fields(self):
+ cfg = ExecutionConfig(criteria="x")
+ out = cfg.to_api_dict()
+ assert out == {"criteria": "x"}
+ assert "executionParams" not in out
+ assert "runResponseGeneration" not in out
+
+ def test_run_response_generation_false_is_kept(self):
+ # Explicit False must be sent (it's meaningfully different from unset).
+ cfg = ExecutionConfig(run_response_generation=False)
+ assert cfg.to_api_dict() == {"runResponseGeneration": False}
+
+ def test_coerce_from_dict(self):
+ cfg = ExecutionConfig.coerce(
+ {
+ "execution_params": {"max_tokens": 100},
+ "criteria": "x",
+ }
+ )
+ assert isinstance(cfg, ExecutionConfig)
+ assert cfg.execution_params == {"max_tokens": 100}
+ assert cfg.criteria == "x"
+
+ def test_coerce_passthrough(self):
+ original = ExecutionConfig(criteria="x")
+ assert ExecutionConfig.coerce(original) is original
+
+ def test_coerce_none(self):
+ assert ExecutionConfig.coerce(None) is None
+
+ def test_coerce_rejects_other_types(self):
+ with pytest.raises(TypeError, match="execution_config must be"):
+ ExecutionConfig.coerce("invalid")
+
+ def test_normalize_execution_params_helper(self):
+ assert _normalize_execution_params(None) is None
+ assert _normalize_execution_params({}) == {}
+ # All known snake_case keys in the map should normalize
+ snake = {k: i for i, k in enumerate(EXECUTION_PARAMS_MAP)}
+ normalized = _normalize_execution_params(snake)
+ assert set(normalized) == set(EXECUTION_PARAMS_MAP.values())
+
+ def test_max_iterations_not_in_execution_params_map(self):
+ # max_iterations is deprecated and must never normalize to a standalone
+ # executionParams.maxIterations — it folds into budget instead.
+ assert "max_iterations" not in EXECUTION_PARAMS_MAP
+
+
+class TestExecutionConfigBudget:
+ """Budget on ExecutionConfig + deprecation of the max_iterations exec param.
+
+ Mirrors the agent run path (see test_agent_budget.py) so sessions and direct
+ runs behave identically: Budget is the single source of truth, the legacy
+ ``execution_params['max_iterations']`` is deprecated and folded into
+ ``executionParams.budget.maxIterations``, and loading a backend session never
+ emits a spurious warning.
+ """
+
+ def test_budget_instance_serializes_into_execution_params(self):
+ cfg = ExecutionConfig(budget=Budget(max_cost=0.5, max_iterations=10))
+ assert cfg.to_api_dict() == {
+ "executionParams": {"budget": {"maxCost": 0.5, "maxIterations": 10}}
+ }
+
+ def test_budget_accepts_dict(self):
+ cfg = ExecutionConfig(budget={"max_duration_seconds": 120})
+ assert isinstance(cfg.budget, Budget)
+ assert cfg.to_api_dict() == {
+ "executionParams": {"budget": {"maxDurationSeconds": 120}}
+ }
+
+ def test_deprecated_max_iterations_warns_and_folds_into_budget(self):
+ cfg = ExecutionConfig(execution_params={"max_iterations": 7})
+ with pytest.warns(DeprecationWarning, match="max_iterations"):
+ out = cfg.to_api_dict()
+ # Folded into budget; no standalone maxIterations is ever emitted.
+ assert out == {"executionParams": {"budget": {"maxIterations": 7}}}
+ assert "maxIterations" not in out["executionParams"]
+
+ def test_deprecated_camelcase_max_iterations_also_folds(self):
+ # A caller passing camelCase maxIterations directly is handled too.
+ cfg = ExecutionConfig(execution_params={"maxIterations": 7})
+ with pytest.warns(DeprecationWarning, match="max_iterations"):
+ out = cfg.to_api_dict()
+ assert out == {"executionParams": {"budget": {"maxIterations": 7}}}
+
+ def test_budget_wins_over_deprecated_max_iterations(self):
+ cfg = ExecutionConfig(
+ execution_params={"max_iterations": 99},
+ budget=Budget(max_iterations=5),
+ )
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ out = cfg.to_api_dict()
+ categories = {w.category for w in caught}
+ assert DeprecationWarning in categories # deprecated exec param
+ assert UserWarning in categories # budget-wins conflict
+ assert out["executionParams"]["budget"]["maxIterations"] == 5
+
+ def test_conflict_warning_points_at_call_site(self):
+ # The budget-wins UserWarning must resolve to the caller's to_api_dict()
+ # line, not into SDK internals (regression: the shared fold helper once
+ # hardcoded the run-path stacklevel, mislocating the session warning).
+ cfg = ExecutionConfig(
+ execution_params={"max_iterations": 99},
+ budget=Budget(max_iterations=5),
+ )
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ cfg.to_api_dict()
+ conflict = [w for w in caught if w.category is UserWarning and "precedence" in str(w.message)]
+ assert conflict, "expected a budget-wins UserWarning"
+ assert conflict[0].filename == __file__, (
+ f"conflict warning resolved to {conflict[0].filename}, not the caller"
+ )
+
+ def test_budget_merges_with_other_execution_params(self):
+ cfg = ExecutionConfig(
+ execution_params={"max_tokens": 256},
+ budget=Budget(max_cost=1.0),
+ )
+ assert cfg.to_api_dict() == {
+ "executionParams": {"maxTokens": 256, "budget": {"maxCost": 1.0}}
+ }
+
+ def test_from_dict_legacy_max_iterations_folds_silently(self):
+ # Backend payload with a legacy standalone maxIterations must NOT warn
+ # on load, and must fold into budget.
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", DeprecationWarning)
+ cfg = ExecutionConfig.from_dict(
+ {"executionParams": {"maxIterations": 4}}
+ )
+ # Round-trips to the budget shape with no standalone maxIterations.
+ out = cfg.to_api_dict()
+ assert out == {"executionParams": {"budget": {"maxIterations": 4}}}
+
+ def test_from_dict_legacy_max_iterations_defers_to_existing_budget(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", DeprecationWarning)
+ cfg = ExecutionConfig.from_dict(
+ {"executionParams": {"maxIterations": 4}, "budget": {"maxIterations": 9}}
+ )
+ assert cfg.to_api_dict()["executionParams"]["budget"]["maxIterations"] == 9
+
+ def test_session_from_dict_legacy_execution_config_no_warning(self):
+ # Loading a Session whose nested executionConfig carries a legacy
+ # maxIterations must not warn and must fold into budget.
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", DeprecationWarning)
+ session = Session.from_dict(
+ {
+ "id": "s1",
+ "agentId": "a1",
+ "executionConfig": {"executionParams": {"maxIterations": 3}},
+ }
+ )
+ out = session.execution_config.to_api_dict()
+ assert out == {"executionParams": {"budget": {"maxIterations": 3}}}
+
+
+class TestSessionExecutionConfigPayload:
+ """Tests that Session save payload carries executionConfig."""
+
+ def test_payload_omits_execution_config_when_unset(self):
+ session = Session(agent_id="a1", name="N")
+ payload = session.build_save_payload()
+ assert "executionConfig" not in payload
+
+ def test_payload_includes_execution_config(self):
+ cfg = ExecutionConfig(
+ execution_params={"output_format": "text", "max_tokens": 512},
+ criteria="be helpful",
+ evolve='{"toEvolve": true}',
+ identifier="abc",
+ run_response_generation=True,
+ )
+ session = Session(agent_id="a1", name="N", execution_config=cfg)
+ payload = session.build_save_payload()
+ assert payload["executionConfig"] == {
+ "executionParams": {"outputFormat": "text", "maxTokens": 512},
+ "criteria": "be helpful",
+ "evolve": '{"toEvolve": true}',
+ "identifier": "abc",
+ "runResponseGeneration": True,
+ }
+
+ def test_session_accepts_dict_execution_config(self):
+ # Session(execution_config={...}) should coerce to ExecutionConfig
+ # so callers don't need to import the dataclass.
+ session = Session(
+ agent_id="a1",
+ execution_config={"criteria": "x", "execution_params": {"max_tokens": 8}},
+ )
+ assert isinstance(session.execution_config, ExecutionConfig)
+ payload = session.build_save_payload()
+ assert payload["executionConfig"] == {
+ "executionParams": {"maxTokens": 8},
+ "criteria": "x",
+ }
+
+ def test_save_sends_execution_config_to_api(self):
+ ctx = _make_mock_context()
+ ctx.client.request.return_value = {**SAMPLE_SESSION_DICT, "id": "new_sess"}
+ BoundSession = _bound_session_class(ctx)
+ # Legacy max_iterations exec param folds into budget on send (deprecated).
+ cfg = ExecutionConfig(
+ execution_params={"max_iterations": 10},
+ run_response_generation=True,
+ )
+ session = BoundSession(agent_id="a1", name="N", execution_config=cfg)
+
+ with pytest.warns(DeprecationWarning, match="max_iterations"):
+ session.save()
+
+ _, kwargs = ctx.client.request.call_args
+ assert kwargs["json"]["executionConfig"] == {
+ "executionParams": {"budget": {"maxIterations": 10}},
+ "runResponseGeneration": True,
+ }
+
+
+class TestSessionExecutionConfigConstruction:
+ """ExecutionConfig is attached to a Session at construction time.
+
+ (The old ``agent.create_session(execution_params=…, criteria=…)`` shortcut
+ plumbing was removed; callers build an ``ExecutionConfig`` and pass it to
+ ``Session(execution_config=…)``, or pass per-run overrides to
+ ``agent.run(query, session=…)``.)
+ """
+
+ def test_shortcut_style_execution_config(self):
+ ctx = _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+
+ cfg = ExecutionConfig(
+ execution_params={"output_format": "text", "max_tokens": 1024},
+ criteria="Be helpful",
+ evolve='{"toEvolve": true}',
+ identifier="user-abc",
+ run_response_generation=True,
+ )
+ session = BoundSession(agent_id="agent_99", name="With Config", execution_config=cfg)
+
+ assert session.execution_config is cfg
+ assert cfg.execution_params == {"output_format": "text", "max_tokens": 1024}
+ assert cfg.criteria == "Be helpful"
+ assert cfg.evolve == '{"toEvolve": true}'
+ assert cfg.identifier == "user-abc"
+ assert cfg.run_response_generation is True
+
+ def test_dict_execution_config_is_coerced(self):
+ ctx = _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+
+ session = BoundSession(agent_id="agent_99", execution_config={"criteria": "x"})
+ assert isinstance(session.execution_config, ExecutionConfig)
+ assert session.execution_config.criteria == "x"
+
+ def test_no_execution_config_by_default(self):
+ ctx = _make_mock_context()
+ BoundSession = _bound_session_class(ctx)
+
+ session = BoundSession(agent_id="agent_99", name="plain")
+ assert session.execution_config is None
+
+ def test_run_response_generation_false_is_kept(self):
+ cfg = ExecutionConfig(run_response_generation=False)
+ assert cfg.run_response_generation is False
diff --git a/tests/unit/v2/test_snake_case_conventions.py b/tests/unit/v2/test_snake_case_conventions.py
index 5775a800..009d4ad7 100644
--- a/tests/unit/v2/test_snake_case_conventions.py
+++ b/tests/unit/v2/test_snake_case_conventions.py
@@ -131,8 +131,11 @@ class TestAgentRunParamsKeys:
def test_snake_case_keys_exist(self):
hints = AgentRunParams.__annotations__
- assert "session_id" in hints
- assert "allow_history_and_session_id" in hints
+ # ``session`` (polymorphic Session/id) replaced the removed id-only
+ # ``session_id`` and the ``allow_history_and_session_id`` flag.
+ assert "session" in hints
+ assert "session_id" not in hints
+ assert "allow_history_and_session_id" not in hints
assert "execution_params" in hints
assert "run_response_generation" in hints
@@ -180,34 +183,29 @@ def _make_agent(self):
agent.expected_output = ""
return agent
- def test_session_id_becomes_camelcase(self):
+ def test_run_response_generation_becomes_camelcase(self):
agent = self._make_agent()
payload = agent.build_run_payload(
query="hello",
- session_id="sess-1",
+ run_response_generation=True,
)
- assert payload["sessionId"] == "sess-1"
- assert "session_id" not in payload
-
- def test_allow_history_becomes_camelcase(self):
- agent = self._make_agent()
- payload = agent.build_run_payload(
- query="hello",
- allow_history_and_session_id=True,
- )
- assert payload["allowHistoryAndSessionId"] is True
- assert "allow_history_and_session_id" not in payload
+ assert payload["runResponseGeneration"] is True
+ assert "run_response_generation" not in payload
def test_execution_params_snake_case_converted(self):
agent = self._make_agent()
- payload = agent.build_run_payload(
- query="hello",
- execution_params={"max_tokens": 100, "max_iterations": 5},
- )
+ with pytest.warns(DeprecationWarning):
+ payload = agent.build_run_payload(
+ query="hello",
+ execution_params={"max_tokens": 100, "max_iterations": 5},
+ )
assert payload["executionParams"]["maxTokens"] == 100
- assert payload["executionParams"]["maxIterations"] == 5
assert "max_tokens" not in payload["executionParams"]
+ # Deprecated ``max_iterations`` now folds into executionParams.budget and
+ # the standalone key is no longer emitted.
+ assert "maxIterations" not in payload["executionParams"]
assert "max_iterations" not in payload["executionParams"]
+ assert payload["executionParams"]["budget"]["maxIterations"] == 5
def test_execution_params_camelcase_still_works(self):
"""Backwards compat: camelCase keys pass through unchanged."""
diff --git a/tests/unit/v2/test_trigger.py b/tests/unit/v2/test_trigger.py
new file mode 100644
index 00000000..a5770643
--- /dev/null
+++ b/tests/unit/v2/test_trigger.py
@@ -0,0 +1,343 @@
+"""Unit tests for the v2 Trigger module (aixplain.v2.trigger)."""
+
+import warnings
+import pytest
+from unittest.mock import Mock
+
+from aixplain.v2.trigger import (
+ Trigger,
+ TriggerConfiguration,
+ TriggerRepeatRule,
+ _normalize_weekdays,
+ _normalize_monthdays,
+ _strip_none,
+)
+from aixplain.v2.integration import TriggerTypeSpec, TriggerEventOption, TriggerTypes
+from aixplain.v2.resource import Page
+
+
+AGENT_ID = "66a060000000000000000000"
+MODEL_URL = "https://models.aixplain.com/api/v2/execute"
+
+
+class FakeAgent:
+ """Minimal agent stand-in with an id."""
+
+ id = AGENT_ID
+
+
+def make_context(client=None):
+ """Build a mock Aixplain context with a client and model_url."""
+ ctx = Mock()
+ ctx.model_url = MODEL_URL
+ ctx.client = client or Mock()
+ return ctx
+
+
+def bound_trigger_class(context):
+ """Return a Trigger subclass bound to a context (as core.init_resources does)."""
+ return type("Trigger", (Trigger,), {"context": context})
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+class TestNormalizers:
+ def test_weekdays_codes_and_full_names(self):
+ assert _normalize_weekdays(["mon", "Thursday", "SUN"]) == ["mon", "thu", "sun"]
+
+ def test_weekdays_single_string(self):
+ assert _normalize_weekdays("fri") == ["fri"]
+
+ def test_weekdays_invalid(self):
+ with pytest.raises(ValueError):
+ _normalize_weekdays(["funday"])
+
+ def test_monthdays_ok(self):
+ assert _normalize_monthdays([1, "15", 31]) == [1, 15, 31]
+
+ def test_monthdays_out_of_range(self):
+ with pytest.raises(ValueError):
+ _normalize_monthdays([0])
+ with pytest.raises(ValueError):
+ _normalize_monthdays([32])
+
+ def test_strip_none_nested(self):
+ assert _strip_none({"a": 1, "b": None, "c": {"d": None, "e": 2}}) == {"a": 1, "c": {"e": 2}}
+
+
+# =============================================================================
+# Schedule mapping -> configuration (the five story examples)
+# =============================================================================
+
+
+class TestScheduleMapping:
+ def _cfg(self, **kw):
+ return Trigger(name="t", agent=FakeAgent(), input="q", **kw).configuration
+
+ def test_once(self):
+ cfg = self._cfg(run_at="2026-01-26T12:00:00Z")
+ assert cfg.type == "once" and cfg.run_at == "2026-01-26T12:00:00Z"
+
+ def test_daily(self):
+ cfg = self._cfg(every="day", at="09:00", timezone="Europe/London")
+ assert cfg.type == "daily" and cfg.time == "09:00" and cfg.timezone == "Europe/London"
+
+ def test_interval_hours(self):
+ cfg = self._cfg(every="hour", interval=2)
+ assert cfg.type == "recurring" and cfg.repeat.every == 2 and cfg.repeat.unit == "hour"
+
+ def test_weekly(self):
+ cfg = self._cfg(every="week", on=["mon", "thu"], at="17:00")
+ assert cfg.type == "weekly" and cfg.days_of_week == ["mon", "thu"] and cfg.time == "17:00"
+
+ def test_monthly(self):
+ cfg = self._cfg(every="month", on=[1, 15], at="09:00")
+ assert cfg.type == "monthly" and cfg.days_of_month == [1, 15] and cfg.time == "09:00"
+
+ def test_day_interval_gt1_uses_recurring(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ cfg = self._cfg(every="day", interval=3, at="09:00")
+ assert cfg.type == "recurring" and cfg.repeat.unit == "day" and cfg.repeat.every == 3
+
+ def test_week_interval_gt1_uses_recurring(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ cfg = self._cfg(every="week", interval=2, on=["mon"])
+ assert cfg.type == "recurring" and cfg.repeat.unit == "week"
+
+ def test_start_at_on_recurring(self):
+ cfg = self._cfg(every="minute", interval=5, start_at="2026-01-01T00:00:00Z")
+ assert cfg.start_at == "2026-01-01T00:00:00Z"
+
+ @pytest.mark.parametrize(
+ "kwargs",
+ [
+ dict(every="month", interval=2, on=[1]),
+ dict(every="week", at="09:00"), # missing on
+ dict(every="month", at="09:00"), # missing on
+ dict(every="fortnight"),
+ ],
+ )
+ def test_invalid_combos_raise(self, kwargs):
+ with pytest.raises(ValueError):
+ Trigger(name="t", agent=FakeAgent(), input="q", **kwargs)
+
+
+# =============================================================================
+# Payload building
+# =============================================================================
+
+
+class TestBuildSavePayload:
+ def test_time_payload_defaults_enabled_true(self):
+ p = Trigger(name="Daily", agent=FakeAgent(), input="Sum", every="day", at="09:00").build_save_payload()
+ assert p["triggerType"] == "time"
+ assert p["assetId"] == AGENT_ID and p["assetType"] == "agent"
+ assert p["enabled"] is True and p["notifications"] is False
+ assert p["configuration"] == {"type": "daily", "time": "09:00"}
+ assert "id" not in p
+
+ def test_notifications_and_enabled_passthrough(self):
+ p = Trigger(
+ name="x", agent=FakeAgent(), input="q", run_at="2026-01-26T12:00:00Z",
+ notifications=True, enabled=False,
+ ).build_save_payload()
+ assert p["notifications"] is True and p["enabled"] is False
+
+ def test_update_payload_includes_id(self):
+ t = Trigger(name="x", agent=FakeAgent(), input="q", run_at="2026-01-26T12:00:00Z")
+ t.id = "abc123"
+ p = t.build_save_payload()
+ assert p["id"] == "abc123"
+
+ def test_external_payload_uses_trigger_id(self):
+ opt = TriggerEventOption(slug="NEW_EMAIL", connection_id="conn1")
+ t = Trigger(name="Triage", agent=FakeAgent(), input="Triage", event=opt)
+ t.trigger_id = "ti_real" # normally filled by activation
+ p = t.build_save_payload()
+ assert p["triggerType"] == "external" and p["triggerId"] == "ti_real"
+ assert "configuration" not in p
+
+ def test_time_trigger_without_schedule_raises(self):
+ t = Trigger(name="x", agent=FakeAgent(), input="q")
+ with pytest.raises(ValueError):
+ t.build_save_payload()
+
+
+# =============================================================================
+# Rehydration (get/search deserialization)
+# =============================================================================
+
+
+class TestRehydration:
+ def _dto(self, **over):
+ dto = {
+ "id": "trig1", "name": "Daily digest", "description": "d",
+ "triggerType": "time", "assetId": AGENT_ID, "assetType": "agent",
+ "type": "daily", "input": "Sum", "enabled": True, "notifications": False,
+ "nextRunAt": "2026-07-14T09:00:00Z", "lastRunAt": None, "failureCount": 0,
+ "enqueued": False, "createdAt": "2026-07-13T00:00:00Z", "updatedAt": "2026-07-13T00:00:00Z",
+ "configuration": {"type": "daily", "time": "09:00", "timezone": "Europe/London", "repeat": None},
+ }
+ dto.update(over)
+ return dto
+
+ def test_from_dict_maps_fields(self):
+ t = Trigger.from_dict(self._dto())
+ assert t.id == "trig1" and t.asset_id == AGENT_ID and t.trigger_type == "time"
+ assert t.schedule_type == "daily" and t.configuration.time == "09:00"
+ assert t.next_run_at == "2026-07-14T09:00:00Z"
+
+ def test_post_init_not_run_on_rehydration(self):
+ # id present -> friendly translation skipped, enabled preserved from DTO (not forced True)
+ t = Trigger.from_dict(self._dto(enabled=False))
+ assert t.enabled is False
+
+ def test_external_rehydration_roundtrips_trigger_id(self):
+ t = Trigger.from_dict(self._dto(triggerType="external", triggerId="ti_x", configuration=None, type=None))
+ p = t.build_save_payload()
+ assert p["triggerType"] == "external" and p["triggerId"] == "ti_x"
+
+
+# =============================================================================
+# search() -> GET /v1/triggers?agentId=
+# =============================================================================
+
+
+class TestSearch:
+ def test_search_by_agent_sends_agentId_param(self):
+ client = Mock()
+ client.get.return_value = [
+ {"id": "t1", "name": "a", "triggerType": "time", "assetId": AGENT_ID,
+ "configuration": {"type": "once", "runAt": "2026-01-01T00:00:00Z"}},
+ {"id": "t2", "name": "b", "triggerType": "time", "assetId": AGENT_ID,
+ "configuration": {"type": "daily", "time": "09:00"}},
+ ]
+ cls = bound_trigger_class(make_context(client))
+ page = cls.search(agent=FakeAgent())
+
+ client.get.assert_called_once_with("v1/triggers", params={"agentId": AGENT_ID})
+ assert isinstance(page, Page)
+ assert page.total == 2 and page.page_total == 1
+ assert [t.id for t in page.results] == ["t1", "t2"]
+
+ def test_search_no_agent_omits_params(self):
+ client = Mock()
+ client.get.return_value = []
+ cls = bound_trigger_class(make_context(client))
+ page = cls.search()
+ client.get.assert_called_once_with("v1/triggers")
+ assert page.total == 0
+
+ def test_search_agent_id_string(self):
+ client = Mock()
+ client.get.return_value = []
+ cls = bound_trigger_class(make_context(client))
+ cls.search(agent_id="deadbeef")
+ client.get.assert_called_once_with("v1/triggers", params={"agentId": "deadbeef"})
+
+
+# =============================================================================
+# External activation on save + deactivation on delete
+# =============================================================================
+
+
+class TestEventLifecycle:
+ def test_activation_called_on_create(self):
+ client = Mock()
+ # 1st call: activate_trigger -> returns composio trigger_id; 2nd: POST /v1/triggers
+ client.request.side_effect = [
+ {"completed": True, "data": {"trigger_id": "ti_activated"}},
+ {"id": "trigNew", "name": "Triage", "triggerType": "external", "triggerId": "ti_activated"},
+ ]
+ ctx = make_context(client)
+ opt = TriggerEventOption(slug="NEW_EMAIL", connection_id="conn1")
+ t = Trigger(name="Triage", agent=FakeAgent(), input="Triage", event=opt)
+ t.context = ctx
+ t.save()
+
+ # First request = activation on the connection's execute URL
+ first = client.request.call_args_list[0]
+ assert first.args[0] == "post"
+ assert first.args[1] == f"{MODEL_URL}/conn1"
+ body = first.kwargs["json"]
+ assert body["action"] == "activate_trigger"
+ assert body["data"] == {"slug": "NEW_EMAIL", "config": {}}
+ assert body["enable"] is True and body["toolkit_versions"] == "latest"
+
+ # Second request = create on /v1/triggers with the real trigger id
+ second = client.request.call_args_list[1]
+ assert second.args[0] == "post" and second.args[1] == "v1/triggers"
+ assert second.kwargs["json"]["triggerId"] == "ti_activated"
+ assert t.trigger_id == "ti_activated" and t.id == "trigNew"
+
+ def test_activation_without_connection_raises(self):
+ opt = TriggerEventOption(slug="NEW_EMAIL") # no connection_id
+ t = Trigger(name="Triage", agent=FakeAgent(), input="Triage", event=opt)
+ t.context = make_context()
+ with pytest.raises(ValueError):
+ t.save()
+
+ def test_delete_deactivates_composio(self):
+ client = Mock()
+ client.request.return_value = {"completed": True, "data": {"deleted": True}}
+ client.request_raw.return_value = Mock()
+ ctx = make_context(client)
+ t = Trigger.from_dict({"id": "trigX", "name": "e", "triggerType": "external", "triggerId": "ti_x"})
+ t.context = ctx
+ t.connection_id = "conn1" # known in-session
+ t.delete()
+
+ deact = client.request.call_args_list[0]
+ assert deact.kwargs["json"]["action"] == "delete_trigger"
+ assert deact.kwargs["json"]["data"] == {"trigger_id": "ti_x"}
+ client.request_raw.assert_called_once()
+ assert client.request_raw.call_args.args[0] == "delete"
+
+ def test_delete_external_without_connection_warns(self):
+ client = Mock()
+ client.request_raw.return_value = Mock()
+ ctx = make_context(client)
+ t = Trigger.from_dict({"id": "trigX", "name": "e", "triggerType": "external", "triggerId": "ti_x"})
+ t.context = ctx
+ with pytest.warns(UserWarning):
+ t.delete()
+ client.request.assert_not_called() # no deactivation attempted
+
+
+# =============================================================================
+# Discovery collection (integration.triggers / tool.triggers)
+# =============================================================================
+
+
+class TestTriggerTypesCollection:
+ def _specs(self):
+ return [
+ TriggerTypeSpec(slug="NEW_EMAIL", name="New email", description="fires on email"),
+ TriggerTypeSpec(slug="NEW_LABEL", name="New label"),
+ ]
+
+ def test_getitem_case_insensitive(self):
+ col = TriggerTypes(self._specs(), connection_id="conn1")
+ opt = col["new_email"]
+ assert isinstance(opt, TriggerEventOption)
+ assert opt.slug == "NEW_EMAIL" and opt.connection_id == "conn1"
+
+ def test_contains_and_len_and_iter(self):
+ col = TriggerTypes(self._specs())
+ assert "NEW_EMAIL" in col and "nope" not in col
+ assert len(col) == 2
+ assert set(iter(col)) == {"NEW_EMAIL", "NEW_LABEL"}
+
+ def test_missing_raises_keyerror(self):
+ col = TriggerTypes(self._specs())
+ with pytest.raises(KeyError):
+ col["DOES_NOT_EXIST"]
+
+ def test_discovery_only_has_no_connection(self):
+ col = TriggerTypes(self._specs()) # from an unconnected integration
+ assert col["NEW_EMAIL"].connection_id is None
diff --git a/tests/unit/v2/test_v2_agent_v3_compat.py b/tests/unit/v2/test_v2_agent_v3_compat.py
index 000dde64..c26f6d60 100644
--- a/tests/unit/v2/test_v2_agent_v3_compat.py
+++ b/tests/unit/v2/test_v2_agent_v3_compat.py
@@ -7,12 +7,12 @@
"""
import json
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
import pytest
from aixplain.v2.agent import Agent
-from aixplain.v2.inspector import Inspector, PrebuiltInspector
+from aixplain.v2.inspector import Inspector
def _agent_from_dict(**overrides):
@@ -88,29 +88,108 @@ def test_inspectors_deserialize_to_inspector_objects(self):
assert isinstance(agent.inspectors[0], Inspector)
assert agent.inspectors[0].name == "Input Gate"
- def test_prebuilt_inspector_payload_deserializes_to_prebuilt_inspector(self):
- # PrebuiltInspector.to_dict() emits a lightweight reference with no
- # "name" key; the backend echoes it back on save. __post_init__ must
- # not route it through Inspector.from_dict (KeyError: 'name').
- prebuilt_payload = PrebuiltInspector.prompt_injection_guard().to_dict()
- assert "name" not in prebuilt_payload # precondition for the regression
-
- agent = _agent_from_dict(inspectors=[prebuilt_payload])
+ def test_prebuilt_guard_payload_deserializes_to_inspector(self):
+ # A prebuilt guard fetched via aix.Inspector.get(...) serializes to the
+ # same inspector shape as a custom inspector — the backend echoes it back
+ # on save, and __post_init__ rehydrates it through Inspector.from_dict.
+ guard_payload = {
+ "name": "Prompt Injection Guard",
+ "targets": ["input"],
+ "action": {"type": "abort"},
+ "evaluator": {"type": "asset", "assetId": "pi-id"},
+ }
+ agent = _agent_from_dict(inspectors=[guard_payload])
- assert isinstance(agent.inspectors[0], PrebuiltInspector)
- assert agent.inspectors[0].preset_id == "prompt_injection_guard"
+ assert isinstance(agent.inspectors[0], Inspector)
+ assert agent.inspectors[0].name == "Prompt Injection Guard"
- def test_mixed_full_and_prebuilt_inspectors_deserialize(self):
- full_payload = {
+ def test_mixed_prebuilt_and_custom_inspectors_deserialize(self):
+ custom_payload = {
"name": "Input Gate",
"targets": ["input"],
"action": {"type": "abort"},
"evaluator": {"type": "asset", "assetId": "model-abc", "prompt": "PASS or FAIL"},
}
- prebuilt_payload = {"presetId": "pii_redaction", "targets": ["output"]}
+ prebuilt_payload = {
+ "name": "PII",
+ "targets": ["output"],
+ "action": {"type": "edit"},
+ "evaluator": {"type": "asset", "assetId": "pii-id"},
+ "editor": {"type": "asset", "assetId": "pii-id"},
+ }
- agent = _agent_from_dict(inspectors=[full_payload, prebuilt_payload])
+ agent = _agent_from_dict(inspectors=[custom_payload, prebuilt_payload])
- assert isinstance(agent.inspectors[0], Inspector)
- assert isinstance(agent.inspectors[1], PrebuiltInspector)
- assert agent.inspectors[1].preset_id == "pii_redaction"
+ assert all(isinstance(i, Inspector) for i in agent.inspectors)
+ assert agent.inspectors[1].name == "PII"
+
+ def test_agent_with_inspector_serializes(self):
+ # Regression: Inspector is a BaseResource, whose `context` field is
+ # init=False. An inspector rehydrated via from_dict (nested in an agent)
+ # has no bound context, and dataclasses_json's _asdict calls
+ # getattr(obj, "context") before honoring `exclude` — which raised
+ # AttributeError on save() until `context` got a None default. Saving a
+ # team agent with a fetched guard hit exactly this path.
+ guard_payload = {
+ "name": "Detect Prompt Attacks Guardrail",
+ "targets": ["input"],
+ "action": {"type": "abort"},
+ "evaluator": {"type": "asset", "assetId": "pi-id"},
+ }
+ agent = _agent_from_dict(inspectors=[guard_payload])
+
+ # The inspector must always expose a context attribute (None when unbound)
+ # and must not leak it into serialized output.
+ assert agent.inspectors[0].context is None
+ serialized = agent.to_dict()
+ assert "context" not in serialized["inspectors"][0]
+
+
+class TestRunPayloadAttachments:
+ """Non-session run path resolves attachments into a structured payload field."""
+
+ def test_url_attachments_not_uploaded_and_type_auto_detected(self):
+ agent = _agent_from_dict()
+ with patch("aixplain.v2.upload_utils.FileUploader") as MockUploader:
+ payload = agent.build_run_payload(
+ query="describe",
+ attachments=[{"url": "https://s3/a.wav", "type": "audio"}, "https://s3/b.png"],
+ )
+ MockUploader.assert_not_called()
+ # Explicit type preserved (+ mime inferred); bare URL string auto-typed as image.
+ assert payload["attachments"][0] == {"url": "https://s3/a.wav", "type": "audio", "mimeType": "audio/wav"}
+ assert payload["attachments"][1]["url"] == "https://s3/b.png"
+ assert payload["attachments"][1]["type"] == "image"
+ assert payload["attachments"][1]["mimeType"] == "image/png"
+
+ def test_local_path_is_uploaded_with_mimetype(self):
+ agent = _agent_from_dict()
+ with (
+ patch("aixplain.v2.upload_utils.FileUploader") as MockUploader,
+ patch("aixplain.v2.upload_utils.MimeTypeDetector") as MockDetector,
+ ):
+ MockUploader.return_value.upload.return_value = "https://cdn/clip.wav"
+ MockDetector.detect_mime_type.return_value = "audio/wav"
+ payload = agent.build_run_payload(query="transcribe", attachments=["/tmp/clip.wav"])
+
+ assert payload["attachments"] == [
+ {"url": "https://cdn/clip.wav", "name": "clip.wav", "type": "audio", "mimeType": "audio/wav"}
+ ]
+
+ def test_files_kwarg_still_works_but_warns(self):
+ agent = _agent_from_dict()
+ with (
+ patch("aixplain.v2.upload_utils.FileUploader") as MockUploader,
+ patch("aixplain.v2.upload_utils.MimeTypeDetector") as MockDetector,
+ pytest.warns(DeprecationWarning, match="files.*deprecated"),
+ ):
+ MockUploader.return_value.upload.return_value = "https://cdn/doc.pdf"
+ MockDetector.detect_mime_type.return_value = "application/pdf"
+ payload = agent.build_run_payload(query="q", files=["/tmp/doc.pdf"])
+
+ assert payload["attachments"][0]["url"] == "https://cdn/doc.pdf"
+
+ def test_no_attachments_means_no_attachments_field(self):
+ agent = _agent_from_dict()
+ payload = agent.build_run_payload(query="hi")
+ assert "attachments" not in payload
diff --git a/uv.lock b/uv.lock
index e91a4f78..a6300268 100644
--- a/uv.lock
+++ b/uv.lock
@@ -20,9 +20,10 @@ resolution-markers = [
[[package]]
name = "aixplain"
-version = "0.2.42"
+version = "0.2.45rc1"
source = { editable = "." }
dependencies = [
+ { name = "babel" },
{ name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
{ name = "click", version = "8.3.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
{ name = "dataclasses-json" },
@@ -54,11 +55,13 @@ test = [
{ name = "pytest-mock" },
{ name = "pytest-rerunfailures", version = "16.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
{ name = "pytest-rerunfailures", version = "16.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
+ { name = "pytest-xdist" },
{ name = "requests-mock" },
]
[package.metadata]
requires-dist = [
+ { name = "babel", specifier = ">=2.12.0" },
{ name = "click", specifier = ">=7.1.2" },
{ name = "dataclasses-json", specifier = ">=0.5.2" },
{ name = "docker", marker = "extra == 'test'", specifier = ">=6.1.3" },
@@ -72,6 +75,7 @@ requires-dist = [
{ name = "pytest", marker = "extra == 'test'", specifier = ">=6.1.0" },
{ name = "pytest-mock", marker = "extra == 'test'", specifier = ">=3.10.0" },
{ name = "pytest-rerunfailures", marker = "extra == 'test'", specifier = ">=16.0" },
+ { name = "pytest-xdist", marker = "extra == 'test'" },
{ name = "python-dotenv", specifier = ">=1.0.0" },
{ name = "pyyaml", specifier = ">=6.0.1" },
{ name = "requests", specifier = ">=2.1.0" },
@@ -115,6 +119,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/38/0e/27be9fdef66e72d64c0cdc3cc2823101b80585f8119b5c112c2e8f5f7dab/anyio-4.12.1-py3-none-any.whl", hash = "sha256:d405828884fc140aa80a3c667b8beed277f1dfedec42ba031bd6ac3db606ab6c", size = 113592, upload-time = "2026-01-06T11:45:19.497Z" },
]
+[[package]]
+name = "babel"
+version = "2.18.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7d/b2/51899539b6ceeeb420d40ed3cd4b7a40519404f9baf3d4ac99dc413a834b/babel-2.18.0.tar.gz", hash = "sha256:b80b99a14bd085fcacfa15c9165f651fbb3406e66cc603abf11c5750937c992d", size = 9959554, upload-time = "2026-02-01T12:30:56.078Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/77/f5/21d2de20e8b8b0408f0681956ca2c69f1320a3848ac50e6e7f39c6159675/babel-2.18.0-py3-none-any.whl", hash = "sha256:e2b422b277c2b9a9630c1d7903c2a00d0830c409c59ac8cae9081c92f1aeba35", size = 10196845, upload-time = "2026-02-01T12:30:53.445Z" },
+]
+
[[package]]
name = "backports-tarfile"
version = "1.2.0"
@@ -535,6 +548,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" },
]
+[[package]]
+name = "execnet"
+version = "2.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/bf/89/780e11f9588d9e7128a3f87788354c7946a9cbb1401ad38a48c4db9a4f07/execnet-2.1.2.tar.gz", hash = "sha256:63d83bfdd9a23e35b9c6a3261412324f964c2ec8dcd8d3c6916ee9373e0befcd", size = 166622, upload-time = "2025-11-12T09:56:37.75Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" },
+]
+
[[package]]
name = "fastapi"
version = "0.128.8"
@@ -2155,6 +2177,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/77/54/60eabb34445e3db3d3d874dc1dfa72751bfec3265bd611cb13c8b290adea/pytest_rerunfailures-16.1-py3-none-any.whl", hash = "sha256:5d11b12c0ca9a1665b5054052fcc1084f8deadd9328962745ef6b04e26382e86", size = 14093, upload-time = "2025-10-10T07:06:00.019Z" },
]
+[[package]]
+name = "pytest-xdist"
+version = "3.8.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "execnet" },
+ { name = "pytest", version = "8.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+ { name = "pytest", version = "9.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" },
+]
+
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"
+ 
+
-
+