mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
496 lines
17 KiB
Python
496 lines
17 KiB
Python
"""
|
|
Tool Quality Manager
|
|
|
|
Core API (called by main flow):
|
|
- record_execution(): Called by the PostToolUse quality_tracking hook
|
|
- adjust_ranking(): Called by ToolPreselector for quality-aware sorting
|
|
|
|
Query API (for inspection/debugging):
|
|
- get_quality_report(), get_tool_insights()
|
|
"""
|
|
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
|
|
|
|
from .types import ToolQualityRecord, ExecutionRecord
|
|
from .store import QualityStore
|
|
from openspace.utils.logging import Logger
|
|
|
|
if TYPE_CHECKING:
|
|
from openspace.grounding.core.tool import BaseTool
|
|
from openspace.grounding.core.types import ToolResult
|
|
|
|
logger = Logger.get_logger(__name__)
|
|
|
|
|
|
class ToolQualityManager:
|
|
"""
|
|
Manages tool quality tracking and quality-aware ranking.
|
|
|
|
Features:
|
|
- Track execution success rate and latency
|
|
- Persistent memory across sessions
|
|
- Quality-integrated tool ranking
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
db_path: Optional[Path] = None,
|
|
enable_persistence: bool = True,
|
|
auto_save: bool = True,
|
|
):
|
|
self._enable_persistence = enable_persistence
|
|
self._auto_save = auto_save
|
|
|
|
# In-memory cache
|
|
self._records: Dict[str, ToolQualityRecord] = {}
|
|
self._global_execution_count: int = 0
|
|
|
|
# Persistent store (SQLite, shares DB file with SkillStore)
|
|
self._store = QualityStore(db_path=db_path) if enable_persistence else None
|
|
|
|
# Load from DB
|
|
if self._store:
|
|
self._records, self._global_execution_count = self._store.load_all()
|
|
|
|
logger.info(
|
|
f"ToolQualityManager initialized "
|
|
f"(persistence={enable_persistence}, records={len(self._records)}, "
|
|
f"global_count={self._global_execution_count})"
|
|
)
|
|
|
|
@property
|
|
def global_execution_count(self) -> int:
|
|
return self._global_execution_count
|
|
|
|
@property
|
|
def store(self) -> QualityStore | None:
|
|
return self._store
|
|
|
|
def get_tool_key(self, tool: "BaseTool") -> str:
|
|
"""Generate unique key for a tool."""
|
|
from openspace.grounding.core.types import BackendType
|
|
|
|
if tool.is_bound:
|
|
backend = tool.runtime_info.backend.value
|
|
server = tool.runtime_info.server_name or "default"
|
|
else:
|
|
backend = tool.backend_type.value if tool.backend_type != BackendType.NOT_SET else "unknown"
|
|
server = "default"
|
|
|
|
return f"{backend}:{server}:{tool.name}"
|
|
|
|
def get_record(self, tool: "BaseTool") -> ToolQualityRecord:
|
|
"""Get or create quality record for a tool."""
|
|
key = self.get_tool_key(tool)
|
|
|
|
if key not in self._records:
|
|
backend, server, name = key.split(":", 2)
|
|
self._records[key] = ToolQualityRecord(
|
|
tool_key=key,
|
|
backend=backend,
|
|
server=server,
|
|
tool_name=name,
|
|
)
|
|
|
|
return self._records[key]
|
|
|
|
def get_quality_score(self, tool: "BaseTool") -> float:
|
|
"""Get quality score for a tool (0-1)."""
|
|
return self.get_record(tool).quality_score
|
|
|
|
# Key-based record access (for cross-system integration)
|
|
def get_or_create_record_by_key(self, tool_key: str) -> ToolQualityRecord:
|
|
"""Get or create a ToolQualityRecord by its canonical key.
|
|
|
|
Used by ExecutionAnalyzer integration where no BaseTool instance
|
|
is available. Parses ``tool_key`` into backend/server/tool_name.
|
|
|
|
Key formats:
|
|
- ``backend:server:tool_name`` → three-part key (canonical for MCP)
|
|
- ``backend:tool_name`` → two-part; tries ``backend:default:tool_name``
|
|
first for matching existing records.
|
|
"""
|
|
# 1. Direct match
|
|
if tool_key in self._records:
|
|
return self._records[tool_key]
|
|
|
|
parts = tool_key.split(":", 2)
|
|
if len(parts) == 3:
|
|
backend, server, name = parts
|
|
elif len(parts) == 2:
|
|
backend, name = parts
|
|
server = "default"
|
|
# Try normalized 3-part key before creating a new record
|
|
canonical = f"{backend}:default:{name}"
|
|
if canonical in self._records:
|
|
return self._records[canonical]
|
|
else:
|
|
backend, server, name = "unknown", "default", tool_key
|
|
|
|
canonical_key = f"{backend}:{server}:{name}"
|
|
if canonical_key in self._records:
|
|
return self._records[canonical_key]
|
|
|
|
record = ToolQualityRecord(
|
|
tool_key=canonical_key,
|
|
backend=backend,
|
|
server=server,
|
|
tool_name=name,
|
|
)
|
|
self._records[canonical_key] = record
|
|
return record
|
|
|
|
def find_record_by_key(self, key: str) -> Optional[ToolQualityRecord]:
|
|
"""Find a record by exact or partial tool key.
|
|
|
|
Tries in order:
|
|
1. Exact match (3-part ``backend:server:tool`` or 2-part)
|
|
2. Normalized 2-part → ``backend:default:tool``
|
|
3. Linear scan matching backend + tool_name (ignoring server)
|
|
"""
|
|
# 1. Exact
|
|
if key in self._records:
|
|
return self._records[key]
|
|
|
|
parts = key.split(":", 2)
|
|
if len(parts) == 3:
|
|
backend, server, tool_name = parts
|
|
# Some LLM analyses emit placeholder keys such as
|
|
# ``backend:default:ls`` instead of the runtime key
|
|
# ``shell:default:ls``. Prefer an existing concrete record for the
|
|
# same tool over creating a parallel aggregate bucket, but only
|
|
# when that match is unambiguous.
|
|
candidates = [
|
|
record for record in self._records.values()
|
|
if (
|
|
record.tool_name == tool_name
|
|
and (
|
|
record.backend == backend
|
|
or backend in {"backend", "tool", "unknown", "default"}
|
|
)
|
|
and (record.server == server or server == "default")
|
|
)
|
|
]
|
|
if len(candidates) == 1:
|
|
return candidates[0]
|
|
elif len(parts) == 2:
|
|
backend, tool_name = parts
|
|
# 2. Normalize
|
|
canonical = f"{backend}:default:{tool_name}"
|
|
if canonical in self._records:
|
|
return self._records[canonical]
|
|
# 3. Scan
|
|
for record in self._records.values():
|
|
if record.backend == backend and record.tool_name == tool_name:
|
|
return record
|
|
return None
|
|
|
|
# Execution Tracking
|
|
async def record_execution(
|
|
self,
|
|
tool: "BaseTool",
|
|
result: "ToolResult",
|
|
execution_time_ms: float,
|
|
) -> ToolQualityRecord:
|
|
"""Record tool execution result and increment global counter."""
|
|
error_message = None
|
|
if result.is_error and result.error:
|
|
error_message = str(result.error)[:500]
|
|
|
|
return await self.record_outcome(
|
|
tool,
|
|
success=result.is_success,
|
|
execution_time_ms=execution_time_ms,
|
|
error_message=error_message,
|
|
)
|
|
|
|
async def record_outcome(
|
|
self,
|
|
tool: "BaseTool",
|
|
*,
|
|
success: bool,
|
|
execution_time_ms: float = 0.0,
|
|
error_message: Optional[str] = None,
|
|
) -> ToolQualityRecord:
|
|
"""Record a normalized pipeline outcome for a tool attempt."""
|
|
record = self.get_record(tool)
|
|
|
|
# Add execution record
|
|
record.add_execution(ExecutionRecord(
|
|
timestamp=datetime.now(),
|
|
success=bool(success),
|
|
execution_time_ms=execution_time_ms,
|
|
error_message=error_message,
|
|
))
|
|
|
|
# Increment global execution count
|
|
self._global_execution_count += 1
|
|
|
|
# Auto-save
|
|
if self._auto_save and self._store:
|
|
await self._store.save_record(record, self._records, self._global_execution_count)
|
|
|
|
logger.debug(
|
|
f"Recorded execution: {record.tool_key} "
|
|
f"success={bool(success)} time={execution_time_ms:.0f}ms "
|
|
f"(global_count={self._global_execution_count})"
|
|
)
|
|
return record
|
|
|
|
# Quality-Aware Ranking
|
|
def adjust_ranking(
|
|
self,
|
|
tools_with_scores: List[Tuple["BaseTool", float]],
|
|
) -> List[Tuple["BaseTool", float]]:
|
|
"""
|
|
Adjust tool ranking using penalty-based approach.
|
|
|
|
Args:
|
|
tools_with_scores: List of (tool, semantic_score) tuples
|
|
"""
|
|
adjusted = []
|
|
for tool, semantic_score in tools_with_scores:
|
|
penalty = self.get_penalty(tool)
|
|
|
|
adjusted_score = semantic_score * penalty
|
|
|
|
adjusted.append((tool, adjusted_score))
|
|
|
|
# Sort by adjusted score (descending)
|
|
adjusted.sort(key=lambda x: x[1], reverse=True)
|
|
|
|
return adjusted
|
|
|
|
def get_penalty(self, tool: "BaseTool") -> float:
|
|
"""Get penalty factor for a tool (0.2-1.0)."""
|
|
return self.get_record(tool).penalty
|
|
|
|
async def save(self) -> None:
|
|
"""
|
|
Manually save all records to disk.
|
|
|
|
Note: Usually not needed - auto_save handles persistence in
|
|
record_execution().
|
|
Provided as public API for explicit save when needed.
|
|
"""
|
|
if self._store:
|
|
await self._store.save_all(self._records)
|
|
|
|
def clear_cache(self) -> None:
|
|
"""Clear all cached data."""
|
|
self._records.clear()
|
|
if self._store:
|
|
self._store.clear()
|
|
|
|
def get_stats(self) -> Dict:
|
|
"""
|
|
Get quality tracking statistics.
|
|
|
|
Note: Query API for inspection, may not be called in main flow.
|
|
"""
|
|
if not self._records:
|
|
return {"total_tools": 0}
|
|
|
|
records = list(self._records.values())
|
|
|
|
return {
|
|
"total_tools": len(records),
|
|
"total_executions": sum(r.total_calls for r in records),
|
|
"avg_success_rate": (
|
|
sum(r.success_rate for r in records) / len(records)
|
|
if records else 0
|
|
),
|
|
"avg_quality_score": (
|
|
sum(r.quality_score for r in records) / len(records)
|
|
if records else 0
|
|
),
|
|
}
|
|
|
|
def get_top_tools(
|
|
self,
|
|
n: int = 10,
|
|
backend: Optional[str] = None,
|
|
min_calls: int = 3,
|
|
) -> List[ToolQualityRecord]:
|
|
"""
|
|
Get top N tools by quality score.
|
|
|
|
Args:
|
|
n: Number of tools to return
|
|
backend: Filter by backend type (optional)
|
|
min_calls: Minimum calls required (to filter untested tools)
|
|
"""
|
|
records = [
|
|
r for r in self._records.values()
|
|
if r.total_calls >= min_calls
|
|
and (backend is None or r.backend == backend)
|
|
]
|
|
|
|
records.sort(key=lambda r: r.quality_score, reverse=True)
|
|
return records[:n]
|
|
|
|
def get_problematic_tools(
|
|
self,
|
|
success_rate_threshold: float = 0.5,
|
|
min_calls: int = 5,
|
|
) -> List[ToolQualityRecord]:
|
|
"""
|
|
Get tools with low success rate (candidates for review/removal).
|
|
|
|
Args:
|
|
success_rate_threshold: Tools below this rate are flagged
|
|
min_calls: Minimum calls required (avoid flagging new tools)
|
|
"""
|
|
return [
|
|
r for r in self._records.values()
|
|
if r.total_calls >= min_calls
|
|
and r.recent_success_rate < success_rate_threshold
|
|
]
|
|
|
|
def get_quality_report(self) -> Dict:
|
|
"""
|
|
Generate comprehensive quality report for upper layer.
|
|
|
|
Returns structured report with:
|
|
- Overall stats
|
|
- Per-backend breakdown
|
|
- Top/problematic tools
|
|
- Improvement suggestions
|
|
"""
|
|
if not self._records:
|
|
return {"status": "no_data", "message": "No quality data collected yet"}
|
|
|
|
records = list(self._records.values())
|
|
tested_records = [r for r in records if r.total_calls >= 3]
|
|
|
|
# Per-backend stats
|
|
backends = {}
|
|
for r in records:
|
|
if r.backend not in backends:
|
|
backends[r.backend] = {
|
|
"tools": 0,
|
|
"total_calls": 0,
|
|
"success_count": 0,
|
|
"servers": set()
|
|
}
|
|
backends[r.backend]["tools"] += 1
|
|
backends[r.backend]["total_calls"] += r.total_calls
|
|
backends[r.backend]["success_count"] += r.success_count
|
|
backends[r.backend]["servers"].add(r.server)
|
|
|
|
# Convert sets to counts
|
|
for b in backends:
|
|
backends[b]["servers"] = len(backends[b]["servers"])
|
|
backends[b]["success_rate"] = (
|
|
backends[b]["success_count"] / backends[b]["total_calls"]
|
|
if backends[b]["total_calls"] > 0 else 0
|
|
)
|
|
|
|
# Top and problematic tools
|
|
top_tools = self.get_top_tools(5)
|
|
problematic = self.get_problematic_tools()
|
|
|
|
return {
|
|
"summary": {
|
|
"total_tools": len(records),
|
|
"tested_tools": len(tested_records),
|
|
"total_executions": sum(r.total_calls for r in records),
|
|
"overall_success_rate": (
|
|
sum(r.success_count for r in records) /
|
|
max(1, sum(r.total_calls for r in records))
|
|
),
|
|
"avg_quality_score": (
|
|
sum(r.quality_score for r in tested_records) / len(tested_records)
|
|
if tested_records else 0
|
|
),
|
|
},
|
|
"by_backend": backends,
|
|
"top_tools": [
|
|
{"key": r.tool_key, "score": r.quality_score, "success_rate": r.success_rate}
|
|
for r in top_tools
|
|
],
|
|
"problematic_tools": [
|
|
{"key": r.tool_key, "success_rate": r.success_rate, "calls": r.total_calls}
|
|
for r in problematic
|
|
],
|
|
"recommendations": self._generate_recommendations(records, problematic),
|
|
}
|
|
|
|
def _generate_recommendations(
|
|
self,
|
|
records: List[ToolQualityRecord],
|
|
problematic: List[ToolQualityRecord],
|
|
) -> List[str]:
|
|
"""Generate actionable recommendations based on quality data."""
|
|
recommendations = []
|
|
|
|
# Check for problematic tools
|
|
if problematic:
|
|
tool_names = [r.tool_name for r in problematic[:3]]
|
|
recommendations.append(
|
|
f"Review low-success tools: {', '.join(tool_names)}"
|
|
)
|
|
|
|
return recommendations
|
|
|
|
def compute_adaptive_quality_weight(self) -> float:
|
|
"""
|
|
Compute adaptive quality weight based on data confidence.
|
|
|
|
Returns higher weight when we have more reliable quality data,
|
|
lower weight when data is sparse.
|
|
"""
|
|
if not self._records:
|
|
return 0.1 # Low weight when no data
|
|
|
|
records = list(self._records.values())
|
|
tested_count = sum(1 for r in records if r.total_calls >= 3)
|
|
|
|
if tested_count == 0:
|
|
return 0.1
|
|
|
|
# More tested tools -> higher confidence -> higher weight
|
|
coverage = tested_count / len(records)
|
|
|
|
# Average calls per tested tool -> data richness
|
|
avg_calls = sum(r.total_calls for r in records) / len(records)
|
|
richness = min(1.0, avg_calls / 20) # Cap at 20 calls average
|
|
|
|
# Combine coverage and richness
|
|
confidence = (coverage * 0.5 + richness * 0.5)
|
|
|
|
# Map to weight range [0.1, 0.5]
|
|
weight = 0.1 + confidence * 0.4
|
|
|
|
return round(weight, 2)
|
|
|
|
def get_tool_insights(self, tool: "BaseTool") -> Dict:
|
|
"""
|
|
Get detailed insights for a specific tool (for debugging/analysis).
|
|
|
|
Returns comprehensive info about tool's quality history.
|
|
"""
|
|
record = self._records.get(self.get_tool_key(tool))
|
|
if not record:
|
|
return {"status": "not_tracked", "tool": tool.name}
|
|
|
|
# Count recent failures
|
|
recent_failures_count = sum(
|
|
1 for e in record.recent_executions[-20:]
|
|
if not e.success
|
|
)
|
|
|
|
return {
|
|
"tool_key": record.tool_key,
|
|
"total_calls": record.total_calls,
|
|
"success_rate": record.success_rate,
|
|
"recent_success_rate": record.recent_success_rate,
|
|
"avg_execution_time_ms": record.avg_execution_time_ms,
|
|
"quality_score": record.quality_score,
|
|
"recent_failures_count": recent_failures_count,
|
|
"first_seen": record.first_seen.isoformat(),
|
|
"last_updated": record.last_updated.isoformat(),
|
|
}
|