This is an automated email from the ASF dual-hosted git repository. xtsong pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 95c1f9c0f08e35dca95e23a9c73c671bbc0f0035 Author: WenjinXie <[email protected]> AuthorDate: Mon Aug 25 15:27:29 2025 +0800 [hotfix] Align class, parameter and decorator name of chat model abstraction. --- python/flink_agents/api/agent.py | 68 ++++++++++++---------- python/flink_agents/api/decorators.py | 12 ++-- python/flink_agents/examples/chat_model_example.py | 44 ++++++++------ python/flink_agents/plan/agent_plan.py | 8 +-- .../python_agent_plan_compatibility_test_agent.py | 4 +- python/flink_agents/plan/tests/test_agent_plan.py | 6 +- .../runtime/tests/test_built_in_actions.py | 8 +-- .../runtime/tests/test_get_resource_in_action.py | 4 +- 8 files changed, 83 insertions(+), 71 deletions(-) diff --git a/python/flink_agents/api/agent.py b/python/flink_agents/api/agent.py index 390df70..7e9af23 100644 --- a/python/flink_agents/api/agent.py +++ b/python/flink_agents/api/agent.py @@ -42,17 +42,19 @@ class Agent(ABC): def my_action(event: Event, ctx: RunnerContext) -> None: action logic - @chat_model_server + @chat_model_connection @staticmethod - def my_server() -> Tuple[Type[BaseChatModelServer], Dict[str, Any]]: - return OllamaChatModelServer, {"name": "my_server", - "model": "qwen2:7b", - "base_url": "http://localhost:11434"} + def my_connection() -> Tuple[Type[BaseChatModelConnection], + Dict[str, Any]]: + return OllamaChatModelConnection, {"name": "my_connection", + "model": "qwen2:7b", + "base_url": "http://localhost:11434"} - @chat_model + @chat_model_setup @staticmethod def my_chat_model() -> Tuple[Type[ChatModel], Dict[str, Any]]: - return OllamaChatModel, {"name": "model", "server": "my_server"} + return OllamaChatModel, {"name": "model", + "connection": "my_connection"} * Add actions and resources to an Agent instance :: @@ -60,27 +62,29 @@ class Agent(ABC): my_agent.add_action(name="my_action", events=[InputEvent], func=action_function) - .add_chat_model_server(name="my_server", - server=OllamaChatModelServer, - arg1=xxx) - .add_chat_model(name="my_chat_model", - chat_model=OllamaChatModel, - server="my_server") + .add_chat_model_connection(name="my_connection", + connection=OllamaChatModelConnection, + arg1=xxx) + .add_chat_model_setup(name="my_chat_model", + chat_model=OllamaChatModelSetup, + connection="my_connection") """ _actions: Dict[str, Tuple[List[Type[Event]], Callable]] _prompts: Dict[str, Prompt] _tools: Dict[str, Callable] - _chat_model_servers: Dict[str, Tuple[Type[BaseChatModelConnection], Dict[str, Any]]] - _chat_models: Dict[str, Tuple[Type[BaseChatModelSetup], Dict[str, Any]]] + _chat_model_connections: Dict[ + str, Tuple[Type[BaseChatModelConnection], Dict[str, Any]] + ] + _chat_model_setups: Dict[str, Tuple[Type[BaseChatModelSetup], Dict[str, Any]]] def __init__(self) -> None: """Init method.""" self._actions = {} self._prompts = {} self._tools = {} - self._chat_model_servers = {} - self._chat_models = {} + self._chat_model_connections = {} + self._chat_model_setups = {} def add_action( self, name: str, events: List[Type[Event]], func: Callable @@ -149,36 +153,36 @@ class Agent(ABC): self._tools[name] = func return self - def add_chat_model_server( - self, name: str, server: Type[BaseChatModelConnection], **kwargs: Any + def add_chat_model_connection( + self, name: str, connection: Type[BaseChatModelConnection], **kwargs: Any ) -> "Agent": - """Add chat model server to agent. + """Add chat model connection to agent. Parameters ---------- name : str - The name of the chat model server, should be unique in the same Agent. - server: Type[BaseChatModelConnection] - The type of chat model server. + The name of the chat model connection, should be unique in the same Agent. + connection: Type[BaseChatModelConnection] + The type of chat model connection. **kwargs: Any - Initialize keyword arguments passed to the chat model server. + Initialize keyword arguments passed to the chat model connection. Returns: ------- Agent The modified Agent instance. """ - if name in self._chat_model_servers: - msg = f"Chat model server {name} already defined" + if name in self._chat_model_connections: + msg = f"Chat model connection {name} already defined" raise ValueError(msg) kwargs["name"] = name - self._chat_model_servers[name] = (server, kwargs) + self._chat_model_connections[name] = (connection, kwargs) return self - def add_chat_model( + def add_chat_model_setup( self, name: str, chat_model: Type[BaseChatModelSetup], **kwargs: Any ) -> "Agent": - """Add chat model to agent. + """Add chat model setup to agent. Parameters ---------- @@ -194,9 +198,9 @@ class Agent(ABC): Agent The modified Agent instance. """ - if name in self._chat_models: - msg = f"Chat model {name} already defined" + if name in self._chat_model_setups: + msg = f"Chat model setup {name} already defined" raise ValueError(msg) kwargs["name"] = name - self._chat_models[name] = (chat_model, kwargs) + self._chat_model_setups[name] = (chat_model, kwargs) return self diff --git a/python/flink_agents/api/decorators.py b/python/flink_agents/api/decorators.py index dbd1cf5..ae59d27 100644 --- a/python/flink_agents/api/decorators.py +++ b/python/flink_agents/api/decorators.py @@ -52,8 +52,8 @@ def action(*listen_events: Type[Event]) -> Callable: return decorator -def chat_model_server(func: Callable) -> Callable: - """Decorator for marking a function declaring a chat model server. +def chat_model_connection(func: Callable) -> Callable: + """Decorator for marking a function declaring a chat model connection. Parameters ---------- @@ -66,12 +66,12 @@ def chat_model_server(func: Callable) -> Callable: Decorator function that marks the target function declare a chat model connection. """ - func._is_chat_model_server = True + func._is_chat_model_connection = True return func -def chat_model(func: Callable) -> Callable: - """Decorator for marking a function declaring a chat model. +def chat_model_setup(func: Callable) -> Callable: + """Decorator for marking a function declaring a chat model setup. Parameters ---------- @@ -83,7 +83,7 @@ def chat_model(func: Callable) -> Callable: Callable Decorator function that marks the target function declare a chat model. """ - func._is_chat_model = True + func._is_chat_model_setup = True return func diff --git a/python/flink_agents/examples/chat_model_example.py b/python/flink_agents/examples/chat_model_example.py index e3b24ca..a46b2b0 100644 --- a/python/flink_agents/examples/chat_model_example.py +++ b/python/flink_agents/examples/chat_model_example.py @@ -24,9 +24,17 @@ from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, BaseChatModelSetup, ) -from flink_agents.api.decorators import action, chat_model, chat_model_server, tool +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 InputEvent, OutputEvent +from flink_agents.api.events.event import ( + InputEvent, + OutputEvent, +) from flink_agents.api.execution_environment import AgentsExecutionEnvironment from flink_agents.api.runner_context import RunnerContext from flink_agents.integrations.chat_models.ollama_chat_model import ( @@ -45,58 +53,58 @@ BACKENDS_TO_RUN: List[str] = ["Tongyi", "Ollama"] class MyAgent(Agent): """Example agent demonstrating the new ChatModel architecture.""" - @chat_model_server + @chat_model_connection @staticmethod - def tongyi_server() -> Tuple[Type[BaseChatModelConnection], Dict[str, Any]]: - """ChatModelServer responsible for tongyi model service connection.""" + def tongyi_connection() -> Tuple[Type[BaseChatModelConnection], Dict[str, Any]]: + """ChatModelConnection responsible for tongyi model service connection.""" if not os.environ.get("DASHSCOPE_API_KEY"): msg = "Please set the 'DASHSCOPE_API_KEY' environment variable." raise ValueError(msg) return TongyiChatModelConnection, { - "name": "tongyi_server", + "name": "tongyi_connection", "model": TONGYI_MODEL, } - @chat_model_server + @chat_model_connection @staticmethod - def ollama_server() -> Tuple[Type[BaseChatModelConnection], Dict[str, Any]]: - """ChatModelServer responsible for ollama model service connection.""" + def ollama_connection() -> Tuple[Type[BaseChatModelConnection], Dict[str, Any]]: + """ChatModelConnection responsible for ollama model service connection.""" return OllamaChatModelConnection, { - "name": "ollama_server", + "name": "ollama_connection", "model": OLLAMA_MODEL, } - @chat_model + @chat_model_setup @staticmethod def math_chat_model() -> Tuple[Type[BaseChatModelSetup], Dict[str, Any]]: - """ChatModel which focus on math, and reuse ChatModelServer.""" + """ChatModel which focus on math, and reuse ChatModelConnection.""" if CURRENT_BACKEND == "Tongyi": return TongyiChatModelSetup, { "name": "math_chat_model", - "connection": "tongyi_server", + "connection": "tongyi_connection", "tools": ["add"], } else: return OllamaChatModelSetup, { "name": "math_chat_model", - "connection": "ollama_server", + "connection": "ollama_connection", "tools": ["add"], "extract_reasoning": True, } - @chat_model + @chat_model_setup @staticmethod def creative_chat_model() -> Tuple[Type[BaseChatModelSetup], Dict[str, Any]]: - """ChatModel which focus on text generate, and reuse ChatModelServer.""" + """ChatModel which focus on text generate, and reuse ChatModelConnection.""" if CURRENT_BACKEND == "Tongyi": return TongyiChatModelSetup, { "name": "creative_chat_model", - "connection": "tongyi_server", + "connection": "tongyi_connection", } else: return OllamaChatModelSetup, { "name": "creative_chat_model", - "connection": "ollama_server", + "connection": "ollama_connection", "extract_reasoning": True, } diff --git a/python/flink_agents/plan/agent_plan.py b/python/flink_agents/plan/agent_plan.py index c697fc3..7118ecd 100644 --- a/python/flink_agents/plan/agent_plan.py +++ b/python/flink_agents/plan/agent_plan.py @@ -235,7 +235,7 @@ def _get_actions(agent: Agent) -> List[Action]: def _get_resource_providers(agent: Agent) -> List[ResourceProvider]: resource_providers = [] for name, value in agent.__class__.__dict__.items(): - if hasattr(value, "_is_chat_model"): + if hasattr(value, "_is_chat_model_setup"): if isinstance(value, staticmethod): value = value.__func__ @@ -249,7 +249,7 @@ def _get_resource_providers(agent: Agent) -> List[ResourceProvider]: kwargs=kwargs, ) resource_providers.append(provider) - elif hasattr(value, "_is_chat_model_server"): + elif hasattr(value, "_is_chat_model_connection"): if isinstance(value, staticmethod): value = value.__func__ @@ -296,7 +296,7 @@ def _get_resource_providers(agent: Agent) -> List[ResourceProvider]: PythonSerializableResourceProvider.from_resource(name=name, resource=tool) ) - for name, chat_model in agent._chat_models.items(): + for name, chat_model in agent._chat_model_setups.items(): clazz, kwargs = chat_model provider = PythonResourceProvider( name=name, @@ -307,7 +307,7 @@ def _get_resource_providers(agent: Agent) -> List[ResourceProvider]: ) resource_providers.append(provider) - for name, connection in agent._chat_model_servers.items(): + for name, connection in agent._chat_model_connections.items(): clazz, kwargs = connection provider = PythonResourceProvider( name=name, diff --git a/python/flink_agents/plan/tests/compatibility/python_agent_plan_compatibility_test_agent.py b/python/flink_agents/plan/tests/compatibility/python_agent_plan_compatibility_test_agent.py index 7c5e2cc..59843f3 100644 --- a/python/flink_agents/plan/tests/compatibility/python_agent_plan_compatibility_test_agent.py +++ b/python/flink_agents/plan/tests/compatibility/python_agent_plan_compatibility_test_agent.py @@ -20,7 +20,7 @@ from typing import Any, Dict, Sequence, Tuple, Type from flink_agents.api.agent import Agent from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.chat_model import BaseChatModelSetup -from flink_agents.api.decorators import action, chat_model, tool +from flink_agents.api.decorators import action, chat_model_setup, tool from flink_agents.api.events.event import Event, InputEvent from flink_agents.api.runner_context import RunnerContext @@ -54,7 +54,7 @@ class PythonAgentPlanCompatibilityTestAgent(Agent): def second_action(event: InputEvent, ctx: RunnerContext) -> None: """Test implementation.""" - @chat_model + @chat_model_setup @staticmethod def chat_model() -> Tuple[Type[BaseChatModelSetup], Dict[str, Any]]: """ChatModel can be used in action.""" diff --git a/python/flink_agents/plan/tests/test_agent_plan.py b/python/flink_agents/plan/tests/test_agent_plan.py index 20d61d8..a3e493c 100644 --- a/python/flink_agents/plan/tests/test_agent_plan.py +++ b/python/flink_agents/plan/tests/test_agent_plan.py @@ -24,7 +24,7 @@ import pytest from flink_agents.api.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 +from flink_agents.api.decorators import action, chat_model_setup from flink_agents.api.events.event import Event, InputEvent, OutputEvent from flink_agents.api.resource import Resource, ResourceType from flink_agents.api.runner_context import RunnerContext @@ -94,7 +94,7 @@ class MockChatModelImpl(BaseChatModelSetup): # noqa: D101 class MyAgent(Agent): # noqa: D101 - @chat_model + @chat_model_setup @staticmethod def mock() -> Tuple[Type[Resource], Dict[str, Any]]: # noqa: D102 return MockChatModelImpl, { @@ -156,7 +156,7 @@ def test_add_action_and_resource_to_agent() -> None: # noqa: D103 my_agent.add_action( name="second_action", events=[InputEvent, MyEvent], func=MyAgent.second_action ) - my_agent.add_chat_model( + my_agent.add_chat_model_setup( name="mock", chat_model=MockChatModelImpl, host="8.8.8.8", diff --git a/python/flink_agents/runtime/tests/test_built_in_actions.py b/python/flink_agents/runtime/tests/test_built_in_actions.py index 9f97e1a..bb20df3 100644 --- a/python/flink_agents/runtime/tests/test_built_in_actions.py +++ b/python/flink_agents/runtime/tests/test_built_in_actions.py @@ -26,8 +26,8 @@ from flink_agents.api.chat_models.chat_model import ( ) from flink_agents.api.decorators import ( action, - chat_model, - chat_model_server, + chat_model_connection, + chat_model_setup, prompt, tool, ) @@ -121,7 +121,7 @@ class MyAgent(Agent): text="Please call the appropriate tool to do the following task: {task}", ) - @chat_model_server + @chat_model_connection @staticmethod def mock_connection() -> Tuple[Type[BaseChatModelConnection], Dict[str, Any]]: """Chat model server can be used by ChatModel.""" @@ -129,7 +129,7 @@ class MyAgent(Agent): "name": "mock_connection", } - @chat_model + @chat_model_setup @staticmethod def mock_chat_model() -> Tuple[Type[BaseChatModelSetup], Dict[str, Any]]: """Chat model can be used in action.""" 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 index 49f4c05..4bec1ec 100644 --- a/python/flink_agents/runtime/tests/test_get_resource_in_action.py +++ b/python/flink_agents/runtime/tests/test_get_resource_in_action.py @@ -20,7 +20,7 @@ from typing import Any, Dict, Sequence, Tuple, Type from flink_agents.api.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, tool +from flink_agents.api.decorators import action, chat_model_setup, tool from flink_agents.api.events.event import InputEvent, OutputEvent from flink_agents.api.execution_environment import AgentsExecutionEnvironment from flink_agents.api.resource import Resource, ResourceType @@ -43,7 +43,7 @@ class MockChatModelImpl(BaseChatModelSetup): # noqa: D101 class MyAgent(Agent): # noqa: D101 - @chat_model + @chat_model_setup @staticmethod def mock_chat_model() -> Tuple[Type[Resource], Dict[str, Any]]: # noqa: D102 return MockChatModelImpl, {
