mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-06-09 06:45:13 +02:00
Feature/tool group (#484)
* Tech spec for tool group * Partial tool group implementation * Tool group tests
This commit is contained in:
cybermaggedon
GitHub
parent
672e358b2f
commit
e74eb5d1ff
9 changed files with 1304 additions and 6 deletions
491
docs/tech-specs/tool-group.md
Normal file
491
docs/tech-specs/tool-group.md
Normal file
|
|
@ -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.
|
||||
|
|
@ -82,8 +82,8 @@ def sample_message_data():
|
|||
},
|
||||
"AgentRequest": {
|
||||
"question": "What is machine learning?",
|
||||
"plan": "",
|
||||
"state": "",
|
||||
"group": [],
|
||||
"history": []
|
||||
},
|
||||
"AgentResponse": {
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
267
tests/integration/test_tool_group_integration.py
Normal file
267
tests/integration/test_tool_group_integration.py
Normal file
|
|
@ -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)
|
||||
321
tests/unit/test_agent/test_tool_filter.py
Normal file
321
tests/unit/test_agent/test_tool_filter.py
Normal file
|
|
@ -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'
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
165
trustgraph-flow/trustgraph/agent/tool_filter.py
Normal file
165
trustgraph-flow/trustgraph/agent/tool_filter.py
Normal file
|
|
@ -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")
|
||||
Loading…
Add table
Add a link
Reference in a new issue