feat(bias): implement data models and storage for tool weighting and bias profiles

This commit is contained in:
2026-03-10 09:27:12 -04:00
parent ee19cc1d2a
commit 77a0b385d5
6 changed files with 264 additions and 154 deletions

46
tests/test_bias_models.py Normal file
View File

@@ -0,0 +1,46 @@
import pytest
from src.models import Tool, ToolPreset, BiasProfile
def test_tool_model():
tool = Tool(name="read_file", weight=5, parameter_bias={"path": "preferred"})
data = tool.to_dict()
assert data["name"] == "read_file"
assert data["weight"] == 5
assert data["parameter_bias"]["path"] == "preferred"
tool2 = Tool.from_dict(data)
assert tool2.name == "read_file"
assert tool2.weight == 5
assert tool2.parameter_bias["path"] == "preferred"
def test_tool_preset_extension():
# Verify that ToolPreset correctly parses and serializes Tool objects
tool_data = {"name": "read_file", "weight": 4, "parameter_bias": {"path": "high"}}
raw_data = {"categories": {"General": [tool_data]}}
# Test parsing via from_dict
preset = ToolPreset.from_dict("test", raw_data)
assert isinstance(preset.categories["General"][0], Tool)
assert preset.categories["General"][0].weight == 4
# Test serialization
data = preset.to_dict()
assert data["categories"]["General"][0]["weight"] == 4
assert data["categories"]["General"][0]["name"] == "read_file"
def test_bias_profile_model():
profile = BiasProfile(
name="Execution-Focused",
tool_weights={"run_powershell": 5},
category_multipliers={"Surgical": 1.5}
)
data = profile.to_dict()
assert data["tool_weights"]["run_powershell"] == 5
assert data["category_multipliers"]["Surgical"] == 1.5
# BiasProfile.from_dict expects 'name' inside the dict as well if coming from load_all_bias_profiles
data["name"] = "Execution-Focused"
profile2 = BiasProfile.from_dict(data)
assert profile2.name == "Execution-Focused"
assert profile2.tool_weights["run_powershell"] == 5
assert profile2.category_multipliers["Surgical"] == 1.5

View File

@@ -2,7 +2,7 @@ import pytest
import tomli_w
from pathlib import Path
from src.tool_presets import ToolPresetManager
from src.models import ToolPreset
from src.models import ToolPreset, BiasProfile, Tool
from src import paths
@pytest.fixture
@@ -25,17 +25,18 @@ def temp_paths(tmp_path, monkeypatch):
"project_presets": project_presets
}
def test_load_all_merged(temp_paths):
def test_load_all_presets_merged(temp_paths):
# Setup global presets
global_data = {
"default": {
"categories": {
"file": {"read": True},
"shell": {"run": False}
"presets": {
"default": {
"categories": {
"General": [{"name": "read_file", "approval": "auto"}]
}
},
"global_only": {
"categories": {"Web": [{"name": "web_search", "approval": "ask"}]}
}
},
"global_only": {
"categories": {"web": {"search": True}}
}
}
with open(temp_paths["global_presets"], "wb") as f:
@@ -43,98 +44,78 @@ def test_load_all_merged(temp_paths):
# Setup project presets (overrides 'default')
project_data = {
"default": {
"categories": {
"file": {"read": True},
"shell": {"run": True} # Override
"presets": {
"default": {
"categories": {
"General": [{"name": "read_file", "approval": "auto", "weight": 5}]
}
}
},
"project_only": {
"categories": {"git": {"commit": True}}
}
}
with open(temp_paths["project_presets"], "wb") as f:
tomli_w.dump(project_data, f)
manager = ToolPresetManager(project_root=temp_paths["project_dir"])
all_presets = manager.load_all()
all_presets = manager.load_all_presets()
assert "default" in all_presets
assert all_presets["default"].categories["shell"]["run"] is True # Overridden
assert isinstance(all_presets["default"].categories["General"][0], Tool)
assert all_presets["default"].categories["General"][0].weight == 5
assert "global_only" in all_presets
assert "project_only" in all_presets
assert all_presets["global_only"].categories["web"]["search"] is True
assert all_presets["project_only"].categories["git"]["commit"] is True
def test_save_preset_global(temp_paths):
manager = ToolPresetManager()
preset = ToolPreset(name="new_global", categories={"test": {"ok": True}})
manager.save_preset(preset, scope="global")
assert temp_paths["global_presets"].exists()
loaded = manager._load_from_path(temp_paths["global_presets"])
assert "new_global" in loaded
assert loaded["new_global"].categories == {"test": {"ok": True}}
def test_save_preset_project(temp_paths):
manager = ToolPresetManager(project_root=temp_paths["project_dir"])
preset = ToolPreset(name="new_project", categories={"test": {"ok": False}})
manager.save_preset(preset, scope="project")
assert temp_paths["project_presets"].exists()
loaded = manager._load_from_path(temp_paths["project_presets"])
assert "new_project" in loaded
assert loaded["new_project"].categories == {"test": {"ok": False}}
def test_delete_preset_global(temp_paths):
# Initial global setup
def test_bias_profiles_merged(temp_paths):
# Setup global biases
global_data = {
"to_delete": {"categories": {}},
"keep": {"categories": {}}
"bias_profiles": {
"Discovery": {
"tool_weights": {"web_search": 5},
"category_multipliers": {"Web": 1.5}
}
}
}
with open(temp_paths["global_presets"], "wb") as f:
tomli_w.dump(global_data, f)
manager = ToolPresetManager()
manager.delete_preset("to_delete", scope="global")
loaded = manager._load_from_path(temp_paths["global_presets"])
assert "to_delete" not in loaded
assert "keep" in loaded
def test_delete_preset_project(temp_paths):
# Initial project setup
# Setup project biases
project_data = {
"to_delete": {"categories": {}},
"keep": {"categories": {}}
"bias_profiles": {
"Execution": {
"tool_weights": {"run_powershell": 5}
}
}
}
with open(temp_paths["project_presets"], "wb") as f:
tomli_w.dump(project_data, f)
manager = ToolPresetManager(project_root=temp_paths["project_dir"])
manager.delete_preset("to_delete", scope="project")
profiles = manager.load_all_bias_profiles()
loaded = manager._load_from_path(temp_paths["project_presets"])
assert "to_delete" not in loaded
assert "keep" in loaded
assert "Discovery" in profiles
assert profiles["Discovery"].category_multipliers["Web"] == 1.5
assert "Execution" in profiles
assert profiles["Execution"].tool_weights["run_powershell"] == 5
def test_save_project_no_root_raises(temp_paths):
manager = ToolPresetManager(project_root=None)
preset = ToolPreset(name="fail", categories={})
with pytest.raises(ValueError, match="Project root not set"):
manager.save_preset(preset, scope="project")
def test_save_bias_profile(temp_paths):
manager = ToolPresetManager(project_root=temp_paths["project_dir"])
profile = BiasProfile(name="Custom", tool_weights={"test": 1})
manager.save_bias_profile(profile, scope="project")
loaded = manager.load_all_bias_profiles()
assert "Custom" in loaded
assert loaded["Custom"].tool_weights["test"] == 1
def test_delete_project_no_root_raises(temp_paths):
manager = ToolPresetManager(project_root=None)
with pytest.raises(ValueError, match="Project root not set"):
manager.delete_preset("any", scope="project")
def test_invalid_scope_raises(temp_paths):
manager = ToolPresetManager()
preset = ToolPreset(name="fail", categories={})
with pytest.raises(ValueError, match="Invalid scope"):
manager.save_preset(preset, scope="invalid")
with pytest.raises(ValueError, match="Invalid scope"):
manager.delete_preset("any", scope="invalid")
def test_delete_bias_profile(temp_paths):
project_data = {
"bias_profiles": {
"to_delete": {"tool_weights": {}}
}
}
with open(temp_paths["project_presets"], "wb") as f:
tomli_w.dump(project_data, f)
manager = ToolPresetManager(project_root=temp_paths["project_dir"])
manager.delete_bias_profile("to_delete", scope="project")
profiles = manager.load_all_bias_profiles()
assert "to_delete" not in profiles

View File

@@ -3,14 +3,14 @@ import asyncio
from src import ai_client
from src import mcp_client
from src import models
from src.models import ToolPreset
from src.models import ToolPreset, Tool
from unittest.mock import MagicMock, patch
@pytest.mark.asyncio
async def test_tool_auto_approval():
# Setup a preset with read_file as auto
preset = ToolPreset(name="AutoTest", categories={
"General": {"read_file": "auto"}
"General": [Tool(name="read_file", approval="auto")]
})
with patch("src.tool_presets.ToolPresetManager.load_all", return_value={"AutoTest": preset}):
@@ -39,7 +39,7 @@ async def test_tool_auto_approval():
async def test_tool_ask_approval():
# Setup a preset with run_powershell as ask
preset = ToolPreset(name="AskTest", categories={
"General": {"run_powershell": "ask"}
"General": [Tool(name="run_powershell", approval="ask")]
})
with patch("src.tool_presets.ToolPresetManager.load_all", return_value={"AskTest": preset}):
@@ -65,7 +65,7 @@ async def test_tool_ask_approval():
async def test_tool_rejection():
# Setup a preset with run_powershell as ask
preset = ToolPreset(name="AskTest", categories={
"General": {"run_powershell": "ask"}
"General": [Tool(name="run_powershell", approval="ask")]
})
with patch("src.tool_presets.ToolPresetManager.load_all", return_value={"AskTest": preset}):