This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new aacc9126 [python] Remove the Python local execution path (#882)
aacc9126 is described below
commit aacc91268aced994dee22c8d8a95aa259d4b3e1f
Author: Weiqing Yang <[email protected]>
AuthorDate: Sun Jul 12 20:12:10 2026 -0700
[python] Remove the Python local execution path (#882)
---
.../docs/development/memory/long_term_memory.md | 14 +-
docs/content/docs/operations/configuration.md | 12 +-
docs/content/docs/operations/deployment.md | 85 +----
python/flink_agents/api/agents/react_agent.py | 5 +-
python/flink_agents/api/execution_environment.py | 113 +++---
python/flink_agents/api/yaml/tests/test_loader.py | 60 ++--
.../chat_model_integration_test.py | 101 ------
.../chat_model_multi_provider_remote_test.py | 7 +-
.../e2e_tests_integration/react_agent_test.py | 79 -----
.../workflow_memory_remote_agent.py | 3 +-
.../e2e_tests_integration/workflow_test.py | 149 --------
python/flink_agents/e2e_tests/test_utils.py | 31 --
.../flink_agents/e2e_tests/test_utils_unit_test.py | 21 --
.../flink_agents/examples/rag/rag_agent_example.py | 43 +--
.../plan/tests/actions/test_chat_model_action.py | 2 +-
.../runtime/local_execution_environment.py | 191 ----------
python/flink_agents/runtime/local_runner.py | 387 ---------------------
.../runtime/remote_execution_environment.py | 18 +-
.../runtime/{ => tests}/local_memory_object.py | 0
.../runtime/tests/test_built_in_actions.py | 218 ------------
.../runtime/tests/test_get_resource_in_action.py | 114 ------
.../tests/test_local_execution_environment.py | 380 --------------------
.../runtime/tests/test_local_memory_object.py | 2 +-
.../tests/test_local_runner_reconcilable.py | 79 -----
.../runtime/tests/test_memory_reference.py | 2 +-
...guage.py => test_plan_java_function_routing.py} | 43 +--
.../tests/test_remote_execution_environment.py | 3 +-
.../runtime/tests/test_runner_context_execute.py | 189 ----------
28 files changed, 125 insertions(+), 2226 deletions(-)
diff --git a/docs/content/docs/development/memory/long_term_memory.md
b/docs/content/docs/development/memory/long_term_memory.md
index 727e5b44..ef0adc8b 100644
--- a/docs/content/docs/development/memory/long_term_memory.md
+++ b/docs/content/docs/development/memory/long_term_memory.md
@@ -56,12 +56,15 @@ When all three options are configured, the framework
automatically creates a Mem
{{< tab "Python" >}}
```python
+from pyflink.datastream import StreamExecutionEnvironment
+
from flink_agents.api.execution_environment import AgentsExecutionEnvironment
from flink_agents.api.core_options import AgentConfigOptions
from flink_agents.api.memory.long_term_memory import LongTermMemoryOptions
-env = AgentsExecutionEnvironment.get_execution_environment()
-agents_config = env.get_config()
+env = StreamExecutionEnvironment.get_execution_environment()
+agents_env = AgentsExecutionEnvironment.get_execution_environment(env)
+agents_config = agents_env.get_config()
# Set job identifier (maps to Mem0 user_id)
agents_config.set(AgentConfigOptions.JOB_IDENTIFIER, "my_job")
@@ -455,6 +458,8 @@ The snippets below show how to read from and write to
Long-Term Memory inside an
{{< tab "Python" >}}
```python
+from pyflink.datastream import StreamExecutionEnvironment
+
from flink_agents.api.decorators import action
from flink_agents.api.execution_environment import AgentsExecutionEnvironment
from flink_agents.api.core_options import AgentConfigOptions
@@ -488,8 +493,9 @@ class PersonalizedAssistant:
ctx.send_event(OutputEvent(output=response))
# Setup
-env = AgentsExecutionEnvironment.get_execution_environment()
-agents_config = env.get_config()
+env = StreamExecutionEnvironment.get_execution_environment()
+agents_env = AgentsExecutionEnvironment.get_execution_environment(env)
+agents_config = agents_env.get_config()
agents_config.set(AgentConfigOptions.JOB_IDENTIFIER, "personalized_assistant")
agents_config.set(LongTermMemoryOptions.Mem0.CHAT_MODEL_SETUP, "my_chat_model")
agents_config.set(LongTermMemoryOptions.Mem0.EMBEDDING_MODEL_SETUP,
"my_embedding_model")
diff --git a/docs/content/docs/operations/configuration.md
b/docs/content/docs/operations/configuration.md
index 44ebbc0f..9ba9e200 100644
--- a/docs/content/docs/operations/configuration.md
+++ b/docs/content/docs/operations/configuration.md
@@ -44,7 +44,8 @@ Users can explicitly modify the configuration when defining
the `AgentsExecution
```python
# Get Flink Agents execution environment
-agents_env = AgentsExecutionEnvironment.get_execution_environment()
+env = StreamExecutionEnvironment.get_execution_environment()
+agents_env = AgentsExecutionEnvironment.get_execution_environment(env)
# Get configuration object from the environment
config = agents_env.get_configuration()
@@ -104,7 +105,7 @@ By default, the configuration is automatically loaded from
`$FLINK_HOME/conf/con
**Special Condition**
-In the following two cases, Flink Agents may not locate the corresponding
configuration file, necessitating manual configuration. If the files are not
set, no configuration files will be loaded, potentially resulting in unexpected
behavior or failures.
+In the following case, Flink Agents may not locate the corresponding
configuration file, necessitating manual configuration. If the file is not set,
no configuration file will be loaded, potentially resulting in unexpected
behavior or failures.
- **For MiniCluster**:
Manual setup is **required** — always export the environment variable before
running the job:
@@ -115,13 +116,6 @@ In the following two cases, Flink Agents may not locate
the corresponding config
This ensures that Flink can locate and load the configuration file correctly.
-- **Local mode**: When [run without flink]({{< ref
"docs/operations/deployment">}}), use the
`AgentsExecutionEnvironment.get_configuration()` API to load the YAML file
directly:
-
- ```python
- config = agents_env.get_configuration("path/to/your/config.yaml")
- ```
-
-
## Built-in configuration options
### Core Options
diff --git a/docs/content/docs/operations/deployment.md
b/docs/content/docs/operations/deployment.md
index bb3cfb30..4178b94d 100644
--- a/docs/content/docs/operations/deployment.md
+++ b/docs/content/docs/operations/deployment.md
@@ -24,88 +24,11 @@ under the License.
## Overview
-We provide two options to run the job:
-
-- **Run without Flink**
- - **Language Support**: Only Python
- - **Input and Output**: Python List
- - **Suitable Use Case**: Local Testing and Debugging
-
-- **Run in Flink**
- - **Language Support**: Python & Java
- - **Input and Output**: DataStream or Table
- - **Suitable Use Case**: Production
-
-These deployment modes differ in supported languages and data formats,
allowing you to choose the one that best fits your use case.
-
-## Run without Flink
-
-After completing the [installation of flink-agents]({{< ref
"docs/get-started/installation" >}}) and building your [ReAct Agent]({{< ref
"docs/development/react_agent" >}}) or [Workflow Agent]({{< ref
"docs/development/workflow_agent" >}}), you can test and execute your agent
locally using a simple Python script. This allows you to validate logic without
requiring a Flink cluster.
-
-### Example for Local Run with Test Data
-
-```python
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from my_module.agents import MyAgent # Replace with your actual agent path
-
-if __name__ == "__main__":
- # 1. Initialize environment
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- # 2. Prepare test data
- input_data = [
- {"key": "0001", "value": "Calculate the sum of 1 and 2."},
- {"key": "0002", "value": "Tell me a joke about cats."}
- ]
-
- # 3. Create agent instance
- agent = MyAgent()
-
- # 4. Build pipeline
- output_data = env.from_list(input_data) \
- .apply(agent) \
- .to_list()
-
- # 5. Execute and show results
- env.execute()
-
- print("\nExecution Results:")
- for record in output_data:
- for key, value in record.items():
- print(f"{key}: {value}")
+Flink Agents jobs run on Flink:
-```
-
-#### Input Data Format
-
-The input data should be a list of dictionaries `List[Dict[str, Any]]` with
the following structure:
-
-```python
-[
- {
- # Optional field: Input key.
- # The key is randomly generated if not provided.
- "key": "key_1",
-
- # Required field: Input content
- # This becomes the `input` field in InputEvent
- "value": "Calculate the sum of 1 and 2.",
- },
- ...
-]
-```
-
-#### Output Data Format
-
-The output data is a list of dictionaries `List[Dict[str, Any]]` where each
dictionary contains a single key-value pair representing the processed result.
The structure is generated from `OutputEvent` objects:
-
-```python
-[
- {key_1: output_1}, # From first OutputEvent
- {key_2: output_2}, # From second OutputEvent
- ...
-]
-```
+- **Language Support**: Python & Java
+- **Input and Output**: DataStream or Table
+- **Suitable Use Case**: Production
## Run in Flink
diff --git a/python/flink_agents/api/agents/react_agent.py
b/python/flink_agents/api/agents/react_agent.py
index c9b310ec..9f5f72df 100644
--- a/python/flink_agents/api/agents/react_agent.py
+++ b/python/flink_agents/api/agents/react_agent.py
@@ -62,11 +62,12 @@ class ReActAgent(Agent):
result: int
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = StreamExecutionEnvironment.get_execution_environment()
+ agents_env =
AgentsExecutionEnvironment.get_execution_environment(env)
# register resource to execution environment
(
- env.add_resource(
+ agents_env.add_resource(
"ollama",
ResourceDescriptor(clazz=OllamaChatModelConnection,
model=model),
)
diff --git a/python/flink_agents/api/execution_environment.py
b/python/flink_agents/api/execution_environment.py
index 115c8616..e9ffac9a 100644
--- a/python/flink_agents/api/execution_environment.py
+++ b/python/flink_agents/api/execution_environment.py
@@ -51,18 +51,6 @@ class AgentBuilder(ABC):
on the environment (e.g. by ``load_yaml``).
"""
- @abstractmethod
- def to_list(self) -> List[Dict[str, Any]]:
- """Get output list of agent execution.
-
- The element in the list is a dict like {'key': output}.
-
- Returns:
- -------
- list
- Outputs of agent execution.
- """
-
@abstractmethod
def to_datastream(self) -> DataStream:
"""Get output datastream of agent execution.
@@ -119,10 +107,14 @@ class AgentsExecutionEnvironment(ABC):
) -> "AgentsExecutionEnvironment":
"""Get agents execution environment.
- Currently, user can run flink agents in ide using
LocalExecutionEnvironment or
- RemoteExecutionEnvironment. To distinguish which environment to use,
when run
- flink agents with pyflink datastream/table, user should pass flink
- StreamExecutionEnvironment when get AgentsExecutionEnvironment.
+ A Flink ``StreamExecutionEnvironment`` is required. When running flink
agents
+ with pyflink datastream/table, pass the ``StreamExecutionEnvironment``
so the
+ agents run on the Flink runtime.
+
+ Parameters
+ ----------
+ env : StreamExecutionEnvironment
+ The Flink stream execution environment the agents run on. Must not
be None.
Returns:
-------
@@ -130,45 +122,42 @@ class AgentsExecutionEnvironment(ABC):
Environment for agent execution.
"""
if env is None:
- return importlib.import_module(
- "flink_agents.runtime.local_execution_environment"
- ).create_instance(env=env, t_env=t_env, **kwargs)
+ err_msg = "A StreamExecutionEnvironment is required."
+ raise ValueError(err_msg)
+
+ major_version = flink_version_manager.major_version
+ if not major_version:
+ err_msg = "Apache Flink is not installed."
+ raise ModuleNotFoundError(err_msg)
+
+ lib_base = files("flink_agents.lib")
+
+ # Load the common JAR (shared dependencies)
+ common_lib = lib_base / "common"
+ if common_lib.is_dir():
+ for jar_file in common_lib.iterdir():
+ if jar_file.is_file() and str(jar_file).endswith(".jar"):
+ env.add_jars(f"file://{jar_file}")
+ else:
+ err_msg = "Flink Agents common JAR not found."
+ raise FileNotFoundError(err_msg)
+
+ # Load the version-specific thin JAR
+ version_dir = f"flink-{major_version}"
+ version_lib = lib_base / version_dir
+
+ # Check if version-specific directory exists
+ if version_lib.is_dir():
+ for jar_file in version_lib.iterdir():
+ if jar_file.is_file() and str(jar_file).endswith(".jar"):
+ env.add_jars(f"file://{jar_file}")
else:
- major_version = flink_version_manager.major_version
- if major_version:
- lib_base = files("flink_agents.lib")
-
- # Load the common JAR (shared dependencies)
- common_lib = lib_base / "common"
- if common_lib.is_dir():
- for jar_file in common_lib.iterdir():
- if jar_file.is_file() and
str(jar_file).endswith(".jar"):
- env.add_jars(f"file://{jar_file}")
- else:
- err_msg = "Flink Agents common JAR not found."
- raise FileNotFoundError(err_msg)
-
- # Load the version-specific thin JAR
- version_dir = f"flink-{major_version}"
- version_lib = lib_base / version_dir
-
- # Check if version-specific directory exists
- if version_lib.is_dir():
- for jar_file in version_lib.iterdir():
- if jar_file.is_file() and
str(jar_file).endswith(".jar"):
- env.add_jars(f"file://{jar_file}")
- else:
- err_msg = (
- f"Flink Agents dist JAR for Flink {major_version} not
found."
- )
- raise FileNotFoundError(err_msg)
-
- return importlib.import_module(
- "flink_agents.runtime.remote_execution_environment"
- ).create_instance(env=env, t_env=t_env, **kwargs)
- else:
- err_msg = "Apache Flink is not installed."
- raise ModuleNotFoundError(err_msg)
+ err_msg = f"Flink Agents dist JAR for Flink {major_version} not
found."
+ raise FileNotFoundError(err_msg)
+
+ return importlib.import_module(
+ "flink_agents.runtime.remote_execution_environment"
+ ).create_instance(env=env, t_env=t_env, **kwargs)
@abstractmethod
def get_config(self, path: str | None = None) -> Configuration:
@@ -180,22 +169,6 @@ class AgentsExecutionEnvironment(ABC):
The configuration for flink agents.
"""
- @abstractmethod
- def from_list(self, input: List[Dict[str, Any]]) -> AgentBuilder:
- """Set input for agents. Used for local execution.
-
- Parameters
- ----------
- input : list
- Receive a list as input. The element in the list should be a dict
like
- {'key': Any, 'value': Any} or {'value': Any} , extra field will be
ignored.
-
- Returns:
- -------
- AgentBuilder
- A new builder to build an agent for specific input.
- """
-
@abstractmethod
def from_datastream(
self, input: DataStream, key_selector: KeySelector | Callable | None =
None
diff --git a/python/flink_agents/api/yaml/tests/test_loader.py
b/python/flink_agents/api/yaml/tests/test_loader.py
index 73a474f7..54986fbf 100644
--- a/python/flink_agents/api/yaml/tests/test_loader.py
+++ b/python/flink_agents/api/yaml/tests/test_loader.py
@@ -16,6 +16,7 @@
# limitations under the License.
#################################################################################
from pathlib import Path
+from unittest.mock import MagicMock, patch
import pytest
@@ -23,7 +24,6 @@ from flink_agents.api.agents.agent import Agent
from flink_agents.api.chat_message import MessageRole
from flink_agents.api.events.chat_event import ChatResponseEvent
from flink_agents.api.events.event import InputEvent
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
from flink_agents.api.function import JavaFunction, PythonFunction
from flink_agents.api.prompts.prompt import LocalPrompt
from flink_agents.api.resource import ResourceDescriptor, ResourceName,
ResourceType
@@ -38,12 +38,28 @@ from flink_agents.api.yaml.loader import (
)
from flink_agents.api.yaml.specs import SkillsSpec
from flink_agents.api.yaml.tests.fixtures import loader_targets
+from flink_agents.runtime.remote_execution_environment import (
+ RemoteExecutionEnvironment,
+)
_FIXTURES = Path(__file__).parent / "fixtures"
_TARGETS_MODULE = "flink_agents.api.yaml.tests.fixtures.loader_targets"
+def _registration_env() -> RemoteExecutionEnvironment:
+ """Return an env for registering and inspecting YAML-loaded agents.
+
+ These tests only call ``load_yaml`` and read back ``_agents`` /
``resources``;
+ they never submit a job, so the Flink environment is mocked and no cluster
or
+ ``flink_agents.lib`` jar is needed.
+ """
+ with patch(
+
"flink_agents.runtime.remote_execution_environment.StreamExecutionEnvironment"
+ ):
+ return RemoteExecutionEnvironment(env=MagicMock())
+
+
def test_resolve_function_python_with_module_attr() -> None:
func = resolve_function(
name="anything", function=f"{_TARGETS_MODULE}:increment"
@@ -206,19 +222,19 @@ def
test_build_agents_handles_shared_resources_and_actions() -> None:
def test_load_yaml_registers_single_agent_on_env() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "single_agent.yaml")
assert "incrementer" in env._agents
def test_load_yaml_registers_multiple_agents() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "multi_agent.yaml")
assert set(env._agents.keys()) == {"a1", "a2"}
def test_load_yaml_merges_shared_action_into_agents() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "with_shared.yaml")
a1 = env._agents["a1"]
a2 = env._agents["a2"]
@@ -235,13 +251,13 @@ def test_load_yaml_merges_shared_action_into_agents() ->
None:
def test_load_yaml_registers_shared_resources_on_env() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "with_shared.yaml")
assert "shared_conn" in env.resources[ResourceType.CHAT_MODEL_CONNECTION]
def test_load_yaml_string_ref_to_missing_shared_action_errors(tmp_path: Path)
-> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
bad = tmp_path / "bad_missing_shared_action.yaml"
bad.write_text("agents:\n - name: a\n actions:\n -
undefined_action\n")
with pytest.raises(ValueError, match="undefined_action"):
@@ -249,7 +265,7 @@ def
test_load_yaml_string_ref_to_missing_shared_action_errors(tmp_path: Path) ->
def test_load_yaml_multi_call_merges() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "multi_file_a.yaml")
load_yaml(env, _FIXTURES / "multi_file_b.yaml")
assert {"file_a_agent", "file_b_agent"} <= set(env._agents.keys())
@@ -258,13 +274,13 @@ def test_load_yaml_multi_call_merges() -> None:
def test_load_yaml_accepts_list_of_paths() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, [_FIXTURES / "multi_file_a.yaml", _FIXTURES /
"multi_file_b.yaml"])
assert {"file_a_agent", "file_b_agent"} <= set(env._agents.keys())
def test_load_yaml_duplicate_agent_across_calls_errors() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "multi_file_a.yaml")
with pytest.raises(ValueError, match="file_a_agent"):
load_yaml(env, _FIXTURES / "multi_file_a.yaml")
@@ -284,7 +300,7 @@ def
test_load_yaml_duplicate_shared_resource_within_file_errors(tmp_path) -> Non
" - name: conn\n"
" clazz: x.Z\n"
)
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
with pytest.raises(ValueError, match="Duplicate shared resource name
'conn'"):
load_yaml(env, bad)
@@ -300,13 +316,13 @@ def
test_load_yaml_duplicate_shared_action_within_file_errors(tmp_path) -> None:
" - name: shared\n"
" trigger_conditions: [input]\n"
)
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
with pytest.raises(ValueError, match="Duplicate shared action name
'shared'"):
load_yaml(env, bad)
def test_load_yaml_duplicate_shared_resource_across_calls_errors(tmp_path) ->
None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "multi_file_a.yaml")
dup = tmp_path / "dup.yaml"
dup.write_text(
@@ -319,24 +335,6 @@ def
test_load_yaml_duplicate_shared_resource_across_calls_errors(tmp_path) -> No
load_yaml(env, dup)
-def test_apply_by_agent_name_runs_yaml_loaded_agent() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
- load_yaml(env, _FIXTURES / "single_agent.yaml")
-
- input_list = []
- output_list = env.from_list(input_list).apply("incrementer").to_list()
- input_list.append({"key": "bob", "value": 1})
- input_list.append({"key": "john", "value": 2})
- env.execute()
- assert output_list == [{"bob": 2}, {"john": 3}]
-
-
-def test_apply_by_unknown_name_errors() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
- with pytest.raises(ValueError, match="ghost"):
- env.from_list([]).apply("ghost")
-
-
def test_build_agents_loads_skills_per_agent_and_shared() -> None:
agents, shared_resources, _ = build_agents(_FIXTURES / "with_skills.yaml")
agent = agents["skills_agent"]
@@ -377,7 +375,7 @@ def test_build_skills_merges_all_schemes() -> None:
def test_load_yaml_registers_shared_skills_on_env() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
+ env = _registration_env()
load_yaml(env, _FIXTURES / "with_skills.yaml")
shared = env.resources[ResourceType.SKILLS]["shared_skills"]
assert isinstance(shared, Skills)
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_test.py
b/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_test.py
deleted file mode 100644
index ddebe7d8..00000000
---
a/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_integration_test.py
+++ /dev/null
@@ -1,101 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-import os
-from pathlib import Path
-
-import pytest
-
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from flink_agents.e2e_tests.e2e_tests_integration.chat_model_integration_agent
import (
- ChatModelTestAgent,
-)
-from flink_agents.e2e_tests.test_utils import pull_model
-
-current_dir = Path(__file__).parent
-
-TONGYI_MODEL = os.environ.get("TONGYI_CHAT_MODEL", "qwen-plus")
-OLLAMA_MODEL = os.environ.get("OLLAMA_CHAT_MODEL", "qwen3:1.7b")
-OPENAI_MODEL = os.environ.get("OPENAI_CHAT_MODEL", "gpt-3.5-turbo")
-AZURE_OPENAI_MODEL = os.environ.get("AZURE_OPENAI_CHAT_MODEL", "gpt-5")
-AZURE_OPENAI_API_VERSION = os.environ.get(
- "AZURE_OPENAI_API_VERSION", "2025-04-01-preview"
-)
-
-DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY")
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY")
-
-client = pull_model(OLLAMA_MODEL)
-
-
[email protected](
- "model_provider",
- [
- pytest.param(
- "Ollama",
- marks=pytest.mark.skipif(
- client is None,
- reason="Ollama client is not available or test model is
missing.",
- ),
- ),
- pytest.param(
- "Tongyi",
- marks=pytest.mark.skipif(
- not DASHSCOPE_API_KEY, reason="Tongyi api key is not set."
- ),
- ),
- pytest.param(
- "OpenAI",
- marks=pytest.mark.skipif(
- not OPENAI_API_KEY, reason="OpenAI api key is not set."
- ),
- ),
- pytest.param(
- "AzureOpenAI",
- marks=pytest.mark.skipif(
- not AZURE_OPENAI_API_KEY, reason="Azure OpenAI api key is not
set."
- ),
- ),
- ],
-)
-def test_chat_model_integration(
- model_provider: str, monkeypatch: pytest.MonkeyPatch
-) -> None:
- monkeypatch.setenv("TONGYI_CHAT_MODEL", TONGYI_MODEL)
- monkeypatch.setenv("OLLAMA_CHAT_MODEL", OLLAMA_MODEL)
- monkeypatch.setenv("OPENAI_CHAT_MODEL", OPENAI_MODEL)
- monkeypatch.setenv("AZURE_OPENAI_CHAT_MODEL", AZURE_OPENAI_MODEL)
- monkeypatch.setenv("AZURE_OPENAI_API_VERSION", AZURE_OPENAI_API_VERSION)
- monkeypatch.setenv("MODEL_PROVIDER", model_provider)
- env = AgentsExecutionEnvironment.get_execution_environment()
- input_list = []
- agent = ChatModelTestAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "0001", "value": "calculate the sum of 1 and
2."})
- input_list.append({"key": "0002", "value": "Tell me a joke about cats."})
-
- env.execute()
-
- for output in output_list:
- for key, value in output.items():
- print(f"{key}: {value}")
-
- assert "3" in output_list[0]["0001"]
- assert "cat" in output_list[1]["0002"]
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_multi_provider_remote_test.py
b/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_multi_provider_remote_test.py
index af243984..374b89ac 100644
---
a/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_multi_provider_remote_test.py
+++
b/python/flink_agents/e2e_tests/e2e_tests_integration/chat_model_multi_provider_remote_test.py
@@ -81,10 +81,9 @@ def test_chat_model_integration_remote(
) -> None:
"""Non-Ollama providers answer math and creative prompts on the remote
path.
- Mirrors the from_list ``chat_model_integration_test`` for the three
- credential-gated providers (Tongyi/OpenAI/AzureOpenAI), but drives the
agent
- through a real StreamExecutionEnvironment with a FileSource and file sink.
- Each parameter skips cleanly when its credential is absent.
+ Covers the three credential-gated providers (Tongyi/OpenAI/AzureOpenAI) and
+ drives the agent through a real StreamExecutionEnvironment with a
FileSource
+ and file sink. Each parameter skips cleanly when its credential is absent.
"""
monkeypatch.setenv("TONGYI_CHAT_MODEL", TONGYI_MODEL)
monkeypatch.setenv("OPENAI_CHAT_MODEL", OPENAI_MODEL)
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py
b/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py
index ad689d05..261378f8 100644
--- a/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py
+++ b/python/flink_agents/e2e_tests/e2e_tests_integration/react_agent_test.py
@@ -21,7 +21,6 @@ import sysconfig
from pathlib import Path
import pytest
-from pydantic import BaseModel
from pyflink.common import Row
from pyflink.common.typeinfo import BasicTypeInfo, ExternalTypeInfo,
RowTypeInfo
from pyflink.datastream import KeySelector, StreamExecutionEnvironment
@@ -48,7 +47,6 @@ from flink_agents.e2e_tests.test_utils import (
assert_tool_invoked,
collect_tool_invocations,
pull_model,
- tool_invocations_from_events,
)
current_dir = Path(__file__).parent
@@ -58,16 +56,6 @@ os.environ["PYTHONPATH"] = sysconfig.get_paths()["purelib"]
OLLAMA_MODEL = os.environ.get("REACT_OLLAMA_MODEL", "qwen3:1.7b")
-class InputData(BaseModel):
- a: int
- b: int
- c: int
-
-
-class OutputData(BaseModel):
- result: int
-
-
class MyKeySelector(KeySelector):
"""KeySelector for extracting key."""
@@ -79,73 +67,6 @@ class MyKeySelector(KeySelector):
client = pull_model(OLLAMA_MODEL)
[email protected](
- client is None, reason="Ollama client is not available or test model is
missing"
-)
-def test_react_agent_on_local_runner(monkeypatch: pytest.MonkeyPatch) -> None:
- monkeypatch.setenv("OLLAMA_CHAT_MODEL", OLLAMA_MODEL)
- env = AgentsExecutionEnvironment.get_execution_environment()
- env.get_config().set(
- AgentExecutionOptions.ERROR_HANDLING_STRATEGY,
ErrorHandlingStrategy.RETRY
- )
- env.get_config().set(AgentExecutionOptions.MAX_RETRIES, 3)
-
- # register resource to execution environment
- (
- env.add_resource(
- "ollama",
- ResourceType.CHAT_MODEL_CONNECTION,
- ResourceDescriptor(
- clazz=ResourceName.ChatModel.OLLAMA_CONNECTION,
request_timeout=240.0
- ),
- )
- .add_resource("add", ResourceType.TOOL, Tool.from_callable(add))
- .add_resource("multiply", ResourceType.TOOL,
Tool.from_callable(multiply))
- )
-
- # prepare prompt
- prompt = Prompt.from_messages(
- messages=[
- ChatMessage(
- role=MessageRole.SYSTEM,
- content='An example of output is {"result": 30.32}.',
- ),
- ChatMessage(role=MessageRole.USER, content="What is ({a} + {b}) *
{c}"),
- ],
- )
-
- # create ReAct agent.
- agent = ReActAgent(
- chat_model=ResourceDescriptor(
- clazz=ResourceName.ChatModel.OLLAMA_SETUP,
- connection="ollama",
- model=OLLAMA_MODEL,
- tools=["add", "multiply"],
- ),
- prompt=prompt,
- output_schema=OutputData,
- )
-
- # execute agent
- input_list = []
-
- output_list = env.from_list(input_list).apply(agent).to_list()
- input_list.append({"key": "0001", "value": InputData(a=2123, b=2321,
c=312)})
-
- env.execute()
-
- assert len(output_list) == 1, (
- "This may be caused by the LLM response does not match the output
schema, you can rerun this case."
- )
- assert int(output_list[0]["0001"].result) == 1386528
-
- # multiply's first arg (4444 = 2123 + 2321) proves the addition was
computed
- # correctly and the multiply tool was used; the model often does the
addition
- # without the add tool, so add is not a reliable signal to assert on.
- invocations = tool_invocations_from_events(env.get_tool_request_events())
- assert_tool_invoked(invocations, "multiply", {"a": 4444, "b": 312})
-
-
@pytest.mark.skipif(
client is None, reason="Ollama client is not available or test model is
missing"
)
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_memory_remote_agent.py
b/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_memory_remote_agent.py
index 840ec131..75803c9d 100644
---
a/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_memory_remote_agent.py
+++
b/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_memory_remote_agent.py
@@ -36,8 +36,7 @@ class WorkflowData(BaseModel):
Payload echoed back in the output.
key : str | None
Routing key. A record that arrives without a key is assigned an
- auto-generated one during deserialization, mirroring the key
- auto-generation the from_list runner performs for keyless records.
+ auto-generated one during deserialization.
"""
value: str
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_test.py
b/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_test.py
deleted file mode 100644
index 477424b6..00000000
--- a/python/flink_agents/e2e_tests/e2e_tests_integration/workflow_test.py
+++ /dev/null
@@ -1,149 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-import json
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar
-
-from pydantic import BaseModel
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.decorators import action
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.events.event_type import EventType
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from flink_agents.api.runner_context import RunnerContext
-
-if TYPE_CHECKING:
- from flink_agents.api.memory_reference import MemoryRef
-
-current_dir = Path(__file__).parent
-
-
-class ProcessedData(BaseModel):
- content: str
- visit_count: int
-
-
-class MyEvent(Event):
- EVENT_TYPE: ClassVar[str] = "_my_event"
-
- def __init__(self, value: Any) -> None:
- """Create a MyEvent with the given value."""
- super().__init__(
- type=MyEvent.EVENT_TYPE,
- attributes={"value": value},
- )
-
- @classmethod
- def from_event(cls, event: "Event") -> "MyEvent":
- """Reconstruct a MyEvent from a generic Event."""
- assert "value" in event.attributes, "Missing 'value' in event
attributes"
- return cls(value=event.attributes["value"])
-
- @property
- def value(self) -> Any:
- """Return the event value."""
- return self.attributes["value"]
-
-
-# TODO: Replace this agent with more practical example.
-class MyAgent(Agent):
- """An example of agent to show the basic usage.
-
- Currently, this agent doesn't really make sense, and it's mainly for
developing
- validation.
- """
-
- @action(EventType.InputEvent)
- @staticmethod
- def first_action(event: Event, ctx: RunnerContext) -> None:
- key = ctx.key
- input_message = InputEvent.from_event(event).input
- memory = ctx.short_term_memory
-
- data_path = f"user_data.{key}"
- stored = memory.get(data_path)
- previous_data = ProcessedData.model_validate(stored) if stored else
None
- current_count = previous_data.visit_count if previous_data else 0
- new_count = current_count + 1
-
- data_to_store = ProcessedData(content=input_message,
visit_count=new_count)
- data_ref = memory.set(data_path, data_to_store.model_dump(mode="json"))
-
- ctx.send_event(MyEvent(value=data_ref))
-
- processed_content = f"{input_message} -> processed_by_first_action"
- key_with_count = f"(visit {new_count} times)"
- ctx.send_event(OutputEvent(output={key_with_count: processed_content}))
-
- @action(MyEvent.EVENT_TYPE)
- @staticmethod
- def second_action(event: Event, ctx: RunnerContext) -> None:
- content_ref: MemoryRef = MyEvent.from_event(event).value
- memory = ctx.short_term_memory
-
- processed_data = ProcessedData.model_validate(memory.get(content_ref))
-
- base_message = processed_data.content
- current_count = processed_data.visit_count
- new_count = current_count + 1
-
- updated_data_to_store = ProcessedData(
- content=base_message, visit_count=new_count
- )
- memory.set(content_ref.path,
updated_data_to_store.model_dump(mode="json"))
-
- final_content = f"{base_message} -> processed by second_action"
- key_with_count = f"(visit {new_count} times)"
- ctx.send_event(OutputEvent(output={key_with_count: final_content}))
-
-
-def test_workflow() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = MyAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "bob", "value": "The message from bob"})
- input_list.append({"k": "john", "v": "The message from john"})
- input_list.append({"key": "john", "value": "Second message from john"})
- input_list.append({"key": "bob", "value": "Second message from bob"})
- input_list.append(
- {"value": "Message from unknown"}
- ) # will automatically generate a new unique key
-
- env.execute()
-
- expected_output = []
- with Path.open(
- Path(f"{current_dir}/../resources/ground_truth/test_workflow.txt")
- ) as f:
- for line in f:
- expected_output.append(json.loads(line)) # noqa:PERF401
-
- assert output_list[:8] == expected_output[:8]
- assert (
- output_list[8][next(iter(output_list[8].keys()))]
- == expected_output[8][next(iter(expected_output[8].keys()))]
- )
- assert (
- output_list[9][next(iter(output_list[9].keys()))]
- == expected_output[9][next(iter(expected_output[9].keys()))]
- )
diff --git a/python/flink_agents/e2e_tests/test_utils.py
b/python/flink_agents/e2e_tests/test_utils.py
index ea1b8832..68ef8009 100644
--- a/python/flink_agents/e2e_tests/test_utils.py
+++ b/python/flink_agents/e2e_tests/test_utils.py
@@ -21,8 +21,6 @@ from pathlib import Path
from ollama import Client
-from flink_agents.api.events.tool_event import ToolRequestEvent
-
current_dir = Path(__file__).parent
@@ -80,35 +78,6 @@ def collect_tool_invocations(log_dir: str | Path) ->
list[dict]:
return invocations
-def tool_invocations_from_events(events: list[ToolRequestEvent]) -> list[dict]:
- """Normalize live ``ToolRequestEvent`` objects to the same invocation
shape.
-
- Adapts the in-memory capture (the ``LocalRunner`` hook) to the same
- ``{name, arguments}`` shape :func:`collect_tool_invocations` returns from
the
- event log, so both sources feed :func:`assert_tool_invoked` identically.
Each
- event's ``tool_calls`` is a list of nested ``{id, type, function:{name,
- arguments}}`` dicts; order is preserved.
-
- Args:
- events: ``ToolRequestEvent`` objects captured during a local run.
-
- Returns:
- Ordered list of ``{"name": str, "arguments": dict | str}``, one per
tool
- call across all events.
- """
- invocations = []
- for event in events:
- for tool_call in event.tool_calls:
- function = tool_call["function"]
- invocations.append(
- {
- "name": function["name"],
- "arguments": function["arguments"],
- }
- )
- return invocations
-
-
def assert_tool_invoked(invocations: list[dict], name: str, arguments: dict)
-> None:
"""Assert some invocation called tool ``name`` with arguments equal to
``arguments``.
diff --git a/python/flink_agents/e2e_tests/test_utils_unit_test.py
b/python/flink_agents/e2e_tests/test_utils_unit_test.py
index 2236abd7..f15089a4 100644
--- a/python/flink_agents/e2e_tests/test_utils_unit_test.py
+++ b/python/flink_agents/e2e_tests/test_utils_unit_test.py
@@ -20,11 +20,9 @@ from pathlib import Path
import pytest
-from flink_agents.api.events.tool_event import ToolRequestEvent
from flink_agents.e2e_tests.test_utils import (
assert_tool_invoked,
collect_tool_invocations,
- tool_invocations_from_events,
)
@@ -129,22 +127,3 @@ def
test_assert_tool_invoked_mismatch_reports_invocations() -> None:
assert_tool_invoked(invocations, "add", {"a": 1, "b": 2})
assert "add" in str(exc_info.value)
assert "9" in str(exc_info.value)
-
-
-def test_tool_invocations_from_events() -> None:
- """Live ToolRequestEvents normalize to the same {name, arguments} shape.
-
- One event carrying two tool calls yields one invocation per call, in order.
- """
- event = ToolRequestEvent(
- model="qwen3:1.7b",
- tool_calls=[
- _function_tool_call("add", {"a": 1, "b": 2}),
- _function_tool_call("multiply", {"a": 4444, "b": 312}),
- ],
- )
-
- assert tool_invocations_from_events([event]) == [
- {"name": "add", "arguments": {"a": 1, "b": 2}},
- {"name": "multiply", "arguments": {"a": 4444, "b": 312}},
- ]
diff --git a/python/flink_agents/examples/rag/rag_agent_example.py
b/python/flink_agents/examples/rag/rag_agent_example.py
index 6972cdf4..00e688b6 100644
--- a/python/flink_agents/examples/rag/rag_agent_example.py
+++ b/python/flink_agents/examples/rag/rag_agent_example.py
@@ -17,6 +17,8 @@
################################################################################
import os
+from pyflink.datastream import StreamExecutionEnvironment
+
from flink_agents.api.agents.agent import Agent
from flink_agents.api.chat_message import ChatMessage, MessageRole
from flink_agents.api.decorators import (
@@ -165,17 +167,9 @@ if __name__ == "__main__":
agent = MyRAGAgent()
- # Prepare example queries
- input_list = []
- test_queries = [
- {"key": "001", "value": "What is Apache Flink?"},
- {"key": "002", "value": "What is Apache Flink Agents?"},
- {"key": "003", "value": "What is Python?"},
- ]
- input_list.extend(test_queries)
-
- # Setup the Agents execution environment
- agents_env = AgentsExecutionEnvironment.get_execution_environment()
+ # Set up the Flink streaming environment and the Agents execution
environment.
+ env = StreamExecutionEnvironment.get_execution_environment()
+ agents_env = AgentsExecutionEnvironment.get_execution_environment(env)
# Setup Ollama embedding and chat model connections
agents_env.add_resource(
@@ -189,15 +183,22 @@ if __name__ == "__main__":
ResourceDescriptor(clazz=ResourceName.ChatModel.OLLAMA_CONNECTION),
)
- output_list = agents_env.from_list(input_list).apply(agent).to_list()
-
- agents_env.execute()
+ # A small stream of example queries, keyed by the query text.
+ query_stream = env.from_collection(
+ [
+ "What is Apache Flink?",
+ "What is Apache Flink Agents?",
+ "What is Python?",
+ ],
+ )
- print("\n" + "=" * 50)
- print("RAG Example Results:")
- print("=" * 50)
+ # Use the RAG agent to answer each query and print the responses to stdout.
+ response_stream = (
+ agents_env.from_datastream(input=query_stream, key_selector=lambda x:
x)
+ .apply(agent)
+ .to_datastream()
+ )
+ response_stream.print()
- for output in output_list:
- for key, value in output.items():
- print(f"\n[{key}] Response: {value}")
- print("-" * 40)
+ # Execute the Flink pipeline.
+ agents_env.execute("RAG Agent Example Job")
diff --git a/python/flink_agents/plan/tests/actions/test_chat_model_action.py
b/python/flink_agents/plan/tests/actions/test_chat_model_action.py
index 94ff8ab9..f0bba58e 100644
--- a/python/flink_agents/plan/tests/actions/test_chat_model_action.py
+++ b/python/flink_agents/plan/tests/actions/test_chat_model_action.py
@@ -31,7 +31,7 @@ from flink_agents.plan.actions.chat_model_action import (
_save_tool_request_event_context,
_update_tool_call_context,
)
-from flink_agents.runtime.local_memory_object import LocalMemoryObject
+from flink_agents.runtime.tests.local_memory_object import LocalMemoryObject
def _memory() -> LocalMemoryObject:
diff --git a/python/flink_agents/runtime/local_execution_environment.py
b/python/flink_agents/runtime/local_execution_environment.py
deleted file mode 100644
index c9768fc3..00000000
--- a/python/flink_agents/runtime/local_execution_environment.py
+++ /dev/null
@@ -1,191 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-from typing import Any, Callable, Dict, List
-
-from pyflink.common import TypeInformation
-from pyflink.datastream import DataStream, KeySelector,
StreamExecutionEnvironment
-from pyflink.table import Schema, StreamTableEnvironment, Table
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.events.tool_event import ToolRequestEvent
-from flink_agents.api.execution_environment import (
- AgentBuilder,
- AgentsExecutionEnvironment,
-)
-from flink_agents.plan.configuration import AgentConfiguration
-from flink_agents.runtime.local_runner import LocalRunner
-
-
-class LocalAgentBuilder(AgentBuilder):
- """LocalAgentBuilder for building agent instance."""
-
- __env: "LocalExecutionEnvironment"
- __input: List[Dict[str, Any]]
- __output: List[Any]
- __runner: LocalRunner = None
- __executed: bool = False
- __config: AgentConfiguration
-
- def __init__(
- self,
- env: "LocalExecutionEnvironment",
- input: List[Dict[str, Any]],
- config: AgentConfiguration,
- ) -> None:
- """Init empty output list."""
- self.__env = env
- self.__input = input
- self.__output = []
- self.__config = config
-
- def apply(self, agent: Agent | str) -> AgentBuilder:
- """Create local runner to execute given agent.
-
- Doesn't support apply multiple Agents.
- """
- if self.__runner is not None:
- err_msg = "LocalAgentBuilder doesn't support apply multiple
agents."
- raise RuntimeError(err_msg)
- if isinstance(agent, str):
- if agent not in self.__env._agents:
- msg = (
- f"No agent named {agent!r} is registered on this "
- "environment. Did you call load_yaml first?"
- )
- raise ValueError(msg)
- agent = self.__env._agents[agent]
- # inspect resources from environment to agent instance.
- registered_resources = self.__env.resources
- for type, name_to_resource in registered_resources.items():
- agent.resources[type] = name_to_resource | agent.resources[type]
- self.__runner = LocalRunner(agent, self.__config)
- self.__env.set_agent(self.__input, self.__output, self.__runner)
- return self
-
- def to_list(self) -> List[Dict[str, Any]]:
- """Get output list of execution environment."""
- return self.__output
-
- def to_datastream(self) -> DataStream:
- """Get output DataStream of agent execution.
-
- This method is not supported for LocalAgentBuilder.
- """
- msg = "LocalAgentBuilder does not support to_datastream."
- raise NotImplementedError(msg)
-
- def to_table(self, schema: Schema, output_type: TypeInformation) -> Table:
- """Get output Table of agent execution.
-
- This method is not supported for LocalAgentBuilder.
- """
- msg = "LocalAgentBuilder does not support to_table."
- raise NotImplementedError(msg)
-
-
-class LocalExecutionEnvironment(AgentsExecutionEnvironment):
- """Implementation of AgentsExecutionEnvironment for local execution
environment."""
-
- __input: List[Dict[str, Any]] = None
- __output: List[Any] = None
- __runner: LocalRunner = None
- __executed: bool = False
- __config: AgentConfiguration = AgentConfiguration()
-
- def get_config(self, path: str | None = None) -> AgentConfiguration:
- """Get configuration of execution environment."""
- if path is not None:
- self.__config.load_from_file(path)
- return self.__config
-
- def from_list(self, input: list) -> LocalAgentBuilder:
- """Set input list of execution environment."""
- if self.__input is not None:
- err_msg = "LocalExecutionEnvironment doesn't support call
from_list multiple times."
- raise RuntimeError(err_msg)
-
- self.__input = input
- return LocalAgentBuilder(env=self, input=input, config=self.__config)
-
- def set_agent(self, input: list, output: list, runner: LocalRunner) ->
None:
- """Set agent input, output and runner."""
- self.__input = input
- self.__runner = runner
- self.__output = output
-
- def execute(self, job_name: str | None = None) -> None:
- """Execute agent individually."""
- if self.__executed:
- err_msg = (
- "LocalExecutionEnvironment doesn't support execute multiple
times."
- )
- raise RuntimeError(err_msg)
- self.__executed = True
- for input in self.__input:
- self.__runner.run(**input)
- outputs = self.__runner.get_outputs()
- for output in outputs:
- self.__output.append(output)
-
- def get_tool_request_events(self) -> List[ToolRequestEvent]:
- """Get the ToolRequestEvents captured by the runner during
execution."""
- return self.__runner.get_tool_request_events()
-
- def from_datastream(
- self, input: DataStream, key_selector: KeySelector | Callable | None =
None
- ) -> AgentBuilder:
- """Set input DataStream of agent execution.
-
- This method is not supported for local execution environments.
- """
- msg = "LocalExecutionEnvironment does not support from_datastream."
- raise NotImplementedError(msg)
-
- def from_table(
- self,
- input: Table,
- key_selector: KeySelector | Callable | None = None,
- ) -> AgentBuilder:
- """Set input Table of agent execution.
-
- This method is not supported for local execution environments.
- """
- msg = "LocalExecutionEnvironment does not support from_table."
- raise NotImplementedError(msg)
-
-
-def create_instance(
- env: StreamExecutionEnvironment, t_env: StreamTableEnvironment, **kwargs:
Any
-) -> AgentsExecutionEnvironment:
- """Factory function to create a remote agents execution environment.
-
- Parameters
- ----------
- env : StreamExecutionEnvironment
- Flink job execution environment.
- t_env: StreamTableEnvironment
- Flink job execution table environment.
- **kwargs : Dict[str, Any]
- The dict of parameters to configure the execution environment.
-
- Returns:
- -------
- AgentsExecutionEnvironment
- A configured agents execution environment instance.
- """
- return LocalExecutionEnvironment()
diff --git a/python/flink_agents/runtime/local_runner.py
b/python/flink_agents/runtime/local_runner.py
deleted file mode 100644
index 38436a2c..00000000
--- a/python/flink_agents/runtime/local_runner.py
+++ /dev/null
@@ -1,387 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-import asyncio
-import logging
-import uuid
-from collections import deque
-from concurrent.futures import Future
-from typing import TYPE_CHECKING, Any, Callable, Dict, List
-
-from typing_extensions import override
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.events.tool_event import ToolRequestEvent
-from flink_agents.api.memory.long_term_memory import BaseLongTermMemory
-from flink_agents.api.memory_object import MemoryObject, MemoryType
-from flink_agents.api.metric_group import MetricGroup
-from flink_agents.api.resource import Resource, ResourceType
-from flink_agents.api.runner_context import AsyncExecutionResult, RunnerContext
-from flink_agents.plan.configuration import AgentConfiguration
-from flink_agents.runtime.agent_runner import AgentRunner
-from flink_agents.runtime.local_memory_object import LocalMemoryObject
-from flink_agents.runtime.resource_cache import ResourceCache
-
-if TYPE_CHECKING:
- from flink_agents.plan.agent_plan import AgentPlan
-
-logging.basicConfig(
- level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
-)
-logger = logging.getLogger(__name__)
-
-
-class LocalRunnerContext(RunnerContext):
- """Implementation of RunnerContext for local agent execution.
-
- Attributes:
- ----------
- __agent_plan : AgentPlan
- Internal agent plan for this context.
- __key : Any
- Unique identifier for the context, correspond to the key in flink
KeyedStream.
- events : deque[Event]
- Queue of events to be processed in this context.
- action_name: str
- Name of the action being executed.
- """
-
- __agent_plan: Any
- __key: Any
- events: deque[Event]
- action_name: str
- _sensory_mem_store: dict[str, Any]
- _short_term_mem_store: dict[str, Any]
- _sensory_memory: MemoryObject
- _short_term_memory: MemoryObject
- _config: AgentConfiguration
-
- def __init__(
- self, agent_plan: "AgentPlan", key: Any, config: AgentConfiguration
- ) -> None:
- """Initialize a new context with the given agent and key.
-
- Parameters
- ----------
- agent_plan : AgentPlan
- Agent plan used for this context.
- key : Any
- Unique context identifier, which is corresponding to the key in
flink
- KeyedStream.
- """
- self.__agent_plan = agent_plan
- self.__resource_cache = ResourceCache(
- agent_plan.resource_providers, agent_plan.config
- )
- self.__key = key
- self.events = deque()
- self._sensory_mem_store = {}
- self._short_term_mem_store = {}
- self._sensory_memory = LocalMemoryObject(
- MemoryType.SENSORY, self._sensory_mem_store,
LocalMemoryObject.ROOT_KEY
- )
- self._short_term_memory = LocalMemoryObject(
- MemoryType.SHORT_TERM,
- self._short_term_mem_store,
- LocalMemoryObject.ROOT_KEY,
- )
- self._config = config
-
- @property
- def key(self) -> Any:
- """Get the unique identifier for this context.
-
- Returns:
- -------
- Any
- The unique identifier for this context.
- """
- return self.__key
-
- @override
- def send_event(self, event: Event) -> None:
- """Send an event to the context's event queue and log it.
-
- Parameters
- ----------
- event : Event
- The event to be added to the queue.
- """
- logger.info("key: %s, send_event: %s", self.__key, event)
- self.events.append(event)
-
- @override
- def get_resource(
- self, name: str, type: ResourceType, metric_group: MetricGroup = None
- ) -> Resource:
- return self.__resource_cache.get_resource(name, type)
-
- @property
- @override
- def action_config(self) -> Dict[str, Any]:
- """Get config of the action."""
- return
self.__agent_plan.get_action_config(action_name=self.action_name)
-
- @override
- def get_action_config_value(self, key: str) -> Any:
- """Get config option value of the key."""
- return self.__agent_plan.get_action_config_value(
- action_name=self.action_name, key=key
- )
-
- @property
- @override
- def sensory_memory(self) -> MemoryObject:
- """Get the sensory memory object associated with this context.
-
- Returns:
- -------
- MemoryObject
- The root object of the short-term memory.
- """
- return self._sensory_memory
-
- @property
- @override
- def short_term_memory(self) -> MemoryObject:
- """Get the short-term memory object associated with this context.
-
- Returns:
- -------
- MemoryObject
- The root object of the short-term memory.
- """
- return self._short_term_memory
-
- @property
- @override
- def long_term_memory(self) -> BaseLongTermMemory:
- err_msg = "Long-Term Memory is not supported for local agent execution
yet."
- raise NotImplementedError(err_msg)
-
- @property
- @override
- def agent_metric_group(self) -> None:
- # TODO: Support metric mechanism for local agent execution.
- err_msg = "Metric mechanism is not supported for local agent execution
yet."
- logger.warning(err_msg)
- return
-
- @property
- @override
- def action_metric_group(self) -> None:
- # TODO: Support metric mechanism for local agent execution.
- err_msg = "Metric mechanism is not supported for local agent execution
yet."
- logger.warning(err_msg)
- return
-
- @override
- def durable_execute(
- self,
- func: Callable[[Any], Any],
- *args: Any,
- reconciler: Callable[[], Any] | None = None,
- **kwargs: Any,
- ) -> Any:
- """Synchronously execute the provided function. Access to memory
- is prohibited within the function.
-
- Note: Local runner does not support durable execution, so recovery
- is not available.
- """
- logger.warning(
- "Local runner does not support durable execution; recovery is not
available."
- )
- return func(*args, **kwargs)
-
- @override
- def durable_execute_async(
- self,
- func: Callable[[Any], Any],
- *args: Any,
- reconciler: Callable[[], Any] | None = None,
- **kwargs: Any,
- ) -> AsyncExecutionResult:
- """Asynchronously execute the provided function. Access to memory
- is prohibited within the function.
-
- Note: Local runner executes synchronously but returns an
AsyncExecutionResult
- for API consistency. Durable execution is not supported.
- """
- logger.warning(
- "Local runner does not support durable execution; recovery is not
available."
- )
- # Execute synchronously and wrap the result in a completed Future
- future: Future = Future()
- try:
- result = func(*args, **kwargs)
- future.set_result(result)
- except Exception as e:
- future.set_exception(e)
-
- # Create a mock executor that returns the pre-completed future
- class _SyncExecutor:
- def __init__(self, completed_future: Future) -> None:
- self._future = completed_future
-
- def submit(self, fn: Callable, *args: Any, **kwargs: Any) ->
Future:
- return self._future
-
- return AsyncExecutionResult(_SyncExecutor(future), func, args, kwargs)
-
- @property
- @override
- def config(self) -> AgentConfiguration:
- return self._config
-
- def clear_sensory_memory(self) -> None:
- """Clean up sensory memory."""
- self._sensory_mem_store.clear()
-
- def close(self) -> None:
- """Cleanup the resource."""
- if self.__resource_cache is not None:
- try:
- self.__resource_cache.close()
- finally:
- self.__resource_cache = None
-
-
-class LocalRunner(AgentRunner):
- """Agent runner implementation for local execution, which is
- convenient for debugging.
-
- Attributes:
- ----------
- __agent_plan : AgentPlan
- Internal agent plan.
- __keyed_contexts : dict[Any, LocalRunnerContext]
- Dictionary of active contexts indexed by key.
- __outputs:
- Outputs generated by agent execution.
- __tool_request_events:
- ToolRequestEvents observed during agent execution.
- __config:
- Internal configration.
- """
-
- __agent_plan: Any
- __keyed_contexts: Dict[Any, LocalRunnerContext]
- __outputs: List[Dict[str, Any]]
- __tool_request_events: List[ToolRequestEvent]
- __config: AgentConfiguration
-
- def __init__(self, agent: Agent, config: AgentConfiguration) -> None:
- """Initialize the runner with the provided agent.
-
- Parameters
- ----------
- agent : Agent
- The agent class to convert and run.
- """
- from flink_agents.plan.agent_plan import AgentPlan
-
- self.__agent_plan = AgentPlan.from_agent(agent, config)
- self.__keyed_contexts = {}
- self.__outputs = []
- self.__tool_request_events = []
- self.__config = config
-
- @override
- def run(self, **data: Dict[str, Any]) -> Any:
- """Execute the agent to process the given data.
-
- Parameters
- ----------
- **data : dict[str, Any]
- input record from upstream.
-
- Returns:
- -------
- key
- The key of the input that was processed.
- """
- if "key" in data:
- key = data["key"]
- elif "k" in data:
- key = data["k"]
- else:
- key = uuid.uuid4()
-
- if key not in self.__keyed_contexts:
- self.__keyed_contexts[key] = LocalRunnerContext(
- self.__agent_plan, key, self.__config
- )
- context = self.__keyed_contexts[key]
- context.clear_sensory_memory()
-
- if "value" in data:
- input_event = InputEvent(input=data["value"])
- elif "v" in data:
- input_event = InputEvent(input=data["v"])
- else:
- msg = "Input data must be dict has 'v' or 'value' field"
- raise RuntimeError(msg)
-
- context.send_event(input_event)
-
- while len(context.events) > 0:
- event = context.events.popleft()
- if isinstance(event, OutputEvent):
- self.__outputs.append({key: event.output})
- continue
- if isinstance(event, ToolRequestEvent):
- self.__tool_request_events.append(event)
- # Fall through: the request must still dispatch to its action.
- event_type = event.get_type()
- for action in self.__agent_plan.get_actions(event_type):
- logger.info("key: %s, performing action: %s", key, action.name)
- context.action_name = action.name
- func_result = action.exec(event, context)
- if asyncio.iscoroutine(func_result):
- # Drive the coroutine to completion using send()
- try:
- while True:
- func_result.send(None)
- except StopIteration:
- # Coroutine completed normally
- pass
- except Exception:
- logger.exception("Error in async execution")
- raise
- return key
-
- def get_outputs(self) -> List[Dict[str, Any]]:
- """Get the outputs generated by agent execution.
-
- Returns:
- -------
- List[Dict[str, Any]]
- The agent execution outputs.
- """
- return self.__outputs
-
- def get_tool_request_events(self) -> List[ToolRequestEvent]:
- """Get the ToolRequestEvents captured during agent execution.
-
- Returns:
- -------
- List[ToolRequestEvent]
- The ToolRequestEvents observed in the run loop, in order.
- """
- return self.__tool_request_events
diff --git a/python/flink_agents/runtime/remote_execution_environment.py
b/python/flink_agents/runtime/remote_execution_environment.py
index c5ac93f8..b14ed0eb 100644
--- a/python/flink_agents/runtime/remote_execution_environment.py
+++ b/python/flink_agents/runtime/remote_execution_environment.py
@@ -18,7 +18,7 @@
import logging
import os
from pathlib import Path
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict
import cloudpickle
from pyflink.common import TypeInformation
@@ -161,14 +161,6 @@ class RemoteAgentBuilder(AgentBuilder):
"""
return self.t_env.from_data_stream(self.to_datastream(output_type),
schema)
- def to_list(self) -> List[Dict[str, Any]]:
- """Get output list of agent execution.
-
- This method is not supported for remote execution environments.
- """
- msg = "RemoteAgentBuilder does not support to_list."
- raise NotImplementedError(msg)
-
class RemoteExecutionEnvironment(AgentsExecutionEnvironment):
"""Implementation of AgentsExecutionEnvironment for execution with
DataStream."""
@@ -271,14 +263,6 @@ class
RemoteExecutionEnvironment(AgentsExecutionEnvironment):
agents=self._agents,
)
- def from_list(self, input: List[Dict[str, Any]]) ->
"AgentsExecutionEnvironment":
- """Set input list of agent execution.
-
- This method is not supported for remote execution environments.
- """
- msg = "RemoteExecutionEnvironment does not support from_list."
- raise NotImplementedError(msg)
-
def execute(self, job_name: str | None = None) -> None:
"""Execute agent."""
self.__env.execute(job_name=job_name)
diff --git a/python/flink_agents/runtime/local_memory_object.py
b/python/flink_agents/runtime/tests/local_memory_object.py
similarity index 100%
rename from python/flink_agents/runtime/local_memory_object.py
rename to python/flink_agents/runtime/tests/local_memory_object.py
diff --git a/python/flink_agents/runtime/tests/test_built_in_actions.py
b/python/flink_agents/runtime/tests/test_built_in_actions.py
deleted file mode 100644
index 057291c1..00000000
--- a/python/flink_agents/runtime/tests/test_built_in_actions.py
+++ /dev/null
@@ -1,218 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-import uuid
-from typing import Any, Dict, List, Sequence
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.chat_message import ChatMessage, MessageRole
-from flink_agents.api.chat_models.chat_model import (
- BaseChatModelConnection,
- BaseChatModelSetup,
-)
-from flink_agents.api.decorators import (
- action,
- chat_model_connection,
- chat_model_setup,
- prompt,
- tool,
-)
-from flink_agents.api.events.chat_event import ChatRequestEvent,
ChatResponseEvent
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.events.event_type import EventType
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from flink_agents.api.prompts.prompt import Prompt
-from flink_agents.api.resource import ResourceDescriptor, ResourceType
-from flink_agents.api.runner_context import RunnerContext
-from flink_agents.api.tools.tool import ToolType
-
-
-class MockChatModelConnection(BaseChatModelConnection):
- """Mock ChatModel for testing integrating prompt and tool."""
-
- def chat(
- self,
- messages: Sequence[ChatMessage],
- tools: List | None = None,
- **kwargs: Any,
- ) -> ChatMessage:
- """Generate tool call or response according to input."""
- # Generate tool call
- if "sum" in messages[-1].content:
- input = messages[-1].content
- # Validate bind_tools
- assert tools[0].name == "add"
- function = {"name": "add", "arguments": {"a": 1, "b": 2}}
- tool_call = {
- "id": uuid.uuid4(),
- "type": ToolType.FUNCTION,
- "function": function,
- }
- return ChatMessage(
- role=MessageRole.ASSISTANT, content=input,
tool_calls=[tool_call]
- )
- # Generate response including tool call context
- else:
- content = "\n".join([message.content for message in messages])
- return ChatMessage(role=MessageRole.ASSISTANT, content=content)
-
-
-class MockChatModel(BaseChatModelSetup):
- """Mock ChatModel for testing integrating prompt and tool."""
-
- def open(self) -> None:
- """Do nothing."""
-
- @property
- def model_kwargs(self) -> Dict[str, Any]:
- """Return model kwargs."""
- return {}
-
- def chat(
- self,
- messages: Sequence[ChatMessage],
- prompt_args: Dict[str, Any] | None = None,
- **kwargs: Any,
- ) -> ChatMessage:
- """Execute chat conversation."""
- # Get model connection
- server = self.resource_context.get_resource(
- self.connection, ResourceType.CHAT_MODEL_CONNECTION
- )
-
- # Apply prompt template
- if self.prompt is not None:
- if isinstance(self.prompt, str):
- # Get prompt resource if it's a string
- prompt = self.resource_context.get_resource(
- self.prompt, ResourceType.PROMPT
- )
- else:
- prompt = self.prompt
-
- if "sum" in messages[-1].content:
- str_prompt_args = (
- {k: str(v) for k, v in prompt_args.items()} if prompt_args
else {}
- )
- messages = prompt.format_messages(**str_prompt_args)
-
- # Bind tools
- tools = None
- if self.tools is not None:
- tools = [
- self.resource_context.get_resource(tool_name,
ResourceType.TOOL)
- for tool_name in self.tools
- ]
-
- # Call server to execute chat
- return server.chat(messages, tools=tools, **kwargs)
-
-
-class MyAgent(Agent):
- """Mock agent for testing built-in actions."""
-
- @prompt
- @staticmethod
- def prompt() -> Prompt:
- """Prompt can be used in action or chat model."""
- return Prompt.from_text(
- text="Please call the appropriate tool to do the following task:
{task}",
- )
-
- @chat_model_connection
- @staticmethod
- def mock_connection() -> ResourceDescriptor:
- """Chat model server can be used by ChatModel."""
- return ResourceDescriptor(
-
clazz=f"{MockChatModelConnection.__module__}.{MockChatModelConnection.__name__}"
- )
-
- @chat_model_setup
- @staticmethod
- def mock_chat_model() -> ResourceDescriptor:
- """Chat model can be used in action."""
- return ResourceDescriptor(
- clazz=f"{MockChatModel.__module__}.{MockChatModel.__name__}",
- connection="mock_connection",
- model="mock-model",
- prompt="prompt",
- tools=["add"],
- )
-
- @tool
- @staticmethod
- def add(a: int, b: int) -> int:
- """Calculate the sum of a and b.
-
- Parameters
- ----------
- a : int
- The first operand
- b : int
- The second operand
-
- Returns:
- -------
- int:
- The sum of a and b
- """
- return a + b
-
- @action(EventType.InputEvent)
- @staticmethod
- def process_input(event: Event, ctx: RunnerContext) -> None:
- """User defined action for processing input.
-
- In this action, we will send ChatRequestEvent to trigger built-in
actions.
- """
- input = InputEvent.from_event(event).input
- ctx.send_event(
- ChatRequestEvent(
- model="mock_chat_model",
- messages=[ChatMessage(role=MessageRole.USER, content=input)],
- prompt_args={"task": input},
- )
- )
-
- @action(EventType.ChatResponseEvent)
- @staticmethod
- def process_chat_response(event: Event, ctx: RunnerContext) -> None:
- """User defined action for processing chat model response."""
- input = ChatResponseEvent.from_event(event).response
- ctx.send_event(OutputEvent(output=input.content))
-
-
-def test_built_in_actions() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = MyAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "0001", "value": "calculate the sum of 1 and
2."})
-
- env.execute()
-
- assert output_list == [
- {
- "0001": "calculate the sum of 1 and 2.\n"
- "Please call the appropriate tool to do the following task: "
- "calculate the sum of 1 and 2.\n"
- "3"
- }
- ]
diff --git a/python/flink_agents/runtime/tests/test_get_resource_in_action.py
b/python/flink_agents/runtime/tests/test_get_resource_in_action.py
deleted file mode 100644
index 8c5988bb..00000000
--- a/python/flink_agents/runtime/tests/test_get_resource_in_action.py
+++ /dev/null
@@ -1,114 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-from typing import Any, Dict, Sequence
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.chat_message import ChatMessage, MessageRole
-from flink_agents.api.chat_models.chat_model import BaseChatModelSetup
-from flink_agents.api.decorators import action, chat_model_setup, tool
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.events.event_type import EventType
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from flink_agents.api.resource import ResourceDescriptor, ResourceType
-from flink_agents.api.runner_context import RunnerContext
-
-
-class MockChatModelImpl(BaseChatModelSetup):
- host: str
- desc: str
-
- def open(self) -> None:
- """Do nothing."""
-
- @property
- def model_kwargs(self) -> Dict[str, Any]:
- return {}
-
- def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) ->
ChatMessage:
- return ChatMessage(
- role=MessageRole.ASSISTANT,
- content=f"{messages[0].content} {self.host} {self.desc}",
- )
-
-
-class MyAgent(Agent):
- @chat_model_setup
- @staticmethod
- def mock_chat_model() -> ResourceDescriptor:
- return ResourceDescriptor(
-
clazz=f"{MockChatModelImpl.__module__}.{MockChatModelImpl.__name__}",
- host="8.8.8.8",
- desc="mock chat model just for testing.",
- connection="mock",
- model="mock-model",
- )
-
- @tool
- @staticmethod
- def mock_tool(input: str) -> str:
- """Mock tool.
-
- Parameters
- ----------
- input : str
- The input message.
-
- Returns:
- -------
- strs
- Response string value.
- """
- return input + " mock tools just for testing."
-
- @action(EventType.InputEvent)
- @staticmethod
- def mock_action(event: Event, ctx: RunnerContext) -> None:
- input = InputEvent.from_event(event).input
- mock_chat_model = ctx.get_resource(
- type=ResourceType.CHAT_MODEL, name="mock_chat_model"
- )
- mock_tool = ctx.get_resource(type=ResourceType.TOOL, name="mock_tool")
- ctx.send_event(
- OutputEvent(
- output=mock_chat_model.chat(
- messages=[ChatMessage(role=MessageRole.USER,
content=input)]
- ).content
- + " "
- + mock_tool.call("call")
- )
- )
-
-
-def test_get_resource_in_action() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = MyAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "bob", "value": "the first message."})
-
- env.execute()
-
- assert output_list == [
- {
- "bob": "the first message. 8.8.8.8 mock chat model "
- "just for testing. call mock tools just for testing."
- }
- ]
diff --git
a/python/flink_agents/runtime/tests/test_local_execution_environment.py
b/python/flink_agents/runtime/tests/test_local_execution_environment.py
deleted file mode 100644
index 294549d2..00000000
--- a/python/flink_agents/runtime/tests/test_local_execution_environment.py
+++ /dev/null
@@ -1,380 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-import time
-import uuid
-from typing import Any, ClassVar, Dict, List, Sequence
-
-import pytest
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.chat_message import ChatMessage, MessageRole
-from flink_agents.api.chat_models.chat_model import (
- BaseChatModelConnection,
- BaseChatModelSetup,
-)
-from flink_agents.api.decorators import (
- action,
- chat_model_connection,
- chat_model_setup,
- tool,
-)
-from flink_agents.api.events.chat_event import ChatRequestEvent,
ChatResponseEvent
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.events.event_type import EventType
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from flink_agents.api.resource import ResourceDescriptor, ResourceType
-from flink_agents.api.runner_context import RunnerContext
-from flink_agents.api.tools.tool import ToolType
-
-
-class Agent1(Agent):
- @action(EventType.InputEvent)
- @staticmethod
- def increment(event: Event, ctx: RunnerContext): # noqa D102
- input = InputEvent.from_event(event).input
- value = input + 1
- ctx.send_event(OutputEvent(output=value))
-
-
-class Agent1WithAsync(Agent):
- @action(EventType.InputEvent)
- @staticmethod
- async def increment(event: Event, ctx: RunnerContext): # noqa D102
- def my_func(value: int) -> int:
- time.sleep(1)
- return value + 1
-
- input = InputEvent.from_event(event).input
- value = await ctx.durable_execute_async(my_func, input)
- ctx.send_event(OutputEvent(output=value))
-
-
-class Agent2(Agent):
- @action(EventType.InputEvent)
- @staticmethod
- def decrease(event: Event, ctx: RunnerContext): # noqa D102
- input = InputEvent.from_event(event).input
- value = input - 1
- ctx.send_event(OutputEvent(output=value))
-
-
-def test_local_execution_environment() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = Agent1()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "bob", "value": 1})
- input_list.append({"k": "john", "v": 2})
-
- env.execute()
-
- assert output_list == [{"bob": 2}, {"john": 3}]
-
-
-def test_local_execution_environment_with_async() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = Agent1WithAsync()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "bob", "value": 1})
- input_list.append({"k": "john", "v": 2})
-
- env.execute()
-
- assert output_list == [{"bob": 2}, {"john": 3}]
-
-
-def test_local_execution_environment_apply_multi_agents() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent1 = Agent1()
- agent2 = Agent2()
-
- with pytest.raises(RuntimeError):
- env.from_list(input_list).apply(agent1).apply(agent2).to_list()
-
-
-def test_local_execution_environment_execute_multi_times() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = Agent1()
-
- env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "bob", "value": 1})
- input_list.append({"k": "john", "v": 2})
-
- env.execute()
- with pytest.raises(RuntimeError):
- env.execute()
-
-
-def test_local_execution_environment_call_from_list_twice() -> None:
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
-
- env.from_list(input_list)
- with pytest.raises(RuntimeError):
- env.from_list(input_list)
-
-
-# ── Unified event E2E tests ──────────────────────────────────────────────
-
-
-class UnifiedEventAgent(Agent):
- @action(EventType.InputEvent)
- @staticmethod
- def on_input(event: Event, ctx: RunnerContext) -> None:
- ctx.send_event(
- Event(
- type="Intermediate",
- attributes={"msg": InputEvent.from_event(event).input},
- )
- )
-
- @action("Intermediate")
- @staticmethod
- def on_intermediate(event: Event, ctx: RunnerContext) -> None:
- ctx.send_event(
- OutputEvent(output=f"processed:{event.get_attr('msg')}")
- )
-
-
-def test_unified_event_workflow() -> None:
- """End-to-end: InputEvent → unified 'Intermediate' → OutputEvent."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = UnifiedEventAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "alice", "value": "hello"})
- env.execute()
-
- assert output_list == [{"alice": "processed:hello"}]
-
-
-class Step1Event(Event):
- """Custom event with a type string."""
-
- EVENT_TYPE: ClassVar[str] = "_step1_event"
-
- def __init__(self, data: str) -> None:
- """Create a Step1Event with the given data."""
- super().__init__(
- type=Step1Event.EVENT_TYPE,
- attributes={"data": data},
- )
-
- @property
- def data(self) -> str:
- """Return the event data."""
- return self.attributes["data"]
-
-
-class MixedEventAgent(Agent):
- """Agent mixing subclassed and string-based event routing."""
-
- @action(EventType.InputEvent)
- @staticmethod
- def start(event: Event, ctx: RunnerContext) -> None:
-
ctx.send_event(Step1Event(data=str(InputEvent.from_event(event).input)))
-
- @action(Step1Event.EVENT_TYPE)
- @staticmethod
- def on_step1(event: Event, ctx: RunnerContext) -> None:
- ctx.send_event(
- Event(type="Step2", attributes={"value": event.get_attr("data")})
- )
-
- @action("Step2")
- @staticmethod
- def on_step2(event: Event, ctx: RunnerContext) -> None:
- ctx.send_event(
- OutputEvent(output=f"done:{event.get_attr('value')}")
- )
-
-
-def test_mixed_event_workflow() -> None:
- """E2E: InputEvent → class Step1Event → unified 'Step2' → OutputEvent."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = MixedEventAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "bob", "value": 42})
- env.execute()
-
- assert output_list == [{"bob": "done:42"}]
-
-
-# ── Tool-request capture hook (Track B) ──────────────────────────────────
-
-
-class _ToolConnection(BaseChatModelConnection):
- """Mock connection emitting a single ``add`` tool call."""
-
- def chat(
- self,
- messages: Sequence[ChatMessage],
- tools: List | None = None,
- **kwargs: Any,
- ) -> ChatMessage:
- """Emit an ``add`` tool call, then echo the tool result as content."""
- last = messages[-1]
- if last.role == MessageRole.TOOL:
- return ChatMessage(role=MessageRole.ASSISTANT,
content=str(last.content))
- tool_call = {
- "id": str(uuid.uuid4()),
- "type": ToolType.FUNCTION,
- "function": {"name": "add", "arguments": {"a": 1, "b": 2}},
- }
- return ChatMessage(
- role=MessageRole.ASSISTANT, content="", tool_calls=[tool_call]
- )
-
-
-class _ToolChatModel(BaseChatModelSetup):
- """Mock setup binding the ``add`` tool to the connection."""
-
- def open(self) -> None:
- """Do nothing."""
-
- @property
- def model_kwargs(self) -> Dict[str, Any]:
- """Return model kwargs."""
- return {}
-
- def chat(
- self,
- messages: Sequence[ChatMessage],
- prompt_args: Dict[str, Any] | None = None,
- **kwargs: Any,
- ) -> ChatMessage:
- """Bind tools and delegate to the connection."""
- server = self.resource_context.get_resource(
- self.connection, ResourceType.CHAT_MODEL_CONNECTION
- )
- tools = [
- self.resource_context.get_resource(name, ResourceType.TOOL)
- for name in (self.tools or [])
- ]
- return server.chat(messages, tools=tools, **kwargs)
-
-
-class ToolRequestAgent(Agent):
- """Agent whose chat model emits a ToolRequestEvent dispatched to ``add``.
-
- The InputEvent action sends a ChatRequestEvent; the mock chat model returns
- an ``add`` tool call, which the built-in chat/tool actions turn into a real
- ToolRequestEvent flowing through the runner. The ToolRequestEvent is
captured
- by the runner AND still dispatched to ``tool_call_action`` — the final
output
- (the tool result) proves capture did not swallow the event.
- """
-
- @chat_model_connection
- @staticmethod
- def conn() -> ResourceDescriptor:
- """Mock chat model connection."""
- return ResourceDescriptor(
- clazz=f"{_ToolConnection.__module__}.{_ToolConnection.__name__}"
- )
-
- @chat_model_setup
- @staticmethod
- def model() -> ResourceDescriptor:
- """Mock chat model bound to the ``add`` tool."""
- return ResourceDescriptor(
- clazz=f"{_ToolChatModel.__module__}.{_ToolChatModel.__name__}",
- connection="conn",
- model="mock-model",
- tools=["add"],
- )
-
- @tool
- @staticmethod
- def add(a: int, b: int) -> int:
- """Return the sum of a and b.
-
- Parameters
- ----------
- a : int
- The first operand.
- b : int
- The second operand.
-
- Returns:
- -------
- int:
- The sum of a and b.
- """
- return a + b
-
- @action(InputEvent.EVENT_TYPE)
- @staticmethod
- def process_input(event: Event, ctx: RunnerContext) -> None:
- """Send a ChatRequestEvent to drive the tool-calling flow."""
- input = InputEvent.from_event(event).input
- ctx.send_event(
- ChatRequestEvent(
- model="model",
- messages=[ChatMessage(role=MessageRole.USER, content=input)],
- )
- )
-
- @action(ChatResponseEvent.EVENT_TYPE)
- @staticmethod
- def process_response(event: Event, ctx: RunnerContext) -> None:
- """Emit the final assistant content as output."""
- response = ChatResponseEvent.from_event(event).response
- ctx.send_event(OutputEvent(output=response.content))
-
-
-def test_local_runner_captures_tool_request_events() -> None:
- """A ToolRequestEvent is captured AND still dispatched to its action."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = ToolRequestAgent()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "0001", "value": "add 1 and 2"})
- env.execute()
-
- captured = env.get_tool_request_events()
- assert len(captured) == 1
- assert captured[0].tool_calls[0]["function"] == {
- "name": "add",
- "arguments": {"a": 1, "b": 2},
- }
- # Dispatch was not swallowed: tool_call_action ran, producing the tool
result
- # that the model echoed back as the final output.
- assert output_list == [{"0001": "3"}]
diff --git a/python/flink_agents/runtime/tests/test_local_memory_object.py
b/python/flink_agents/runtime/tests/test_local_memory_object.py
index 18b56587..bc8444d7 100644
--- a/python/flink_agents/runtime/tests/test_local_memory_object.py
+++ b/python/flink_agents/runtime/tests/test_local_memory_object.py
@@ -18,7 +18,7 @@
from typing import Dict, List
from flink_agents.api.memory_object import MemoryType
-from flink_agents.runtime.local_memory_object import LocalMemoryObject
+from flink_agents.runtime.tests.local_memory_object import LocalMemoryObject
def create_memory() -> LocalMemoryObject:
diff --git
a/python/flink_agents/runtime/tests/test_local_runner_reconcilable.py
b/python/flink_agents/runtime/tests/test_local_runner_reconcilable.py
deleted file mode 100644
index a213423f..00000000
--- a/python/flink_agents/runtime/tests/test_local_runner_reconcilable.py
+++ /dev/null
@@ -1,79 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-import asyncio
-from typing import Any
-
-from flink_agents.runtime.local_runner import LocalRunnerContext
-
-
-def reconciled_add(x: int, y: int) -> int:
- """Return a simple deterministic value for local-runner tests."""
- return x + y
-
-
-def _create_local_runner_context() -> LocalRunnerContext:
- return LocalRunnerContext.__new__(LocalRunnerContext)
-
-
-def test_local_runner_context_reconciler_durable_execute_degrades() -> None:
- """Keep sync local execution on the existing non-durable path."""
- ctx = _create_local_runner_context()
- reconciler_called = False
-
- def reconciler() -> int:
- nonlocal reconciler_called
- reconciler_called = True
- return 999
-
- result = ctx.durable_execute(reconciled_add, 5, 10, reconciler=reconciler)
-
- assert result == 15
- assert reconciler_called is False
-
-
-def test_local_runner_context_reconciler_durable_execute_async_degrades() ->
None:
- """Keep async local execution on the existing non-durable path."""
- ctx = _create_local_runner_context()
- reconciler_called = False
-
- def reconciler() -> int:
- nonlocal reconciler_called
- reconciler_called = True
- return 999
-
- async_result = ctx.durable_execute_async(
- reconciled_add, 5, 10, reconciler=reconciler
- )
-
- async def _await_result() -> Any:
- return await async_result
-
- assert asyncio.run(_await_result()) == 15
- assert reconciler_called is False
-
-
-def test_local_runner_context_reconciler_kwarg_is_not_forwarded() -> None:
- """Do not forward the reserved reconciler kwarg to the user function."""
- ctx = _create_local_runner_context()
-
- def collect_kwargs(**kwargs: Any) -> dict[str, Any]:
- return kwargs
-
- result = ctx.durable_execute(collect_kwargs, reconciler=lambda: "unused")
-
- assert result == {}
diff --git a/python/flink_agents/runtime/tests/test_memory_reference.py
b/python/flink_agents/runtime/tests/test_memory_reference.py
index 86428d2a..62a44e2e 100644
--- a/python/flink_agents/runtime/tests/test_memory_reference.py
+++ b/python/flink_agents/runtime/tests/test_memory_reference.py
@@ -17,7 +17,7 @@
#################################################################################
from flink_agents.api.memory_object import MemoryType
from flink_agents.api.memory_reference import MemoryRef
-from flink_agents.runtime.local_memory_object import LocalMemoryObject
+from flink_agents.runtime.tests.local_memory_object import LocalMemoryObject
class MockRunnerContext:
diff --git
a/python/flink_agents/runtime/tests/test_local_runner_cross_language.py
b/python/flink_agents/runtime/tests/test_plan_java_function_routing.py
similarity index 72%
rename from
python/flink_agents/runtime/tests/test_local_runner_cross_language.py
rename to python/flink_agents/runtime/tests/test_plan_java_function_routing.py
index 3e9c0412..27972c8e 100644
--- a/python/flink_agents/runtime/tests/test_local_runner_cross_language.py
+++ b/python/flink_agents/runtime/tests/test_plan_java_function_routing.py
@@ -15,53 +15,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
-"""Action-dispatch tests via Python's local runner."""
+"""Routing tests for the plan-side ``JavaFunction`` dispatch."""
from typing import Any, List, Tuple
import pytest
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.function import JavaFunction as ApiJavaFunction
-from flink_agents.api.runner_context import RunnerContext
-from flink_agents.plan.configuration import AgentConfiguration
+from flink_agents.api.events.event import InputEvent
from flink_agents.plan.function import JavaFunction as PlanJavaFunction
-from flink_agents.runtime.local_runner import LocalRunner
-
-
-def echo_action(event: Event, ctx: RunnerContext) -> None:
- value = InputEvent.from_event(event).input
- ctx.send_event(OutputEvent(output=value))
-
-
-def _make_java_function_descriptor() -> ApiJavaFunction:
- return ApiJavaFunction.for_action("com.example.Handlers", "handle")
-
-
-def test_local_runner_dispatches_python_function_action() -> None:
- agent = Agent()
- agent.add_action(
- name="echo", trigger_conditions=[InputEvent.EVENT_TYPE],
func=echo_action
- )
-
- runner = LocalRunner(agent, AgentConfiguration())
- runner.run(key="k1", value="hello")
-
- assert runner.get_outputs() == [{"k1": "hello"}]
-
-
-def
test_local_runner_dispatch_of_java_function_action_fails_without_jvm_bridge()
-> None:
- agent = Agent()
- agent.add_action(
- name="handle",
- trigger_conditions=[InputEvent.EVENT_TYPE],
- func=_make_java_function_descriptor(),
- )
-
- runner = LocalRunner(agent, AgentConfiguration())
- with pytest.raises(RuntimeError, match="JVM resource adapter"):
- runner.run(key="k1", value="hello")
class _RecordingJavaAdapter:
diff --git
a/python/flink_agents/runtime/tests/test_remote_execution_environment.py
b/python/flink_agents/runtime/tests/test_remote_execution_environment.py
index 28eca85e..3a58cc22 100644
--- a/python/flink_agents/runtime/tests/test_remote_execution_environment.py
+++ b/python/flink_agents/runtime/tests/test_remote_execution_environment.py
@@ -205,8 +205,7 @@ def test_apply_by_unknown_name_errors() -> None:
The guard fires at apply() time on the remote builder, so no cluster is
started and no job is submitted. The Flink environment and input datastream
- are never exercised to reach the guard, so both are mocked. Mirrors the
same
- contract on the from_list builder.
+ are never exercised to reach the guard, so both are mocked.
"""
with patch(
"flink_agents.runtime.remote_execution_environment.StreamExecutionEnvironment"
diff --git a/python/flink_agents/runtime/tests/test_runner_context_execute.py
b/python/flink_agents/runtime/tests/test_runner_context_execute.py
deleted file mode 100644
index abde15a6..00000000
--- a/python/flink_agents/runtime/tests/test_runner_context_execute.py
+++ /dev/null
@@ -1,189 +0,0 @@
-################################################################################
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#################################################################################
-"""Tests for RunnerContext durable_execute() method."""
-
-from flink_agents.api.agents.agent import Agent
-from flink_agents.api.decorators import action
-from flink_agents.api.events.event import Event, InputEvent, OutputEvent
-from flink_agents.api.events.event_type import EventType
-from flink_agents.api.execution_environment import AgentsExecutionEnvironment
-from flink_agents.api.runner_context import RunnerContext
-
-
-def slow_computation(x: int, y: int) -> int:
- """A sample function that simulates slow computation."""
- return x + y
-
-
-def multiply(x: int, y: int) -> int:
- """A sample function that multiplies two numbers."""
- return x * y
-
-
-def raise_exception(msg: str) -> None:
- """A sample function that raises an exception."""
- raise ValueError(msg)
-
-
-class AgentWithDurableExecute(Agent):
- """Agent that uses synchronous durable_execute() method."""
-
- @action(EventType.InputEvent)
- @staticmethod
- def process(event: Event, ctx: RunnerContext) -> None:
- """Process an event using durable_execute()."""
- input_val = InputEvent.from_event(event).input
- result = ctx.durable_execute(slow_computation, input_val, 10)
- ctx.send_event(OutputEvent(output=result))
-
-
-class AgentWithMultipleDurableExecute(Agent):
- """Agent that makes multiple durable_execute() calls."""
-
- @action(EventType.InputEvent)
- @staticmethod
- def process(event: Event, ctx: RunnerContext) -> None:
- """Process an event with multiple durable_execute() calls."""
- input_val = InputEvent.from_event(event).input
- result1 = ctx.durable_execute(slow_computation, input_val, 5)
- result2 = ctx.durable_execute(multiply, result1, 2)
- ctx.send_event(OutputEvent(output=result2))
-
-
-class AgentWithDurableExecuteAndAsync(Agent):
- """Agent that uses both durable_execute() and durable_execute_async()."""
-
- @action(EventType.InputEvent)
- @staticmethod
- async def process(event: Event, ctx: RunnerContext) -> None:
- """Process an event using both durable_execute() and
durable_execute_async()."""
- input_val = InputEvent.from_event(event).input
- # Use synchronous durable execute
- sync_result = ctx.durable_execute(slow_computation, input_val, 5)
- # Use async durable execute
- async_result = await ctx.durable_execute_async(multiply, sync_result,
3)
- ctx.send_event(OutputEvent(output=async_result))
-
-
-class AgentWithDurableExecuteException(Agent):
- """Agent that uses durable_execute() with a function that raises an
exception."""
-
- @action(EventType.InputEvent)
- @staticmethod
- def process(event: Event, ctx: RunnerContext) -> None:
- """Process an event where durable_execute() raises an exception."""
- input_val = InputEvent.from_event(event).input
- try:
- ctx.durable_execute(raise_exception, f"Test error: {input_val}")
- except ValueError as e:
- ctx.send_event(OutputEvent(output=f"Caught: {e}"))
-
-
-class AgentWithKwargs(Agent):
- """Agent that uses durable_execute() with keyword arguments."""
-
- @action(EventType.InputEvent)
- @staticmethod
- def process(event: Event, ctx: RunnerContext) -> None:
- """Process an event using durable_execute() with kwargs."""
- input_val = InputEvent.from_event(event).input
- result = ctx.durable_execute(slow_computation, x=input_val, y=20)
- ctx.send_event(OutputEvent(output=result))
-
-
-def test_durable_execute_basic() -> None:
- """Test basic synchronous durable_execute() functionality."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = AgentWithDurableExecute()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "alice", "value": 5})
- input_list.append({"key": "bob", "value": 15})
-
- env.execute()
-
- assert output_list == [{"alice": 15}, {"bob": 25}]
-
-
-def test_durable_execute_multiple_calls() -> None:
- """Test multiple durable_execute() calls in a single action."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = AgentWithMultipleDurableExecute()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "alice", "value": 10})
-
- env.execute()
-
- # (10 + 5) * 2 = 30
- assert output_list == [{"alice": 30}]
-
-
-def test_durable_execute_with_async() -> None:
- """Test durable_execute() and durable_execute_async() in the same
action."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = AgentWithDurableExecuteAndAsync()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "alice", "value": 7})
-
- env.execute()
-
- # (7 + 5) * 3 = 36
- assert output_list == [{"alice": 36}]
-
-
-def test_durable_execute_exception_handling() -> None:
- """Test that exceptions from durable_execute() can be caught."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = AgentWithDurableExecuteException()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "alice", "value": "test"})
-
- env.execute()
-
- assert output_list == [{"alice": "Caught: Test error: test"}]
-
-
-def test_durable_execute_with_kwargs() -> None:
- """Test durable_execute() with keyword arguments."""
- env = AgentsExecutionEnvironment.get_execution_environment()
-
- input_list = []
- agent = AgentWithKwargs()
-
- output_list = env.from_list(input_list).apply(agent).to_list()
-
- input_list.append({"key": "alice", "value": 5})
-
- env.execute()
-
- assert output_list == [{"alice": 25}]
