> ## Documentation Index
> Fetch the complete documentation index at: https://docs.cognite.com/llms.txt
> Use this file to discover all available pages before exploring further.
# Run Python code
> Reference information for using the Run Python code tool in Atlas AI agents, including supported parameter types, code structure requirements, and best practices
The **Run Python code** tool enables Atlas AI agents to run custom Python scripts with runtime arguments. The tool automatically validates your code, generates contracts for the agent from Python type annotations, and provides testing capabilities within the Atlas AI interface.
## Required code structure
To ensure your script works with the Atlas AI agent, you must define specific entry points and follow the required function signatures.
### The handle() function
Every Python script must include a `handle()` function as the entry point. Name the function `handle`, annotate all parameters with their types, and include a docstring following the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods).
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle(param1: str, param2: int) -> str:
"""
Brief description of what the function does.
Args:
param1 (str): Description of the first parameter.
param2 (int): Description of the second parameter.
Returns:
str: Description of what is returned.
"""
# Your implementation here
return f"Processed {param1} with value {param2}"
```
**Requirements:**
* Name the function `handle`.
* Annotate all parameters with their types. See [Supported parameter types](#supported-parameter-types).
* Include a docstring following the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods).
* The return type annotation is recommended but not required.
* The function can be synchronous (`def handle()`) or asynchronous (`async def handle()`).
The AI agent uses only the parameter descriptions from the `Args:` section of the docstring. It does not use the function description or return value description. To provide descriptions for the agent on how to use the tool, use the Instructions field in the agent configuration interface.
### Async handle() function
You can define `handle()` as an async function using `async def` to perform asynchronous operations.
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
import asyncio
async def handle(value: int) -> int:
"""Async handle function."""
await asyncio.sleep(0.1)
return value * 2
```
**When to use async handle:**
* **Eventual consistency**: When you create a resource and need to wait before retrieving it to ensure it's available.
* **Delayed operations**: When you need to introduce delays between operations.
* **Async SDK operations**: When using async methods from libraries that support async/await patterns.
**Requirements for async handle:**
* Use `async def` instead of `def` for the function definition.
* Use `await` for asynchronous operations like `asyncio.sleep()`.
* All the same requirements apply as for synchronous `handle()` functions (type annotations, docstrings, etc.).
* The same parameter types and return types are supported for both sync and async functions.
### The test() function
Include a `test()` function for validation and testing. Use the Test code button in the Atlas AI agent builder to trigger the function.
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def test():
"""Test function that runs when you click 'Test code' in the UI."""
# You can use assertions
assert handle("test", 42) == "Processed test with value 42"
# Or return a value to display in the console
return handle("example", 100)
```
**Requirements:**
* Name the function `test`.
* No parameters are allowed.
* The function can be synchronous (`def test()`) or asynchronous (`async def test()`).
* The function can use assertions or return values for display.
* The function runs when you click the Test code button in the UI.
### Async test() function
You can define `test()` as an async function to test async operations or await async `handle()` functions:
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
import asyncio
async def handle(value: int) -> int:
"""Async handle function."""
await asyncio.sleep(0.1)
return value * 2
async def test():
"""Async test function that tests async handle."""
# Await the async handle function directly
result = await handle(21)
assert result == 42, f"Expected 42, got {result}"
return result
```
**When to use async test:**
* When testing async `handle()` functions you can await them directly.
* When you need to perform async operations in your test setup or teardown.
* When testing scenarios that involve delays or eventual consistency.
## Supported parameter types
Parameters can be of primitive types, `DateTime`, `NodeId` and lists of these. Primitive types can have a default value.
### Primitive types
| Python type | JSON schema type | Example usage | Runtime value |
| ----------- | ---------------- | -------------------- | ------------- |
| `str` | `"string"` | `name: str` | `"John Doe"` |
| `int` | `"integer"` | `age: int` | `25` |
| `float` | `"number"` | `temperature: float` | `23.5` |
| `bool` | `"boolean"` | `is_active: bool` | `true` |
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle(
name: str,
age: int,
temperature: float,
is_active: bool
) -> dict:
"""
Example with primitive types.
Args:
name (str): User's name.
age (int): User's age in years.
temperature (float): Temperature in Celsius.
is_active (bool): Whether user is active.
Returns:
dict: Processed user data.
"""
return {
"name": name,
"age": age,
"temperature": temperature,
"is_active": is_active
}
```
### Default values
Primitive types support default values:
```python theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle(
message: str,
count: int = 1,
multiplier: float = 1.0,
enabled: bool = True
) -> str:
"""
Example with default values.
Args:
message (str): The message to process.
count (int): Number of repetitions (default=1).
multiplier (float): Multiplier factor (default=1.0).
enabled (bool): Whether processing is enabled (default=True).
Returns:
str: Processed message.
"""
if not enabled:
return "Processing disabled"
return (message * count) * int(multiplier)
```
### List types
We support lists of primitive types:
```python theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle(
names: list[str],
values: list[int],
temperatures: list[float]
) -> dict:
"""
Example with list parameters.
Args:
names (list[str]): List of names.
values (list[int]): List of integer values.
temperatures (list[float]): List of temperature readings.
Returns:
dict: Summary of the lists.
"""
return {
"name_count": len(names),
"total_value": sum(values),
"avg_temperature": sum(temperatures) / len(temperatures) if temperatures else 0
}
```
We don't support nested lists (e.g., `list[list[str]]`).
### DateTime type
For temporal data use `datetime`:
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
import datetime
def handle(
start_time: datetime.datetime,
end_time: datetime.datetime,
timestamps: list[datetime.datetime]
) -> dict:
"""
Example with datetime parameters.
Args:
start_time (datetime.datetime): Start of time range.
end_time (datetime.datetime): End of time range.
timestamps (list[datetime.datetime]): List of timestamps.
Returns:
dict: Time range information.
"""
duration = end_time - start_time
return {
"duration_seconds": duration.total_seconds(),
"timestamp_count": len(timestamps)
}
```
### NodeId type
To work with Cognite Data Fusion (CDF) instances from data models, use [`NodeId`](https://cognite-sdk-python.readthedocs-hosted.com/en/latest/data_modeling.html) to identify the instances:
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
from cognite.client.data_classes.data_modeling import NodeId
def handle(asset_id: NodeId, related_asset_ids: list[NodeId]) -> dict:
"""
Example with NodeId parameters.
Args:
asset_id (NodeId): Primary asset identifier with space and external_id.
related_asset_ids (list[NodeId]): List of related asset identifiers.
Returns:
dict: Asset information.
"""
# NodeId objects have .space and .external_id properties
return {
"primary_asset": {
"space": asset_id.space,
"external_id": asset_id.external_id
},
"related_count": len(related_asset_ids)
}
```
**Runtime format:** The tool passes `NodeId` parameters to the language model as objects with `space` and `externalId` properties:
```json theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
{
"space": "my-space",
"externalId": "asset-123"
}
```
The `NodeId` class uniquely identifies nodes in Cognite's data modeling framework, combining a `space` (workspace identifier) and `external_id` to form a unique node reference. See the [Cognite Python SDK data modeling documentation](https://cognite-sdk-python.readthedocs-hosted.com/en/latest/data_modeling.html) for details.
### Limitations and restrictions
The following types are **not supported** as arguments to `handle()`:
* Custom classes and objects
* Nested lists (`list[list[str]]`)
* Complex generics beyond `list[T]`
* Union types (`str | int`)
* Optional types (`Optional[str]`)
* Tuples, sets, and other collections
* String literal types (e.g., `Literal["A"]`, `Literal["A", "B"]`)
The following restrictions apply for default argument values for `handle()`:
* Default values are only supported for primitive types (`str`, `int`, `float`, `bool`).
* Lists cannot have default values.
* `NodeId` and `datetime` parameters cannot have default values.
## Supported return types
Your `handle()` function can return any JSON-serializable data. The following return types are supported:
### Primitive return types
```python theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle() -> str:
return "Success message"
def handle() -> int:
return 42
def handle() -> float:
return 3.14159
def handle() -> bool:
return True
```
### Dictionary returns
```python theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle() -> dict:
return {
"status": "completed",
"count": 10,
"results": ["item1", "item2"]
}
```
### List returns
```python theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle() -> list[str]:
return ["result1", "result2", "result3"]
def handle() -> list[dict]:
return [
{"id": 1, "name": "Item 1"},
{"id": 2, "name": "Item 2"}
]
```
### None/null returns
```python theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle() -> None:
# Perform an action and return None
print("Operation completed")
return None
```
### Unsupported return types
The following return types are **not supported**:
* Raw Cognite SDK objects
* Custom class instances
* Functions or callable objects
* File handles or binary data
* Complex objects that can't be serialized to JSON
## Cognite SDK and authentication
The Cognite Python SDK comes pre-installed and configured with your session credentials. The tool uses your current credentials and requires no additional configuration.
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteAsset
def handle() -> dict:
"""
List assets.
Returns:
dict: Asset information including ID and metadata.
"""
# Client is automatically configured - no setup required
client = CogniteClient()
# List assets
assets = client.data_modeling.instances.list(sources=[CogniteAsset.get_source()])
if assets:
return assets[0].dump() # Use .dump() for JSON serialization
return {"error": "No assets found"}
```
For complete SDK usage and API reference, see the [Cognite Python SDK documentation](https://cognite-sdk-python.readthedocs-hosted.com/en/latest/).
### Cognite SDK return types
We don't support direct serialization of Cognite SDK objects. To reduce token usage and make the data readable to the agent, extract relevant values and reshape the data as needed.
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteAsset
def handle() -> list[dict]:
"""
Get simplified asset information.
Returns:
list[dict]: List of assets with only essential fields.
"""
client = CogniteClient()
assets = client.data_modeling.instances.list(
sources=[CogniteAsset.get_source()],
limit=10
)
# Extract only relevant fields
return [
{
"external_id": asset.external_id,
"name": asset.name,
"description": asset.description
}
for asset in assets
]
```
As a last resort, you can use the `.dump()` method to serialize all information from Cognite SDK objects, but this may include unnecessary data that increases token usage.
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle() -> dict:
"""Get complete asset data using dump()."""
from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteAsset
client = CogniteClient()
assets = client.data_modeling.instances.list(sources=[CogniteAsset.get_source()])
return assets[0].dump() # ✅ Correct - JSON serializable
def handle() -> list[dict]:
"""Get all assets using dump()."""
from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteAsset
client = CogniteClient()
assets = client.data_modeling.instances.list(sources=[CogniteAsset.get_source()])
return assets.dump() # ✅ Correct - JSON serializable
```
## When to use Run Python code vs Call Function
Both the **Run Python code** tool and the **Call Function** tool let your agent run custom Python logic. They differ in where the code runs and what capabilities are available.
| Capability | Run Python code | Call Function |
| ----------------------- | ------------------------------------------- | ----------------------------------------- |
| Code location | Written and maintained in the Agent builder | Deployed separately as a Cognite Function |
| Execution environment | WebAssembly sandbox (Pyodide, Python 3.11) | Cognite cloud infrastructure |
| Custom packages | Not supported | Supported via `requirements.txt` |
| External network access | Not supported (CDF APIs only) | Supported |
| Execution time | Limited | Up to the configured max polling time |
| Memory | Limited by browser memory | More headroom for larger datasets |
Use Run Python code when you need a focused calculation or data transformation that works within the WebAssembly sandbox. The code is stored in the agent configuration, making it easy to write and test. This tool works well for domain-specific logic like unit conversions, threshold checks, or formatting results from Query tool output.
Use Call Function when the logic requires custom packages, external network calls, or longer execution times. A Cognite Function runs on Cognite's cloud infrastructure, so it can access third-party APIs, use packages from PyPI, and handle larger datasets. This tool requires the function to be deployed separately before the agent can call it.
If you are not sure which to use, start with Run Python code. Switch to Call Function if you encounter package or environment limitations.
## Best practices
We recommend the following best practices when writing Python code for the **Run Python code** tool.
### Error handling
Use descriptive error messages:
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle(value: float) -> float:
"""
Process a value with validation.
Args:
value (float): Input value to process.
Returns:
float: Processed value.
"""
if value < 0:
raise ValueError("The value must be non-negative.") # Clear, actionable message
if value > 1000:
raise ValueError("The value exceeds the maximum allowed limit of 1000.")
return value * 2
```
### Documentation
* Always include docstrings following the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods).
* **Parameter descriptions are critical**: Only the `Args:` section descriptions are passed to the AI agent as part of the JSON schema.
* Write clear, specific parameter descriptions that help the agent understand what values to provide.
* Include examples in the parameter descriptions when helpful.
**Example of parameter documentation for the agent:**
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def handle(asset_id: NodeId, threshold: float) -> dict:
"""
Analyze asset performance against threshold.
Args:
asset_id (NodeId): The specific asset to analyze. Provide both space and externalId.
threshold (float): The performance threshold value in percentage (0-100).
Returns:
dict: The analysis results with performance metrics.
"""
# Implementation here
pass
```
### Testing
Write comprehensive test functions. You can use both synchronous and asynchronous test functions:
```python wrap theme={"languages":{"custom":["/_languages/kuiper.json","../_languages/kuiper.json"]}}
def test():
"""Comprehensive test coverage."""
# Test normal cases
result1 = handle("test", 42)
assert result1 == expected_value, "Normal case failed"
# Test edge cases
result2 = handle("", 0)
assert result2 is not None, "Edge case failed"
# Test error conditions
try:
handle("invalid", -1)
assert False, "Should have raised ValueError"
except ValueError:
pass # Expected
# Return a sample result for display
return handle("sample", 100)
```
## Troubleshooting
These are common errors when using the **Run Python code** tool, and how to fix them.
### "Function 'handle' not found"
* Ensure your function is named `handle`.
* Check for indentation errors.
### "Missing type annotation for parameter(s)"
* Make sure all parameters have type annotations.
* Use supported types only.
### "Duplicate argument name(s) found"
* Make sure the parameter names are unique.
* Check for copy-paste errors.
### "Unsupported type for parameter"
* Use only [supported types](#supported-parameter-types).
* Check for typos in type annotations.
### Serialization errors
* Use the `.dump()` method on Cognite SDK objects.
* Make sure the return values are serializable to JSON.
## Code execution environment capabilities and limitations
The **Run Python code** tool runs in a controlled WebAssembly environment with specific capabilities and limitations:
* **Runtime**: Pyodide (Python 3.11 in WebAssembly).
* **Execution context**: Runs in a web worker for isolation.
* **Performance**: WebAssembly execution may be slower than native Python.
* **Memory**: Limited by browser memory constraints.
## Pre-installed packages
The following packages are automatically available:
* **Cognite Python SDK** (`cognite-sdk`): Pre-configured with authentication.
* **Standard library**: Full Python 3.11 standard library.
* **Core scientific packages** (via Pyodide):
* `numpy`: Numerical computing
* `pandas`: Data analysis and manipulation
* `json`: JSON encoding/decoding
* `datetime`: Date and time handling
* `uuid`: UUID generation
We don't support installing custom packages. This includes:
* `pip install` commands
* Dynamic package installation
* Custom or third-party packages not included in Pyodide
If you need custom packages or more advanced functionality, consider using [Cognite Functions](/cdf/functions/) instead. Functions run on Cognite's cloud infrastructure and support:
* Custom package installation via `requirements.txt`
* Longer execution times
* More memory and compute resources
* Full Python ecosystem access
## System limitations
The execution environment has certain constraints that you should be aware of when writing your scripts.
* **File system**: No access to local file system or file I/O operations. The scripts have access to a virtual memory-based file system.
* **Network access**: Only API calls to Cognite Data Fusion APIs are supported. Network calls to other servers might get blocked.
* **Process control**: No subprocess or system command execution.
* **Threading**: Limited threading capabilities in WebAssembly environment.
* **Graphics**: Matplotlib backend set to 'AGG' (non-interactive).
## Use via API or SDK
When you call an agent with this tool enabled through the REST API or Python SDK, the agent will request confirmation before running the Python code. Your application receives a `toolConfirmation` action in the response and must reply with `ALLOW` to proceed or `DENY` to cancel.
See [Use Atlas AI agents via API and SDK](/cdf/atlas_ai/guides/atlas_ai_agent_api_sdk) for a complete example.
## How the AI agent uses your code
The Atlas AI agent processes your Python code through these steps to understand and execute your tool functions:
1. **JSON schema as tool descriptor**: The generated JSON schema is passed to the AI agent's language model as a tool descriptor. This tells the agent exactly what parameters are available, their types, and what they're for.
2. **Parameter descriptions guide the agent**: The agent uses the parameter descriptions from your docstring's `Args:` section to understand what values to provide.
3. **Type validation**: The schema ensures the language model provides correctly typed arguments that match your function signature.
4. **Runtime execution**: When the agent calls your tool, it validates the arguments against the schema and transforms them (for example, creating NodeId objects from dictionaries) before passing them to your `handle()` function.
Your parameter docstrings are critical. They directly influence how the agent understands and uses your tool.