diff --git a/tests/integration/test_agent_structured_query_integration.py b/tests/integration/test_agent_structured_query_integration.py index 2cd646c4..2d556f93 100644 --- a/tests/integration/test_agent_structured_query_integration.py +++ b/tests/integration/test_agent_structured_query_integration.py @@ -120,7 +120,8 @@ Args: { # Verify structured query was called mock_structured_client.structured_query.assert_called_once() call_args = mock_structured_client.structured_query.call_args - question_arg = call_args[0][0] # positional argument + # Check keyword arguments + question_arg = call_args.kwargs.get("question") or call_args[1].get("question") assert "customers" in question_arg.lower() assert "new york" in question_arg.lower() @@ -202,8 +203,9 @@ Args: { # Agent should handle the error gracefully assert any(isinstance(resp, AgentResponse) for resp in responses) # The tool should have returned an error response that contains error info - structured_query_call_args = mock_structured_client.structured_query.call_args[0] - assert "table" in structured_query_call_args[0].lower() or "exist" in structured_query_call_args[0].lower() + call_args = mock_structured_client.structured_query.call_args + question_arg = call_args.kwargs.get("question") or call_args[1].get("question") + assert "table" in question_arg.lower() or "exist" in question_arg.lower() @pytest.mark.asyncio async def test_agent_multi_step_structured_query_reasoning(self, agent_processor, structured_query_tool_config): @@ -278,8 +280,9 @@ Args: { assert any(isinstance(resp, AgentResponse) for resp in responses) # Verify the structured query was called with customer-related question - call_args = mock_structured_client.structured_query.call_args[0] - assert "california" in call_args[0].lower() + call_args = mock_structured_client.structured_query.call_args + question_arg = call_args.kwargs.get("question") or call_args[1].get("question") + assert "california" in question_arg.lower() @pytest.mark.asyncio async def test_agent_structured_query_with_collection_parameter(self, agent_processor): @@ -367,8 +370,9 @@ Args: { assert any(isinstance(resp, AgentResponse) for resp in responses) # Check the query was about sales/transactions - call_args = mock_structured_client.structured_query.call_args[0] - assert "sales" in call_args[0].lower() or "transactions" in call_args[0].lower() + call_args = mock_structured_client.structured_query.call_args + question_arg = call_args.kwargs.get("question") or call_args[1].get("question") + assert "sales" in question_arg.lower() or "transactions" in question_arg.lower() @pytest.mark.asyncio async def test_agent_structured_query_tool_argument_validation(self, agent_processor, structured_query_tool_config): @@ -473,5 +477,6 @@ Args: { assert any(isinstance(resp, AgentResponse) for resp in responses) # Check that the query was about customer information - call_args = mock_structured_client.structured_query.call_args[0] - assert "customer" in call_args[0].lower() \ No newline at end of file + call_args = mock_structured_client.structured_query.call_args + question_arg = call_args.kwargs.get("question") or call_args[1].get("question") + assert "customer" in question_arg.lower() \ No newline at end of file diff --git a/tests/integration/test_structured_query_integration.py b/tests/integration/test_structured_query_integration.py index 2e836b33..cf8037d0 100644 --- a/tests/integration/test_structured_query_integration.py +++ b/tests/integration/test_structured_query_integration.py @@ -41,7 +41,9 @@ class TestStructuredQueryServiceIntegration: """Test complete structured query processing pipeline""" # Arrange - Create realistic query request request = StructuredQueryRequest( - question="Show me all customers from California who have made purchases over $500" + question="Show me all customers from California who have made purchases over $500", + user="trustgraph", + collection="default" ) msg = MagicMock() diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 63c16d30..06bf7610 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -148,7 +148,8 @@ class Processor(AgentService): elif impl_id == "structured-query": impl = functools.partial( StructuredQueryImpl, - collection=data.get("collection") + collection=data.get("collection"), + user=None # User will be provided dynamically via context ) arguments = StructuredQueryImpl.get_arguments() else: diff --git a/trustgraph-flow/trustgraph/agent/tool_filter.py b/trustgraph-flow/trustgraph/agent/tool_filter.py index 0d66b990..f27d947c 100644 --- a/trustgraph-flow/trustgraph/agent/tool_filter.py +++ b/trustgraph-flow/trustgraph/agent/tool_filter.py @@ -29,9 +29,9 @@ def filter_tools_by_group_and_state( """ # Apply defaults as specified in tech spec - if requested_groups is None: + if requested_groups is None or requested_groups == []: requested_groups = ["default"] - if current_state is None: + if current_state is None or current_state == "": current_state = "undefined" logger.info(f"Filtering tools with groups={requested_groups}, state={current_state}")