From a5c116f7b473c300ccb4fa15f74551cf1b9a2b9c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Feb 2026 17:48:06 +0000 Subject: [PATCH] feat: add agentSchema and agentCode sync actions to ComponentBase Co-Authored-By: zdenek.srotyr@keboola.com --- src/keboola/component/base.py | 99 +++++++ .../data_examples/data_agent_code/config.json | 1 + .../data_agent_schema/config.json | 12 + tests/test_agent.py | 266 ++++++++++++++++++ 4 files changed, 378 insertions(+) create mode 100644 tests/data_examples/data_agent_code/config.json create mode 100644 tests/data_examples/data_agent_schema/config.json create mode 100644 tests/test_agent.py diff --git a/src/keboola/component/base.py b/src/keboola/component/base.py index af2f5a7..e4f90af 100644 --- a/src/keboola/component/base.py +++ b/src/keboola/component/base.py @@ -1,4 +1,5 @@ import contextlib +import inspect import json import logging import os @@ -12,6 +13,7 @@ from . import dao from . import table_schema as ts +from .exceptions import UserException from .interface import CommonInterface from .sync_actions import SyncActionResult, process_sync_action_result @@ -249,6 +251,103 @@ def execute_action(self): raise AttributeError(f"The defined action {action} is not implemented!") from e return action_method() + def get_agent_schema(self) -> Optional[dict]: + """Override in component to provide a custom agent schema. + + If not overridden, the schema is auto-generated via introspection. + Return ``None`` to explicitly disable agent mode for this component. + """ + return self._generate_agent_schema() + + def get_agent_context(self) -> Optional[dict]: + """Override in component to provide execution context for agentCode. + + Should return a dict of variable names to initialized objects that + will be available in the agent code namespace. + For example: ``{"sf": authenticated_salesforce_client}`` + + Returns ``None`` by default (agent code runs with ``comp`` only). + """ + return None + + @sync_action('agentSchema') + def agent_schema(self): + schema = self.get_agent_schema() + if schema is None: + raise UserException("This component does not support agent mode.") + return schema + + @sync_action('agentCode') + def agent_code(self): + code = self.configuration.parameters.get('agent_code') + if not code: + raise UserException("Parameter 'agent_code' is required.") + context = self.get_agent_context() + local_ns = {'comp': self} + if context: + local_ns.update(context) + exec(code, {'__builtins__': __builtins__}, local_ns) # noqa: S102 + raw_result = local_ns.get('result') + if raw_result is None: + return None + return {'result': raw_result} + + def _generate_agent_schema(self) -> dict: + comp_class = type(self) + comp_module = inspect.getmodule(comp_class) + schema = { + 'component_id': self.environment_variables.component_id, + 'modules': {}, + 'sync_actions': [], + 'context_variables': {}, + 'installed_packages': self._get_installed_packages(), + } + try: + schema['modules']['component'] = inspect.getsource(comp_class) + except (TypeError, OSError): + pass + if comp_module: + for name, obj in vars(comp_module).items(): + if name.startswith('_'): + continue + if inspect.isclass(obj) and obj is not comp_class: + try: + schema['modules'][name] = inspect.getsource(obj) + except (TypeError, OSError): + pass + elif inspect.ismodule(obj): + try: + schema['modules'][name] = inspect.getsource(obj) + except (TypeError, OSError): + pass + for action_name, method_name in _SYNC_ACTION_MAPPING.items(): + if action_name in ('run', 'agentSchema', 'agentCode'): + continue + method = getattr(self, method_name, None) + if method: + schema['sync_actions'].append({ + 'action': action_name, + 'method': method_name, + 'doc': inspect.getdoc(method) or '', + }) + context = self.get_agent_context() + if context: + for var_name, obj in context.items(): + schema['context_variables'][var_name] = { + 'type': type(obj).__qualname__, + 'module': type(obj).__module__, + 'methods': [m for m in dir(obj) if not m.startswith('_')], + } + return schema + + @staticmethod + def _get_installed_packages() -> List[str]: + try: + from importlib.metadata import distributions + return sorted([f"{d.metadata['Name']}=={d.metadata['Version']}" for d in distributions()]) + except Exception: + return [] + def _generate_table_metadata_legacy(self, table_schema: ts.TableSchema) -> dao.TableMetadata: """ Generates a TableMetadata object for the table definition using a TableSchema object. diff --git a/tests/data_examples/data_agent_code/config.json b/tests/data_examples/data_agent_code/config.json new file mode 100644 index 0000000..d1f2ebf --- /dev/null +++ b/tests/data_examples/data_agent_code/config.json @@ -0,0 +1 @@ +{"storage": {"input": {"tables": []}, "output": {"tables": []}}, "parameters": {"agent_code": "result = 1 + 1"}, "action": "agentCode"} \ No newline at end of file diff --git a/tests/data_examples/data_agent_schema/config.json b/tests/data_examples/data_agent_schema/config.json new file mode 100644 index 0000000..7751dc5 --- /dev/null +++ b/tests/data_examples/data_agent_schema/config.json @@ -0,0 +1,12 @@ +{ + "storage": { + "input": { + "tables": [] + }, + "output": { + "tables": [] + } + }, + "parameters": {}, + "action": "agentSchema" +} diff --git a/tests/test_agent.py b/tests/test_agent.py new file mode 100644 index 0000000..c318430 --- /dev/null +++ b/tests/test_agent.py @@ -0,0 +1,266 @@ +import json +import os +import unittest +from io import StringIO +from unittest.mock import patch + +from keboola.component.base import ComponentBase, sync_action + + +class SimpleAgentComponent(ComponentBase): + def run(self): + pass + + def get_agent_context(self): + return {'calculator': {'add': lambda a, b: a + b}} + + +class AgentDisabledComponent(ComponentBase): + def run(self): + pass + + def get_agent_schema(self): + return None + + +class CustomSchemaComponent(ComponentBase): + def run(self): + pass + + def get_agent_schema(self): + return { + 'component_name': 'test.component', + 'description': 'A test component', + 'methods': [ + {'name': 'do_thing', 'args': {}, 'returns': 'str'} + ], + } + + def get_agent_context(self): + return {'greeting': 'hello'} + + +class ComponentWithSyncActions(ComponentBase): + def run(self): + pass + + @sync_action('testConnection') + def test_connection(self): + pass + + @sync_action('loadData') + def load_data(self): + """Loads data from source.""" + pass + + +class TestAgentSchema(unittest.TestCase): + + def setUp(self): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_schema') + os.environ['KBC_DATADIR'] = path + + def test_agent_schema_returns_auto_generated_schema(self): + comp = SimpleAgentComponent() + schema = comp.get_agent_schema() + self.assertIsInstance(schema, dict) + self.assertIn('modules', schema) + self.assertIn('sync_actions', schema) + self.assertIn('installed_packages', schema) + self.assertIn('context_variables', schema) + + def test_agent_schema_contains_component_source(self): + comp = SimpleAgentComponent() + schema = comp.get_agent_schema() + self.assertIn('component', schema['modules']) + self.assertIn('SimpleAgentComponent', schema['modules']['component']) + + def test_agent_schema_disabled_returns_none(self): + comp = AgentDisabledComponent() + schema = comp.get_agent_schema() + self.assertIsNone(schema) + + def test_agent_schema_custom_schema(self): + comp = CustomSchemaComponent() + schema = comp.get_agent_schema() + self.assertEqual(schema['component_name'], 'test.component') + self.assertEqual(schema['description'], 'A test component') + self.assertIn('methods', schema) + + def test_agent_schema_includes_context_variables(self): + comp = SimpleAgentComponent() + schema = comp.get_agent_schema() + self.assertIn('calculator', schema['context_variables']) + self.assertEqual(schema['context_variables']['calculator']['type'], 'dict') + + def test_agent_schema_includes_installed_packages(self): + comp = SimpleAgentComponent() + schema = comp.get_agent_schema() + self.assertIsInstance(schema['installed_packages'], list) + + @patch('sys.stdout', new_callable=StringIO) + def test_agent_schema_sync_action_outputs_json(self, stdout): + comp = SimpleAgentComponent() + comp.execute_action() + output = stdout.getvalue() + parsed = json.loads(output) + self.assertIn('modules', parsed) + + def test_agent_schema_sync_action_disabled_exits(self): + with self.assertRaises(SystemExit): + AgentDisabledComponent().execute_action() + + def test_agent_schema_lists_sync_actions(self): + comp = ComponentWithSyncActions() + schema = comp.get_agent_schema() + action_names = [a['action'] for a in schema['sync_actions']] + self.assertIn('testConnection', action_names) + self.assertIn('loadData', action_names) + self.assertNotIn('agentSchema', action_names) + self.assertNotIn('agentCode', action_names) + self.assertNotIn('run', action_names) + + def test_agent_schema_sync_action_has_doc(self): + comp = ComponentWithSyncActions() + schema = comp.get_agent_schema() + load_data_action = next(a for a in schema['sync_actions'] if a['action'] == 'loadData') + self.assertEqual(load_data_action['doc'], 'Loads data from source.') + + +class TestAgentCode(unittest.TestCase): + + def setUp(self): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_code') + os.environ['KBC_DATADIR'] = path + + @patch('sys.stdout', new_callable=StringIO) + def test_agent_code_executes_simple_expression(self, stdout): + comp = SimpleAgentComponent() + comp.execute_action() + output = stdout.getvalue() + self.assertEqual(json.loads(output), {'result': 2}) + + @patch('sys.stdout', new_callable=StringIO) + def test_agent_code_has_comp_in_context(self, stdout): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_code') + os.environ['KBC_DATADIR'] = path + + config_path = os.path.join(path, 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + original_code = config['parameters']['agent_code'] + + config['parameters']['agent_code'] = 'result = type(comp).__name__' + with open(config_path, 'w') as f: + json.dump(config, f) + + try: + comp = SimpleAgentComponent() + comp.execute_action() + output = stdout.getvalue() + self.assertEqual(json.loads(output), {'result': 'SimpleAgentComponent'}) + finally: + config['parameters']['agent_code'] = original_code + with open(config_path, 'w') as f: + json.dump(config, f) + + @patch('sys.stdout', new_callable=StringIO) + def test_agent_code_has_custom_context(self, stdout): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_code') + os.environ['KBC_DATADIR'] = path + + config_path = os.path.join(path, 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + original_code = config['parameters']['agent_code'] + + config['parameters']['agent_code'] = "result = 'context_has_calculator' if 'add' in calculator else 'no'" + with open(config_path, 'w') as f: + json.dump(config, f) + + try: + comp = SimpleAgentComponent() + comp.execute_action() + output = stdout.getvalue() + self.assertEqual(json.loads(output), {'result': 'context_has_calculator'}) + finally: + config['parameters']['agent_code'] = original_code + with open(config_path, 'w') as f: + json.dump(config, f) + + def test_agent_code_missing_code_param_exits(self): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_code') + os.environ['KBC_DATADIR'] = path + + config_path = os.path.join(path, 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + original_code = config['parameters']['agent_code'] + + config['parameters']['agent_code'] = '' + with open(config_path, 'w') as f: + json.dump(config, f) + + try: + with self.assertRaises(SystemExit): + SimpleAgentComponent().execute_action() + finally: + config['parameters']['agent_code'] = original_code + with open(config_path, 'w') as f: + json.dump(config, f) + + @patch('sys.stdout', new_callable=StringIO) + def test_agent_code_returns_none_when_no_result(self, stdout): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_code') + os.environ['KBC_DATADIR'] = path + + config_path = os.path.join(path, 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + original_code = config['parameters']['agent_code'] + + config['parameters']['agent_code'] = 'x = 42' + with open(config_path, 'w') as f: + json.dump(config, f) + + try: + comp = SimpleAgentComponent() + comp.execute_action() + output = stdout.getvalue() + self.assertEqual(json.loads(output), {"status": "success"}) + finally: + config['parameters']['agent_code'] = original_code + with open(config_path, 'w') as f: + json.dump(config, f) + + def test_agent_code_syntax_error_exits(self): + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), + 'data_examples', 'data_agent_code') + os.environ['KBC_DATADIR'] = path + + config_path = os.path.join(path, 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + original_code = config['parameters']['agent_code'] + + config['parameters']['agent_code'] = 'def broken(' + with open(config_path, 'w') as f: + json.dump(config, f) + + try: + with self.assertRaises(SystemExit): + SimpleAgentComponent().execute_action() + finally: + config['parameters']['agent_code'] = original_code + with open(config_path, 'w') as f: + json.dump(config, f) + + +if __name__ == '__main__': + unittest.main()