diff --git a/TOOLS.md b/TOOLS.md index a6b8e53c..6e5df5bd 100644 --- a/TOOLS.md +++ b/TOOLS.md @@ -20,6 +20,7 @@ The Flo AI tools system allows you to create and configure tools that can be use - **Entity-attached tools**: Tools that require specific context or configuration (e.g., BigQuery with datasource ID, email with SMTP settings) The system supports: + - Basic tool creation and usage - Partial tools with pre-filled parameters - YAML-based tool configuration @@ -55,6 +56,7 @@ async def calculate(expression: str, precision: int = 2) -> str: ### Tool Properties Every tool has these properties: + - `name`: Tool name (defaults to function name) - `description`: Tool description (defaults to docstring) - `parameters`: Dictionary of parameter definitions @@ -81,6 +83,7 @@ Partial tools allow you to pre-fill some parameters during agent building, hidin ### Why Use Partial Tools? **Problem**: Some tools require context-specific parameters that shouldn't be provided by the AI: + - BigQuery tools need `datasource_id` and `project_id` - Email tools need SMTP server configuration - Database tools need connection strings @@ -130,7 +133,7 @@ agent = (AgentBuilder() .with_tools([ { "tool": bigquery_query.tool, - "pre_filled_params": { + "prefilled_params": { "datasource_id": "ds_production_123", "project_id": "my-company-prod", "dataset": "analytics" @@ -194,7 +197,7 @@ from flo_ai.tool.tool_config import ToolConfig tool_config = ToolConfig( tool=my_tool, - pre_filled_params={"param1": "value1", "param2": "value2"}, + prefilled_params={"param1": "value1", "param2": "value2"}, name_override="custom_tool_name", description_override="Custom tool description" ) @@ -207,7 +210,7 @@ configured_tool = tool_config.to_tool() - `is_partial()`: Check if the tool has pre-filled parameters - `to_tool()`: Convert configuration to a Tool object -- `get_pre_filled_params()`: Get pre-filled parameters +- `get_prefilled_params()`: Get pre-filled parameters ## YAML Configuration @@ -225,19 +228,19 @@ agent: tools: # Simple tool reference - "calculate" - + # Tool with pre-filled parameters - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_production_123" project_id: "my-company-prod" dataset: "analytics" name_override: "query_production_data" description_override: "Query production BigQuery data" - + # Tool with different configuration - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 5 language: "en" name_override: "search_web" @@ -257,15 +260,15 @@ agent: name: "gpt-4" tools: - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_production_123" project_id: "my-company-prod" dataset: "analytics" name_override: "query_production_data" description_override: "Query production BigQuery data" - + - name: "send_email" - pre_filled_params: + prefilled_params: smtp_server: "smtp.company.com" smtp_port: 587 name_override: "send_notification" @@ -283,15 +286,15 @@ agent: name: "gpt-4" tools: - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_dev_456" project_id: "my-company-dev" dataset: "test_data" name_override: "query_dev_data" description_override: "Query development BigQuery data" - + - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 3 language: "en" name_override: "search_web" @@ -394,14 +397,14 @@ agent: tools: - "calculate" - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_production_123" project_id: "my-company-prod" dataset: "analytics" name_override: "query_production_data" description_override: "Query production BigQuery data" - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 5 language: "en" name_override: "search_web" @@ -442,6 +445,7 @@ agent = AgentBuilder.from_yaml( ``` **Parameters:** + - `name`: Custom name for the tool (defaults to function name) - `description`: Tool description (defaults to function docstring) - `parameter_descriptions`: Dict mapping parameter names to descriptions @@ -453,16 +457,17 @@ class ToolConfig: def __init__( self, tool: Tool, - pre_filled_params: Optional[Dict[str, Any]] = None, + prefilled_params: Optional[Dict[str, Any]] = None, name_override: Optional[str] = None, description_override: Optional[str] = None, ) ``` **Methods:** + - `is_partial() -> bool`: Check if tool has pre-filled parameters - `to_tool() -> Tool`: Convert to Tool object -- `get_pre_filled_params() -> Dict[str, Any]`: Get pre-filled parameters +- `get_prefilled_params() -> Dict[str, Any]`: Get pre-filled parameters ### PartialTool Class @@ -471,23 +476,24 @@ class PartialTool(Tool): def __init__( self, base_tool: Tool, - pre_filled_params: Dict[str, Any], + prefilled_params: Dict[str, Any], name_override: Optional[str] = None, description_override: Optional[str] = None, ) ``` **Methods:** + - `get_original_tool() -> Tool`: Get the original tool -- `get_pre_filled_params() -> Dict[str, Any]`: Get pre-filled parameters -- `add_pre_filled_param(key: str, value: Any) -> PartialTool`: Add parameter -- `remove_pre_filled_param(key: str) -> PartialTool`: Remove parameter +- `get_prefilled_params() -> Dict[str, Any]`: Get pre-filled parameters +- `add_prefilled_param(key: str, value: Any) -> PartialTool`: Add parameter +- `remove_prefilled_param(key: str) -> PartialTool`: Remove parameter ### AgentBuilder Methods ```python # Add single tool with optional pre-filled parameters -def add_tool(self, tool: Tool, **pre_filled_params) -> 'AgentBuilder' +def add_tool(self, tool: Tool, **prefilled_params) -> 'AgentBuilder' # Add multiple tools (supports Tool, ToolConfig, or dict) def with_tools(self, tools: Union[List[Tool], List[ToolConfig], List[Dict[str, Any]]]) -> 'AgentBuilder' @@ -551,7 +557,7 @@ def test_tool_execution(): function=mock_function, parameters={"param1": {"type": "string", "description": "Param 1", "required": True}} ) - + result = await tool.execute(param1="test_value") assert result == "test_result" mock_function.assert_called_once_with(param1="test_value") @@ -569,7 +575,7 @@ def test_tool_execution(): ### Debug Tips 1. **Check tool parameters**: Use `tool.parameters` to see what the AI sees -2. **Verify pre-filled params**: Use `partial_tool.get_pre_filled_params()` +2. **Verify pre-filled params**: Use `partial_tool.get_prefilled_params()` 3. **Test tool execution**: Test tools independently before using in agents 4. **Check YAML structure**: Validate YAML with online validators diff --git a/flo_ai/examples/partial_tool_example.py b/flo_ai/examples/partial_tool_example.py index 11a36386..cf559795 100644 --- a/flo_ai/examples/partial_tool_example.py +++ b/flo_ai/examples/partial_tool_example.py @@ -119,7 +119,7 @@ async def main(): # Tool with pre-filled parameters using dictionary format { 'tool': bigquery_query.tool, - 'pre_filled_params': { + 'prefilled_params': { 'datasource_id': 'ds_production_456', 'project_id': 'company-prod', 'dataset': 'analytics', @@ -128,7 +128,7 @@ async def main(): # Tool with pre-filled parameters using dictionary format { 'tool': web_search.tool, - 'pre_filled_params': {'max_results': 3, 'language': 'en'}, + 'prefilled_params': {'max_results': 3, 'language': 'en'}, }, # Regular tool without pre-filling calculate.tool, @@ -171,12 +171,12 @@ async def main(): print('5. Demonstrating parameter management...') # Add a new pre-filled parameter - bigquery_partial.add_pre_filled_param('timeout', 30) - print(f'Added timeout parameter: {bigquery_partial.get_pre_filled_params()}') + bigquery_partial.add_prefilled_param('timeout', 30) + print(f'Added timeout parameter: {bigquery_partial.get_prefilled_params()}') # Remove a parameter - bigquery_partial.remove_pre_filled_param('timeout') - print(f'Removed timeout parameter: {bigquery_partial.get_pre_filled_params()}') + bigquery_partial.remove_prefilled_param('timeout') + print(f'Removed timeout parameter: {bigquery_partial.get_prefilled_params()}') print() # Show the difference between original and partial tools @@ -191,7 +191,7 @@ async def main(): print(f" - {param}: {info['type']} (required: {info['required']})") print( - f'\nPre-filled parameters (hidden from AI): {bigquery_partial.get_pre_filled_params()}' + f'\nPre-filled parameters (hidden from AI): {bigquery_partial.get_prefilled_params()}' ) diff --git a/flo_ai/examples/tools_quickstart.py b/flo_ai/examples/tools_quickstart.py index 08ee7972..02561bef 100644 --- a/flo_ai/examples/tools_quickstart.py +++ b/flo_ai/examples/tools_quickstart.py @@ -125,9 +125,9 @@ async def main(): print(f" - {param}: {info['type']} (required: {info['required']})") # Show pre-filled parameters - if hasattr(partial_tool, 'get_pre_filled_params'): + if hasattr(partial_tool, 'get_prefilled_params'): print( - f'\nPre-filled parameters (hidden from AI): {partial_tool.get_pre_filled_params()}' + f'\nPre-filled parameters (hidden from AI): {partial_tool.get_prefilled_params()}' ) print() @@ -157,14 +157,14 @@ async def main(): tools: - "calculate" - name: "query_database" - pre_filled_params: + prefilled_params: database_url: "postgresql://yaml-db.company.com:5432/data" timeout: 45 max_rows: 2000 name_override: "query_yaml_database" description_override: "Query YAML-configured database" - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 3 language: "en" name_override: "search_web" @@ -189,14 +189,14 @@ async def main(): # 7. Tool parameter management print('7. Tool parameter management...') - if hasattr(partial_tool, 'add_pre_filled_param'): + if hasattr(partial_tool, 'add_prefilled_param'): # Add a new pre-filled parameter - partial_tool.add_pre_filled_param('retry_count', 3) - print(f'Added retry_count parameter: {partial_tool.get_pre_filled_params()}') + partial_tool.add_prefilled_param('retry_count', 3) + print(f'Added retry_count parameter: {partial_tool.get_prefilled_params()}') # Remove a parameter - partial_tool.remove_pre_filled_param('retry_count') - print(f'Removed retry_count parameter: {partial_tool.get_pre_filled_params()}') + partial_tool.remove_prefilled_param('retry_count') + print(f'Removed retry_count parameter: {partial_tool.get_prefilled_params()}') print() print('=== Quick Start Complete ===') diff --git a/flo_ai/examples/yaml_tool_config_example.py b/flo_ai/examples/yaml_tool_config_example.py index 9a51bb21..d2ff0168 100644 --- a/flo_ai/examples/yaml_tool_config_example.py +++ b/flo_ai/examples/yaml_tool_config_example.py @@ -115,7 +115,7 @@ async def main(): name: "gpt-4" tools: - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_production_123" project_id: "my-company-prod" dataset: "analytics" @@ -123,14 +123,14 @@ async def main(): description_override: "Query production BigQuery data" - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 5 language: "en" name_override: "search_web" description_override: "Search the web for information" - name: "send_email" - pre_filled_params: + prefilled_params: smtp_server: "smtp.company.com" smtp_port: 587 name_override: "send_notification" @@ -160,7 +160,7 @@ async def main(): name: "gpt-4" tools: - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_dev_456" project_id: "my-company-dev" dataset: "test_data" @@ -168,7 +168,7 @@ async def main(): description_override: "Query development BigQuery data" - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 3 language: "en" name_override: "search_web" @@ -198,13 +198,13 @@ async def main(): - "calculate" # Simple reference - name: "bigquery_query" - pre_filled_params: + prefilled_params: datasource_id: "ds_mixed_789" project_id: "mixed-project" dataset: "mixed_data" - name: "web_search" - pre_filled_params: + prefilled_params: max_results: 10 language: "en" name_override: "search_web" @@ -238,7 +238,7 @@ async def main(): tool_config = ToolConfig( tool=configured_tool, - pre_filled_params={ + prefilled_params={ 'datasource_id': 'ds_production_123', 'project_id': 'my-company-prod', 'dataset': 'analytics', @@ -249,7 +249,7 @@ async def main(): for param, info in configured_tool_for_ai.parameters.items(): print(f" - {param}: {info['type']} (required: {info['required']})") - print(f'\nPre-filled parameters (hidden from AI): {tool_config.pre_filled_params}') + print(f'\nPre-filled parameters (hidden from AI): {tool_config.prefilled_params}') if __name__ == '__main__': diff --git a/flo_ai/flo_ai/__init__.py b/flo_ai/flo_ai/__init__.py index 176c6b09..900a6d1c 100644 --- a/flo_ai/flo_ai/__init__.py +++ b/flo_ai/flo_ai/__init__.py @@ -34,7 +34,6 @@ from .arium import ( Arium, BaseArium, - create_arium, MessageMemory, BaseMemory, StartNode, @@ -43,7 +42,9 @@ AriumBuilder, AriumEvent, AriumEventType, + MessageMemoryItem, default_event_callback, + create_arium, ) # Utils package - Utility functions @@ -99,6 +100,7 @@ 'create_arium', 'MessageMemory', 'BaseMemory', + 'MessageMemoryItem', 'StartNode', 'EndNode', 'Edge', diff --git a/flo_ai/flo_ai/arium/__init__.py b/flo_ai/flo_ai/arium/__init__.py index fa372477..d005161f 100644 --- a/flo_ai/flo_ai/arium/__init__.py +++ b/flo_ai/flo_ai/arium/__init__.py @@ -1,7 +1,7 @@ from .arium import Arium from .base import BaseArium from .builder import AriumBuilder, create_arium -from .memory import MessageMemory, BaseMemory +from .memory import MessageMemory, BaseMemory, MessageMemoryItem from .models import StartNode, EndNode, Edge from .events import AriumEventType, AriumEvent, default_event_callback from .llm_router import ( @@ -20,6 +20,7 @@ 'create_arium', 'MessageMemory', 'BaseMemory', + 'MessageMemoryItem', 'StartNode', 'EndNode', 'Edge', diff --git a/flo_ai/flo_ai/arium/builder.py b/flo_ai/flo_ai/arium/builder.py index 8ee349da..e6d867fb 100644 --- a/flo_ai/flo_ai/arium/builder.py +++ b/flo_ai/flo_ai/arium/builder.py @@ -350,7 +350,9 @@ def from_yaml( function_name: function2 description: "Function 2" input_filter: ["input1", "input2"] - + prefilled_params: + param1: "value1" + param2: "value2" # LLM Router definitions (NEW) routers: - name: content_router @@ -529,6 +531,9 @@ def from_yaml( for function_node_config in function_nodes_config: function_node_name = function_node_config['name'] function_name = function_node_config['function_name'] + prefilled_params = function_node_config.get('prefilled_params', None) + description = function_node_config.get('description', None) + input_filter = function_node_config.get('input_filter', None) function = function_registry.get(function_name) if function is None: @@ -540,9 +545,10 @@ def from_yaml( function_node = FunctionNode( name=function_node_name, - description=function_node_config.get('description', ''), + description=description, function=function, - input_filter=function_node_config.get('input_filter', None), + input_filter=input_filter, + prefilled_params=prefilled_params, ) function_nodes_dict[function_node_name] = function_node diff --git a/flo_ai/flo_ai/arium/nodes.py b/flo_ai/flo_ai/arium/nodes.py index e463cb27..79489dc6 100644 --- a/flo_ai/flo_ai/arium/nodes.py +++ b/flo_ai/flo_ai/arium/nodes.py @@ -169,11 +169,13 @@ def __init__( name: str, description: str, function: Callable[..., Any], + prefilled_params: Optional[Dict[str, Any]] = None, input_filter: Optional[List[str]] = None, ) -> None: self.name = name self.description = description self.function = function + self.prefilled_params = prefilled_params or {} self.input_filter: Optional[List[str]] = input_filter async def run( @@ -188,14 +190,19 @@ async def run( if asyncio.iscoroutinefunction(self.function): logger.info(f"Executing FunctionNode '{self.name}' as a coroutine function") - result = await self.function(inputs=inputs, variables=variables, **kwargs) + result = await self.function( + inputs=inputs, variables=variables, **self.prefilled_params, **kwargs + ) return UserMessage(content=result) + logger.info(f"Executing FunctionNode '{self.name}' as a regular function") + result = self.function( + inputs=inputs, variables=variables, **self.prefilled_params, **kwargs + ) + if asyncio.iscoroutine(result): logger.info(f"Executing FunctionNode '{self.name}' as a coroutine") content = await result return UserMessage(content=content) - logger.info(f"Executing FunctionNode '{self.name}' as a regular function") - result = self.function(inputs=inputs, variables=variables, **kwargs) return UserMessage(content=result) diff --git a/flo_ai/flo_ai/builder/agent_builder.py b/flo_ai/flo_ai/builder/agent_builder.py index 31b8fa37..d3e916c1 100644 --- a/flo_ai/flo_ai/builder/agent_builder.py +++ b/flo_ai/flo_ai/builder/agent_builder.py @@ -61,7 +61,7 @@ def with_tools( tools: List of tools, tool configurations, or tool dictionaries. Each tool dictionary should have: - 'tool': The Tool object - - 'pre_filled_params': Optional dict of pre-filled parameters + - 'prefilled_params': Optional dict of pre-filled parameters - 'name_override': Optional custom name - 'description_override': Optional custom description @@ -71,14 +71,14 @@ def with_tools( # Tool configurations builder.with_tools([ - ToolConfig(tool1, pre_filled_params={"param1": "value1"}), - ToolConfig(tool2, pre_filled_params={"param2": "value2"}) + ToolConfig(tool1, prefilled_params={"param1": "value1"}), + ToolConfig(tool2, prefilled_params={"param2": "value2"}) ]) # Tool dictionaries builder.with_tools([ - {"tool": tool1, "pre_filled_params": {"param1": "value1"}}, - {"tool": tool2, "pre_filled_params": {"param2": "value2"}} + {"tool": tool1, "prefilled_params": {"param1": "value1"}}, + {"tool": tool2, "prefilled_params": {"param2": "value2"}} ]) """ processed_tools = [] @@ -93,13 +93,13 @@ def with_tools( elif isinstance(tool_item, dict): # Tool dictionary - convert to ToolConfig then to tool tool = tool_item['tool'] - pre_filled_params = tool_item.get('pre_filled_params', {}) + prefilled_params = tool_item.get('prefilled_params', {}) name_override = tool_item.get('name_override') description_override = tool_item.get('description_override') tool_config = ToolConfig( tool=tool, - pre_filled_params=pre_filled_params, + prefilled_params=prefilled_params, name_override=name_override, description_override=description_override, ) @@ -110,13 +110,13 @@ def with_tools( self._tools = processed_tools return self - def add_tool(self, tool: Tool, **pre_filled_params) -> 'AgentBuilder': + def add_tool(self, tool: Tool, **prefilled_params) -> 'AgentBuilder': """ Add a single tool with optional pre-filled parameters. Args: tool: The tool to add - **pre_filled_params: Pre-filled parameters for the tool + **prefilled_params: Pre-filled parameters for the tool Example: builder.add_tool( @@ -125,8 +125,8 @@ def add_tool(self, tool: Tool, **pre_filled_params) -> 'AgentBuilder': project_id="my-project" ) """ - if pre_filled_params: - tool_config = create_tool_config(tool, **pre_filled_params) + if prefilled_params: + tool_config = create_tool_config(tool, **prefilled_params) self._tools.append(tool_config.to_tool()) else: self._tools.append(tool) @@ -299,21 +299,21 @@ def _process_yaml_tools( raise ValueError(f"Tool '{tool_name}' not found in tool registry") # Extract configuration - pre_filled_params = tool_config.get('pre_filled_params', {}) + prefilled_params = tool_config.get('prefilled_params', {}) name_override = tool_config.get('name_override') description_override = tool_config.get('description_override') # Create tool configuration tool_config_obj = ToolConfig( tool=base_tool, - pre_filled_params=pre_filled_params, + prefilled_params=prefilled_params, name_override=name_override, description_override=description_override, ) # If there are pre-filled parameters or custom name/description, convert to tool if ( - pre_filled_params + prefilled_params or name_override is not None or description_override is not None ): diff --git a/flo_ai/flo_ai/tool/partial_tool.py b/flo_ai/flo_ai/tool/partial_tool.py index d0bc9333..2a25baef 100644 --- a/flo_ai/flo_ai/tool/partial_tool.py +++ b/flo_ai/flo_ai/tool/partial_tool.py @@ -12,7 +12,7 @@ class PartialTool(Tool): def __init__( self, base_tool: Tool, - pre_filled_params: Dict[str, Any], + prefilled_params: Dict[str, Any], name_override: Optional[str] = None, description_override: Optional[str] = None, ): @@ -21,17 +21,17 @@ def __init__( Args: base_tool: The original tool to wrap - pre_filled_params: Parameters to pre-fill (datasource_id, etc.) + prefilled_params: Parameters to pre-fill (datasource_id, etc.) name_override: Optional custom name for the partial tool description_override: Optional custom description """ self.base_tool = base_tool - self.pre_filled_params = pre_filled_params.copy() + self.prefilled_params = prefilled_params.copy() # Create filtered parameters (remove pre-filled ones from AI's view) filtered_parameters = {} for param_name, param_info in base_tool.parameters.items(): - if param_name not in pre_filled_params: + if param_name not in prefilled_params: filtered_parameters[param_name] = param_info.copy() super().__init__( @@ -47,7 +47,7 @@ async def execute(self, **kwargs) -> Any: try: # Merge pre-filled params with AI-provided params # AI params take precedence over pre-filled ones - merged_params = {**self.pre_filled_params, **kwargs} + merged_params = {**self.prefilled_params, **kwargs} logger.info( f'Executing partial tool {self.name} with merged params: {merged_params}' @@ -64,29 +64,29 @@ def get_original_tool(self) -> Tool: """Get the original tool without pre-filled parameters.""" return self.base_tool - def get_pre_filled_params(self) -> Dict[str, Any]: + def get_prefilled_params(self) -> Dict[str, Any]: """Get the pre-filled parameters.""" - return self.pre_filled_params.copy() + return self.prefilled_params.copy() - def add_pre_filled_param(self, key: str, value: Any) -> 'PartialTool': + def add_prefilled_param(self, key: str, value: Any) -> 'PartialTool': """Add or update a pre-filled parameter.""" - self.pre_filled_params[key] = value + self.prefilled_params[key] = value return self - def remove_pre_filled_param(self, key: str) -> 'PartialTool': + def remove_prefilled_param(self, key: str) -> 'PartialTool': """Remove a pre-filled parameter.""" - if key in self.pre_filled_params: - del self.pre_filled_params[key] + if key in self.prefilled_params: + del self.prefilled_params[key] return self -def create_partial_tool(tool: Tool, **pre_filled_params) -> PartialTool: +def create_partial_tool(tool: Tool, **prefilled_params) -> PartialTool: """ Create a partial tool with pre-filled parameters. Args: tool: The original tool - **pre_filled_params: Parameters to pre-fill + **prefilled_params: Parameters to pre-fill Returns: PartialTool: A tool with pre-filled parameters @@ -108,4 +108,4 @@ async def bigquery_query(query: str, datasource_id: str, project_id: str): # AI only needs to provide the query result = await partial_tool.execute(query="SELECT * FROM users") """ - return PartialTool(tool, pre_filled_params) + return PartialTool(tool, prefilled_params) diff --git a/flo_ai/flo_ai/tool/tool_config.py b/flo_ai/flo_ai/tool/tool_config.py index 0aa5b38d..de039b7b 100644 --- a/flo_ai/flo_ai/tool/tool_config.py +++ b/flo_ai/flo_ai/tool/tool_config.py @@ -12,7 +12,7 @@ class ToolConfig: def __init__( self, tool: Tool, - pre_filled_params: Optional[Dict[str, Any]] = None, + prefilled_params: Optional[Dict[str, Any]] = None, name_override: Optional[str] = None, description_override: Optional[str] = None, ): @@ -21,24 +21,24 @@ def __init__( Args: tool: The base tool - pre_filled_params: Optional pre-filled parameters + prefilled_params: Optional pre-filled parameters name_override: Optional custom name description_override: Optional custom description """ self.tool = tool - self.pre_filled_params = pre_filled_params or {} + self.prefilled_params = prefilled_params or {} self.name_override = name_override self.description_override = description_override def is_partial(self) -> bool: """Check if this tool configuration has pre-filled parameters.""" - return bool(self.pre_filled_params) + return bool(self.prefilled_params) def to_tool(self) -> Tool: """Convert this configuration to a Tool (either original or partial).""" # If there are pre-filled parameters or custom name/description, create partial tool if ( - self.pre_filled_params + self.prefilled_params or self.name_override is not None or self.description_override is not None ): @@ -47,7 +47,7 @@ def to_tool(self) -> Tool: return PartialTool( base_tool=self.tool, - pre_filled_params=self.pre_filled_params, + prefilled_params=self.prefilled_params, name_override=self.name_override, description_override=self.description_override, ) @@ -56,15 +56,15 @@ def to_tool(self) -> Tool: return self.tool -def create_tool_config(tool: Tool, **pre_filled_params) -> ToolConfig: +def create_tool_config(tool: Tool, **prefilled_params) -> ToolConfig: """ Create a tool configuration with pre-filled parameters. Args: tool: The base tool - **pre_filled_params: Pre-filled parameters + **prefilled_params: Pre-filled parameters Returns: ToolConfig: A tool configuration """ - return ToolConfig(tool, pre_filled_params) + return ToolConfig(tool, prefilled_params) diff --git a/flo_ai/pyproject.toml b/flo_ai/pyproject.toml index 7334a634..453cc861 100644 --- a/flo_ai/pyproject.toml +++ b/flo_ai/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "flo_ai" -version = "1.0.8-rc3" +version = "1.1.0-rc2" description = "A easy way to create structured AI agents" authors = [{ name = "rootflo", email = "*@rootflo.ai" }] requires-python = ">=3.10,<4.0" diff --git a/flo_ai/tests/unit-tests/test_agent_builder_tools.py b/flo_ai/tests/unit-tests/test_agent_builder_tools.py index 9d8f013f..737a36a6 100644 --- a/flo_ai/tests/unit-tests/test_agent_builder_tools.py +++ b/flo_ai/tests/unit-tests/test_agent_builder_tools.py @@ -47,7 +47,7 @@ def test_with_tools_tool_configs(self): """Test adding tool configurations to agent builder.""" tool_config = ToolConfig( tool=self.base_tool, - pre_filled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, + prefilled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, ) builder = AgentBuilder() @@ -62,7 +62,7 @@ def test_with_tools_tool_dictionaries(self): """Test adding tool dictionaries to agent builder.""" tool_dict = { 'tool': self.base_tool, - 'pre_filled_params': { + 'prefilled_params': { 'datasource_id': 'ds_123', 'project_id': 'my-project', }, @@ -81,12 +81,12 @@ def test_with_tools_tool_dictionaries(self): def test_with_tools_mixed_types(self): """Test adding mixed tool types to agent builder.""" tool_config = ToolConfig( - tool=self.base_tool, pre_filled_params={'datasource_id': 'ds_123'} + tool=self.base_tool, prefilled_params={'datasource_id': 'ds_123'} ) tool_dict = { 'tool': self.base_tool, - 'pre_filled_params': {'project_id': 'my-project'}, + 'prefilled_params': {'project_id': 'my-project'}, } builder = AgentBuilder() diff --git a/flo_ai/tests/unit-tests/test_arium_builder.py b/flo_ai/tests/unit-tests/test_arium_builder.py index edd85c59..bac36c7d 100644 --- a/flo_ai/tests/unit-tests/test_arium_builder.py +++ b/flo_ai/tests/unit-tests/test_arium_builder.py @@ -71,6 +71,32 @@ def test_add_function_nodes(self): function_node in builder._function_nodes for function_node in function_nodes ) + def test_add_function_node_with_prefilled_params(self): + """Test adding a function node with prefilled_params""" + builder = AriumBuilder() + + def test_function(inputs, variables=None, param1=None, param2=None): + return f'Processed with param1={param1}, param2={param2}' + + prefilled_params = {'param1': 'value1', 'param2': 'value2'} + + function_node = FunctionNode( + name='test_function_node', + description='Test function node with prefilled params', + function=test_function, + prefilled_params=prefilled_params, + ) + + result = builder.add_function_node(function_node) + + assert result is builder + + assert function_node in builder._function_nodes + + assert function_node.prefilled_params == prefilled_params + assert function_node.prefilled_params['param1'] == 'value1' + assert function_node.prefilled_params['param2'] == 'value2' + def test_with_memory(self): """Test setting custom memory""" builder = AriumBuilder() diff --git a/flo_ai/tests/unit-tests/test_arium_yaml.py b/flo_ai/tests/unit-tests/test_arium_yaml.py index d47b011e..442e338e 100644 --- a/flo_ai/tests/unit-tests/test_arium_yaml.py +++ b/flo_ai/tests/unit-tests/test_arium_yaml.py @@ -207,6 +207,72 @@ async def test_function(inputs): assert builder._function_nodes[0].name == 'test_function_node' assert builder._function_nodes[0].function == test_function + def test_from_yaml_with_function_nodes_prefilled_params(self): + """Test YAML configuration with function nodes using prefilled_params.""" + yaml_config = """ + arium: + agents: + - name: test_agent + yaml_config: | + agent: + name: test_agent + job: "Test agent" + model: + provider: openai + name: gpt-4o-mini + + function_nodes: + - name: test_function_node + function_name: test_function + description: "Test function node with prefilled params" + prefilled_params: + param1: value1 + param2: value2 + param3: 42 + + workflow: + start: test_agent + edges: + - from: test_agent + to: [test_function_node] + - from: test_function_node + to: [end] + end: [test_function_node] + """ + + # Create mock function + async def test_function( + inputs, variables=None, param1=None, param2=None, param3=None + ): + return f'processed with param1={param1}, param2={param2}, param3={param3}' + + function_registry = {'test_function': test_function} + + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: + mock_agent = Mock(spec=Agent) + mock_agent.name = 'test_agent' + + mock_builder_instance = Mock() + mock_builder_instance.build.return_value = mock_agent + mock_agent_builder.from_yaml.return_value = mock_builder_instance + + builder = AriumBuilder.from_yaml( + yaml_str=yaml_config, function_registry=function_registry + ) + + # Verify function nodes were added with prefilled_params + assert len(builder._function_nodes) == 1 + function_node = builder._function_nodes[0] + assert function_node.name == 'test_function_node' + assert function_node.function == test_function + assert ( + function_node.description == 'Test function node with prefilled params' + ) + assert function_node.prefilled_params is not None + assert function_node.prefilled_params['param1'] == 'value1' + assert function_node.prefilled_params['param2'] == 'value2' + assert function_node.prefilled_params['param3'] == 42 + def test_from_yaml_with_routers(self): """Test YAML configuration with custom routers.""" yaml_config = """ diff --git a/flo_ai/tests/unit-tests/test_partial_tool.py b/flo_ai/tests/unit-tests/test_partial_tool.py index 06ef510a..13965e35 100644 --- a/flo_ai/tests/unit-tests/test_partial_tool.py +++ b/flo_ai/tests/unit-tests/test_partial_tool.py @@ -37,14 +37,14 @@ def test_partial_tool_creation(self): # Create partial tool with pre-filled datasource_id and project_id partial_tool = PartialTool( base_tool=base_tool, - pre_filled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, + prefilled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, ) # Verify partial tool properties assert partial_tool.name == 'test_tool_partial' assert 'pre-configured parameters' in partial_tool.description assert partial_tool.base_tool == base_tool - assert partial_tool.get_pre_filled_params() == { + assert partial_tool.get_prefilled_params() == { 'datasource_id': 'ds_123', 'project_id': 'my-project', } @@ -69,7 +69,7 @@ def test_partial_tool_with_custom_name_and_description(self): partial_tool = PartialTool( base_tool=base_tool, - pre_filled_params={'param1': 'value1'}, + prefilled_params={'param1': 'value1'}, name_override='custom_name', description_override='Custom description', ) @@ -106,7 +106,7 @@ async def test_partial_tool_execution(self): partial_tool = PartialTool( base_tool=base_tool, - pre_filled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, + prefilled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, ) # Execute with AI-provided query @@ -141,7 +141,7 @@ async def test_partial_tool_execution_ai_params_override_prefilled(self): ) partial_tool = PartialTool( - base_tool=base_tool, pre_filled_params={'datasource_id': 'ds_123'} + base_tool=base_tool, prefilled_params={'datasource_id': 'ds_123'} ) # Execute with AI-provided datasource_id that should override pre-filled one @@ -169,23 +169,23 @@ def test_partial_tool_parameter_management(self): ) partial_tool = PartialTool( - base_tool=base_tool, pre_filled_params={'param1': 'value1'} + base_tool=base_tool, prefilled_params={'param1': 'value1'} ) # Test adding a parameter - partial_tool.add_pre_filled_param('param2', 'value2') - assert partial_tool.get_pre_filled_params() == { + partial_tool.add_prefilled_param('param2', 'value2') + assert partial_tool.get_prefilled_params() == { 'param1': 'value1', 'param2': 'value2', } # Test removing a parameter - partial_tool.remove_pre_filled_param('param1') - assert partial_tool.get_pre_filled_params() == {'param2': 'value2'} + partial_tool.remove_prefilled_param('param1') + assert partial_tool.get_prefilled_params() == {'param2': 'value2'} # Test removing non-existent parameter (should not raise error) - partial_tool.remove_pre_filled_param('non_existent') - assert partial_tool.get_pre_filled_params() == {'param2': 'value2'} + partial_tool.remove_prefilled_param('non_existent') + assert partial_tool.get_prefilled_params() == {'param2': 'value2'} def test_create_partial_tool_helper_function(self): """Test the create_partial_tool helper function.""" @@ -214,7 +214,7 @@ def test_create_partial_tool_helper_function(self): ) assert isinstance(partial_tool, PartialTool) - assert partial_tool.get_pre_filled_params() == { + assert partial_tool.get_prefilled_params() == { 'datasource_id': 'ds_123', 'project_id': 'my-project', } @@ -234,7 +234,7 @@ async def test_partial_tool_error_handling(self): ) partial_tool = PartialTool( - base_tool=base_tool, pre_filled_params={'param1': 'value1'} + base_tool=base_tool, prefilled_params={'param1': 'value1'} ) # Execute and expect ToolExecutionError @@ -278,7 +278,7 @@ def test_partial_tool_parameter_filtering(self): # Pre-fill some required and some optional parameters partial_tool = PartialTool( base_tool=base_tool, - pre_filled_params={ + prefilled_params={ 'datasource_id': 'ds_123', 'project_id': 'my-project', 'optional_param': 'default_value', diff --git a/flo_ai/tests/unit-tests/test_tool_config.py b/flo_ai/tests/unit-tests/test_tool_config.py index 1add0c71..ad3baf85 100644 --- a/flo_ai/tests/unit-tests/test_tool_config.py +++ b/flo_ai/tests/unit-tests/test_tool_config.py @@ -30,13 +30,13 @@ def test_tool_config_creation(self): # Create tool config with pre-filled parameters tool_config = ToolConfig( tool=base_tool, - pre_filled_params={'datasource_id': 'ds_123'}, + prefilled_params={'datasource_id': 'ds_123'}, name_override='custom_name', description_override='Custom description', ) assert tool_config.tool == base_tool - assert tool_config.pre_filled_params == {'datasource_id': 'ds_123'} + assert tool_config.prefilled_params == {'datasource_id': 'ds_123'} assert tool_config.name_override == 'custom_name' assert tool_config.description_override == 'Custom description' assert tool_config.is_partial() is True @@ -55,7 +55,7 @@ def test_tool_config_without_prefilled_params(self): tool_config = ToolConfig(tool=base_tool) assert tool_config.tool == base_tool - assert tool_config.pre_filled_params == {} + assert tool_config.prefilled_params == {} assert tool_config.is_partial() is False def test_tool_config_to_tool_conversion(self): @@ -86,7 +86,7 @@ def test_tool_config_to_tool_conversion(self): # Test with pre-filled params (should return partial tool) tool_config_with_params = ToolConfig( - tool=base_tool, pre_filled_params={'datasource_id': 'ds_123'} + tool=base_tool, prefilled_params={'datasource_id': 'ds_123'} ) converted_tool = tool_config_with_params.to_tool() assert converted_tool.name == 'test_tool_partial' @@ -119,7 +119,7 @@ def test_create_tool_config_helper_function(self): assert isinstance(tool_config, ToolConfig) assert tool_config.tool == base_tool - assert tool_config.pre_filled_params == { + assert tool_config.prefilled_params == { 'datasource_id': 'ds_123', 'project_id': 'my-project', } diff --git a/flo_ai/tests/unit-tests/test_yaml_tool_config.py b/flo_ai/tests/unit-tests/test_yaml_tool_config.py index f9c2615b..44590eb2 100644 --- a/flo_ai/tests/unit-tests/test_yaml_tool_config.py +++ b/flo_ai/tests/unit-tests/test_yaml_tool_config.py @@ -77,7 +77,7 @@ def test_tool_configuration_with_prefilled_params(self): name: "gpt-4" tools: - name: "test_tool" - pre_filled_params: + prefilled_params: datasource_id: "ds_123" project_id: "my-project" name_override: "custom_tool_name" @@ -117,7 +117,7 @@ def test_mixed_tool_types(self): tools: - "test_tool" # Simple reference - name: "another_tool" - pre_filled_params: + prefilled_params: datasource_id: "ds_456" project_id: "another-project" """ @@ -170,7 +170,7 @@ def test_tool_configuration_missing_name(self): provider: "openai" name: "gpt-4" tools: - - pre_filled_params: + - prefilled_params: datasource_id: "ds_123" """ @@ -263,7 +263,7 @@ def test_process_yaml_tools_method(self): 'test_tool', # Simple reference { 'name': 'another_tool', - 'pre_filled_params': {'datasource_id': 'ds_123'}, + 'prefilled_params': {'datasource_id': 'ds_123'}, 'name_override': 'custom_name', }, ]