diff --git a/docs/tech-specs/tool-group.md b/docs/tech-specs/tool-group.md new file mode 100644 index 00000000..e4816de5 --- /dev/null +++ b/docs/tech-specs/tool-group.md @@ -0,0 +1,491 @@ +# TrustGraph Tool Group System +## Technical Specification v1.0 + +### Executive Summary + +This specification defines a tool grouping system for TrustGraph agents that allows fine-grained control over which tools are available for specific requests. The system introduces group-based tool filtering through configuration and request-level specification, enabling better security boundaries, resource management, and functional partitioning of agent capabilities. + +### 1. Overview + +#### 1.1 Problem Statement + +Currently, TrustGraph agents have access to all configured tools regardless of request context or security requirements. This creates several challenges: + +- **Security Risk**: Sensitive tools (e.g., data modification) are available even for read-only queries +- **Resource Waste**: Complex tools are loaded even when simple queries don't require them +- **Functional Confusion**: Agents may select inappropriate tools when simpler alternatives exist +- **Multi-tenant Isolation**: Different user groups need access to different tool sets + +#### 1.2 Solution Overview + +The tool group system introduces: + +1. **Group Classification**: Tools are tagged with group memberships during configuration +2. **Request-level Filtering**: AgentRequest specifies which tool groups are permitted +3. **Runtime Enforcement**: Agents only have access to tools matching the requested groups +4. **Flexible Grouping**: Tools can belong to multiple groups for complex scenarios + +### 2. Schema Changes + +#### 2.1 Tool Configuration Schema Enhancement + +The existing tool configuration is enhanced with a `group` field: + +**Before:** +```json +{ + "name": "knowledge-query", + "type": "knowledge-query", + "description": "Query the knowledge graph" +} +``` + +**After:** +```json +{ + "name": "knowledge-query", + "type": "knowledge-query", + "description": "Query the knowledge graph", + "group": ["read-only", "knowledge", "basic"] +} +``` + +**Group Field Specification:** +- `group`: Array(String) - List of groups this tool belongs to +- **Optional**: Tools without group field belong to "default" group +- **Multi-membership**: Tools can belong to multiple groups +- **Case-sensitive**: Group names are exact string matches + +#### 2.1.2 Tool State Transition Enhancement + +Tools can optionally specify state transitions and state-based availability: + +```json +{ + "name": "knowledge-query", + "type": "knowledge-query", + "description": "Query the knowledge graph", + "group": ["read-only", "knowledge", "basic"], + "state": "analysis", + "available_in_states": ["undefined", "research"] +} +``` + +**State Field Specification:** +- `state`: String - **Optional** - State to transition to after successful tool execution +- `available_in_states`: Array(String) - **Optional** - States in which this tool is available +- **Default behavior**: Tools without `available_in_states` are available in all states +- **State transition**: Only occurs after successful tool execution + +#### 2.2 AgentRequest Schema Enhancement + +The `AgentRequest` schema in `trustgraph-base/trustgraph/schema/services/agent.py` is enhanced: + +**Current AgentRequest:** +- `question`: String - User query +- `plan`: String - Execution plan (can be removed) +- `state`: String - Agent state +- `history`: Array(AgentStep) - Execution history + +**Enhanced AgentRequest:** +- `question`: String - User query +- `state`: String - Agent execution state (now actively used for tool filtering) +- `history`: Array(AgentStep) - Execution history +- `group`: Array(String) - **NEW** - Tool groups allowed for this request + +**Schema Changes:** +- **Removed**: `plan` field is no longer needed and can be removed (was originally intended for tool specification) +- **Added**: `group` field for tool group specification +- **Enhanced**: `state` field now controls tool availability during execution + +**Field Behaviors:** + +**Group Field:** +- **Optional**: If not specified, defaults to ["default"] +- **Intersection**: Only tools matching at least one specified group are available +- **Empty array**: No tools available (agent can only use internal reasoning) +- **Wildcard**: Special group "*" grants access to all tools + +**State Field:** +- **Optional**: If not specified, defaults to "undefined" +- **State-based filtering**: Only tools available in current state are eligible +- **Default state**: "undefined" state allows all tools (subject to group filtering) +- **State transitions**: Tools can change state after successful execution + +### 3. Custom Group Examples + +Organizations can define domain-specific groups: + +```json +{ + "financial-tools": ["stock-query", "portfolio-analysis"], + "medical-tools": ["diagnosis-assist", "drug-interaction"], + "legal-tools": ["contract-analysis", "case-search"] +} +``` + +### 4. Implementation Details + +#### 4.1 Tool Loading and Filtering + +**Configuration Phase:** +1. All tools are loaded from configuration with their group assignments +2. Tools without explicit groups are assigned to "default" group +3. Group membership is validated and stored in tool registry + +**Request Processing Phase:** +1. AgentRequest arrives with optional group specification +2. Agent filters available tools based on group intersection +3. Only matching tools are passed to agent execution context +4. Agent operates with filtered tool set throughout request lifecycle + +#### 4.2 Tool Filtering Logic + +**Combined Group and State Filtering:** + +``` +For each configured tool: + tool_groups = tool.group || ["default"] + tool_states = tool.available_in_states || ["*"] // Available in all states + +For each request: + requested_groups = request.group || ["default"] + current_state = request.state || "undefined" + +Tool is available if: + // Group filtering + (intersection(tool_groups, requested_groups) is not empty OR "*" in requested_groups) + AND + // State filtering + (current_state in tool_states OR "*" in tool_states) +``` + +**State Transition Logic:** + +``` +After successful tool execution: + if tool.state is defined: + next_request.state = tool.state + else: + next_request.state = current_request.state // No change +``` + +#### 4.3 Agent Integration Points + +**ReAct Agent:** +- Tool filtering occurs in agent_manager.py during tool registry creation +- Available tools list is filtered by both group and state before plan generation +- State transitions update AgentRequest.state field after successful tool execution +- Next iteration uses updated state for tool filtering + +**Confidence-Based Agent:** +- Tool filtering occurs in planner.py during plan generation +- ExecutionStep validation ensures only group+state eligible tools are used +- Flow controller enforces tool availability at runtime +- State transitions managed by Flow Controller between steps + +### 5. Configuration Examples + +#### 5.1 Tool Configuration with Groups and States + +```yaml +tool: + knowledge-query: + type: knowledge-query + name: "Knowledge Graph Query" + description: "Query the knowledge graph for entities and relationships" + group: ["read-only", "knowledge", "basic"] + state: "analysis" + available_in_states: ["undefined", "research"] + + graph-update: + type: graph-update + name: "Graph Update" + description: "Add or modify entities in the knowledge graph" + group: ["write", "knowledge", "admin"] + available_in_states: ["analysis", "modification"] + + text-completion: + type: text-completion + name: "Text Completion" + description: "Generate text using language models" + group: ["read-only", "text", "basic"] + state: "undefined" + # No available_in_states = available in all states + + complex-analysis: + type: mcp-tool + name: "Complex Analysis Tool" + description: "Perform complex data analysis" + group: ["advanced", "compute", "expensive"] + state: "results" + available_in_states: ["analysis"] + mcp_tool_id: "analysis-server" + + reset-workflow: + type: mcp-tool + name: "Reset Workflow" + description: "Reset to initial state" + group: ["admin"] + state: "undefined" + available_in_states: ["analysis", "results"] +``` + +#### 5.2 Request Examples with State Workflows + +**Initial Research Request:** +```json +{ + "question": "What entities are connected to Company X?", + "group": ["read-only", "knowledge"], + "state": "undefined" +} +``` +*Available tools: knowledge-query, text-completion* +*After knowledge-query: state → "analysis"* + +**Analysis Phase:** +```json +{ + "question": "Continue analysis based on previous results", + "group": ["advanced", "compute", "write"], + "state": "analysis" +} +``` +*Available tools: complex-analysis, graph-update, reset-workflow* +*After complex-analysis: state → "results"* + +**Results Phase:** +```json +{ + "question": "What should I do with these results?", + "group": ["admin"], + "state": "results" +} +``` +*Available tools: reset-workflow only* +*After reset-workflow: state → "undefined"* + +**Workflow Example - Complete Flow:** +1. **Start (undefined)**: Use knowledge-query → transitions to "analysis" +2. **Analysis state**: Use complex-analysis → transitions to "results" +3. **Results state**: Use reset-workflow → transitions back to "undefined" +4. **Back to start**: All initial tools available again + +### 6. Security Considerations + +#### 6.1 Access Control Integration + +**Gateway-Level Filtering:** +- Gateway can enforce group restrictions based on user permissions +- Prevent elevation of privileges through request manipulation +- Audit trail includes requested and granted tool groups + +**Example Gateway Logic:** +``` +user_permissions = get_user_permissions(request.user_id) +allowed_groups = user_permissions.tool_groups +requested_groups = request.group + +# Validate request doesn't exceed permissions +if not is_subset(requested_groups, allowed_groups): + reject_request("Insufficient permissions for requested tool groups") +``` + +#### 6.2 Audit and Monitoring + +**Enhanced Audit Trail:** +- Log requested tool groups and initial state per request +- Track state transitions and tool usage by group membership +- Monitor unauthorized group access attempts and invalid state transitions +- Alert on unusual group usage patterns or suspicious state workflows + +### 7. Migration Strategy + +#### 7.1 Backward Compatibility + +**Phase 1: Additive Changes** +- Add optional `group` field to tool configurations +- Add optional `group` field to AgentRequest schema +- Default behavior: All existing tools belong to "default" group +- Existing requests without group field use "default" group + +**Existing Behavior Preserved:** +- Tools without group configuration continue to work (default group) +- Tools without state configuration are available in all states +- Requests without group specification access all tools (default group) +- Requests without state specification use "undefined" state (all tools available) +- No breaking changes to existing deployments + +### 8. Monitoring and Observability + +#### 8.1 New Metrics + +**Tool Group Usage:** +- `agent_tool_group_requests_total` - Counter of requests by group +- `agent_tool_group_availability` - Gauge of tools available per group +- `agent_filtered_tools_count` - Histogram of tool count after group+state filtering + +**State Workflow Metrics:** +- `agent_state_transitions_total` - Counter of state transitions by tool +- `agent_workflow_duration_seconds` - Histogram of time spent in each state +- `agent_state_availability` - Gauge of tools available per state + +**Security Metrics:** +- `agent_group_access_denied_total` - Counter of unauthorized group access +- `agent_invalid_state_transition_total` - Counter of invalid state transitions +- `agent_privilege_escalation_attempts_total` - Counter of suspicious requests + +#### 8.2 Logging Enhancements + +**Request Logging:** +```json +{ + "request_id": "req-123", + "requested_groups": ["read-only", "knowledge"], + "initial_state": "undefined", + "state_transitions": [ + {"tool": "knowledge-query", "from": "undefined", "to": "analysis", "timestamp": "2024-01-01T10:00:01Z"} + ], + "available_tools": ["knowledge-query", "text-completion"], + "filtered_by_group": ["graph-update", "admin-tool"], + "filtered_by_state": [], + "execution_time": "1.2s" +} +``` + +### 9. Testing Strategy + +#### 9.1 Unit Tests + +**Tool Filtering Logic:** +- Test group intersection calculations +- Test state-based filtering logic +- Verify default group and state assignment +- Test wildcard group behavior +- Validate empty group handling +- Test combined group+state filtering scenarios + +**Configuration Validation:** +- Test tool loading with various group and state configurations +- Verify schema validation for invalid group and state specifications +- Test backward compatibility with existing configurations +- Validate state transition definitions and cycles + +#### 9.2 Integration Tests + +**Agent Behavior:** +- Verify agents only see group+state filtered tools +- Test request execution with various group combinations +- Test state transitions during agent execution +- Validate error handling when no tools are available +- Test workflow progression through multiple states + +**Security Testing:** +- Test privilege escalation prevention +- Verify audit trail accuracy +- Test gateway integration with user permissions + +#### 9.3 End-to-End Scenarios + +**Multi-tenant Usage with State Workflows:** +``` +Scenario: Different users with different tool access and workflow states +Given: User A has "read-only" permissions, state "undefined" + And: User B has "write" permissions, state "analysis" +When: Both request knowledge operations +Then: User A gets read-only tools available in "undefined" state + And: User B gets write tools available in "analysis" state + And: State transitions are tracked per user session + And: All usage and transitions are properly audited +``` + +**Workflow State Progression:** +``` +Scenario: Complete workflow execution +Given: Request with groups ["knowledge", "compute"] and state "undefined" +When: Agent executes knowledge-query tool (transitions to "analysis") + And: Agent executes complex-analysis tool (transitions to "results") + And: Agent executes reset-workflow tool (transitions to "undefined") +Then: Each step has correctly filtered available tools + And: State transitions are logged with timestamps + And: Final state allows initial workflow to repeat +``` + +### 10. Performance Considerations + +#### 10.1 Tool Loading Impact + +**Configuration Loading:** +- Group and state metadata loaded once at startup +- Minimal memory overhead per tool (additional fields) +- No impact on tool initialization time + +**Request Processing:** +- Combined group+state filtering occurs once per request +- O(n) complexity where n = number of configured tools +- State transitions add minimal overhead (string assignment) +- Negligible impact for typical tool counts (< 100) + +#### 10.2 Optimization Strategies + +**Pre-computed Tool Sets:** +- Cache tool sets by group+state combination +- Avoid repeated filtering for common group/state patterns +- Memory vs computation tradeoff for frequently used combinations + +**Lazy Loading:** +- Load tool implementations only when needed +- Reduce startup time for deployments with many tools +- Dynamic tool registration based on group requirements + +### 11. Future Enhancements + +#### 11.1 Dynamic Group Assignment + +**Context-Aware Grouping:** +- Assign tools to groups based on request context +- Time-based group availability (business hours only) +- Load-based group restrictions (expensive tools during low usage) + +#### 11.2 Group Hierarchies + +**Nested Group Structure:** +```json +{ + "knowledge": { + "read": ["knowledge-query", "entity-search"], + "write": ["graph-update", "entity-create"] + } +} +``` + +#### 11.3 Tool Recommendations + +**Group-Based Suggestions:** +- Suggest optimal tool groups for request types +- Learn from usage patterns to improve recommendations +- Provide fallback groups when preferred tools are unavailable + +### 12. Open Questions + +1. **Group Validation**: Should invalid group names in requests cause hard failures or warnings? + +2. **Group Discovery**: Should the system provide an API to list available groups and their tools? + +3. **Dynamic Groups**: Should groups be configurable at runtime or only at startup? + +4. **Group Inheritance**: Should tools inherit groups from their parent categories or implementations? + +5. **Performance Monitoring**: What additional metrics are needed to track group-based tool usage effectively? + +### 13. Conclusion + +The tool group system provides: + +- **Security**: Fine-grained access control over agent capabilities +- **Performance**: Reduced tool loading and selection overhead +- **Flexibility**: Multi-dimensional tool classification +- **Compatibility**: Seamless integration with existing agent architectures + +This system enables TrustGraph deployments to better manage tool access, improve security boundaries, and optimize resource usage while maintaining full backward compatibility with existing configurations and requests. diff --git a/tests/contract/conftest.py b/tests/contract/conftest.py index 5c5b82cb..3d184d3d 100644 --- a/tests/contract/conftest.py +++ b/tests/contract/conftest.py @@ -82,8 +82,8 @@ def sample_message_data(): }, "AgentRequest": { "question": "What is machine learning?", - "plan": "", "state": "", + "group": [], "history": [] }, "AgentResponse": { diff --git a/tests/contract/test_message_contracts.py b/tests/contract/test_message_contracts.py index 861e5368..e4779d8b 100644 --- a/tests/contract/test_message_contracts.py +++ b/tests/contract/test_message_contracts.py @@ -198,8 +198,8 @@ class TestAgentMessageContracts: # Test required fields request = AgentRequest(**request_data) assert hasattr(request, 'question') - assert hasattr(request, 'plan') assert hasattr(request, 'state') + assert hasattr(request, 'group') assert hasattr(request, 'history') def test_agent_response_schema_contract(self, sample_message_data): diff --git a/tests/integration/test_tool_group_integration.py b/tests/integration/test_tool_group_integration.py new file mode 100644 index 00000000..2c01cb61 --- /dev/null +++ b/tests/integration/test_tool_group_integration.py @@ -0,0 +1,267 @@ +""" +Integration tests for the tool group system. + +Tests the complete workflow of tool filtering and execution logic. +""" + +import pytest +import json +import sys +import os +from unittest.mock import Mock, AsyncMock, patch + +# Add trustgraph paths for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'trustgraph-base')) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'trustgraph-flow')) + +from trustgraph.agent.tool_filter import filter_tools_by_group_and_state, get_next_state, validate_tool_config + + +@pytest.fixture +def sample_tools(): + """Sample tools with different groups and states for testing.""" + return { + 'knowledge_query': Mock(config={ + 'group': ['read-only', 'knowledge', 'basic'], + 'state': 'analysis', + 'applicable-states': ['undefined', 'research'] + }), + 'graph_update': Mock(config={ + 'group': ['write', 'knowledge', 'admin'], + 'applicable-states': ['analysis', 'modification'] + }), + 'text_completion': Mock(config={ + 'group': ['read-only', 'text', 'basic'], + 'state': 'undefined' + # No applicable-states = available in all states + }), + 'complex_analysis': Mock(config={ + 'group': ['advanced', 'compute', 'expensive'], + 'state': 'results', + 'applicable-states': ['analysis'] + }) + } + + +class TestToolGroupFiltering: + """Test tool group filtering integration scenarios.""" + + def test_basic_group_filtering(self, sample_tools): + """Test that filtering only returns tools matching requested groups.""" + + # Filter for read-only and knowledge tools + filtered = filter_tools_by_group_and_state( + sample_tools, + ['read-only', 'knowledge'], + 'undefined' + ) + + # Should include tools with matching groups and correct state + assert 'knowledge_query' in filtered # Has read-only + knowledge, available in undefined + assert 'text_completion' in filtered # Has read-only, available in all states + assert 'graph_update' not in filtered # Has knowledge but no read-only + assert 'complex_analysis' not in filtered # Wrong groups and state + + def test_state_based_filtering(self, sample_tools): + """Test filtering based on current state.""" + + # Filter for analysis state with advanced tools + filtered = filter_tools_by_group_and_state( + sample_tools, + ['advanced', 'compute'], + 'analysis' + ) + + # Should only include tools available in analysis state + assert 'complex_analysis' in filtered # Available in analysis state + assert 'knowledge_query' not in filtered # Not available in analysis state + assert 'graph_update' not in filtered # Wrong group (no advanced/compute) + assert 'text_completion' not in filtered # Wrong group + + def test_state_transition_handling(self, sample_tools): + """Test state transitions after tool execution.""" + + # Get knowledge_query tool and test state transition + knowledge_tool = sample_tools['knowledge_query'] + + # Test state transition + next_state = get_next_state(knowledge_tool, 'undefined') + assert next_state == 'analysis' # knowledge_query should transition to analysis + + # Test tool with no state transition + text_tool = sample_tools['text_completion'] + next_state = get_next_state(text_tool, 'research') + assert next_state == 'undefined' # text_completion transitions to undefined + + def test_wildcard_group_access(self, sample_tools): + """Test wildcard group grants access to all tools.""" + + # Filter with wildcard group access + filtered = filter_tools_by_group_and_state( + sample_tools, + ['*'], # Wildcard access + 'undefined' + ) + + # Should include all tools that are available in undefined state + assert 'knowledge_query' in filtered # Available in undefined + assert 'text_completion' in filtered # Available in all states + assert 'graph_update' not in filtered # Not available in undefined + assert 'complex_analysis' not in filtered # Not available in undefined + + def test_no_matching_tools(self, sample_tools): + """Test behavior when no tools match the requested groups.""" + + # Filter with non-matching group + filtered = filter_tools_by_group_and_state( + sample_tools, + ['nonexistent-group'], + 'undefined' + ) + + # Should return empty dictionary + assert len(filtered) == 0 + + def test_default_group_behavior(self): + """Test default group behavior when no group is specified.""" + + # Create tools with and without explicit groups + tools = { + 'default_tool': Mock(config={}), # No group = default group + 'admin_tool': Mock(config={'group': ['admin']}) + } + + # Filter with no group specified (should default to ["default"]) + filtered = filter_tools_by_group_and_state(tools, None, 'undefined') + + # Only default_tool should be available + assert 'default_tool' in filtered + assert 'admin_tool' not in filtered + + +class TestToolConfigurationValidation: + """Test tool configuration validation with group metadata.""" + + def test_tool_config_validation_invalid(self): + """Test that invalid tool configurations are rejected.""" + + # Test invalid group field (should be list) + invalid_config = { + "name": "invalid_tool", + "description": "Invalid tool", + "type": "text-completion", + "group": "not-a-list" # Should be list + } + + # Should raise validation error + with pytest.raises(ValueError, match="'group' field must be a list"): + validate_tool_config(invalid_config) + + def test_tool_config_validation_valid(self): + """Test that valid tool configurations are accepted.""" + + valid_config = { + "name": "valid_tool", + "description": "Valid tool", + "type": "text-completion", + "group": ["read-only", "text"], + "state": "analysis", + "applicable-states": ["undefined", "research"] + } + + # Should not raise any exception + validate_tool_config(valid_config) + + def test_kebab_case_field_names(self): + """Test that kebab-case field names are properly handled.""" + + config = { + "name": "test_tool", + "group": ["basic"], + "applicable-states": ["undefined", "analysis"] # kebab-case + } + + # Should validate without error + validate_tool_config(config) + + # Create mock tool and test filtering + tool = Mock(config=config) + + # Test that kebab-case field is properly read + filtered = filter_tools_by_group_and_state( + {'test_tool': tool}, + ['basic'], + 'analysis' + ) + + assert 'test_tool' in filtered + + +class TestCompleteWorkflow: + """Test complete multi-step workflows with state transitions.""" + + def test_research_analysis_workflow(self, sample_tools): + """Test complete research -> analysis -> results workflow.""" + + # Step 1: Initial research phase (undefined state) + step1_filtered = filter_tools_by_group_and_state( + sample_tools, + ['read-only', 'knowledge'], + 'undefined' + ) + + # Should have access to knowledge_query and text_completion + assert 'knowledge_query' in step1_filtered + assert 'text_completion' in step1_filtered + assert 'complex_analysis' not in step1_filtered # Not available in undefined + + # Simulate executing knowledge_query tool + knowledge_tool = step1_filtered['knowledge_query'] + next_state = get_next_state(knowledge_tool, 'undefined') + assert next_state == 'analysis' # Transition to analysis state + + # Step 2: Analysis phase + step2_filtered = filter_tools_by_group_and_state( + sample_tools, + ['advanced', 'compute', 'text'], # Include text for text_completion + 'analysis' + ) + + # Should have access to complex_analysis and text_completion + assert 'complex_analysis' in step2_filtered + assert 'text_completion' in step2_filtered # Available in all states + assert 'knowledge_query' not in step2_filtered # Not available in analysis + + # Simulate executing complex_analysis tool + analysis_tool = step2_filtered['complex_analysis'] + final_state = get_next_state(analysis_tool, 'analysis') + assert final_state == 'results' # Transition to results state + + def test_multi_tenant_scenario(self, sample_tools): + """Test different users with different permissions.""" + + # User A: Read-only permissions in undefined state + user_a_tools = filter_tools_by_group_and_state( + sample_tools, + ['read-only'], + 'undefined' + ) + + # Should only have access to read-only tools in undefined state + assert 'knowledge_query' in user_a_tools # read-only + available in undefined + assert 'text_completion' in user_a_tools # read-only + available in all states + assert 'graph_update' not in user_a_tools # write permissions required + assert 'complex_analysis' not in user_a_tools # advanced permissions required + + # User B: Admin permissions in analysis state + user_b_tools = filter_tools_by_group_and_state( + sample_tools, + ['write', 'admin'], + 'analysis' + ) + + # Should have access to admin tools available in analysis state + assert 'graph_update' in user_b_tools # admin + available in analysis + assert 'complex_analysis' not in user_b_tools # wrong group (needs advanced/compute) + assert 'knowledge_query' not in user_b_tools # not available in analysis state + assert 'text_completion' not in user_b_tools # wrong group (no admin) \ No newline at end of file diff --git a/tests/unit/test_agent/test_tool_filter.py b/tests/unit/test_agent/test_tool_filter.py new file mode 100644 index 00000000..c7e7cf3e --- /dev/null +++ b/tests/unit/test_agent/test_tool_filter.py @@ -0,0 +1,321 @@ +""" +Unit tests for the tool filtering logic in the tool group system. +""" + +import pytest +import sys +import os +from unittest.mock import Mock + +# Add trustgraph-flow to path for imports +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', '..', 'trustgraph-flow')) + +from trustgraph.agent.tool_filter import ( + filter_tools_by_group_and_state, + get_next_state, + validate_tool_config, + _is_tool_available +) + + +class TestToolFiltering: + """Test tool filtering based on groups and states.""" + + def test_filter_tools_default_group(self): + """Tools without groups should belong to 'default' group.""" + tools = { + 'tool1': Mock(config={}), + 'tool2': Mock(config={'group': ['read-only']}) + } + + # Request default group (implicit) + filtered = filter_tools_by_group_and_state(tools, None, None) + + # Only tool1 should be available (no group = default group) + assert 'tool1' in filtered + assert 'tool2' not in filtered + + def test_filter_tools_explicit_groups(self): + """Test filtering with explicit group membership.""" + tools = { + 'read_tool': Mock(config={'group': ['read-only', 'basic']}), + 'write_tool': Mock(config={'group': ['write', 'admin']}), + 'mixed_tool': Mock(config={'group': ['read-only', 'write']}) + } + + # Request read-only tools + filtered = filter_tools_by_group_and_state(tools, ['read-only'], None) + + assert 'read_tool' in filtered + assert 'write_tool' not in filtered + assert 'mixed_tool' in filtered # Has read-only in its groups + + def test_filter_tools_multiple_requested_groups(self): + """Test filtering with multiple requested groups.""" + tools = { + 'tool1': Mock(config={'group': ['read-only']}), + 'tool2': Mock(config={'group': ['write']}), + 'tool3': Mock(config={'group': ['admin']}) + } + + # Request read-only and write tools + filtered = filter_tools_by_group_and_state(tools, ['read-only', 'write'], None) + + assert 'tool1' in filtered + assert 'tool2' in filtered + assert 'tool3' not in filtered + + def test_filter_tools_wildcard_group(self): + """Test wildcard group grants access to all tools.""" + tools = { + 'tool1': Mock(config={'group': ['read-only']}), + 'tool2': Mock(config={'group': ['admin']}), + 'tool3': Mock(config={}) # default group + } + + # Request wildcard access + filtered = filter_tools_by_group_and_state(tools, ['*'], None) + + assert len(filtered) == 3 + assert all(tool in filtered for tool in tools) + + def test_filter_tools_by_state(self): + """Test filtering based on applicable-states.""" + tools = { + 'init_tool': Mock(config={'applicable-states': ['undefined']}), + 'analysis_tool': Mock(config={'applicable-states': ['analysis']}), + 'any_state_tool': Mock(config={}) # available in all states + } + + # Filter for 'analysis' state + filtered = filter_tools_by_group_and_state(tools, ['default'], 'analysis') + + assert 'init_tool' not in filtered + assert 'analysis_tool' in filtered + assert 'any_state_tool' in filtered + + def test_filter_tools_state_wildcard(self): + """Test tools with '*' in applicable-states are always available.""" + tools = { + 'wildcard_tool': Mock(config={'applicable-states': ['*']}), + 'specific_tool': Mock(config={'applicable-states': ['research']}) + } + + # Filter for 'analysis' state + filtered = filter_tools_by_group_and_state(tools, ['default'], 'analysis') + + assert 'wildcard_tool' in filtered + assert 'specific_tool' not in filtered + + def test_filter_tools_combined_group_and_state(self): + """Test combined group and state filtering.""" + tools = { + 'valid_tool': Mock(config={ + 'group': ['read-only'], + 'applicable-states': ['analysis'] + }), + 'wrong_group': Mock(config={ + 'group': ['admin'], + 'applicable-states': ['analysis'] + }), + 'wrong_state': Mock(config={ + 'group': ['read-only'], + 'applicable-states': ['research'] + }), + 'wrong_both': Mock(config={ + 'group': ['admin'], + 'applicable-states': ['research'] + }) + } + + filtered = filter_tools_by_group_and_state( + tools, ['read-only'], 'analysis' + ) + + assert 'valid_tool' in filtered + assert 'wrong_group' not in filtered + assert 'wrong_state' not in filtered + assert 'wrong_both' not in filtered + + def test_filter_tools_empty_request_groups(self): + """Test that empty group list results in no available tools.""" + tools = { + 'tool1': Mock(config={'group': ['read-only']}), + 'tool2': Mock(config={}) + } + + filtered = filter_tools_by_group_and_state(tools, [], None) + + assert len(filtered) == 0 + + +class TestStateTransitions: + """Test state transition logic.""" + + def test_get_next_state_with_transition(self): + """Test state transition when tool defines next state.""" + tool = Mock(config={'state': 'analysis'}) + + next_state = get_next_state(tool, 'undefined') + + assert next_state == 'analysis' + + def test_get_next_state_no_transition(self): + """Test no state change when tool doesn't define next state.""" + tool = Mock(config={}) + + next_state = get_next_state(tool, 'research') + + assert next_state == 'research' + + def test_get_next_state_empty_config(self): + """Test with tool that has no config.""" + tool = Mock(config=None) + tool.config = None + + next_state = get_next_state(tool, 'initial') + + assert next_state == 'initial' + + +class TestConfigValidation: + """Test tool configuration validation.""" + + def test_validate_valid_config(self): + """Test validation of valid configuration.""" + config = { + 'group': ['read-only', 'basic'], + 'state': 'analysis', + 'applicable-states': ['undefined', 'research'] + } + + # Should not raise an exception + validate_tool_config(config) + + def test_validate_group_not_list(self): + """Test validation fails when group is not a list.""" + config = {'group': 'read-only'} # Should be list + + with pytest.raises(ValueError, match="'group' field must be a list"): + validate_tool_config(config) + + def test_validate_group_non_string_elements(self): + """Test validation fails when group contains non-strings.""" + config = {'group': ['read-only', 123]} # 123 is not string + + with pytest.raises(ValueError, match="All group names must be strings"): + validate_tool_config(config) + + def test_validate_state_not_string(self): + """Test validation fails when state is not a string.""" + config = {'state': 123} # Should be string + + with pytest.raises(ValueError, match="'state' field must be a string"): + validate_tool_config(config) + + def test_validate_applicable_states_not_list(self): + """Test validation fails when applicable-states is not a list.""" + config = {'applicable-states': 'undefined'} # Should be list + + with pytest.raises(ValueError, match="'applicable-states' field must be a list"): + validate_tool_config(config) + + def test_validate_applicable_states_non_string_elements(self): + """Test validation fails when applicable-states contains non-strings.""" + config = {'applicable-states': ['undefined', 123]} + + with pytest.raises(ValueError, match="All state names must be strings"): + validate_tool_config(config) + + def test_validate_minimal_config(self): + """Test validation of minimal valid configuration.""" + config = {'name': 'test', 'description': 'Test tool'} + + # Should not raise an exception + validate_tool_config(config) + + +class TestToolAvailability: + """Test the internal _is_tool_available function.""" + + def test_tool_available_default_groups_and_states(self): + """Test tool with default groups and states.""" + tool = Mock(config={}) + + # Default group request, default state + assert _is_tool_available(tool, ['default'], 'undefined') + + # Non-default group request should fail + assert not _is_tool_available(tool, ['admin'], 'undefined') + + def test_tool_available_string_group_conversion(self): + """Test that single group string is converted to list.""" + tool = Mock(config={'group': 'read-only'}) # Single string + + assert _is_tool_available(tool, ['read-only'], 'undefined') + assert not _is_tool_available(tool, ['admin'], 'undefined') + + def test_tool_available_string_state_conversion(self): + """Test that single state string is converted to list.""" + tool = Mock(config={'applicable-states': 'analysis'}) # Single string + + assert _is_tool_available(tool, ['default'], 'analysis') + assert not _is_tool_available(tool, ['default'], 'research') + + def test_tool_no_config_attribute(self): + """Test tool without config attribute.""" + tool = Mock() + del tool.config # Remove config attribute + + # Should use defaults and be available for default group/state + assert _is_tool_available(tool, ['default'], 'undefined') + assert not _is_tool_available(tool, ['admin'], 'undefined') + + +class TestWorkflowScenarios: + """Test complete workflow scenarios from the tech spec.""" + + def test_research_to_analysis_workflow(self): + """Test the research -> analysis workflow from tech spec.""" + tools = { + 'knowledge_query': Mock(config={ + 'group': ['read-only', 'knowledge'], + 'state': 'analysis', + 'applicable-states': ['undefined', 'research'] + }), + 'complex_analysis': Mock(config={ + 'group': ['advanced', 'compute'], + 'state': 'results', + 'applicable-states': ['analysis'] + }), + 'text_completion': Mock(config={ + 'group': ['read-only', 'text', 'basic'] + # No applicable-states = available in all states + }) + } + + # Phase 1: Initial research (undefined state) + phase1_filtered = filter_tools_by_group_and_state( + tools, ['read-only', 'knowledge'], 'undefined' + ) + assert 'knowledge_query' in phase1_filtered + assert 'text_completion' in phase1_filtered + assert 'complex_analysis' not in phase1_filtered + + # Simulate tool execution and state transition + executed_tool = phase1_filtered['knowledge_query'] + next_state = get_next_state(executed_tool, 'undefined') + assert next_state == 'analysis' + + # Phase 2: Analysis state (include basic group for text_completion) + phase2_filtered = filter_tools_by_group_and_state( + tools, ['advanced', 'compute', 'basic'], 'analysis' + ) + assert 'knowledge_query' not in phase2_filtered # Not available in analysis + assert 'complex_analysis' in phase2_filtered + assert 'text_completion' in phase2_filtered # Always available + + # Simulate complex analysis execution + executed_tool = phase2_filtered['complex_analysis'] + final_state = get_next_state(executed_tool, 'analysis') + assert final_state == 'results' \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/agent.py b/trustgraph-base/trustgraph/schema/services/agent.py index 21d2fe1f..55f2ae0f 100644 --- a/trustgraph-base/trustgraph/schema/services/agent.py +++ b/trustgraph-base/trustgraph/schema/services/agent.py @@ -16,8 +16,8 @@ class AgentStep(Record): class AgentRequest(Record): question = String() - plan = String() state = String() + group = Array(String()) history = Array(AgentStep()) class AgentResponse(Record): diff --git a/trustgraph-cli/trustgraph/cli/set_tool.py b/trustgraph-cli/trustgraph/cli/set_tool.py index e39dfad7..2b7e2093 100644 --- a/trustgraph-cli/trustgraph/cli/set_tool.py +++ b/trustgraph-cli/trustgraph/cli/set_tool.py @@ -63,6 +63,9 @@ def set_tool( collection : str, template : str, arguments : List[Argument], + group : List[str], + state : str, + applicable_states : List[str], ): api = Api(url).config() @@ -93,6 +96,12 @@ def set_tool( for a in arguments ] + if group: object["group"] = group + + if state: object["state"] = state + + if applicable_states: object["applicable-states"] = applicable_states + values = api.put([ ConfigValue( type="tool", key=f"{id}", value=json.dumps(object) @@ -179,6 +188,23 @@ def main(): help=f'Tool arguments in the form: name:type:description (can specify multiple)', ) + parser.add_argument( + '--group', + nargs="*", + help=f'Tool groups (e.g., read-only, knowledge, admin)', + ) + + parser.add_argument( + '--state', + help=f'State to transition to after successful execution', + ) + + parser.add_argument( + '--applicable-states', + nargs="*", + help=f'States in which this tool is available', + ) + args = parser.parse_args() try: @@ -219,6 +245,9 @@ def main(): collection=args.collection, template=args.template, arguments=arguments, + group=args.group or [], + state=args.state, + applicable_states=getattr(args, 'applicable_states', None) or [], ) except Exception as e: diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 74b89a1e..fbdb3f96 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -18,6 +18,7 @@ from ... schema import AgentRequest, AgentResponse, AgentStep, Error from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl from . agent_manager import AgentManager +from ..tool_filter import validate_tool_config, filter_tools_by_group_and_state, get_next_state from . types import Final, Action, Tool, Argument @@ -142,6 +143,9 @@ class Processor(AgentService): f"Tool type {impl_id} not known" ) + # Validate tool configuration + validate_tool_config(data) + tools[name] = Tool( name=name, description=data.get("description"), @@ -219,9 +223,24 @@ class Processor(AgentService): await respond(r) + # Apply tool filtering based on request groups and state + filtered_tools = filter_tools_by_group_and_state( + tools=self.agent.tools, + requested_groups=getattr(request, 'group', None), + current_state=getattr(request, 'state', None) + ) + + logger.info(f"Filtered from {len(self.agent.tools)} to {len(filtered_tools)} available tools") + + # Create temporary agent with filtered tools + temp_agent = AgentManager( + tools=filtered_tools, + additional_context=self.agent.additional_context + ) + logger.debug("Call React") - act = await self.agent.react( + act = await temp_agent.react( question = request.question, history = history, think = think, @@ -255,11 +274,17 @@ class Processor(AgentService): logger.debug("Send next...") history.append(act) + + # Handle state transitions if tool execution was successful + next_state = request.state + if act.name in filtered_tools: + executed_tool = filtered_tools[act.name] + next_state = get_next_state(executed_tool, request.state or "undefined") r = AgentRequest( question=request.question, - plan=request.plan, - state=request.state, + state=next_state, + group=getattr(request, 'group', []), history=[ AgentStep( thought=h.thought, diff --git a/trustgraph-flow/trustgraph/agent/tool_filter.py b/trustgraph-flow/trustgraph/agent/tool_filter.py new file mode 100644 index 00000000..0d66b990 --- /dev/null +++ b/trustgraph-flow/trustgraph/agent/tool_filter.py @@ -0,0 +1,165 @@ +""" +Tool filtering logic for the TrustGraph tool group system. + +Provides functions to filter available tools based on group membership +and execution state as defined in the tool-group tech spec. +""" + +import logging +from typing import Dict, List, Optional, Any + +logger = logging.getLogger(__name__) + + +def filter_tools_by_group_and_state( + tools: Dict[str, Any], + requested_groups: Optional[List[str]] = None, + current_state: Optional[str] = None +) -> Dict[str, Any]: + """ + Filter tools based on group membership and execution state. + + Args: + tools: Dictionary of tool_name -> tool_object + requested_groups: List of groups requested (defaults to ["default"]) + current_state: Current execution state (defaults to "undefined") + + Returns: + Dictionary of filtered tools that match group and state criteria + """ + + # Apply defaults as specified in tech spec + if requested_groups is None: + requested_groups = ["default"] + if current_state is None: + current_state = "undefined" + + logger.info(f"Filtering tools with groups={requested_groups}, state={current_state}") + + filtered_tools = {} + + for tool_name, tool in tools.items(): + if _is_tool_available(tool, requested_groups, current_state): + filtered_tools[tool_name] = tool + else: + logger.debug(f"Tool {tool_name} filtered out") + + logger.info(f"Filtered {len(tools)} tools to {len(filtered_tools)} available tools") + return filtered_tools + + +def _is_tool_available( + tool: Any, + requested_groups: List[str], + current_state: str +) -> bool: + """ + Check if a tool is available based on group and state criteria. + + Args: + tool: Tool object with config attribute containing group/state metadata + requested_groups: List of requested groups + current_state: Current execution state + + Returns: + True if tool should be available, False otherwise + """ + + # Extract tool configuration + config = getattr(tool, 'config', {}) + + # Get tool groups (default to ["default"] if not specified) + tool_groups = config.get('group', ["default"]) + if not isinstance(tool_groups, list): + tool_groups = [tool_groups] + + # Get tool applicable states (default to all states if not specified) + applicable_states = config.get('applicable-states', ["*"]) + if not isinstance(applicable_states, list): + applicable_states = [applicable_states] + + # Apply group filtering logic from tech spec: + # Tool is available if intersection(tool_groups, requested_groups) is not empty + # OR "*" is in requested_groups (wildcard access) + group_match = ( + "*" in requested_groups or + bool(set(tool_groups) & set(requested_groups)) + ) + + # Apply state filtering logic from tech spec: + # Tool is available if current_state is in applicable_states + # OR "*" is in applicable_states (available in all states) + state_match = ( + "*" in applicable_states or + current_state in applicable_states + ) + + is_available = group_match and state_match + + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + f"Tool availability check: tool_groups={tool_groups}, " + f"requested_groups={requested_groups}, applicable_states={applicable_states}, " + f"current_state={current_state}, group_match={group_match}, " + f"state_match={state_match}, is_available={is_available}" + ) + + return is_available + + +def get_next_state(tool: Any, current_state: str) -> str: + """ + Get the next state after successful tool execution. + + Args: + tool: Tool object with config attribute + current_state: Current execution state + + Returns: + Next state, or current_state if no transition is defined + """ + config = getattr(tool, 'config', {}) + if config is None: + config = {} + next_state = config.get('state') + + if next_state: + logger.debug(f"State transition: {current_state} -> {next_state}") + return next_state + else: + logger.debug(f"No state transition defined, staying in {current_state}") + return current_state + + +def validate_tool_config(config: Dict[str, Any]) -> None: + """ + Validate tool configuration for group and state fields. + + Args: + config: Tool configuration dictionary + + Raises: + ValueError: If configuration is invalid + """ + + # Validate group field + if 'group' in config: + groups = config['group'] + if not isinstance(groups, list): + raise ValueError("Tool 'group' field must be a list of strings") + if not all(isinstance(g, str) for g in groups): + raise ValueError("All group names must be strings") + + # Validate state field + if 'state' in config: + state = config['state'] + if not isinstance(state, str): + raise ValueError("Tool 'state' field must be a string") + + # Validate applicable-states field + if 'applicable-states' in config: + states = config['applicable-states'] + if not isinstance(states, list): + raise ValueError("Tool 'applicable-states' field must be a list of strings") + if not all(isinstance(s, str) for s in states): + raise ValueError("All state names must be strings") \ No newline at end of file