diff --git a/README.md b/README.md index aa316766..86a35e5e 100644 --- a/README.md +++ b/README.md @@ -959,7 +959,7 @@ async def get_weather(city: str, country: Optional[str] = None) -> str: return f"Weather in {city}: sunny" ``` -> 📖 **For detailed documentation on the `@flo_tool` decorator, see [README_flo_tool.md](flo_ai/README_flo_tool.md)** +> 📖 **For detailed documentation on the `@flo_tool` decorator, see [README_flo_tool.md](TOOLS.md)** ## 🧠 Reasoning Patterns diff --git a/TOOLS.md b/TOOLS.md new file mode 100644 index 00000000..a6b8e53c --- /dev/null +++ b/TOOLS.md @@ -0,0 +1,607 @@ +# Tools Documentation + +This document provides comprehensive documentation for the Flo AI tools system, including basic tools, partial tools, and YAML configuration. + +## Table of Contents + +- [Overview](#overview) +- [Basic Tools](#basic-tools) +- [Partial Tools](#partial-tools) +- [Tool Configuration](#tool-configuration) +- [YAML Configuration](#yaml-configuration) +- [Examples](#examples) +- [API Reference](#api-reference) + +## Overview + +The Flo AI tools system allows you to create and configure tools that can be used by AI agents. Tools can be: + +- **Standalone tools**: Simple functions that don't require external context (e.g., calculator, web search) +- **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 +- Tool registries for easy management + +## Basic Tools + +### Creating Tools with @flo_tool Decorator + +The `@flo_tool` decorator is the easiest way to create tools: + +```python +from flo_ai.tool.flo_tool import flo_tool + +@flo_tool( + description="Calculate mathematical expressions", + parameter_descriptions={ + "expression": "Mathematical expression to evaluate", + "precision": "Number of decimal places for the result" + } +) +async def calculate(expression: str, precision: int = 2) -> str: + """Calculate mathematical expressions.""" + try: + result = eval(expression) + return f"Result: {result:.{precision}f}" + except Exception as e: + return f"Error: {str(e)}" + +# The tool is automatically available as calculate.tool +``` + +### 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 +- `function`: The actual function to execute + +### Using Tools in Agents + +```python +from flo_ai.builder.agent_builder import AgentBuilder +from flo_ai.llm import OpenAI + +agent = (AgentBuilder() + .with_name("Math Assistant") + .with_prompt("You are a math assistant. Use the calculator tool.") + .with_llm(OpenAI(model="gpt-4")) + .add_tool(calculate.tool) + .build()) +``` + +## Partial Tools + +Partial tools allow you to pre-fill some parameters during agent building, hiding them from the AI while still allowing the AI to provide additional parameters. + +### 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 + +**Solution**: Pre-fill these parameters during agent building, so the AI only sees what it needs to provide. + +### Creating Partial Tools + +#### Method 1: Using AgentBuilder.add_tool() + +```python +from flo_ai.builder.agent_builder import AgentBuilder + +# BigQuery tool with multiple parameters +@flo_tool(description="Query BigQuery database") +async def bigquery_query( + query: str, + datasource_id: str, + project_id: str, + dataset: str +) -> str: + return f"Executed: {query} on {project_id}.{dataset}" + +# Create agent with pre-filled parameters +agent = (AgentBuilder() + .with_name("Data Analyst") + .with_prompt("You are a data analyst.") + .with_llm(OpenAI(model="gpt-4")) + .add_tool( + bigquery_query.tool, + datasource_id="ds_production_123", + project_id="my-company-prod", + dataset="analytics" + ) + .build()) + +# AI only needs to provide the query parameter +``` + +#### Method 2: Using Tool Dictionaries + +```python +agent = (AgentBuilder() + .with_name("Data Analyst") + .with_prompt("You are a data analyst.") + .with_llm(OpenAI(model="gpt-4")) + .with_tools([ + { + "tool": bigquery_query.tool, + "pre_filled_params": { + "datasource_id": "ds_production_123", + "project_id": "my-company-prod", + "dataset": "analytics" + }, + "name_override": "query_production_data", + "description_override": "Query production BigQuery data" + } + ]) + .build()) +``` + +#### Method 3: Using ToolConfig Class + +```python +from flo_ai.tool.tool_config import ToolConfig, create_tool_config + +# Create partial tool configuration +partial_tool = create_tool_config( + bigquery_query.tool, + datasource_id="ds_production_123", + project_id="my-company-prod", + dataset="analytics" +) + +agent = (AgentBuilder() + .with_name("Data Analyst") + .with_prompt("You are a data analyst.") + .with_llm(OpenAI(model="gpt-4")) + .with_tools([partial_tool]) + .build()) +``` + +### How Partial Tools Work + +1. **Parameter Filtering**: Pre-filled parameters are completely hidden from the AI +2. **Automatic Merging**: During execution, pre-filled parameters are merged with AI-provided parameters +3. **AI Override**: If the AI provides a parameter that's pre-filled, the AI's value takes precedence + +```python +# Original tool parameters (AI sees all) +bigquery_query.tool.parameters +# {'query': {...}, 'datasource_id': {...}, 'project_id': {...}, 'dataset': {...}} + +# Partial tool parameters (AI only sees non-pre-filled) +partial_tool.parameters +# {'query': {...}} # Only query is visible to AI + +# During execution, parameters are merged: +# AI provides: {"query": "SELECT * FROM users"} +# Final call: {"query": "SELECT * FROM users", "datasource_id": "ds_123", ...} +``` + +## Tool Configuration + +### ToolConfig Class + +The `ToolConfig` class provides a flexible way to configure tools: + +```python +from flo_ai.tool.tool_config import ToolConfig + +tool_config = ToolConfig( + tool=my_tool, + pre_filled_params={"param1": "value1", "param2": "value2"}, + name_override="custom_tool_name", + description_override="Custom tool description" +) + +# Convert to tool +configured_tool = tool_config.to_tool() +``` + +### ToolConfig Methods + +- `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 + +## YAML Configuration + +YAML configuration allows you to define tools and their configurations in YAML files, making it easy to manage different environments and configurations. + +### Basic YAML Structure + +```yaml +agent: + name: "Data Analyst Assistant" + job: "You are a data analyst with access to BigQuery and web search tools." + model: + provider: "openai" + name: "gpt-4" + tools: + # Simple tool reference + - "calculate" + + # Tool with pre-filled parameters + - name: "bigquery_query" + pre_filled_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: + max_results: 5 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" +``` + +### Environment-Specific Configurations + +#### Production Configuration (`agent_config.yaml`) + +```yaml +agent: + name: "Production Data Analyst" + job: "You are a data analyst working with production data." + model: + provider: "openai" + name: "gpt-4" + tools: + - name: "bigquery_query" + pre_filled_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: + smtp_server: "smtp.company.com" + smtp_port: 587 + name_override: "send_notification" + description_override: "Send email notifications" +``` + +#### Development Configuration (`agent_config_dev.yaml`) + +```yaml +agent: + name: "Development Data Analyst" + job: "You are a data analyst working with test data." + model: + provider: "openai" + name: "gpt-4" + tools: + - name: "bigquery_query" + pre_filled_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: + max_results: 3 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" +``` + +### Using YAML Configuration + +```python +from flo_ai.builder.agent_builder import AgentBuilder + +# Create tool registry +tool_registry = { + "bigquery_query": bigquery_query.tool, + "web_search": web_search.tool, + "calculate": calculate.tool, + "send_email": send_email.tool, +} + +# Load from YAML +with open("agent_config.yaml", "r") as f: + yaml_str = f.read() + +agent = AgentBuilder.from_yaml( + yaml_str, + tool_registry=tool_registry +) +``` + +## Examples + +### Complete Example: Data Analyst Agent + +```python +import asyncio +from flo_ai.tool.flo_tool import flo_tool +from flo_ai.builder.agent_builder import AgentBuilder +from flo_ai.llm import OpenAI + +# Define tools +@flo_tool(description="Query BigQuery database") +async def bigquery_query( + query: str, + datasource_id: str, + project_id: str, + dataset: str +) -> str: + return f"Executed: {query} on {project_id}.{dataset}" + +@flo_tool(description="Search the web") +async def web_search(query: str, max_results: int = 10) -> str: + return f"Found {max_results} results for: {query}" + +@flo_tool(description="Calculate mathematical expressions") +async def calculate(expression: str) -> str: + try: + result = eval(expression) + return f"Result: {result}" + except Exception as e: + return f"Error: {str(e)}" + +# Create agent with partial tools +agent = (AgentBuilder() + .with_name("Data Analyst") + .with_prompt("You are a data analyst. Use the provided tools to analyze data.") + .with_llm(OpenAI(model="gpt-4")) + .add_tool(calculate.tool) # Regular tool + .add_tool( + bigquery_query.tool, + datasource_id="ds_production_123", + project_id="my-company-prod", + dataset="analytics" + ) # Partial tool + .add_tool( + web_search.tool, + max_results=5 + ) # Partial tool + .build()) + +# Test the agent +async def test_agent(): + # The AI only needs to provide the query for BigQuery + # datasource_id, project_id, and dataset are pre-filled + result = await agent.tools[1].execute(query="SELECT COUNT(*) FROM users") + print(result) # "Executed: SELECT COUNT(*) FROM users on my-company-prod.analytics" + +asyncio.run(test_agent()) +``` + +### YAML Example + +```yaml +# agent_config.yaml +agent: + name: "Data Analyst Assistant" + job: "You are a data analyst with access to BigQuery and web search tools." + model: + provider: "openai" + name: "gpt-4" + tools: + - "calculate" + - name: "bigquery_query" + pre_filled_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: + max_results: 5 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" + settings: + temperature: 0.7 + max_retries: 3 + reasoning_pattern: "DIRECT" +``` + +```python +# Load from YAML +tool_registry = { + "bigquery_query": bigquery_query.tool, + "web_search": web_search.tool, + "calculate": calculate.tool, +} + +with open("agent_config.yaml", "r") as f: + yaml_str = f.read() + +agent = AgentBuilder.from_yaml( + yaml_str, + tool_registry=tool_registry +) +``` + +## API Reference + +### @flo_tool Decorator + +```python +@flo_tool( + name: Optional[str] = None, + description: Optional[str] = None, + parameter_descriptions: Optional[Dict[str, str]] = None +) +``` + +**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 + +### ToolConfig Class + +```python +class ToolConfig: + def __init__( + self, + tool: Tool, + pre_filled_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 + +### PartialTool Class + +```python +class PartialTool(Tool): + def __init__( + self, + base_tool: Tool, + pre_filled_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 + +### AgentBuilder Methods + +```python +# Add single tool with optional pre-filled parameters +def add_tool(self, tool: Tool, **pre_filled_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' + +# Load from YAML +@classmethod +def from_yaml( + cls, + yaml_str: str, + tools: Optional[List[Tool]] = None, + base_llm: Optional[BaseLLM] = None, + tool_registry: Optional[Dict[str, Tool]] = None, +) -> 'AgentBuilder' +``` + +## Best Practices + +### 1. Tool Design + +- **Single Responsibility**: Each tool should have one clear purpose +- **Clear Parameters**: Use descriptive parameter names and descriptions +- **Error Handling**: Always handle errors gracefully +- **Type Hints**: Use type hints for better documentation + +### 2. Partial Tools + +- **Pre-fill Context**: Pre-fill parameters that are environment-specific +- **Hide Complexity**: Hide technical details from the AI +- **Consistent Naming**: Use consistent naming for pre-filled parameters + +### 3. YAML Configuration + +- **Environment Separation**: Use different YAML files for different environments +- **Tool Registry**: Maintain a centralized tool registry +- **Version Control**: Keep YAML files in version control +- **Documentation**: Document tool configurations + +### 4. Error Handling + +```python +@flo_tool(description="Query database") +async def query_database(query: str, connection_string: str) -> str: + try: + # Database query logic + return result + except Exception as e: + return f"Database error: {str(e)}" +``` + +### 5. Testing + +```python +import pytest +from unittest.mock import Mock, AsyncMock + +def test_tool_execution(): + mock_function = AsyncMock(return_value="test_result") + tool = Tool( + name="test_tool", + description="A test tool", + 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") +``` + +## Troubleshooting + +### Common Issues + +1. **Tool not found in registry**: Ensure the tool is registered with the correct name +2. **Parameter validation errors**: Check that all required parameters are provided +3. **YAML parsing errors**: Validate YAML syntax and structure +4. **Import errors**: Ensure all required modules are imported + +### 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()` +3. **Test tool execution**: Test tools independently before using in agents +4. **Check YAML structure**: Validate YAML with online validators + +## Migration Guide + +### From Basic Tools to Partial Tools + +```python +# Before: AI needs to provide all parameters +agent.add_tool(bigquery_query.tool) + +# After: Pre-fill context-specific parameters +agent.add_tool( + bigquery_query.tool, + datasource_id="ds_123", + project_id="my-project" +) +``` + +### From Code to YAML Configuration + +```python +# Before: Hard-coded in code +agent = (AgentBuilder() + .with_name("Data Analyst") + .with_prompt("You are a data analyst.") + .with_llm(OpenAI(model="gpt-4")) + .add_tool(bigquery_query.tool) + .build()) + +# After: YAML configuration +agent = AgentBuilder.from_yaml(yaml_str, tool_registry=tool_registry) +``` + +This documentation covers all aspects of the Flo AI tools system. For more specific examples, see the `examples/` directory in the repository. diff --git a/flo_ai/README.md b/flo_ai/README.md index 19968686..5c8e15e3 100644 --- a/flo_ai/README.md +++ b/flo_ai/README.md @@ -1,263 +1 @@ -# @flo_tool Decorator - -The `@flo_tool` decorator is a powerful utility that automatically converts any Python function into a `Tool` object for use with Flo AI agents. It extracts function parameters, type hints, and descriptions to create a fully functional tool with minimal boilerplate code. - -## Features - -- **Automatic parameter extraction**: Uses Python's `inspect` module to extract function parameters and type hints -- **Flexible descriptions**: Supports custom descriptions, docstring extraction, and parameter-specific descriptions -- **Type conversion**: Automatically converts Python types to JSON schema types -- **Dual functionality**: Functions can be called normally AND used as tools -- **Async support**: Works seamlessly with both sync and async functions - -## Basic Usage - -### Simple Decorator - -```python -from flo_ai.tool import flo_tool - -@flo_tool() -async def calculate(operation: str, x: float, y: float) -> float: - """Calculate mathematical operations between two numbers.""" - operations = { - 'add': lambda: x + y, - 'subtract': lambda: x - y, - 'multiply': lambda: x * y, - 'divide': lambda: x / y if y != 0 else 'Cannot divide by zero', - } - if operation not in operations: - raise ValueError(f'Unknown operation: {operation}') - return operations[operation]() - -# Function can be called normally -result = await calculate("add", 5, 3) # Returns 8 - -# Tool object is accessible via .tool attribute -tool = calculate.tool -print(tool.name) # "calculate" -print(tool.description) # Uses function docstring -print(tool.parameters) # Automatically extracted from type hints -``` - -### With Custom Descriptions - -```python -@flo_tool( - name="weather_checker", - description="Get current weather information for a city", - parameter_descriptions={ - "city": "The city to get weather for", - "country": "The country (optional)", - } -) -async def get_weather(city: str, country: str = None) -> str: - """Get weather information for a specific city.""" - # Implementation here - return f"Weather in {city}: sunny" -``` - -### Using Docstrings for Descriptions - -```python -@flo_tool() -async def convert_units(value: float, from_unit: str, to_unit: str) -> str: - """ - Convert between different units (km/miles, kg/lbs, celsius/fahrenheit). - - Args: - value: The value to convert - from_unit: The unit to convert from - to_unit: The unit to convert to - """ - # Implementation here - return f"{value} {from_unit} = {result} {to_unit}" -``` - -## Advanced Usage - -### Creating Tools from Existing Functions - -If you have existing functions that you want to convert to tools without modifying them: - -```python -from flo_ai.tool import create_tool_from_function - -async def existing_function(text: str, style: str = "normal") -> str: - """Format text in different styles.""" - styles = { - "uppercase": text.upper(), - "lowercase": text.lower(), - "title": text.title(), - "normal": text - } - return styles.get(style, text) - -# Convert to tool -format_tool = create_tool_from_function( - existing_function, - name="text_formatter", - description="Format text in different styles", - parameter_descriptions={ - "text": "The text to format", - "style": "The formatting style (uppercase, lowercase, title, normal)" - } -) -``` - -### Using with Agents - -```python -from flo_ai.builder.agent_builder import AgentBuilder -from flo_ai.models.base_agent import ReasoningPattern - -# Create tools from decorated functions -tools = [ - calculate.tool, - get_weather.tool, - convert_units.tool, - format_tool # From create_tool_from_function -] - -# Build agent with tools -agent = ( - AgentBuilder() - .with_name("Multi-Tool Agent") - .with_prompt("You are a helpful assistant with access to various tools.") - .with_llm(llm) - .with_tools(tools) - .with_reasoning(ReasoningPattern.REACT) - .build() -) - -# Use the agent -response = await agent.run("Calculate 5 + 3 and then convert 10 km to miles") -``` - -## Parameter Types - -The decorator automatically converts Python types to JSON schema types: - -| Python Type | JSON Schema Type | -|-------------|------------------| -| `str` | `string` | -| `int` | `integer` | -| `float` | `number` | -| `bool` | `boolean` | -| `list` | `array` | -| `dict` | `object` | -| `Optional[T]` | `T` (required: false) | -| No annotation | `string` (default) | - -## Examples - -### Complete Example - -```python -import asyncio -from flo_ai.tool import flo_tool -from flo_ai.builder.agent_builder import AgentBuilder -from flo_ai.models.base_agent import ReasoningPattern -from flo_ai.llm.openai_llm import OpenAI - -# Define tools with decorator -@flo_tool( - description="Perform basic calculations", - parameter_descriptions={ - "operation": "The operation to perform (add, subtract, multiply, divide)", - "x": "First number", - "y": "Second number" - } -) -async def calculate(operation: str, x: float, y: float) -> float: - operations = { - 'add': lambda: x + y, - 'subtract': lambda: x - y, - 'multiply': lambda: x * y, - 'divide': lambda: x / y if y != 0 else 'Cannot divide by zero', - } - if operation not in operations: - raise ValueError(f'Unknown operation: {operation}') - return operations[operation]() - -@flo_tool() -async def convert_units(value: float, from_unit: str, to_unit: str) -> str: - """Convert between different units.""" - # Implementation here - return f"{value} {from_unit} = {result} {to_unit}" - -async def main(): - # Create agent with tools - llm = OpenAI(model='gpt-4-turbo-preview') - - agent = ( - AgentBuilder() - .with_name("Calculator Agent") - .with_prompt("You can perform calculations and unit conversions.") - .with_llm(llm) - .with_tools([calculate.tool, convert_units.tool]) - .with_reasoning(ReasoningPattern.REACT) - .build() - ) - - # Test the agent - response = await agent.run("Calculate 10 + 5 and convert 20 km to miles") - print(response) - -if __name__ == '__main__': - asyncio.run(main()) -``` - -## Benefits - -1. **Reduced Boilerplate**: No need to manually create Tool objects with parameter definitions -2. **Type Safety**: Leverages Python's type hints for automatic parameter type detection -3. **Documentation**: Uses docstrings and parameter descriptions for better tool documentation -4. **Flexibility**: Functions can be used both as regular functions and as tools -5. **Maintainability**: Changes to function signatures automatically update the tool definition - -## Migration from Manual Tool Creation - -### Before (Manual) -```python -from flo_ai.tool.base_tool import Tool - -async def calculate(operation: str, x: float, y: float) -> float: - # Implementation - pass - -calculator_tool = Tool( - name='calculate', - description='Perform basic calculations', - function=calculate, - parameters={ - 'operation': { - 'type': 'string', - 'description': 'The operation to perform', - }, - 'x': {'type': 'number', 'description': 'First number'}, - 'y': {'type': 'number', 'description': 'Second number'}, - }, -) -``` - -### After (With @flo_tool) -```python -from flo_ai.tool import flo_tool - -@flo_tool( - description="Perform basic calculations", - parameter_descriptions={ - "operation": "The operation to perform", - "x": "First number", - "y": "Second number" - } -) -async def calculate(operation: str, x: float, y: float) -> float: - # Implementation - pass - -# Tool is automatically available as calculate.tool -``` - -The `@flo_tool` decorator significantly reduces the amount of code needed to create tools while maintaining all the functionality and flexibility of the original Tool class. \ No newline at end of file +Please refer to [Main README.md](../README.md)** \ No newline at end of file diff --git a/flo_ai/examples/partial_tool_example.py b/flo_ai/examples/partial_tool_example.py new file mode 100644 index 00000000..11a36386 --- /dev/null +++ b/flo_ai/examples/partial_tool_example.py @@ -0,0 +1,199 @@ +""" +Example demonstrating partial tool functionality. + +This example shows how to create tools with pre-filled parameters +that are hidden from the AI, allowing for cleaner agent interactions. +""" + +import asyncio +from flo_ai.tool.flo_tool import flo_tool +from flo_ai.tool.partial_tool import create_partial_tool +from flo_ai.builder.agent_builder import AgentBuilder +from flo_ai.llm import OpenAI + + +# Example 1: BigQuery tool with datasource configuration +@flo_tool( + description='Query BigQuery database', + parameter_descriptions={ + 'query': 'SQL query to execute', + 'datasource_id': 'ID of the data source', + 'project_id': 'Google Cloud project ID', + 'dataset': 'BigQuery dataset name', + }, +) +async def bigquery_query( + query: str, datasource_id: str, project_id: str, dataset: str +) -> str: + """Execute a BigQuery query.""" + # Simulate BigQuery execution + return f"Executed query '{query}' on {project_id}.{dataset} using datasource {datasource_id}" + + +# Example 2: Web search tool with configuration +@flo_tool( + description='Search the web for information', + parameter_descriptions={ + 'query': 'Search query', + 'max_results': 'Maximum number of results to return', + 'language': 'Language for search results', + }, +) +async def web_search(query: str, max_results: int = 10, language: str = 'en') -> str: + """Search the web for information.""" + # Simulate web search + return f"Found {max_results} results for '{query}' in {language}" + + +# Example 3: Calculator tool (standalone, no pre-filling needed) +@flo_tool( + description='Perform mathematical calculations', + parameter_descriptions={'expression': 'Mathematical expression to evaluate'}, +) +async def calculate(expression: str) -> str: + """Calculate mathematical expressions.""" + try: + # Simple evaluation (in production, use a safer method) + result = eval(expression) + return f'Result: {result}' + except Exception as e: + return f'Error: {str(e)}' + + +async def main(): + """Demonstrate partial tool usage.""" + print('=== Partial Tool Example ===\n') + + # Create partial tools with pre-filled parameters + print('1. Creating partial tools with pre-filled parameters...') + + # BigQuery tool with pre-filled datasource info + bigquery_partial = create_partial_tool( + bigquery_query.tool, + datasource_id='ds_analytics_123', + project_id='my-company-project', + dataset='user_analytics', + ) + + # Web search tool with pre-filled configuration + web_search_partial = create_partial_tool( + web_search.tool, max_results=5, language='en' + ) + + print( + f'BigQuery partial tool parameters: {list(bigquery_partial.parameters.keys())}' + ) + print( + f'Web search partial tool parameters: {list(web_search_partial.parameters.keys())}' + ) + print() + + # Test the partial tools + print('2. Testing partial tools...') + + # Test BigQuery partial tool + bigquery_result = await bigquery_partial.execute( + query='SELECT COUNT(*) FROM users WHERE active = true' + ) + print(f'BigQuery result: {bigquery_result}') + + # Test web search partial tool + web_result = await web_search_partial.execute( + query='artificial intelligence trends 2024' + ) + print(f'Web search result: {web_result}') + print() + + # Create an agent with partial tools using the simplified interface + print('3. Creating agent with partial tools...') + + agent = ( + AgentBuilder() + .with_name('Data Analyst Assistant') + .with_prompt( + 'You are a data analyst. Use the provided tools to analyze data and search for information.' + ) + .with_llm(OpenAI(model='gpt-4')) + .with_tools( + [ + # Tool with pre-filled parameters using dictionary format + { + 'tool': bigquery_query.tool, + 'pre_filled_params': { + 'datasource_id': 'ds_production_456', + 'project_id': 'company-prod', + 'dataset': 'analytics', + }, + }, + # Tool with pre-filled parameters using dictionary format + { + 'tool': web_search.tool, + 'pre_filled_params': {'max_results': 3, 'language': 'en'}, + }, + # Regular tool without pre-filling + calculate.tool, + ] + ) + .build() + ) + + print(f'Agent created with {len(agent.tools)} tools') + print('Tools available to AI:') + for tool in agent.tools: + print(f' - {tool.name}: {list(tool.parameters.keys())}') + print() + + # Alternative: Using add_tool method for individual tools + print('4. Alternative: Using add_tool method...') + + agent2 = ( + AgentBuilder() + .with_name('Simple Assistant') + .with_prompt('You are a simple assistant.') + .with_llm(OpenAI(model='gpt-4')) + .add_tool(calculate.tool) # Regular tool + .add_tool( + bigquery_query.tool, + datasource_id='ds_simple_789', + project_id='simple-project', + dataset='data', + ) # Tool with pre-filled parameters + .build() + ) + + print(f'Alternative agent created with {len(agent2.tools)} tools') + print('Tools available to AI:') + for tool in agent2.tools: + print(f' - {tool.name}: {list(tool.parameters.keys())}') + print() + + # Demonstrate tool parameter management + 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()}') + + # Remove a parameter + bigquery_partial.remove_pre_filled_param('timeout') + print(f'Removed timeout parameter: {bigquery_partial.get_pre_filled_params()}') + print() + + # Show the difference between original and partial tools + print('6. Comparing original vs partial tools...') + + print('Original BigQuery tool parameters:') + for param, info in bigquery_query.tool.parameters.items(): + print(f" - {param}: {info['type']} (required: {info['required']})") + + print('\nPartial BigQuery tool parameters (AI view):') + for param, info in bigquery_partial.parameters.items(): + print(f" - {param}: {info['type']} (required: {info['required']})") + + print( + f'\nPre-filled parameters (hidden from AI): {bigquery_partial.get_pre_filled_params()}' + ) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/tools_quickstart.py b/flo_ai/examples/tools_quickstart.py new file mode 100644 index 00000000..08ee7972 --- /dev/null +++ b/flo_ai/examples/tools_quickstart.py @@ -0,0 +1,207 @@ +""" +Quick start example demonstrating Flo AI tools functionality. + +This example shows: +1. Basic tool creation with @flo_tool decorator +2. Partial tools with pre-filled parameters +3. YAML configuration +4. Tool execution +""" + +import asyncio +from flo_ai.tool.flo_tool import flo_tool +from flo_ai.builder.agent_builder import AgentBuilder +from flo_ai.llm import OpenAI + + +# Example 1: Basic tool creation +@flo_tool( + description='Calculate mathematical expressions', + parameter_descriptions={ + 'expression': 'Mathematical expression to evaluate', + 'precision': 'Number of decimal places for the result', + }, +) +async def calculate(expression: str, precision: int = 2) -> str: + """Calculate mathematical expressions safely.""" + try: + # Simple evaluation (in production, use a safer method) + result = eval(expression) + return f'Result: {result:.{precision}f}' + except Exception as e: + return f'Error: {str(e)}' + + +# Example 2: Tool with multiple parameters (for partial tool demo) +@flo_tool( + description='Query a database', + parameter_descriptions={ + 'query': 'SQL query to execute', + 'database_url': 'Database connection URL', + 'timeout': 'Query timeout in seconds', + 'max_rows': 'Maximum number of rows to return', + }, +) +async def query_database( + query: str, database_url: str, timeout: int = 30, max_rows: int = 1000 +) -> str: + """Execute a database query.""" + return f"Executed query '{query}' on {database_url} (timeout: {timeout}s, max_rows: {max_rows})" + + +# Example 3: Web search tool +@flo_tool( + description='Search the web for information', + parameter_descriptions={ + 'query': 'Search query', + 'max_results': 'Maximum number of results', + 'language': 'Language for search results', + }, +) +async def web_search(query: str, max_results: int = 10, language: str = 'en') -> str: + """Search the web for information.""" + return f"Found {max_results} results for '{query}' in {language}" + + +async def main(): + """Demonstrate tools functionality.""" + print('=== Flo AI Tools Quick Start ===\n') + + # 1. Basic tool usage + print('1. Basic tool usage...') + result = await calculate('2 + 3 * 4', precision=1) + print(f'Calculation result: {result}') + print() + + # 2. Show tool properties + print('2. Tool properties...') + print(f'Tool name: {calculate.tool.name}') + print(f'Tool description: {calculate.tool.description}') + print(f'Tool parameters: {list(calculate.tool.parameters.keys())}') + for param, info in calculate.tool.parameters.items(): + print(f" - {param}: {info['type']} (required: {info['required']})") + print() + + # 3. Partial tool demonstration + print('3. Partial tool demonstration...') + + # Create agent with partial tools + agent = ( + AgentBuilder() + .with_name('Data Analyst Assistant') + .with_prompt('You are a data analyst. Use the provided tools to analyze data.') + .with_llm(OpenAI(model='gpt-4')) + .add_tool(calculate.tool) # Regular tool + .add_tool( + query_database.tool, + database_url='postgresql://prod-db.company.com:5432/analytics', + timeout=60, + max_rows=5000, + ) # Partial tool with pre-filled parameters + .add_tool( + web_search.tool, max_results=5, language='en' + ) # Partial tool with pre-filled parameters + .build() + ) + + print(f'Agent created with {len(agent.tools)} tools') + print('Tools available to AI:') + for i, tool in enumerate(agent.tools): + print(f' {i+1}. {tool.name}: {list(tool.parameters.keys())}') + print() + + # 4. Demonstrate parameter filtering + print('4. Parameter filtering demonstration...') + + # Original tool parameters (all visible) + print('Original database tool parameters:') + for param, info in query_database.tool.parameters.items(): + print(f" - {param}: {info['type']} (required: {info['required']})") + + # Partial tool parameters (filtered) + partial_tool = agent.tools[1] # The query_database tool + print('\nPartial database tool parameters (AI view):') + for param, info in partial_tool.parameters.items(): + print(f" - {param}: {info['type']} (required: {info['required']})") + + # Show pre-filled parameters + if hasattr(partial_tool, 'get_pre_filled_params'): + print( + f'\nPre-filled parameters (hidden from AI): {partial_tool.get_pre_filled_params()}' + ) + print() + + # 5. Tool execution demonstration + print('5. Tool execution demonstration...') + + # Execute the partial tool (AI only needs to provide query) + try: + result = await partial_tool.execute( + query='SELECT COUNT(*) FROM users WHERE active = true' + ) + print(f'Database query result: {result}') + except Exception as e: + print(f'Error executing tool: {e}') + print() + + # 6. YAML configuration example + print('6. YAML configuration example...') + + yaml_config = """ +agent: + name: "YAML Configured Agent" + job: "You are an agent configured via YAML." + model: + provider: "openai" + name: "gpt-4" + tools: + - "calculate" + - name: "query_database" + pre_filled_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: + max_results: 3 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" +""" + + # Create tool registry + tool_registry = { + 'calculate': calculate.tool, + 'query_database': query_database.tool, + 'web_search': web_search.tool, + } + + try: + yaml_agent = AgentBuilder.from_yaml(yaml_config, tool_registry=tool_registry) + print(f'YAML agent created: {yaml_agent._name}') + print(f'YAML agent tools: {[tool.name for tool in yaml_agent._tools]}') + except Exception as e: + print(f'Error creating YAML agent (expected - no API key): {e}') + print() + + # 7. Tool parameter management + print('7. Tool parameter management...') + + if hasattr(partial_tool, 'add_pre_filled_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()}') + + # Remove a parameter + partial_tool.remove_pre_filled_param('retry_count') + print(f'Removed retry_count parameter: {partial_tool.get_pre_filled_params()}') + print() + + print('=== Quick Start Complete ===') + print('For more detailed documentation, see TOOLS.md') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/yaml_tool_config_example.py b/flo_ai/examples/yaml_tool_config_example.py new file mode 100644 index 00000000..9a51bb21 --- /dev/null +++ b/flo_ai/examples/yaml_tool_config_example.py @@ -0,0 +1,256 @@ +""" +Example demonstrating YAML tool configuration with pre-filled parameters. + +This example shows how to configure agents with tools that have pre-filled +parameters using YAML configuration files. +""" + +import asyncio +from flo_ai.tool.flo_tool import flo_tool +from flo_ai.builder.agent_builder import AgentBuilder + + +# Define some example tools +@flo_tool( + description='Query BigQuery database', + parameter_descriptions={ + 'query': 'SQL query to execute', + 'datasource_id': 'ID of the data source', + 'project_id': 'Google Cloud project ID', + 'dataset': 'BigQuery dataset name', + }, +) +async def bigquery_query( + query: str, datasource_id: str, project_id: str, dataset: str +) -> str: + """Execute a BigQuery query.""" + return f"Executed query '{query}' on {project_id}.{dataset} using datasource {datasource_id}" + + +@flo_tool( + description='Search the web for information', + parameter_descriptions={ + 'query': 'Search query', + 'max_results': 'Maximum number of results to return', + 'language': 'Language for search results', + }, +) +async def web_search(query: str, max_results: int = 10, language: str = 'en') -> str: + """Search the web for information.""" + return f"Found {max_results} results for '{query}' in {language}" + + +@flo_tool( + description='Perform mathematical calculations', + parameter_descriptions={'expression': 'Mathematical expression to evaluate'}, +) +async def calculate(expression: str) -> str: + """Calculate mathematical expressions.""" + try: + result = eval(expression) + return f'Result: {result}' + except Exception as e: + return f'Error: {str(e)}' + + +@flo_tool( + description='Send email notifications', + parameter_descriptions={ + 'to': 'Recipient email address', + 'subject': 'Email subject', + 'body': 'Email body content', + 'smtp_server': 'SMTP server address', + 'smtp_port': 'SMTP server port', + }, +) +async def send_email( + to: str, subject: str, body: str, smtp_server: str, smtp_port: int +) -> str: + """Send an email notification.""" + return f"Email sent to {to} with subject '{subject}' via {smtp_server}:{smtp_port}" + + +async def main(): + """Demonstrate YAML tool configuration.""" + print('=== YAML Tool Configuration Example ===\n') + + # Create a tool registry + tool_registry = { + 'bigquery_query': bigquery_query.tool, + 'web_search': web_search.tool, + 'calculate': calculate.tool, + 'send_email': send_email.tool, + } + + # Example 1: Simple YAML with tool references + print('1. Simple YAML with tool references...') + simple_yaml = """ +agent: + name: "Simple Data Analyst" + job: "You are a data analyst. Use the provided tools to analyze data." + model: + provider: "openai" + name: "gpt-4" + tools: + - "calculate" + - "web_search" +""" + + try: + agent1 = AgentBuilder.from_yaml(simple_yaml, tool_registry=tool_registry) + print(f'Created agent: {agent1._name}') + print(f'Tools: {[tool.name for tool in agent1._tools]}') + except Exception as e: + print(f'Error (expected - no API key): {e}') + print() + + # Example 2: YAML with tool configurations and pre-filled parameters + print('2. YAML with tool configurations and pre-filled parameters...') + config_yaml = """ +agent: + name: "Advanced Data Analyst" + job: "You are an advanced data analyst with access to BigQuery and web search." + model: + provider: "openai" + name: "gpt-4" + tools: + - name: "bigquery_query" + pre_filled_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: + max_results: 5 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" + + - name: "send_email" + pre_filled_params: + smtp_server: "smtp.company.com" + smtp_port: 587 + name_override: "send_notification" + description_override: "Send email notifications" + + - "calculate" # Regular tool without pre-filling +""" + + try: + agent2 = AgentBuilder.from_yaml(config_yaml, tool_registry=tool_registry) + print(f'Created agent: {agent2._name}') + print('Tools available to AI:') + for tool in agent2._tools: + print(f' - {tool.name}: {list(tool.parameters.keys())}') + except Exception as e: + print(f'Error (expected - no API key): {e}') + print() + + # Example 3: Environment-specific configurations + print('3. Environment-specific configurations...') + dev_yaml = """ +agent: + name: "Development Data Analyst" + job: "You are a data analyst working in the development environment." + model: + provider: "openai" + name: "gpt-4" + tools: + - name: "bigquery_query" + pre_filled_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: + max_results: 3 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" +""" + + try: + agent3 = AgentBuilder.from_yaml(dev_yaml, tool_registry=tool_registry) + print(f'Created agent: {agent3._name}') + print('Tools available to AI:') + for tool in agent3._tools: + print(f' - {tool.name}: {list(tool.parameters.keys())}') + except Exception as e: + print(f'Error (expected - no API key): {e}') + print() + + # Example 4: Mixed tool types + print('4. Mixed tool types...') + mixed_yaml = """ +agent: + name: "Mixed Tool Agent" + job: "You are an agent with mixed tool configurations." + model: + provider: "openai" + name: "gpt-4" + tools: + - "calculate" # Simple reference + + - name: "bigquery_query" + pre_filled_params: + datasource_id: "ds_mixed_789" + project_id: "mixed-project" + dataset: "mixed_data" + + - name: "web_search" + pre_filled_params: + max_results: 10 + language: "en" + name_override: "search_web" + description_override: "Search the web for information" +""" + + try: + agent4 = AgentBuilder.from_yaml(mixed_yaml, tool_registry=tool_registry) + print(f'Created agent: {agent4._name}') + print('Tools available to AI:') + for tool in agent4._tools: + print(f' - {tool.name}: {list(tool.parameters.keys())}') + except Exception as e: + print(f'Error (expected - no API key): {e}') + print() + + # Example 5: Show the difference between original and configured tools + print('5. Comparing original vs configured tools...') + + # Create a tool with pre-filled parameters + configured_tool = tool_registry['bigquery_query'] + + print('Original BigQuery tool parameters:') + for param, info in configured_tool.parameters.items(): + print(f" - {param}: {info['type']} (required: {info['required']})") + + print('\nConfigured BigQuery tool parameters (AI view):') + # This would be the tool as seen by the AI after YAML configuration + # In practice, this would be created by the YAML processing + from flo_ai.tool.tool_config import ToolConfig + + tool_config = ToolConfig( + tool=configured_tool, + pre_filled_params={ + 'datasource_id': 'ds_production_123', + 'project_id': 'my-company-prod', + 'dataset': 'analytics', + }, + ) + configured_tool_for_ai = tool_config.to_tool() + + 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}') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/flo_ai/builder/agent_builder.py b/flo_ai/flo_ai/builder/agent_builder.py index 9b08ada8..f6db72a9 100644 --- a/flo_ai/flo_ai/builder/agent_builder.py +++ b/flo_ai/flo_ai/builder/agent_builder.py @@ -4,6 +4,7 @@ from flo_ai.models.base_agent import ReasoningPattern from flo_ai.llm import BaseLLM, OpenAI, Anthropic, Gemini, OllamaLLM, VertexAI from flo_ai.tool.base_tool import Tool +from flo_ai.tool.tool_config import ToolConfig, create_tool_config from flo_ai.formatter.yaml_format_parser import FloYamlParser from pydantic import BaseModel @@ -42,9 +43,85 @@ def with_llm(self, llm: BaseLLM) -> 'AgentBuilder': self._llm = llm return self - def with_tools(self, tools: List[Tool]) -> 'AgentBuilder': - """Add tools to the agent""" - self._tools = tools + def with_tools( + self, tools: Union[List[Tool], List[ToolConfig], List[Dict[str, Any]]] + ) -> 'AgentBuilder': + """ + Add tools to the agent. + + Args: + 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 + - 'name_override': Optional custom name + - 'description_override': Optional custom description + + Examples: + # Regular tools + builder.with_tools([tool1, tool2]) + + # Tool configurations + builder.with_tools([ + ToolConfig(tool1, pre_filled_params={"param1": "value1"}), + ToolConfig(tool2, pre_filled_params={"param2": "value2"}) + ]) + + # Tool dictionaries + builder.with_tools([ + {"tool": tool1, "pre_filled_params": {"param1": "value1"}}, + {"tool": tool2, "pre_filled_params": {"param2": "value2"}} + ]) + """ + processed_tools = [] + + for tool_item in tools: + if isinstance(tool_item, Tool): + # Regular tool - add as is + processed_tools.append(tool_item) + elif isinstance(tool_item, ToolConfig): + # Tool configuration - convert to tool + processed_tools.append(tool_item.to_tool()) + 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', {}) + 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, + name_override=name_override, + description_override=description_override, + ) + processed_tools.append(tool_config.to_tool()) + else: + raise ValueError(f'Unsupported tool type: {type(tool_item)}') + + self._tools = processed_tools + return self + + def add_tool(self, tool: Tool, **pre_filled_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 + + Example: + builder.add_tool( + bigquery_tool, + datasource_id="ds_123", + project_id="my-project" + ) + """ + if pre_filled_params: + tool_config = create_tool_config(tool, **pre_filled_params) + self._tools.append(tool_config.to_tool()) + else: + self._tools.append(tool) return self def with_reasoning(self, pattern: ReasoningPattern) -> 'AgentBuilder': @@ -98,12 +175,16 @@ def from_yaml( yaml_str: str, tools: Optional[List[Tool]] = None, base_llm: Optional[BaseLLM] = None, + tool_registry: Optional[Dict[str, Tool]] = None, ) -> 'AgentBuilder': """Create an agent builder from a YAML configuration string Args: yaml_str: YAML string containing agent configuration tools: Optional list of tools to use with the agent + base_llm: Optional base LLM to use + tool_registry: Optional dictionary mapping tool names to Tool objects + Used to resolve tool references in YAML Returns: AgentBuilder: Configured agent builder instance @@ -159,8 +240,13 @@ def from_yaml( ) builder.with_llm(base_llm) - # Set tools if provided - if tools: + # Handle tools configuration + if 'tools' in agent_config: + # Process tools from YAML configuration + yaml_tools = cls._process_yaml_tools(agent_config['tools'], tool_registry) + builder.with_tools(yaml_tools) + elif tools: + # Use provided tools builder.with_tools(tools) # Set parser if present @@ -179,3 +265,69 @@ def from_yaml( builder.with_reasoning(ReasoningPattern[settings['reasoning_pattern']]) return builder + + @classmethod + def _process_yaml_tools( + cls, + tools_config: List[Dict[str, Any]], + tool_registry: Optional[Dict[str, Tool]] = None, + ) -> List[Tool]: + """Process tools configuration from YAML. + + Args: + tools_config: List of tool configurations from YAML + tool_registry: Optional dictionary mapping tool names to Tool objects + + Returns: + List[Tool]: Processed tools + """ + processed_tools = [] + + for tool_config in tools_config: + if isinstance(tool_config, str): + # Simple string reference - look up in registry + if tool_registry and tool_config in tool_registry: + processed_tools.append(tool_registry[tool_config]) + else: + raise ValueError(f"Tool '{tool_config}' not found in tool registry") + elif isinstance(tool_config, dict): + # Tool configuration dictionary + tool_name = tool_config.get('name') + if not tool_name: + raise ValueError("Tool configuration must have a 'name' field") + + # Look up tool in registry + if tool_registry and tool_name in tool_registry: + base_tool = tool_registry[tool_name] + else: + raise ValueError(f"Tool '{tool_name}' not found in tool registry") + + # Extract configuration + pre_filled_params = tool_config.get('pre_filled_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, + 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 + or name_override is not None + or description_override is not None + ): + processed_tools.append(tool_config_obj.to_tool()) + else: + # No pre-filled params and no custom name/description, use original tool + processed_tools.append(base_tool) + else: + raise ValueError( + f'Invalid tool configuration type: {type(tool_config)}' + ) + + return processed_tools diff --git a/flo_ai/flo_ai/tool/__init__.py b/flo_ai/flo_ai/tool/__init__.py index 9d985ea5..144436f9 100644 --- a/flo_ai/flo_ai/tool/__init__.py +++ b/flo_ai/flo_ai/tool/__init__.py @@ -1,4 +1,15 @@ from .base_tool import Tool, ToolExecutionError from .flo_tool import flo_tool, create_tool_from_function +from .partial_tool import PartialTool, create_partial_tool +from .tool_config import ToolConfig, create_tool_config -__all__ = ['Tool', 'ToolExecutionError', 'flo_tool', 'create_tool_from_function'] +__all__ = [ + 'Tool', + 'ToolExecutionError', + 'flo_tool', + 'create_tool_from_function', + 'PartialTool', + 'create_partial_tool', + 'ToolConfig', + 'create_tool_config', +] diff --git a/flo_ai/flo_ai/tool/partial_tool.py b/flo_ai/flo_ai/tool/partial_tool.py new file mode 100644 index 00000000..d0bc9333 --- /dev/null +++ b/flo_ai/flo_ai/tool/partial_tool.py @@ -0,0 +1,111 @@ +from typing import Dict, Any, Optional +from .base_tool import Tool, ToolExecutionError +from flo_ai.utils.logger import logger + + +class PartialTool(Tool): + """ + A tool that has some parameters pre-filled during agent building. + The AI can still provide additional parameters during execution. + """ + + def __init__( + self, + base_tool: Tool, + pre_filled_params: Dict[str, Any], + name_override: Optional[str] = None, + description_override: Optional[str] = None, + ): + """ + Create a partial tool with pre-filled parameters. + + Args: + base_tool: The original tool to wrap + pre_filled_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() + + # 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: + filtered_parameters[param_name] = param_info.copy() + + super().__init__( + name=name_override or f'{base_tool.name}_partial', + description=description_override + or f'{base_tool.description} (with pre-configured parameters)', + function=base_tool.function, + parameters=filtered_parameters, + ) + + async def execute(self, **kwargs) -> Any: + """Execute the tool with pre-filled parameters merged with AI-provided ones.""" + try: + # Merge pre-filled params with AI-provided params + # AI params take precedence over pre-filled ones + merged_params = {**self.pre_filled_params, **kwargs} + + logger.info( + f'Executing partial tool {self.name} with merged params: {merged_params}' + ) + tool_result = await self.base_tool.function(**merged_params) + logger.info(f'Partial tool {self.name} returned: {tool_result}') + return tool_result + except Exception as e: + raise ToolExecutionError( + f'Error executing partial tool {self.name}: {str(e)}', original_error=e + ) + + 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]: + """Get the pre-filled parameters.""" + return self.pre_filled_params.copy() + + def add_pre_filled_param(self, key: str, value: Any) -> 'PartialTool': + """Add or update a pre-filled parameter.""" + self.pre_filled_params[key] = value + return self + + def remove_pre_filled_param(self, key: str) -> 'PartialTool': + """Remove a pre-filled parameter.""" + if key in self.pre_filled_params: + del self.pre_filled_params[key] + return self + + +def create_partial_tool(tool: Tool, **pre_filled_params) -> PartialTool: + """ + Create a partial tool with pre-filled parameters. + + Args: + tool: The original tool + **pre_filled_params: Parameters to pre-fill + + Returns: + PartialTool: A tool with pre-filled parameters + + Example: + # Original tool + @flo_tool(description="Query BigQuery") + async def bigquery_query(query: str, datasource_id: str, project_id: str): + # implementation + pass + + # Create partial tool with pre-filled datasource_id + partial_tool = create_partial_tool( + bigquery_query.tool, + datasource_id="ds_123", + project_id="my-project" + ) + + # AI only needs to provide the query + result = await partial_tool.execute(query="SELECT * FROM users") + """ + return PartialTool(tool, pre_filled_params) diff --git a/flo_ai/flo_ai/tool/tool_config.py b/flo_ai/flo_ai/tool/tool_config.py new file mode 100644 index 00000000..0aa5b38d --- /dev/null +++ b/flo_ai/flo_ai/tool/tool_config.py @@ -0,0 +1,70 @@ +from typing import Dict, Any, Optional +from .base_tool import Tool + + +class ToolConfig: + """ + Represents a tool with optional pre-filled parameters. + This allows the AgentBuilder to handle both regular tools and tools with configurations + transparently through the same interface. + """ + + def __init__( + self, + tool: Tool, + pre_filled_params: Optional[Dict[str, Any]] = None, + name_override: Optional[str] = None, + description_override: Optional[str] = None, + ): + """ + Create a tool configuration. + + Args: + tool: The base tool + pre_filled_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.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) + + 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 + or self.name_override is not None + or self.description_override is not None + ): + # Import here to avoid circular imports + from .partial_tool import PartialTool + + return PartialTool( + base_tool=self.tool, + pre_filled_params=self.pre_filled_params, + name_override=self.name_override, + description_override=self.description_override, + ) + else: + # No customizations, return original tool + return self.tool + + +def create_tool_config(tool: Tool, **pre_filled_params) -> ToolConfig: + """ + Create a tool configuration with pre-filled parameters. + + Args: + tool: The base tool + **pre_filled_params: Pre-filled parameters + + Returns: + ToolConfig: A tool configuration + """ + return ToolConfig(tool, pre_filled_params) diff --git a/flo_ai/tests/test_agent_builder_tools.py b/flo_ai/tests/test_agent_builder_tools.py new file mode 100644 index 00000000..9d8f013f --- /dev/null +++ b/flo_ai/tests/test_agent_builder_tools.py @@ -0,0 +1,196 @@ +import pytest +from unittest.mock import Mock, AsyncMock +from flo_ai.tool.base_tool import Tool +from flo_ai.tool.tool_config import ToolConfig +from flo_ai.builder.agent_builder import AgentBuilder +from flo_ai.llm import OpenAI + + +class TestAgentBuilderTools: + """Test cases for AgentBuilder tool functionality.""" + + def setup_method(self): + """Set up test fixtures.""" + self.mock_function = AsyncMock(return_value='test_result') + self.base_tool = Tool( + name='test_tool', + description='A test tool', + function=self.mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + 'project_id': { + 'type': 'string', + 'description': 'Project ID', + 'required': True, + }, + }, + ) + + def test_with_tools_regular_tools(self): + """Test adding regular tools to agent builder.""" + builder = AgentBuilder() + builder.with_tools([self.base_tool]) + + assert len(builder._tools) == 1 + assert builder._tools[0] == self.base_tool + + 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'}, + ) + + builder = AgentBuilder() + builder.with_tools([tool_config]) + + assert len(builder._tools) == 1 + # Should be converted to a partial tool + assert builder._tools[0].name == 'test_tool_partial' + assert 'pre-configured parameters' in builder._tools[0].description + + def test_with_tools_tool_dictionaries(self): + """Test adding tool dictionaries to agent builder.""" + tool_dict = { + 'tool': self.base_tool, + 'pre_filled_params': { + 'datasource_id': 'ds_123', + 'project_id': 'my-project', + }, + 'name_override': 'custom_tool_name', + 'description_override': 'Custom tool description', + } + + builder = AgentBuilder() + builder.with_tools([tool_dict]) + + assert len(builder._tools) == 1 + # Should be converted to a partial tool with custom name and description + assert builder._tools[0].name == 'custom_tool_name' + assert builder._tools[0].description == 'Custom tool description' + + 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_dict = { + 'tool': self.base_tool, + 'pre_filled_params': {'project_id': 'my-project'}, + } + + builder = AgentBuilder() + builder.with_tools([self.base_tool, tool_config, tool_dict]) + + assert len(builder._tools) == 3 + # First tool should be the original + assert builder._tools[0] == self.base_tool + # Second and third should be partial tools + assert 'partial' in builder._tools[1].name + assert 'partial' in builder._tools[2].name + + def test_with_tools_invalid_type(self): + """Test that invalid tool types raise an error.""" + builder = AgentBuilder() + + with pytest.raises(ValueError, match='Unsupported tool type'): + builder.with_tools(['invalid_tool']) + + def test_add_tool_regular_tool(self): + """Test adding a regular tool using add_tool method.""" + builder = AgentBuilder() + builder.add_tool(self.base_tool) + + assert len(builder._tools) == 1 + assert builder._tools[0] == self.base_tool + + def test_add_tool_with_prefilled_params(self): + """Test adding a tool with pre-filled parameters using add_tool method.""" + builder = AgentBuilder() + builder.add_tool( + self.base_tool, datasource_id='ds_123', project_id='my-project' + ) + + assert len(builder._tools) == 1 + # Should be converted to a partial tool + assert builder._tools[0].name == 'test_tool_partial' + assert 'pre-configured parameters' in builder._tools[0].description + + def test_add_tool_multiple_tools(self): + """Test adding multiple tools using add_tool method.""" + builder = AgentBuilder() + builder.add_tool(self.base_tool) # Regular tool + builder.add_tool( + self.base_tool, datasource_id='ds_123' + ) # Tool with pre-filled params + + assert len(builder._tools) == 2 + assert builder._tools[0] == self.base_tool # Original tool + assert 'partial' in builder._tools[1].name # Partial tool + + def test_build_agent_with_tools(self): + """Test building an agent with tools.""" + # Mock LLM + mock_llm = Mock(spec=OpenAI) + + builder = AgentBuilder() + agent = ( + builder.with_name('Test Agent') + .with_prompt('You are a test agent') + .with_llm(mock_llm) + .add_tool(self.base_tool) + .add_tool(self.base_tool, datasource_id='ds_123', project_id='my-project') + .build() + ) + + assert agent.name == 'Test Agent' + assert agent.system_prompt == 'You are a test agent' + assert agent.llm == mock_llm + assert len(agent.tools) == 2 + assert agent.tools[0] == self.base_tool + assert 'partial' in agent.tools[1].name + + def test_tool_parameter_filtering(self): + """Test that pre-filled parameters are properly filtered from AI view.""" + builder = AgentBuilder() + builder.add_tool( + self.base_tool, datasource_id='ds_123', project_id='my-project' + ) + + # The tool should have only non-pre-filled parameters visible to AI + partial_tool = builder._tools[0] + ai_visible_params = partial_tool.parameters + + assert 'query' in ai_visible_params + assert 'datasource_id' not in ai_visible_params + assert 'project_id' not in ai_visible_params + + def test_tool_execution_with_merged_params(self): + """Test that tool execution merges pre-filled and AI-provided parameters.""" + builder = AgentBuilder() + builder.add_tool( + self.base_tool, datasource_id='ds_123', project_id='my-project' + ) + + partial_tool = builder._tools[0] + + # Execute with AI-provided query + import asyncio + + result = asyncio.run(partial_tool.execute(query='SELECT * FROM users')) + + # Verify the function was called with merged parameters + self.mock_function.assert_called_once_with( + query='SELECT * FROM users', datasource_id='ds_123', project_id='my-project' + ) + assert result == 'test_result' diff --git a/flo_ai/tests/test_partial_tool.py b/flo_ai/tests/test_partial_tool.py new file mode 100644 index 00000000..06ef510a --- /dev/null +++ b/flo_ai/tests/test_partial_tool.py @@ -0,0 +1,297 @@ +import pytest +from unittest.mock import AsyncMock +from flo_ai.tool.base_tool import Tool +from flo_ai.tool.partial_tool import PartialTool, create_partial_tool + + +class TestPartialTool: + """Test cases for PartialTool functionality.""" + + def test_partial_tool_creation(self): + """Test creating a partial tool with pre-filled parameters.""" + # Create a mock base tool + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + 'project_id': { + 'type': 'string', + 'description': 'Project ID', + 'required': True, + }, + }, + ) + + # 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'}, + ) + + # 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() == { + 'datasource_id': 'ds_123', + 'project_id': 'my-project', + } + + # Verify filtered parameters (only query should be visible to AI) + assert 'query' in partial_tool.parameters + assert 'datasource_id' not in partial_tool.parameters + assert 'project_id' not in partial_tool.parameters + assert partial_tool.parameters['query']['required'] is True + + def test_partial_tool_with_custom_name_and_description(self): + """Test creating a partial tool with custom name and description.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'param1': {'type': 'string', 'description': 'Param 1', 'required': True} + }, + ) + + partial_tool = PartialTool( + base_tool=base_tool, + pre_filled_params={'param1': 'value1'}, + name_override='custom_name', + description_override='Custom description', + ) + + assert partial_tool.name == 'custom_name' + assert partial_tool.description == 'Custom description' + + @pytest.mark.asyncio + async def test_partial_tool_execution(self): + """Test executing a partial tool with merged parameters.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + 'project_id': { + 'type': 'string', + 'description': 'Project ID', + 'required': True, + }, + }, + ) + + partial_tool = PartialTool( + base_tool=base_tool, + pre_filled_params={'datasource_id': 'ds_123', 'project_id': 'my-project'}, + ) + + # Execute with AI-provided query + result = await partial_tool.execute(query='SELECT * FROM users') + + # Verify the function was called with merged parameters + mock_function.assert_called_once_with( + query='SELECT * FROM users', datasource_id='ds_123', project_id='my-project' + ) + assert result == 'test_result' + + @pytest.mark.asyncio + async def test_partial_tool_execution_ai_params_override_prefilled(self): + """Test that AI-provided parameters override pre-filled ones.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + }, + ) + + partial_tool = PartialTool( + base_tool=base_tool, pre_filled_params={'datasource_id': 'ds_123'} + ) + + # Execute with AI-provided datasource_id that should override pre-filled one + await partial_tool.execute( + query='SELECT * FROM users', + datasource_id='ds_456', # This should override the pre-filled "ds_123" + ) + + # Verify the function was called with AI-provided datasource_id + mock_function.assert_called_once_with( + query='SELECT * FROM users', + datasource_id='ds_456', # AI-provided value should take precedence + ) + + def test_partial_tool_parameter_management(self): + """Test adding and removing pre-filled parameters.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'param1': {'type': 'string', 'description': 'Param 1', 'required': True} + }, + ) + + partial_tool = PartialTool( + base_tool=base_tool, pre_filled_params={'param1': 'value1'} + ) + + # Test adding a parameter + partial_tool.add_pre_filled_param('param2', 'value2') + assert partial_tool.get_pre_filled_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'} + + # 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'} + + def test_create_partial_tool_helper_function(self): + """Test the create_partial_tool helper function.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + }, + ) + + # Use helper function + partial_tool = create_partial_tool( + base_tool, datasource_id='ds_123', project_id='my-project' + ) + + assert isinstance(partial_tool, PartialTool) + assert partial_tool.get_pre_filled_params() == { + 'datasource_id': 'ds_123', + 'project_id': 'my-project', + } + + @pytest.mark.asyncio + async def test_partial_tool_error_handling(self): + """Test error handling in partial tool execution.""" + # Create a mock function that raises an exception + mock_function = AsyncMock(side_effect=Exception('Test error')) + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'param1': {'type': 'string', 'description': 'Param 1', 'required': True} + }, + ) + + partial_tool = PartialTool( + base_tool=base_tool, pre_filled_params={'param1': 'value1'} + ) + + # Execute and expect ToolExecutionError + with pytest.raises(Exception) as exc_info: + await partial_tool.execute() + + assert 'Error executing partial tool' in str(exc_info.value) + assert 'Test error' in str(exc_info.value) + + def test_partial_tool_parameter_filtering(self): + """Test that pre-filled parameters are properly filtered from AI view.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + 'project_id': { + 'type': 'string', + 'description': 'Project ID', + 'required': True, + }, + 'optional_param': { + 'type': 'string', + 'description': 'Optional param', + 'required': False, + }, + }, + ) + + # Pre-fill some required and some optional parameters + partial_tool = PartialTool( + base_tool=base_tool, + pre_filled_params={ + 'datasource_id': 'ds_123', + 'project_id': 'my-project', + 'optional_param': 'default_value', + }, + ) + + # Only non-pre-filled parameters should be visible to AI + ai_visible_params = partial_tool.parameters + assert 'query' in ai_visible_params + assert 'datasource_id' not in ai_visible_params + assert 'project_id' not in ai_visible_params + assert 'optional_param' not in ai_visible_params + + # Verify the parameter info is preserved + assert ai_visible_params['query']['type'] == 'string' + assert ai_visible_params['query']['required'] is True diff --git a/flo_ai/tests/test_tool_config.py b/flo_ai/tests/test_tool_config.py new file mode 100644 index 00000000..1add0c71 --- /dev/null +++ b/flo_ai/tests/test_tool_config.py @@ -0,0 +1,126 @@ +from unittest.mock import AsyncMock +from flo_ai.tool.base_tool import Tool +from flo_ai.tool.tool_config import ToolConfig, create_tool_config + + +class TestToolConfig: + """Test cases for ToolConfig functionality.""" + + def test_tool_config_creation(self): + """Test creating a tool configuration.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + }, + ) + + # Create tool config with pre-filled parameters + tool_config = ToolConfig( + tool=base_tool, + pre_filled_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.name_override == 'custom_name' + assert tool_config.description_override == 'Custom description' + assert tool_config.is_partial() is True + + def test_tool_config_without_prefilled_params(self): + """Test creating a tool configuration without pre-filled parameters.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'param1': {'type': 'string', 'description': 'Param 1', 'required': True} + }, + ) + + tool_config = ToolConfig(tool=base_tool) + assert tool_config.tool == base_tool + assert tool_config.pre_filled_params == {} + assert tool_config.is_partial() is False + + def test_tool_config_to_tool_conversion(self): + """Test converting tool configuration to tool.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + }, + ) + + # Test without pre-filled params (should return original tool) + tool_config = ToolConfig(tool=base_tool) + converted_tool = tool_config.to_tool() + assert converted_tool == base_tool + + # Test with pre-filled params (should return partial tool) + tool_config_with_params = ToolConfig( + tool=base_tool, pre_filled_params={'datasource_id': 'ds_123'} + ) + converted_tool = tool_config_with_params.to_tool() + assert converted_tool.name == 'test_tool_partial' + assert 'pre-configured parameters' in converted_tool.description + + def test_create_tool_config_helper_function(self): + """Test the create_tool_config helper function.""" + mock_function = AsyncMock(return_value='test_result') + base_tool = Tool( + name='test_tool', + description='A test tool', + function=mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + }, + ) + + tool_config = create_tool_config( + base_tool, datasource_id='ds_123', project_id='my-project' + ) + + assert isinstance(tool_config, ToolConfig) + assert tool_config.tool == base_tool + assert tool_config.pre_filled_params == { + 'datasource_id': 'ds_123', + 'project_id': 'my-project', + } + assert tool_config.is_partial() is True diff --git a/flo_ai/tests/test_yaml_tool_config.py b/flo_ai/tests/test_yaml_tool_config.py new file mode 100644 index 00000000..f9c2615b --- /dev/null +++ b/flo_ai/tests/test_yaml_tool_config.py @@ -0,0 +1,280 @@ +import pytest +from unittest.mock import Mock, AsyncMock +from flo_ai.tool.base_tool import Tool +from flo_ai.builder.agent_builder import AgentBuilder +from flo_ai.llm import OpenAI + + +class TestYamlToolConfig: + """Test cases for YAML tool configuration functionality.""" + + def setup_method(self): + """Set up test fixtures.""" + self.mock_function = AsyncMock(return_value='test_result') + self.base_tool = Tool( + name='test_tool', + description='A test tool', + function=self.mock_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'Query string', + 'required': True, + }, + 'datasource_id': { + 'type': 'string', + 'description': 'Data source ID', + 'required': True, + }, + 'project_id': { + 'type': 'string', + 'description': 'Project ID', + 'required': True, + }, + }, + ) + + self.tool_registry = { + 'test_tool': self.base_tool, + 'another_tool': self.base_tool, + } + + def test_simple_tool_reference(self): + """Test YAML with simple tool references.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - "test_tool" + - "another_tool" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + agent = AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + assert agent._name == 'Test Agent' + assert agent._system_prompt == 'You are a test agent' + assert len(agent._tools) == 2 + assert agent._tools[0] == self.base_tool + assert agent._tools[1] == self.base_tool + + def test_tool_configuration_with_prefilled_params(self): + """Test YAML with tool configuration and pre-filled parameters.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - name: "test_tool" + pre_filled_params: + datasource_id: "ds_123" + project_id: "my-project" + name_override: "custom_tool_name" + description_override: "Custom tool description" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + agent = AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + assert agent._name == 'Test Agent' + assert len(agent._tools) == 1 + + # Check that the tool is a partial tool with custom name and description + tool = agent._tools[0] + assert tool.name == 'custom_tool_name' + assert tool.description == 'Custom tool description' + + # Check that pre-filled parameters are hidden from AI + ai_visible_params = tool.parameters + assert 'query' in ai_visible_params + assert 'datasource_id' not in ai_visible_params + assert 'project_id' not in ai_visible_params + + def test_mixed_tool_types(self): + """Test YAML with mixed tool types (simple references and configurations).""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - "test_tool" # Simple reference + - name: "another_tool" + pre_filled_params: + datasource_id: "ds_456" + project_id: "another-project" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + agent = AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + assert len(agent._tools) == 2 + + # First tool should be the original tool + assert agent._tools[0] == self.base_tool + + # Second tool should be a partial tool + assert 'partial' in agent._tools[1].name + + def test_tool_not_found_in_registry(self): + """Test error when tool is not found in registry.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - "nonexistent_tool" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + with pytest.raises( + ValueError, match="Tool 'nonexistent_tool' not found in tool registry" + ): + AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + def test_tool_configuration_missing_name(self): + """Test error when tool configuration is missing name field.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - pre_filled_params: + datasource_id: "ds_123" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + with pytest.raises( + ValueError, match="Tool configuration must have a 'name' field" + ): + AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + def test_invalid_tool_configuration_type(self): + """Test error when tool configuration has invalid type.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - 123 # Invalid type +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + with pytest.raises(ValueError, match='Invalid tool configuration type'): + AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + def test_no_tool_registry_provided(self): + """Test error when no tool registry is provided but tools are referenced.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - "test_tool" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + with pytest.raises( + ValueError, match="Tool 'test_tool' not found in tool registry" + ): + AgentBuilder.from_yaml(yaml_str, tool_registry=None, base_llm=mock_llm) + + def test_tool_configuration_without_prefilled_params(self): + """Test tool configuration without pre-filled parameters.""" + yaml_str = """ +agent: + name: "Test Agent" + job: "You are a test agent" + model: + provider: "openai" + name: "gpt-4" + tools: + - name: "test_tool" + name_override: "custom_tool_name" + description_override: "Custom tool description" +""" + + # Mock LLM to avoid API key requirement + mock_llm = Mock(spec=OpenAI) + + agent = AgentBuilder.from_yaml( + yaml_str, tool_registry=self.tool_registry, base_llm=mock_llm + ) + + assert len(agent._tools) == 1 + tool = agent._tools[0] + + # Should be a partial tool with custom name and description + assert tool.name == 'custom_tool_name' + assert tool.description == 'Custom tool description' + # Should be a partial tool because we have name_override and description_override + assert hasattr(tool, 'base_tool') # Partial tool has base_tool attribute + + def test_process_yaml_tools_method(self): + """Test the _process_yaml_tools method directly.""" + tools_config = [ + 'test_tool', # Simple reference + { + 'name': 'another_tool', + 'pre_filled_params': {'datasource_id': 'ds_123'}, + 'name_override': 'custom_name', + }, + ] + + processed_tools = AgentBuilder._process_yaml_tools( + tools_config, self.tool_registry + ) + + assert len(processed_tools) == 2 + assert processed_tools[0] == self.base_tool # Simple reference + assert ( + processed_tools[1].name == 'custom_name' + ) # Configuration with name_override + assert hasattr(processed_tools[1], 'base_tool') # Should be a partial tool