diff --git a/concordia/document/interactive_document_tools.py b/concordia/document/interactive_document_tools.py index c9274c00f..5eb455c34 100644 --- a/concordia/document/interactive_document_tools.py +++ b/concordia/document/interactive_document_tools.py @@ -22,11 +22,12 @@ from collections.abc import Collection, Iterable, Sequence import json import re -from typing import Any +from typing import Any, Literal from concordia.document import document from concordia.document import interactive_document from concordia.document import tool as tool_module +from concordia.document import tool_policy from concordia.language_model import language_model import numpy as np @@ -38,6 +39,18 @@ # New tags for tool-related content TOOL_CALL_TAG = 'tool_call' TOOL_RESULT_TAG = 'tool_result' +TOOL_POLICY_TAG = 'tool_policy' +TOOL_POLICY_ALLOW_TAG = 'tool_policy_allow' +TOOL_POLICY_DENY_OBSERVED_TAG = 'tool_policy_deny_observed' +TOOL_POLICY_EDIT_OBSERVED_TAG = 'tool_policy_edit_observed' +TOOL_POLICY_DENY_ENFORCED_TAG = 'tool_policy_deny_enforced' +TOOL_POLICY_EDIT_ENFORCED_TAG = 'tool_policy_edit_enforced' +TOOL_POLICY_ERROR_OBSERVED_TAG = 'tool_policy_error_observed' +TOOL_POLICY_ERROR_ENFORCED_TAG = 'tool_policy_error_enforced' + +_VALID_ENFORCEMENT_MODES = frozenset({'observe', 'enforce'}) + +EnforcementMode = Literal['observe', 'enforce'] # Regex pattern to find JSON tool calls in LLM output # Matches {"tool": "name", "args": {...}} @@ -83,6 +96,8 @@ def __init__( rng: np.random.Generator | None = None, max_tool_calls_per_question: int = DEFAULT_MAX_TOOL_CALLS_PER_QUESTION, max_tool_result_length: int = DEFAULT_MAX_TOOL_RESULT_LENGTH, + policy: tool_policy.ToolPolicy | None = None, + enforcement_mode: EnforcementMode = 'observe', ) -> None: """Initializes the instance. @@ -95,11 +110,20 @@ def __init__( question before forcing a final answer. max_tool_result_length: Maximum character length for tool results. Results exceeding this will be truncated. + policy: Optional policy used to evaluate tool calls. + enforcement_mode: Whether policy decisions are observed or enforced. """ super().__init__(model=model, contents=contents, rng=rng) + if enforcement_mode not in _VALID_ENFORCEMENT_MODES: + raise ValueError( + 'enforcement_mode must be "observe" or "enforce", ' + f'got "{enforcement_mode}".' + ) self._tools = {t.name: t for t in tools} self._max_tool_calls = max_tool_calls_per_question self._max_result_length = max_tool_result_length + self._policy = policy + self._enforcement_mode = enforcement_mode def copy(self) -> 'InteractiveDocumentWithTools': """See base class.""" @@ -110,6 +134,8 @@ def copy(self) -> 'InteractiveDocumentWithTools': rng=self._rng, max_tool_calls_per_question=self._max_tool_calls, max_tool_result_length=self._max_result_length, + policy=self._policy, + enforcement_mode=self._enforcement_mode, ) def _format_tool_descriptions(self) -> str: @@ -188,6 +214,252 @@ def _tool_result( """Appends a tool result record to the document.""" self.append(text + end, tags=[TOOL_RESULT_TAG, *tags]) + def _tool_policy( + self, text: str, *, tags: Collection[str] = (), end: str = '' + ) -> None: + """Appends a tool policy record to the document.""" + self.append(text + end, tags=[TOOL_POLICY_TAG, *tags]) + + @staticmethod + def _merge_tags(*tag_groups: Collection[str]) -> tuple[str, ...]: + """Merges tags while preserving order and removing duplicates.""" + merged: list[str] = [] + seen: set[str] = set() + for tag_group in tag_groups: + for tag in tag_group: + if tag not in seen: + merged.append(tag) + seen.add(tag) + return tuple(merged) + + def _serialize_args(self, args: dict[str, Any]) -> str: + """Converts args to a stable string for logs.""" + try: + return json.dumps(args, sort_keys=True) + except TypeError: + return str(args) + + def _format_policy_note( + self, + *, + action: str, + reason: str = '', + edited_args: dict[str, Any] | None = None, + ) -> str: + """Formats a policy note for structured logging.""" + segments = [f'action={action}'] + if reason: + segments.append(f'reason={reason}') + if edited_args is not None: + segments.append(f'edited_args={self._serialize_args(edited_args)}') + return '[Tool Policy: ' + '; '.join(segments) + ']' + + def _evaluate_policy_decision( + self, + *, + tool_name: str, + args: dict[str, Any], + raw_response: str, + attempt_index: int, + ) -> tool_policy.PolicyDecision: + """Evaluates the configured policy for one tool call.""" + if self._policy is None: + raise RuntimeError('Policy is not configured.') + return self._policy.evaluate( + tool_policy.ToolCall( + tool_name=tool_name, + args=dict(args), + raw_response=raw_response, + attempt_index=attempt_index, + ), + self._tools, + ) + + def _coerce_policy_decision( + self, decision: Any + ) -> tool_policy.PolicyDecision: + """Validates and normalizes the policy decision shape.""" + if not isinstance(decision, tool_policy.PolicyDecision): + raise ValueError( + 'Policy must return tool_policy.PolicyDecision, got ' + f'{type(decision).__name__}.' + ) + if not isinstance(decision.action, tool_policy.PolicyAction): + raise ValueError( + 'Policy decision action must be PolicyAction, got ' + f'{type(decision.action).__name__}.' + ) + + if isinstance(decision.tags, str): + raise ValueError( + 'Policy decision tags must be an iterable of strings, not str.' + ) + try: + tags = tuple(decision.tags) + except TypeError as error: + raise ValueError( + 'Policy decision tags must be an iterable of strings.' + ) from error + if any(not isinstance(tag, str) for tag in tags): + raise ValueError('Policy decision tags must contain only strings.') + + reason = str(decision.reason) + edited_args = decision.edited_args + if ( + decision.action is tool_policy.PolicyAction.EDIT + and not isinstance(edited_args, dict) + ): + raise ValueError( + 'Policy decision action EDIT requires edited_args as dict.' + ) + return tool_policy.PolicyDecision( + action=decision.action, + reason=reason, + edited_args=edited_args, + tags=tags, + ) + + def _policy_error_outcome( + self, *, tool_name: str, args: dict[str, Any], error: Exception + ) -> tuple[dict[str, Any], str, tuple[str, ...], str]: + """Converts policy failures into observe/enforce runtime outcomes.""" + if self._enforcement_mode == 'observe': + tags = (TOOL_POLICY_ERROR_OBSERVED_TAG,) + note = self._format_policy_note( + action='error_observed', reason=str(error) + ) + return args, self._execute_tool(tool_name, args), tags, note + + tags = (TOOL_POLICY_ERROR_ENFORCED_TAG,) + note = self._format_policy_note(action='error_enforced', reason=str(error)) + result = f'Error: Tool call "{tool_name}" blocked due to policy error.' + return args, result, tags, note + + def _apply_policy_decision( + self, + *, + tool_name: str, + args: dict[str, Any], + decision: tool_policy.PolicyDecision, + ) -> tuple[dict[str, Any], str, tuple[str, ...], str]: + """Applies a validated policy decision to produce execution outcome.""" + tags = decision.tags + reason = decision.reason + + if decision.action is tool_policy.PolicyAction.ALLOW: + merged_tags = self._merge_tags((TOOL_POLICY_ALLOW_TAG,), tags) + note = self._format_policy_note(action='allow', reason=reason) + return args, self._execute_tool(tool_name, args), merged_tags, note + + if decision.action is tool_policy.PolicyAction.DENY: + if self._enforcement_mode == 'observe': + merged_tags = self._merge_tags((TOOL_POLICY_DENY_OBSERVED_TAG,), tags) + note = self._format_policy_note(action='deny_observed', reason=reason) + result = self._execute_tool(tool_name, args) + return args, result, merged_tags, note + merged_tags = self._merge_tags((TOOL_POLICY_DENY_ENFORCED_TAG,), tags) + note = self._format_policy_note(action='deny_enforced', reason=reason) + result = f'Error: Tool call "{tool_name}" denied by policy.' + return args, result, merged_tags, note + + if decision.action is tool_policy.PolicyAction.EDIT: + if self._enforcement_mode == 'observe': + merged_tags = self._merge_tags((TOOL_POLICY_EDIT_OBSERVED_TAG,), tags) + note = self._format_policy_note( + action='edit_observed', + reason=reason, + edited_args=decision.edited_args, + ) + return args, self._execute_tool(tool_name, args), merged_tags, note + + edited_args = decision.edited_args + if not isinstance(edited_args, dict): + return self._policy_error_outcome( + tool_name=tool_name, + args=args, + error=ValueError( + 'Policy decision action EDIT requires edited_args as dict.' + ), + ) + + merged_tags = self._merge_tags((TOOL_POLICY_EDIT_ENFORCED_TAG,), tags) + note = self._format_policy_note( + action='edit_enforced', reason=reason, edited_args=edited_args + ) + return ( + edited_args, + self._execute_tool(tool_name, edited_args), + merged_tags, + note, + ) + + return self._policy_error_outcome( + tool_name=tool_name, + args=args, + error=ValueError( + 'Policy decision action must be one of ALLOW, DENY, EDIT.' + ), + ) + + def _run_policy( + self, + *, + tool_name: str, + args: dict[str, Any], + raw_response: str, + attempt_index: int, + ) -> tuple[dict[str, Any], str, tuple[str, ...], str | None]: + """Evaluates policy and executes tool call according to mode.""" + if self._policy is None: + return args, self._execute_tool(tool_name, args), (), None + + try: + decision = self._evaluate_policy_decision( + tool_name=tool_name, + args=args, + raw_response=raw_response, + attempt_index=attempt_index, + ) + decision = self._coerce_policy_decision(decision) + except Exception as error: # pylint: disable=broad-exception-caught + return self._policy_error_outcome( + tool_name=tool_name, args=args, error=error + ) + + return self._apply_policy_decision( + tool_name=tool_name, args=args, decision=decision + ) + + def _record_tool_interaction( + self, + *, + raw_response: str, + tool_name: str, + original_args: dict[str, Any], + executed_args: dict[str, Any], + result: str, + policy_tags: tuple[str, ...], + policy_note: str | None, + ) -> None: + """Records tool call, policy decision, and tool result in the document.""" + self._model_response(raw_response) + self._response('\n') + + call_text = ( + f'[Tool Call: {tool_name}({self._serialize_args(original_args)})]' + ) + if executed_args != original_args: + call_text += ( + ' [Executed with args: ' + f'{self._serialize_args(executed_args)}]' + ) + self._tool_call(call_text + '\n', tags=policy_tags) + + if policy_note: + self._tool_policy(policy_note + '\n', tags=policy_tags) + + self._tool_result(f'[Tool Result: {result}]\n', tags=policy_tags) + def open_question( self, question: str, @@ -280,17 +552,21 @@ def open_question( tool_name, args = tool_call tool_calls_made += 1 - # Record the raw LLM output (contains tool call) - self._model_response(response) - self._response('\n') - - # Record the tool call - args_json = json.dumps(args) - self._tool_call(f'[Tool Call: {tool_name}({args_json})]\n') - - # Execute and record result - result = self._execute_tool(tool_name, args) - self._tool_result(f'[Tool Result: {result}]\n') + executed_args, result, policy_tags, policy_note = self._run_policy( + tool_name=tool_name, + args=args, + raw_response=response, + attempt_index=tool_calls_made, + ) + self._record_tool_interaction( + raw_response=response, + tool_name=tool_name, + original_args=args, + executed_args=executed_args, + result=result, + policy_tags=policy_tags, + policy_note=policy_note, + ) # Prepare for next iteration self._response(f'{answer_label}: ') @@ -406,17 +682,21 @@ def multiple_choice_question( tool_name, args = tool_call tool_calls_made += 1 - # Record the raw LLM output (contains tool call) - self._model_response(response) - self._response('\n') - - # Record the tool call - args_json = json.dumps(args) - self._tool_call(f'[Tool Call: {tool_name}({args_json})]\n') - - # Execute and record result - result = self._execute_tool(tool_name, args) - self._tool_result(f'[Tool Result: {result}]\n') + executed_args, result, policy_tags, policy_note = self._run_policy( + tool_name=tool_name, + args=args, + raw_response=response, + attempt_index=tool_calls_made, + ) + self._record_tool_interaction( + raw_response=response, + tool_name=tool_name, + original_args=args, + executed_args=executed_args, + result=result, + policy_tags=policy_tags, + policy_note=policy_note, + ) # Prepare for next iteration self._response('Answer: ') diff --git a/concordia/document/interactive_document_tools_test.py b/concordia/document/interactive_document_tools_test.py index a4953b46c..208cf6e6f 100644 --- a/concordia/document/interactive_document_tools_test.py +++ b/concordia/document/interactive_document_tools_test.py @@ -15,6 +15,7 @@ """Unit tests for InteractiveDocumentWithTools tool-calling functionality.""" import functools +from typing import cast from unittest import mock from absl.testing import absltest @@ -22,6 +23,7 @@ from concordia.document import document from concordia.document import interactive_document_tools from concordia.document import tool as tool_module +from concordia.document import tool_policy from concordia.language_model import language_model @@ -63,6 +65,73 @@ def execute(self, **kwargs) -> str: return self._return_value +class MockPolicy: + """Policy stub with configurable decisions.""" + + def __init__( + self, + decision: tool_policy.PolicyDecision | None = None, + *, + raise_error: bool = False, + ): + self._decision = decision or tool_policy.PolicyDecision() + self._raise_error = raise_error + self.calls: list[tool_policy.ToolCall] = [] + + def evaluate( + self, + call: tool_policy.ToolCall, + available_tools: dict[str, tool_module.Tool], + ) -> tool_policy.PolicyDecision: + del available_tools + self.calls.append(call) + if self._raise_error: + raise ValueError('policy failure') + return self._decision + + +class MockNonDecisionPolicy: + """Policy stub that returns an invalid decision shape.""" + + def evaluate( + self, + call: tool_policy.ToolCall, + available_tools: dict[str, tool_module.Tool], + ) -> object: + del call, available_tools + return {'action': 'allow'} + + +class MockInvalidActionPolicy: + """Policy stub that returns a decision with malformed action.""" + + def evaluate( + self, + call: tool_policy.ToolCall, + available_tools: dict[str, tool_module.Tool], + ) -> tool_policy.PolicyDecision: + del call, available_tools + return tool_policy.PolicyDecision( + action=cast(tool_policy.PolicyAction, 'allow'), + reason='invalid action type', + ) + + +class MockStringTagsPolicy: + """Policy stub with malformed string tags payload.""" + + def evaluate( + self, + call: tool_policy.ToolCall, + available_tools: dict[str, tool_module.Tool], + ) -> tool_policy.PolicyDecision: + del call, available_tools + return tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.ALLOW, + tags=cast(tuple[str, ...], 'invalid'), + ) + + class InteractiveDocumentWithToolsTest(parameterized.TestCase): def test_open_question_no_tools(self): @@ -264,9 +333,16 @@ def test_copy_preserves_tools(self): language_model.LanguageModel, instance=True, spec_set=True ) tool = MockTool('search', 'Search', 'result') + policy = MockPolicy( + tool_policy.PolicyDecision(action=tool_policy.PolicyAction.DENY) + ) doc = interactive_document_tools.InteractiveDocumentWithTools( - model, tools=[tool], max_tool_calls_per_question=10 + model, + tools=[tool], + max_tool_calls_per_question=10, + policy=policy, + enforcement_mode='enforce', ) doc.statement('Some content') @@ -274,6 +350,343 @@ def test_copy_preserves_tools(self): self.assertIn('search', copied._tools) self.assertEqual(copied._max_tool_calls, 10) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + copied.open_question('Question?') + self.assertEqual(tool.call_count, 0) + + def test_invalid_enforcement_mode_raises(self): + """Invalid policy enforcement mode should raise ValueError.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + invalid_mode = cast( + interactive_document_tools.EnforcementMode, 'invalid' + ) + with self.assertRaises(ValueError): + interactive_document_tools.InteractiveDocumentWithTools( + model, enforcement_mode=invalid_mode + ) + + def test_policy_observe_mode_runs_tool_and_logs_allow(self): + """Observe mode should execute tool and log allow decision.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy( + tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.ALLOW, reason='safe' + ) + ) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='observe' + ) + + response = doc.open_question('Question?') + + self.assertEqual(response, 'Final answer') + self.assertEqual(tool.call_count, 1) + self.assertLen(policy.calls, 1) + self.assertEqual(policy.calls[0].attempt_index, 1) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_allow', tags_found) + + def test_policy_observe_mode_denied_call_still_executes_tool(self): + """Observe mode should record deny decision without blocking execution.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy( + tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.DENY, reason='blocked in observe' + ) + ) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='observe' + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_deny_observed', tags_found) + + def test_policy_observe_mode_edit_keeps_original_args(self): + """Observe mode should execute original args for edit decisions.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "original"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy( + tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.EDIT, + reason='suggest edit', + edited_args={'query': 'edited'}, + ) + ) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='observe' + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + self.assertEqual(tool.last_args, {'query': 'original'}) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_edit_observed', tags_found) + + def test_policy_enforce_mode_denied_call_blocks_execution(self): + """Enforce mode deny should block tool execution.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy( + tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.DENY, reason='blocked' + ) + ) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='enforce' + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 0) + doc_text = doc.text() + self.assertIn('Error: Tool call "search" denied by policy.', doc_text) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_deny_enforced', tags_found) + + def test_policy_enforce_mode_edit_executes_edited_args(self): + """Enforce mode edit should execute with edited args.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "original"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy( + tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.EDIT, + reason='rewrite args', + edited_args={'query': 'edited'}, + ) + ) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='enforce' + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + self.assertEqual(tool.last_args, {'query': 'edited'}) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_edit_enforced', tags_found) + + def test_policy_error_observe_mode_executes_and_logs_error(self): + """Observe mode should fail open when policy raises an error.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy(raise_error=True) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='observe' + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_observed', tags_found) + + def test_policy_error_enforce_mode_blocks_and_logs_error(self): + """Enforce mode should fail closed when policy raises an error.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + policy = MockPolicy(raise_error=True) + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, tools=[tool], policy=policy, enforcement_mode='enforce' + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 0) + doc_text = doc.text() + self.assertIn('blocked due to policy error', doc_text) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_enforced', tags_found) + + def test_policy_observe_mode_invalid_action_fails_open(self): + """Observe mode executes tool if policy returns invalid action.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, + tools=[tool], + policy=MockInvalidActionPolicy(), + enforcement_mode='observe', + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_observed', tags_found) + + def test_policy_observe_mode_string_tags_fails_open(self): + """Observe mode executes tool if policy returns string tags.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, + tools=[tool], + policy=MockStringTagsPolicy(), + enforcement_mode='observe', + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_observed', tags_found) + + def test_policy_enforce_mode_invalid_action_blocks(self): + """Enforce mode blocks tool if policy returns invalid action.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, + tools=[tool], + policy=MockInvalidActionPolicy(), + enforcement_mode='enforce', + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 0) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_enforced', tags_found) + + def test_policy_enforce_mode_string_tags_blocks(self): + """Enforce mode blocks tool if policy returns string tags.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, + tools=[tool], + policy=MockStringTagsPolicy(), + enforcement_mode='enforce', + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 0) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_enforced', tags_found) + + def test_policy_observe_mode_non_decision_fails_open(self): + """Observe mode executes tool if policy returns non-decision object.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, + tools=[tool], + policy=MockNonDecisionPolicy(), + enforcement_mode='observe', + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 1) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_observed', tags_found) + + def test_policy_enforce_mode_non_decision_blocks(self): + """Enforce mode blocks tool if policy returns non-decision object.""" + model = mock.create_autospec( + language_model.LanguageModel, instance=True, spec_set=True + ) + model.sample_text.side_effect = [ + '{"tool": "search", "args": {"query": "q1"}}', + 'Final answer', + ] + tool = MockTool('search', 'Search', 'result') + doc = interactive_document_tools.InteractiveDocumentWithTools( + model, + tools=[tool], + policy=MockNonDecisionPolicy(), + enforcement_mode='enforce', + ) + + doc.open_question('Question?') + + self.assertEqual(tool.call_count, 0) + tags_found = {tag for c in doc.contents() for tag in c.tags} + self.assertIn('tool_policy_error_enforced', tags_found) def test_multiple_choice_with_tool_call_then_answer(self): """LLM uses a tool then answers multiple choice.""" diff --git a/concordia/document/tool.py b/concordia/document/tool.py index 74c99fc13..754d32f8f 100644 --- a/concordia/document/tool.py +++ b/concordia/document/tool.py @@ -15,6 +15,7 @@ """Tool interface for interactive documents with tool use.""" import abc +from collections.abc import Mapping from typing import Any @@ -72,3 +73,20 @@ def execute(self, **kwargs: Any) -> str: truncated by the calling code to fit context limits. """ raise NotImplementedError + + @property + def input_schema(self) -> Mapping[str, Any] | None: + """Schema describing tool arguments. + + This is optional metadata that policy implementations can use to validate + arguments before execution. + """ + return None + + @property + def risk_level(self) -> str: + """Risk category for tool execution. + + Recommended values are "read", "sensitive", and "destructive". + """ + return 'read' diff --git a/concordia/document/tool_policy.py b/concordia/document/tool_policy.py new file mode 100644 index 000000000..0778b70b0 --- /dev/null +++ b/concordia/document/tool_policy.py @@ -0,0 +1,62 @@ +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Policy interfaces for tool execution in interactive documents.""" + +from collections.abc import Mapping +import dataclasses +import enum +from typing import Any, Protocol + +from concordia.document import tool as tool_module + + +class PolicyAction(enum.Enum): + """Action selected by a tool policy.""" + + ALLOW = 'allow' + DENY = 'deny' + EDIT = 'edit' + + +@dataclasses.dataclass(frozen=True) +class PolicyDecision: + """Result of a policy evaluation.""" + + action: PolicyAction = PolicyAction.ALLOW + reason: str = '' + edited_args: dict[str, Any] | None = None + tags: tuple[str, ...] = () + + +@dataclasses.dataclass(frozen=True) +class ToolCall: + """Structured tool call data passed to policy implementations.""" + + tool_name: str + args: dict[str, Any] + raw_response: str + attempt_index: int + + +class ToolPolicy(Protocol): + """Protocol for policy implementations used by tool-enabled documents.""" + + def evaluate( + self, + call: ToolCall, + available_tools: Mapping[str, tool_module.Tool], + ) -> PolicyDecision: + """Evaluates a tool call and returns a decision.""" + raise NotImplementedError diff --git a/concordia/document/tool_policy_test.py b/concordia/document/tool_policy_test.py new file mode 100644 index 000000000..7f47e70c6 --- /dev/null +++ b/concordia/document/tool_policy_test.py @@ -0,0 +1,94 @@ +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for tool policy primitives.""" + +from collections.abc import Mapping +from typing import Any + +from absl.testing import absltest +from concordia.document import tool as tool_module +from concordia.document import tool_policy + + +class _StubTool(tool_module.Tool): + """Minimal tool implementation used by tests.""" + + @property + def name(self) -> str: + return 'stub' + + @property + def description(self) -> str: + return 'Stub tool.' + + def execute(self, **kwargs: Any) -> str: + del kwargs + return 'ok' + + +class _DenyPolicy: + """Policy stub that always denies tool calls.""" + + def evaluate( + self, + call: tool_policy.ToolCall, + available_tools: Mapping[str, tool_module.Tool], + ) -> tool_policy.PolicyDecision: + del available_tools + return tool_policy.PolicyDecision( + action=tool_policy.PolicyAction.DENY, + reason=f'Denied {call.tool_name}.', + tags=('test_tag',), + ) + + +class ToolPolicyTest(absltest.TestCase): + + def test_policy_decision_defaults_to_allow(self): + decision = tool_policy.PolicyDecision() + self.assertEqual(decision.action, tool_policy.PolicyAction.ALLOW) + self.assertEqual(decision.reason, '') + self.assertIsNone(decision.edited_args) + self.assertEqual(decision.tags, ()) + + def test_tool_call_captures_input(self): + call = tool_policy.ToolCall( + tool_name='search', + args={'query': 'weather'}, + raw_response='{"tool":"search"}', + attempt_index=2, + ) + self.assertEqual(call.tool_name, 'search') + self.assertEqual(call.args, {'query': 'weather'}) + self.assertEqual(call.raw_response, '{"tool":"search"}') + self.assertEqual(call.attempt_index, 2) + + def test_tool_policy_protocol_shape(self): + policy = _DenyPolicy() + decision = policy.evaluate( + tool_policy.ToolCall( + tool_name='stub', + args={}, + raw_response='{}', + attempt_index=1, + ), + {'stub': _StubTool()}, + ) + self.assertEqual(decision.action, tool_policy.PolicyAction.DENY) + self.assertEqual(decision.tags, ('test_tag',)) + + +if __name__ == '__main__': + absltest.main()