Implement Phase 1: core tracking and sidebar UI
- Add model usage tracking via PromptExecutor hook - Support 12+ loader node types (checkpoints, LoRAs, VAEs, etc.) - Pattern matching for custom node loaders - Cross-platform file locking (Windows/Unix) - JSON storage with schema migrations - REST API endpoints for usage data - Sidebar panel with timeframe filtering and sorting - Stale model highlighting (>30 days unused)
This commit is contained in:
Setmaster
@@ -1 +1,4 @@
|
||||
.claude/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
# ModelPulse
|
||||
|
||||
A ComfyUI extension that tracks model usage frequency, helping you identify abandoned or underutilized models to manage storage efficiently.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## Features
|
||||
|
||||
- **Automatic Tracking** - Monitors usage of all model types: checkpoints, LoRAs, VAEs, ControlNets, CLIP, UNETs, upscalers, and more
|
||||
- **Sidebar Panel** - View usage statistics directly in ComfyUI's sidebar
|
||||
- **Time Filtering** - View usage for all time, this month, or this week
|
||||
- **Smart Sorting** - Sort by last used, usage count, or name
|
||||
- **Stale Detection** - Models unused for 30+ days are highlighted
|
||||
- **Category Grouping** - Models organized by type with collapsible sections
|
||||
- **Zero Dependencies** - Uses only Python standard library
|
||||
|
||||
## Installation
|
||||
|
||||
### Option 1: ComfyUI Manager (Recommended)
|
||||
|
||||
Search for "ModelPulse" in ComfyUI Manager and click Install.
|
||||
|
||||
### Option 2: Manual Installation
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/your-username/ComfyUI-ModelPulse.git
|
||||
```
|
||||
|
||||
Restart ComfyUI after installation.
|
||||
|
||||
## Usage
|
||||
|
||||
1. Click the **chart icon** in the ComfyUI sidebar to open ModelPulse
|
||||
2. Run workflows as usual - model usage is tracked automatically
|
||||
3. Use the tabs to filter by time period (All Time / Month / Week)
|
||||
4. Use the dropdown to sort by Last Used, Usage Count, or Name
|
||||
5. Click category headers to collapse/expand sections
|
||||
|
||||
## Tracked Model Types
|
||||
|
||||
| Category | Loader Nodes |
|
||||
|----------|-------------|
|
||||
| Checkpoints | CheckpointLoaderSimple, CheckpointLoader |
|
||||
| LoRAs | LoraLoader, LoraLoaderModelOnly |
|
||||
| VAEs | VAELoader |
|
||||
| ControlNets | ControlNetLoader |
|
||||
| CLIP | CLIPLoader, DualCLIPLoader |
|
||||
| UNETs | UNETLoader |
|
||||
| Upscalers | UpscaleModelLoader |
|
||||
| Style Models | StyleModelLoader |
|
||||
| GLIGEN | GLIGENLoader |
|
||||
|
||||
Custom node loaders (Impact Pack, Efficiency Nodes, etc.) are also detected via pattern matching.
|
||||
|
||||
## API Endpoints
|
||||
|
||||
ModelPulse exposes REST endpoints for programmatic access:
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
|----------|--------|-------------|
|
||||
| `/modelpulse/usage` | GET | Get all usage data (supports `?timeframe=` and `?sort=` params) |
|
||||
| `/modelpulse/model/{id}` | GET | Get detailed data for a specific model |
|
||||
| `/modelpulse/categories` | GET | Get category list with counts |
|
||||
| `/modelpulse/reset` | POST | Reset all tracking data (requires `{"confirm": true}`) |
|
||||
|
||||
## Data Storage
|
||||
|
||||
Usage data is stored in:
|
||||
```
|
||||
ComfyUI/user/default/modelpulse/usage_data.json
|
||||
```
|
||||
|
||||
The data persists across ComfyUI restarts and can be backed up or edited manually.
|
||||
|
||||
## Requirements
|
||||
|
||||
- ComfyUI 0.3.0 or later (Vue 3 frontend)
|
||||
- Python 3.10+
|
||||
- No external Python dependencies
|
||||
|
||||
## License
|
||||
|
||||
MIT License - see [LICENSE](LICENSE) for details.
|
||||
|
||||
## Contributing
|
||||
|
||||
Contributions are welcome! Please feel free to submit issues and pull requests.
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
"""
|
||||
ModelPulse - ComfyUI Model Usage Tracker
|
||||
|
||||
Tracks model usage frequency to help identify abandoned or underutilized models.
|
||||
"""
|
||||
|
||||
from .py.tracking import ModelTracker
|
||||
from .py.routes import setup_routes
|
||||
|
||||
# Global tracker instance
|
||||
tracker = ModelTracker()
|
||||
|
||||
# Store original execute for restoration if needed
|
||||
_original_execute = None
|
||||
|
||||
|
||||
def init_tracking():
|
||||
"""Initialize the tracking hook on PromptExecutor."""
|
||||
global _original_execute
|
||||
|
||||
try:
|
||||
from execution import PromptExecutor
|
||||
except ImportError:
|
||||
print("[ModelPulse] Warning: Could not import PromptExecutor. Tracking disabled.")
|
||||
return
|
||||
|
||||
_original_execute = PromptExecutor.execute
|
||||
|
||||
def tracked_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
|
||||
"""Wrapper that tracks model usage before execution."""
|
||||
try:
|
||||
# Extract and record models from the prompt
|
||||
tracker.extract_and_record_models(prompt)
|
||||
except Exception as e:
|
||||
# Don't let tracking errors break execution
|
||||
print(f"[ModelPulse] Error tracking models: {e}")
|
||||
|
||||
# Call original execute (synchronous)
|
||||
return _original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
|
||||
|
||||
PromptExecutor.execute = tracked_execute
|
||||
print("[ModelPulse] Tracking initialized")
|
||||
|
||||
|
||||
def init_routes():
|
||||
"""Initialize API routes."""
|
||||
try:
|
||||
from server import PromptServer
|
||||
|
||||
if PromptServer.instance is None:
|
||||
print("[ModelPulse] Warning: PromptServer not initialized. Routes disabled.")
|
||||
return
|
||||
|
||||
setup_routes(PromptServer.instance.routes, tracker)
|
||||
print("[ModelPulse] API routes registered")
|
||||
|
||||
except ImportError:
|
||||
print("[ModelPulse] Warning: Could not import PromptServer. Routes disabled.")
|
||||
except Exception as e:
|
||||
print(f"[ModelPulse] Error registering routes: {e}")
|
||||
|
||||
|
||||
# Initialize on module load
|
||||
init_tracking()
|
||||
init_routes()
|
||||
|
||||
|
||||
# ComfyUI node registration (empty - we don't add any nodes)
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
# Tell ComfyUI where to find our JavaScript files
|
||||
WEB_DIRECTORY = "./js"
|
||||
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
"WEB_DIRECTORY",
|
||||
"tracker",
|
||||
]
|
||||
@@ -0,0 +1,511 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
|
||||
/**
|
||||
* ModelPulse - Usage Tracker Extension
|
||||
*
|
||||
* Displays model usage statistics in a sidebar panel.
|
||||
*/
|
||||
|
||||
// State management
|
||||
let currentTimeframe = "all";
|
||||
let currentSort = "last_used";
|
||||
let usageData = null;
|
||||
let isLoading = false;
|
||||
|
||||
/**
|
||||
* Fetch usage data from the API
|
||||
*/
|
||||
async function fetchUsageData(timeframe = "all", sortBy = "last_used") {
|
||||
try {
|
||||
const response = await fetch(
|
||||
`/modelpulse/usage?timeframe=${timeframe}&sort=${sortBy}`
|
||||
);
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}`);
|
||||
}
|
||||
return await response.json();
|
||||
} catch (error) {
|
||||
console.error("[ModelPulse] Error fetching usage data:", error);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Format ISO date to relative time string
|
||||
*/
|
||||
function formatLastUsed(isoDate) {
|
||||
if (!isoDate) return "Never";
|
||||
|
||||
const date = new Date(isoDate);
|
||||
const now = new Date();
|
||||
const diffMs = now - date;
|
||||
const diffMins = Math.floor(diffMs / (1000 * 60));
|
||||
const diffHours = Math.floor(diffMs / (1000 * 60 * 60));
|
||||
const diffDays = Math.floor(diffHours / 24);
|
||||
|
||||
if (diffMins < 1) return "just now";
|
||||
if (diffMins < 60) return `${diffMins}m ago`;
|
||||
if (diffHours < 24) return `${diffHours}h ago`;
|
||||
if (diffDays < 30) return `${diffDays}d ago`;
|
||||
|
||||
const diffMonths = Math.floor(diffDays / 30);
|
||||
if (diffMonths < 12) return `${diffMonths}mo ago`;
|
||||
|
||||
const diffYears = Math.floor(diffDays / 365);
|
||||
return `${diffYears}y ago`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get category display info
|
||||
*/
|
||||
function getCategoryInfo(category) {
|
||||
const categories = {
|
||||
checkpoint: { name: "Checkpoints", icon: "📦" },
|
||||
lora: { name: "LoRAs", icon: "✨" },
|
||||
vae: { name: "VAEs", icon: "⚙️" },
|
||||
controlnet: { name: "ControlNets", icon: "🎛️" },
|
||||
clip: { name: "CLIP", icon: "💬" },
|
||||
unet: { name: "UNETs", icon: "🖥️" },
|
||||
upscaler: { name: "Upscalers", icon: "🔍" },
|
||||
style_model: { name: "Style Models", icon: "🎨" },
|
||||
gligen: { name: "GLIGEN", icon: "📐" },
|
||||
};
|
||||
return categories[category] || { name: category, icon: "📁" };
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a model is stale (unused for > 30 days)
|
||||
*/
|
||||
function isStale(lastUsed) {
|
||||
if (!lastUsed) return true;
|
||||
const date = new Date(lastUsed);
|
||||
const now = new Date();
|
||||
const diffDays = Math.floor((now - date) / (1000 * 60 * 60 * 24));
|
||||
return diffDays > 30;
|
||||
}
|
||||
|
||||
/**
|
||||
* Group models by category
|
||||
*/
|
||||
function groupByCategory(models) {
|
||||
const groups = {};
|
||||
for (const model of models) {
|
||||
const cat = model.category;
|
||||
if (!groups[cat]) {
|
||||
groups[cat] = [];
|
||||
}
|
||||
groups[cat].push(model);
|
||||
}
|
||||
return groups;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create the styles for the UI
|
||||
*/
|
||||
function createStyles() {
|
||||
const style = document.createElement("style");
|
||||
style.textContent = `
|
||||
.modelpulse-container {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
height: 100%;
|
||||
font-family: var(--comfy-font, sans-serif);
|
||||
font-size: 13px;
|
||||
color: var(--input-text, #ddd);
|
||||
}
|
||||
|
||||
.modelpulse-header {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
padding: 12px;
|
||||
border-bottom: 1px solid var(--border-color, #444);
|
||||
}
|
||||
|
||||
.modelpulse-title {
|
||||
font-size: 15px;
|
||||
font-weight: 600;
|
||||
}
|
||||
|
||||
.modelpulse-refresh {
|
||||
background: transparent;
|
||||
border: none;
|
||||
color: var(--input-text, #ddd);
|
||||
cursor: pointer;
|
||||
padding: 4px 8px;
|
||||
border-radius: 4px;
|
||||
font-size: 14px;
|
||||
}
|
||||
|
||||
.modelpulse-refresh:hover {
|
||||
background: var(--comfy-menu-bg-hover, #333);
|
||||
}
|
||||
|
||||
.modelpulse-refresh.loading {
|
||||
opacity: 0.5;
|
||||
cursor: wait;
|
||||
}
|
||||
|
||||
.modelpulse-tabs {
|
||||
display: flex;
|
||||
border-bottom: 1px solid var(--border-color, #444);
|
||||
}
|
||||
|
||||
.modelpulse-tab {
|
||||
flex: 1;
|
||||
padding: 10px;
|
||||
text-align: center;
|
||||
cursor: pointer;
|
||||
background: transparent;
|
||||
border: none;
|
||||
color: var(--input-text, #ddd);
|
||||
border-bottom: 2px solid transparent;
|
||||
transition: all 0.15s ease;
|
||||
}
|
||||
|
||||
.modelpulse-tab:hover {
|
||||
background: var(--comfy-menu-bg-hover, #333);
|
||||
}
|
||||
|
||||
.modelpulse-tab.active {
|
||||
border-bottom-color: var(--primary-color, #4a9eff);
|
||||
color: var(--primary-color, #4a9eff);
|
||||
}
|
||||
|
||||
.modelpulse-controls {
|
||||
display: flex;
|
||||
padding: 8px 12px;
|
||||
gap: 8px;
|
||||
align-items: center;
|
||||
border-bottom: 1px solid var(--border-color, #444);
|
||||
}
|
||||
|
||||
.modelpulse-sort-label {
|
||||
font-size: 12px;
|
||||
color: var(--descrip-text, #999);
|
||||
}
|
||||
|
||||
.modelpulse-sort {
|
||||
flex: 1;
|
||||
padding: 6px 8px;
|
||||
background: var(--comfy-input-bg, #222);
|
||||
border: 1px solid var(--border-color, #444);
|
||||
border-radius: 4px;
|
||||
color: var(--input-text, #ddd);
|
||||
font-size: 12px;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.modelpulse-list {
|
||||
flex: 1;
|
||||
overflow-y: auto;
|
||||
padding: 8px 0;
|
||||
}
|
||||
|
||||
.modelpulse-category {
|
||||
margin-bottom: 4px;
|
||||
}
|
||||
|
||||
.modelpulse-category-header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
padding: 8px 12px;
|
||||
background: var(--comfy-menu-bg, #2a2a2a);
|
||||
cursor: pointer;
|
||||
user-select: none;
|
||||
}
|
||||
|
||||
.modelpulse-category-header:hover {
|
||||
background: var(--comfy-menu-bg-hover, #333);
|
||||
}
|
||||
|
||||
.modelpulse-category-icon {
|
||||
font-size: 14px;
|
||||
}
|
||||
|
||||
.modelpulse-category-name {
|
||||
flex: 1;
|
||||
font-weight: 500;
|
||||
}
|
||||
|
||||
.modelpulse-category-count {
|
||||
font-size: 11px;
|
||||
color: var(--descrip-text, #999);
|
||||
background: var(--comfy-input-bg, #222);
|
||||
padding: 2px 6px;
|
||||
border-radius: 10px;
|
||||
}
|
||||
|
||||
.modelpulse-category-arrow {
|
||||
font-size: 10px;
|
||||
transition: transform 0.15s ease;
|
||||
}
|
||||
|
||||
.modelpulse-category.collapsed .modelpulse-category-arrow {
|
||||
transform: rotate(-90deg);
|
||||
}
|
||||
|
||||
.modelpulse-category-models {
|
||||
display: block;
|
||||
}
|
||||
|
||||
.modelpulse-category.collapsed .modelpulse-category-models {
|
||||
display: none;
|
||||
}
|
||||
|
||||
.modelpulse-model {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
padding: 10px 12px 10px 32px;
|
||||
border-bottom: 1px solid var(--border-color, #333);
|
||||
cursor: default;
|
||||
}
|
||||
|
||||
.modelpulse-model:hover {
|
||||
background: var(--comfy-menu-bg-hover, #333);
|
||||
}
|
||||
|
||||
.modelpulse-model-name {
|
||||
font-weight: 500;
|
||||
margin-bottom: 4px;
|
||||
word-break: break-all;
|
||||
}
|
||||
|
||||
.modelpulse-model-stats {
|
||||
display: flex;
|
||||
gap: 16px;
|
||||
font-size: 11px;
|
||||
color: var(--descrip-text, #999);
|
||||
}
|
||||
|
||||
.modelpulse-model.stale .modelpulse-model-name {
|
||||
color: var(--error-text, #ff6b6b);
|
||||
}
|
||||
|
||||
.modelpulse-model.stale .modelpulse-last-used {
|
||||
color: var(--error-text, #ff6b6b);
|
||||
}
|
||||
|
||||
.modelpulse-empty {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
height: 200px;
|
||||
color: var(--descrip-text, #999);
|
||||
text-align: center;
|
||||
padding: 20px;
|
||||
}
|
||||
|
||||
.modelpulse-empty-icon {
|
||||
font-size: 48px;
|
||||
margin-bottom: 12px;
|
||||
opacity: 0.5;
|
||||
}
|
||||
|
||||
.modelpulse-loading {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
height: 100px;
|
||||
color: var(--descrip-text, #999);
|
||||
}
|
||||
`;
|
||||
return style;
|
||||
}
|
||||
|
||||
/**
|
||||
* Render the ModelPulse UI into the container element
|
||||
*/
|
||||
async function renderModelPulseUI(container) {
|
||||
// Add styles
|
||||
container.appendChild(createStyles());
|
||||
|
||||
// Create main container
|
||||
const main = document.createElement("div");
|
||||
main.className = "modelpulse-container";
|
||||
container.appendChild(main);
|
||||
|
||||
// Header
|
||||
const header = document.createElement("div");
|
||||
header.className = "modelpulse-header";
|
||||
header.innerHTML = `
|
||||
<span class="modelpulse-title">Usage</span>
|
||||
<button class="modelpulse-refresh" title="Refresh">↻</button>
|
||||
`;
|
||||
main.appendChild(header);
|
||||
|
||||
// Timeframe tabs
|
||||
const tabs = document.createElement("div");
|
||||
tabs.className = "modelpulse-tabs";
|
||||
tabs.innerHTML = `
|
||||
<button class="modelpulse-tab ${currentTimeframe === 'all' ? 'active' : ''}" data-timeframe="all">All Time</button>
|
||||
<button class="modelpulse-tab ${currentTimeframe === 'month' ? 'active' : ''}" data-timeframe="month">Month</button>
|
||||
<button class="modelpulse-tab ${currentTimeframe === 'week' ? 'active' : ''}" data-timeframe="week">Week</button>
|
||||
`;
|
||||
main.appendChild(tabs);
|
||||
|
||||
// Controls (sort dropdown)
|
||||
const controls = document.createElement("div");
|
||||
controls.className = "modelpulse-controls";
|
||||
controls.innerHTML = `
|
||||
<span class="modelpulse-sort-label">Sort:</span>
|
||||
<select class="modelpulse-sort">
|
||||
<option value="last_used" ${currentSort === 'last_used' ? 'selected' : ''}>Last Used</option>
|
||||
<option value="usage_count" ${currentSort === 'usage_count' ? 'selected' : ''}>Usage Count</option>
|
||||
<option value="name" ${currentSort === 'name' ? 'selected' : ''}>Name</option>
|
||||
</select>
|
||||
`;
|
||||
main.appendChild(controls);
|
||||
|
||||
// Model list container
|
||||
const list = document.createElement("div");
|
||||
list.className = "modelpulse-list";
|
||||
main.appendChild(list);
|
||||
|
||||
// Event handlers
|
||||
const refreshBtn = header.querySelector(".modelpulse-refresh");
|
||||
refreshBtn.addEventListener("click", () => refreshData(list, refreshBtn));
|
||||
|
||||
tabs.addEventListener("click", (e) => {
|
||||
const tab = e.target.closest(".modelpulse-tab");
|
||||
if (tab) {
|
||||
currentTimeframe = tab.dataset.timeframe;
|
||||
tabs.querySelectorAll(".modelpulse-tab").forEach(t => t.classList.remove("active"));
|
||||
tab.classList.add("active");
|
||||
refreshData(list, refreshBtn);
|
||||
}
|
||||
});
|
||||
|
||||
const sortSelect = controls.querySelector(".modelpulse-sort");
|
||||
sortSelect.addEventListener("change", () => {
|
||||
currentSort = sortSelect.value;
|
||||
refreshData(list, refreshBtn);
|
||||
});
|
||||
|
||||
// Initial load
|
||||
await refreshData(list, refreshBtn);
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh the data and re-render the list
|
||||
*/
|
||||
async function refreshData(listContainer, refreshBtn) {
|
||||
if (isLoading) return;
|
||||
|
||||
isLoading = true;
|
||||
refreshBtn.classList.add("loading");
|
||||
listContainer.innerHTML = '<div class="modelpulse-loading">Loading...</div>';
|
||||
|
||||
try {
|
||||
usageData = await fetchUsageData(currentTimeframe, currentSort);
|
||||
renderModelList(listContainer);
|
||||
} catch (error) {
|
||||
listContainer.innerHTML = `
|
||||
<div class="modelpulse-empty">
|
||||
<div class="modelpulse-empty-icon">⚠️</div>
|
||||
<div>Error loading data</div>
|
||||
</div>
|
||||
`;
|
||||
} finally {
|
||||
isLoading = false;
|
||||
refreshBtn.classList.remove("loading");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Render the model list grouped by category
|
||||
*/
|
||||
function renderModelList(container) {
|
||||
container.innerHTML = "";
|
||||
|
||||
if (!usageData || !usageData.models || usageData.models.length === 0) {
|
||||
container.innerHTML = `
|
||||
<div class="modelpulse-empty">
|
||||
<div class="modelpulse-empty-icon">📊</div>
|
||||
<div>No model usage recorded yet</div>
|
||||
<div style="margin-top: 8px; font-size: 11px;">Run some workflows to start tracking</div>
|
||||
</div>
|
||||
`;
|
||||
return;
|
||||
}
|
||||
|
||||
const grouped = groupByCategory(usageData.models);
|
||||
|
||||
// Sort categories by total usage
|
||||
const sortedCategories = Object.keys(grouped).sort((a, b) => {
|
||||
const aTotal = grouped[a].reduce((sum, m) => sum + m.timeframe_count, 0);
|
||||
const bTotal = grouped[b].reduce((sum, m) => sum + m.timeframe_count, 0);
|
||||
return bTotal - aTotal;
|
||||
});
|
||||
|
||||
for (const category of sortedCategories) {
|
||||
const models = grouped[category];
|
||||
const catInfo = getCategoryInfo(category);
|
||||
|
||||
const catDiv = document.createElement("div");
|
||||
catDiv.className = "modelpulse-category";
|
||||
|
||||
// Category header
|
||||
const catHeader = document.createElement("div");
|
||||
catHeader.className = "modelpulse-category-header";
|
||||
catHeader.innerHTML = `
|
||||
<span class="modelpulse-category-icon">${catInfo.icon}</span>
|
||||
<span class="modelpulse-category-name">${catInfo.name}</span>
|
||||
<span class="modelpulse-category-count">${models.length}</span>
|
||||
<span class="modelpulse-category-arrow">▼</span>
|
||||
`;
|
||||
catHeader.addEventListener("click", () => {
|
||||
catDiv.classList.toggle("collapsed");
|
||||
});
|
||||
catDiv.appendChild(catHeader);
|
||||
|
||||
// Models list
|
||||
const modelsDiv = document.createElement("div");
|
||||
modelsDiv.className = "modelpulse-category-models";
|
||||
|
||||
for (const model of models) {
|
||||
const modelDiv = document.createElement("div");
|
||||
modelDiv.className = "modelpulse-model";
|
||||
if (isStale(model.last_used)) {
|
||||
modelDiv.classList.add("stale");
|
||||
}
|
||||
|
||||
const countLabel = currentTimeframe === "all" ? "Uses" : "Uses (period)";
|
||||
const count = currentTimeframe === "all" ? model.usage_count : model.timeframe_count;
|
||||
|
||||
modelDiv.innerHTML = `
|
||||
<div class="modelpulse-model-name" title="${model.model_id}">${model.name}</div>
|
||||
<div class="modelpulse-model-stats">
|
||||
<span class="modelpulse-usage">${countLabel}: ${count}</span>
|
||||
<span class="modelpulse-last-used">Last: ${formatLastUsed(model.last_used)}</span>
|
||||
</div>
|
||||
`;
|
||||
modelsDiv.appendChild(modelDiv);
|
||||
}
|
||||
|
||||
catDiv.appendChild(modelsDiv);
|
||||
container.appendChild(catDiv);
|
||||
}
|
||||
}
|
||||
|
||||
// Register the extension
|
||||
app.registerExtension({
|
||||
name: "ModelPulse",
|
||||
async setup() {
|
||||
// Register sidebar tab
|
||||
app.extensionManager.registerSidebarTab({
|
||||
id: "modelpulse",
|
||||
icon: "pi pi-chart-bar",
|
||||
title: "Usage",
|
||||
tooltip: "Track model usage frequency",
|
||||
type: "custom",
|
||||
render: (el) => {
|
||||
renderModelPulseUI(el);
|
||||
}
|
||||
});
|
||||
|
||||
console.log("[ModelPulse] Extension loaded");
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1 @@
|
||||
"""ModelPulse Python package."""
|
||||
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
ModelPulse - Model Type Definitions
|
||||
|
||||
Maps ComfyUI loader node types to their model categories and input keys.
|
||||
"""
|
||||
|
||||
# Standard ComfyUI model loaders
|
||||
# Format: "NodeClassName": ("category", "input_key" or ["input_key1", "input_key2"])
|
||||
MODEL_LOADERS: dict[str, tuple[str, str | list[str]]] = {
|
||||
# Checkpoint loaders
|
||||
"CheckpointLoaderSimple": ("checkpoint", "ckpt_name"),
|
||||
"CheckpointLoader": ("checkpoint", "ckpt_name"),
|
||||
|
||||
# LoRA loaders
|
||||
"LoraLoader": ("lora", "lora_name"),
|
||||
"LoraLoaderModelOnly": ("lora", "lora_name"),
|
||||
|
||||
# VAE loaders
|
||||
"VAELoader": ("vae", "vae_name"),
|
||||
|
||||
# ControlNet loaders
|
||||
"ControlNetLoader": ("controlnet", "control_net_name"),
|
||||
|
||||
# CLIP loaders
|
||||
"CLIPLoader": ("clip", "clip_name"),
|
||||
"DualCLIPLoader": ("clip", ["clip_name1", "clip_name2"]),
|
||||
|
||||
# UNET loaders
|
||||
"UNETLoader": ("unet", "unet_name"),
|
||||
|
||||
# Upscaler loaders
|
||||
"UpscaleModelLoader": ("upscaler", "model_name"),
|
||||
|
||||
# Style model loaders
|
||||
"StyleModelLoader": ("style_model", "style_model_name"),
|
||||
|
||||
# GLIGEN loaders
|
||||
"GLIGENLoader": ("gligen", "gligen_name"),
|
||||
}
|
||||
|
||||
# Patterns for detecting custom node loaders (Impact Pack, Efficiency Nodes, etc.)
|
||||
# These are checked if the node class name isn't in MODEL_LOADERS
|
||||
LOADER_PATTERNS: list[tuple[str, str, list[str]]] = [
|
||||
# (pattern_in_class_name, category, possible_input_keys)
|
||||
("Checkpoint", "checkpoint", ["ckpt_name", "checkpoint", "model_name"]),
|
||||
("Lora", "lora", ["lora_name", "lora"]),
|
||||
("VAE", "vae", ["vae_name", "vae"]),
|
||||
("ControlNet", "controlnet", ["control_net_name", "controlnet", "control_net"]),
|
||||
("CLIP", "clip", ["clip_name", "clip"]),
|
||||
("UNET", "unet", ["unet_name", "unet"]),
|
||||
("Upscale", "upscaler", ["model_name", "upscale_model"]),
|
||||
]
|
||||
|
||||
# Model categories for display purposes
|
||||
MODEL_CATEGORIES: dict[str, dict[str, str]] = {
|
||||
"checkpoint": {
|
||||
"name": "Checkpoints",
|
||||
"icon": "pi-box",
|
||||
},
|
||||
"lora": {
|
||||
"name": "LoRAs",
|
||||
"icon": "pi-sparkles",
|
||||
},
|
||||
"vae": {
|
||||
"name": "VAEs",
|
||||
"icon": "pi-cog",
|
||||
},
|
||||
"controlnet": {
|
||||
"name": "ControlNets",
|
||||
"icon": "pi-sitemap",
|
||||
},
|
||||
"clip": {
|
||||
"name": "CLIP Models",
|
||||
"icon": "pi-comments",
|
||||
},
|
||||
"unet": {
|
||||
"name": "UNETs",
|
||||
"icon": "pi-server",
|
||||
},
|
||||
"upscaler": {
|
||||
"name": "Upscalers",
|
||||
"icon": "pi-expand",
|
||||
},
|
||||
"style_model": {
|
||||
"name": "Style Models",
|
||||
"icon": "pi-palette",
|
||||
},
|
||||
"gligen": {
|
||||
"name": "GLIGEN",
|
||||
"icon": "pi-th-large",
|
||||
},
|
||||
}
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
"""
|
||||
ModelPulse - API Routes
|
||||
|
||||
Custom HTTP routes for the ModelPulse extension.
|
||||
"""
|
||||
|
||||
from aiohttp import web
|
||||
from urllib.parse import unquote
|
||||
|
||||
from .tracking import ModelTracker
|
||||
|
||||
|
||||
def setup_routes(routes, tracker: ModelTracker) -> None:
|
||||
"""
|
||||
Register ModelPulse API routes.
|
||||
|
||||
Args:
|
||||
routes: PromptServer.instance.routes
|
||||
tracker: The ModelTracker instance
|
||||
"""
|
||||
|
||||
@routes.get("/modelpulse/usage")
|
||||
async def get_usage_data(request: web.Request) -> web.Response:
|
||||
"""
|
||||
Get all model usage data.
|
||||
|
||||
Query parameters:
|
||||
timeframe: "all" (default), "month", or "week"
|
||||
sort: "last_used" (default), "usage_count", or "name"
|
||||
category: Optional category filter
|
||||
"""
|
||||
timeframe = request.query.get("timeframe", "all")
|
||||
sort_by = request.query.get("sort", "last_used")
|
||||
category = request.query.get("category")
|
||||
|
||||
# Validate parameters
|
||||
if timeframe not in ("all", "month", "week"):
|
||||
timeframe = "all"
|
||||
if sort_by not in ("last_used", "usage_count", "name"):
|
||||
sort_by = "last_used"
|
||||
|
||||
data = tracker.get_usage_data(
|
||||
timeframe=timeframe,
|
||||
sort_by=sort_by,
|
||||
category=category,
|
||||
)
|
||||
return web.json_response(data)
|
||||
|
||||
@routes.get("/modelpulse/model/{model_id:.*}")
|
||||
async def get_model_detail(request: web.Request) -> web.Response:
|
||||
"""
|
||||
Get detailed usage data for a specific model.
|
||||
|
||||
Path parameters:
|
||||
model_id: The model identifier (category/filename)
|
||||
"""
|
||||
model_id = unquote(request.match_info["model_id"])
|
||||
data = tracker.get_model_detail(model_id)
|
||||
|
||||
if data is None:
|
||||
return web.json_response(
|
||||
{"error": "Model not found"},
|
||||
status=404,
|
||||
)
|
||||
|
||||
return web.json_response(data)
|
||||
|
||||
@routes.get("/modelpulse/categories")
|
||||
async def get_categories(request: web.Request) -> web.Response:
|
||||
"""Get list of model categories with counts."""
|
||||
from .model_types import MODEL_CATEGORIES
|
||||
|
||||
usage_data = tracker.get_usage_data()
|
||||
category_counts: dict[str, int] = {}
|
||||
|
||||
for model in usage_data["models"]:
|
||||
cat = model["category"]
|
||||
category_counts[cat] = category_counts.get(cat, 0) + 1
|
||||
|
||||
categories = []
|
||||
for cat_id, cat_info in MODEL_CATEGORIES.items():
|
||||
categories.append({
|
||||
"id": cat_id,
|
||||
"name": cat_info["name"],
|
||||
"icon": cat_info["icon"],
|
||||
"count": category_counts.get(cat_id, 0),
|
||||
})
|
||||
|
||||
# Sort by count descending
|
||||
categories.sort(key=lambda x: x["count"], reverse=True)
|
||||
|
||||
return web.json_response({"categories": categories})
|
||||
|
||||
@routes.post("/modelpulse/reset")
|
||||
async def reset_tracking(request: web.Request) -> web.Response:
|
||||
"""
|
||||
Reset all tracking data.
|
||||
|
||||
Requires confirmation in request body: {"confirm": true}
|
||||
"""
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return web.json_response(
|
||||
{"error": "Invalid JSON body"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
if body.get("confirm") is not True:
|
||||
return web.json_response(
|
||||
{"error": "Confirmation required. Send {\"confirm\": true}"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
tracker.reset()
|
||||
return web.json_response({"status": "ok", "message": "Tracking data reset"})
|
||||
|
||||
@routes.post("/modelpulse/cleanup")
|
||||
async def cleanup_old_data(request: web.Request) -> web.Response:
|
||||
"""
|
||||
Clean up old usage log entries.
|
||||
|
||||
Optional body parameter: {"max_days": 365}
|
||||
"""
|
||||
try:
|
||||
body = await request.json()
|
||||
max_days = body.get("max_days", 365)
|
||||
except Exception:
|
||||
max_days = 365
|
||||
|
||||
if not isinstance(max_days, int) or max_days < 1:
|
||||
max_days = 365
|
||||
|
||||
tracker.cleanup(max_days=max_days)
|
||||
return web.json_response({
|
||||
"status": "ok",
|
||||
"message": f"Cleaned up entries older than {max_days} days",
|
||||
})
|
||||
+258
@@ -0,0 +1,258 @@
|
||||
"""
|
||||
ModelPulse - Storage Layer
|
||||
|
||||
JSON-based persistence for model usage data with cross-platform file locking.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# Cross-platform file locking
|
||||
if sys.platform == "win32":
|
||||
import msvcrt
|
||||
|
||||
def _lock_file(f, exclusive: bool = False) -> None:
|
||||
"""Lock file on Windows using msvcrt."""
|
||||
msvcrt.locking(f.fileno(), msvcrt.LK_NBLCK if exclusive else msvcrt.LK_LOCK, 1)
|
||||
|
||||
def _unlock_file(f) -> None:
|
||||
"""Unlock file on Windows using msvcrt."""
|
||||
try:
|
||||
f.seek(0)
|
||||
msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1)
|
||||
except OSError:
|
||||
pass # Already unlocked or file closed
|
||||
else:
|
||||
import fcntl
|
||||
|
||||
def _lock_file(f, exclusive: bool = False) -> None:
|
||||
"""Lock file on Unix using fcntl."""
|
||||
fcntl.flock(f.fileno(), fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
|
||||
|
||||
def _unlock_file(f) -> None:
|
||||
"""Unlock file on Unix using fcntl."""
|
||||
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
class FileLock:
|
||||
"""Cross-platform file lock context manager using a .lock file."""
|
||||
|
||||
def __init__(self, path: Path, timeout: float = 10.0):
|
||||
self.lock_path = path.with_suffix(path.suffix + ".lock")
|
||||
self.timeout = timeout
|
||||
self._lock_file = None
|
||||
|
||||
def __enter__(self):
|
||||
start = time.time()
|
||||
while True:
|
||||
try:
|
||||
# Create lock file exclusively
|
||||
self._lock_file = open(self.lock_path, "x")
|
||||
return self
|
||||
except FileExistsError:
|
||||
# Check if lock is stale (older than timeout)
|
||||
try:
|
||||
lock_age = time.time() - self.lock_path.stat().st_mtime
|
||||
if lock_age > self.timeout:
|
||||
# Stale lock, remove it
|
||||
self.lock_path.unlink()
|
||||
continue
|
||||
except FileNotFoundError:
|
||||
continue # Lock was released, retry
|
||||
|
||||
if time.time() - start > self.timeout:
|
||||
raise TimeoutError(f"Could not acquire lock on {self.lock_path}")
|
||||
time.sleep(0.1)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if self._lock_file:
|
||||
self._lock_file.close()
|
||||
try:
|
||||
self.lock_path.unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
|
||||
def get_storage_path() -> Path:
|
||||
"""Get the path to the usage data file."""
|
||||
# ComfyUI stores user data in user/default/
|
||||
# We need to find the ComfyUI root directory
|
||||
# This file is in custom_nodes/ComfyUI-ModelPulse/py/
|
||||
current_file = Path(__file__).resolve()
|
||||
custom_nodes_dir = current_file.parent.parent.parent # Go up to custom_nodes
|
||||
comfyui_root = custom_nodes_dir.parent # Go up to ComfyUI root
|
||||
|
||||
storage_dir = comfyui_root / "user" / "default" / "modelpulse"
|
||||
storage_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return storage_dir / "usage_data.json"
|
||||
|
||||
|
||||
def get_backup_path() -> Path:
|
||||
"""Get the path for backup files."""
|
||||
storage_path = get_storage_path()
|
||||
return storage_path.with_suffix(".backup.json")
|
||||
|
||||
|
||||
def create_empty_data() -> dict[str, Any]:
|
||||
"""Create an empty data structure."""
|
||||
now = datetime.utcnow().isoformat() + "Z"
|
||||
return {
|
||||
"version": SCHEMA_VERSION,
|
||||
"models": {},
|
||||
"metadata": {
|
||||
"tracking_started": now,
|
||||
"last_updated": now,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_data() -> dict[str, Any]:
|
||||
"""
|
||||
Load usage data from JSON file.
|
||||
|
||||
Creates the file with empty data if it doesn't exist.
|
||||
Handles schema migrations if needed.
|
||||
"""
|
||||
storage_path = get_storage_path()
|
||||
|
||||
if not storage_path.exists():
|
||||
data = create_empty_data()
|
||||
save_data(data)
|
||||
return data
|
||||
|
||||
try:
|
||||
with FileLock(storage_path):
|
||||
with open(storage_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Handle schema migrations
|
||||
data = migrate_schema(data)
|
||||
return data
|
||||
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
# Corrupted file - backup and create fresh
|
||||
print(f"[ModelPulse] Warning: Corrupted data file, creating backup. Error: {e}")
|
||||
backup_corrupted_file(storage_path)
|
||||
data = create_empty_data()
|
||||
save_data(data)
|
||||
return data
|
||||
|
||||
|
||||
def save_data(data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Save usage data to JSON file with file locking.
|
||||
|
||||
Uses exclusive lock to prevent concurrent writes.
|
||||
"""
|
||||
storage_path = get_storage_path()
|
||||
|
||||
# Update last_updated timestamp
|
||||
data["metadata"]["last_updated"] = datetime.utcnow().isoformat() + "Z"
|
||||
|
||||
# Write to temp file first, then rename (atomic operation)
|
||||
temp_path = storage_path.with_suffix(".tmp")
|
||||
|
||||
try:
|
||||
with FileLock(storage_path):
|
||||
with open(temp_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
|
||||
# Atomic rename
|
||||
temp_path.replace(storage_path)
|
||||
|
||||
except Exception as e:
|
||||
# Clean up temp file if it exists
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
raise e
|
||||
|
||||
|
||||
def migrate_schema(data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Migrate data to the current schema version.
|
||||
|
||||
Creates a backup before migration.
|
||||
"""
|
||||
current_version = data.get("version", 0)
|
||||
|
||||
if current_version == SCHEMA_VERSION:
|
||||
return data
|
||||
|
||||
# Backup before migration
|
||||
backup_path = get_backup_path()
|
||||
with open(backup_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
print(f"[ModelPulse] Created backup at {backup_path} before schema migration")
|
||||
|
||||
# Apply migrations sequentially
|
||||
if current_version < 1:
|
||||
data = migrate_to_v1(data)
|
||||
|
||||
# Future migrations would go here:
|
||||
# if current_version < 2:
|
||||
# data = migrate_to_v2(data)
|
||||
|
||||
data["version"] = SCHEMA_VERSION
|
||||
return data
|
||||
|
||||
|
||||
def migrate_to_v1(data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Migrate from pre-v1 (or no version) to v1 schema."""
|
||||
# Ensure required structure exists
|
||||
if "models" not in data:
|
||||
data["models"] = {}
|
||||
|
||||
if "metadata" not in data:
|
||||
now = datetime.utcnow().isoformat() + "Z"
|
||||
data["metadata"] = {
|
||||
"tracking_started": now,
|
||||
"last_updated": now,
|
||||
}
|
||||
|
||||
# Ensure all models have required fields
|
||||
for model_id, model_data in data.get("models", {}).items():
|
||||
if "usage_log" not in model_data:
|
||||
model_data["usage_log"] = []
|
||||
if "first_used" not in model_data:
|
||||
model_data["first_used"] = model_data.get(
|
||||
"last_used", datetime.utcnow().isoformat() + "Z"
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def backup_corrupted_file(storage_path: Path) -> None:
|
||||
"""Backup a corrupted file for manual recovery."""
|
||||
if storage_path.exists():
|
||||
timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = storage_path.with_suffix(f".corrupted_{timestamp}.json")
|
||||
storage_path.rename(backup_path)
|
||||
print(f"[ModelPulse] Backed up corrupted file to {backup_path}")
|
||||
|
||||
|
||||
def cleanup_old_usage_logs(data: dict[str, Any], max_days: int = 365) -> dict[str, Any]:
|
||||
"""
|
||||
Clean up old daily usage logs to prevent unbounded growth.
|
||||
|
||||
Keeps the last `max_days` days of daily data.
|
||||
"""
|
||||
cutoff_date = datetime.utcnow().date()
|
||||
from datetime import timedelta
|
||||
|
||||
cutoff_str = (cutoff_date - timedelta(days=max_days)).isoformat()
|
||||
|
||||
for model_data in data.get("models", {}).values():
|
||||
usage_log = model_data.get("usage_log", [])
|
||||
model_data["usage_log"] = [
|
||||
entry for entry in usage_log if entry.get("date", "") >= cutoff_str
|
||||
]
|
||||
|
||||
return data
|
||||
+263
@@ -0,0 +1,263 @@
|
||||
"""
|
||||
ModelPulse - Core Tracking Logic
|
||||
|
||||
Extracts models from prompts and records usage statistics.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from .model_types import MODEL_LOADERS, LOADER_PATTERNS
|
||||
from .storage import load_data, save_data, cleanup_old_usage_logs
|
||||
|
||||
|
||||
class ModelTracker:
|
||||
"""Tracks model usage across ComfyUI workflow executions."""
|
||||
|
||||
def __init__(self):
|
||||
self._data: dict[str, Any] | None = None
|
||||
|
||||
@property
|
||||
def data(self) -> dict[str, Any]:
|
||||
"""Lazy-load data from storage."""
|
||||
if self._data is None:
|
||||
self._data = load_data()
|
||||
return self._data
|
||||
|
||||
def _save(self) -> None:
|
||||
"""Save current data to storage."""
|
||||
if self._data is not None:
|
||||
save_data(self._data)
|
||||
|
||||
def extract_models_from_prompt(self, prompt: dict[str, Any]) -> list[dict[str, str]]:
|
||||
"""
|
||||
Extract all model references from a ComfyUI prompt.
|
||||
|
||||
Args:
|
||||
prompt: The prompt dictionary from ComfyUI execution
|
||||
|
||||
Returns:
|
||||
List of dicts with 'category', 'name', and 'model_id' keys
|
||||
"""
|
||||
models = []
|
||||
seen = set() # Avoid duplicates within same prompt
|
||||
|
||||
for node_id, node_data in prompt.items():
|
||||
class_type = node_data.get("class_type", "")
|
||||
inputs = node_data.get("inputs", {})
|
||||
|
||||
# Check standard loaders first
|
||||
if class_type in MODEL_LOADERS:
|
||||
category, input_keys = MODEL_LOADERS[class_type]
|
||||
extracted = self._extract_from_inputs(category, input_keys, inputs)
|
||||
for model in extracted:
|
||||
if model["model_id"] not in seen:
|
||||
seen.add(model["model_id"])
|
||||
models.append(model)
|
||||
|
||||
# Check pattern-based loaders for custom nodes
|
||||
elif "Loader" in class_type:
|
||||
extracted = self._extract_from_patterns(class_type, inputs)
|
||||
for model in extracted:
|
||||
if model["model_id"] not in seen:
|
||||
seen.add(model["model_id"])
|
||||
models.append(model)
|
||||
|
||||
return models
|
||||
|
||||
def _extract_from_inputs(
|
||||
self, category: str, input_keys: str | list[str], inputs: dict[str, Any]
|
||||
) -> list[dict[str, str]]:
|
||||
"""Extract model names from node inputs."""
|
||||
models = []
|
||||
|
||||
if isinstance(input_keys, str):
|
||||
input_keys = [input_keys]
|
||||
|
||||
for key in input_keys:
|
||||
value = inputs.get(key)
|
||||
if value and isinstance(value, str):
|
||||
model_id = f"{category}/{value}"
|
||||
models.append({
|
||||
"category": category,
|
||||
"name": value,
|
||||
"model_id": model_id,
|
||||
})
|
||||
|
||||
return models
|
||||
|
||||
def _extract_from_patterns(
|
||||
self, class_type: str, inputs: dict[str, Any]
|
||||
) -> list[dict[str, str]]:
|
||||
"""Extract models using pattern matching for custom node loaders."""
|
||||
models = []
|
||||
|
||||
for pattern, category, possible_keys in LOADER_PATTERNS:
|
||||
if pattern.lower() in class_type.lower():
|
||||
for key in possible_keys:
|
||||
value = inputs.get(key)
|
||||
if value and isinstance(value, str):
|
||||
model_id = f"{category}/{value}"
|
||||
models.append({
|
||||
"category": category,
|
||||
"name": value,
|
||||
"model_id": model_id,
|
||||
})
|
||||
break # Found a match for this pattern
|
||||
break # Only match first pattern
|
||||
|
||||
return models
|
||||
|
||||
def record_usage(self, models: list[dict[str, str]]) -> None:
|
||||
"""
|
||||
Record usage for a list of models.
|
||||
|
||||
Args:
|
||||
models: List of model dicts from extract_models_from_prompt
|
||||
"""
|
||||
if not models:
|
||||
return
|
||||
|
||||
now = datetime.utcnow()
|
||||
now_iso = now.isoformat() + "Z"
|
||||
today = now.date().isoformat()
|
||||
|
||||
for model in models:
|
||||
model_id = model["model_id"]
|
||||
|
||||
if model_id not in self.data["models"]:
|
||||
# New model - create entry
|
||||
self.data["models"][model_id] = {
|
||||
"category": model["category"],
|
||||
"name": model["name"],
|
||||
"path": model_id, # We don't have full path, use model_id
|
||||
"first_used": now_iso,
|
||||
"last_used": now_iso,
|
||||
"usage_count": 0,
|
||||
"usage_log": [],
|
||||
}
|
||||
|
||||
model_data = self.data["models"][model_id]
|
||||
|
||||
# Update timestamps and count
|
||||
model_data["last_used"] = now_iso
|
||||
model_data["usage_count"] += 1
|
||||
|
||||
# Update daily usage log
|
||||
usage_log = model_data["usage_log"]
|
||||
if usage_log and usage_log[-1]["date"] == today:
|
||||
usage_log[-1]["count"] += 1
|
||||
else:
|
||||
usage_log.append({"date": today, "count": 1})
|
||||
|
||||
self._save()
|
||||
|
||||
def extract_and_record_models(self, prompt: dict[str, Any]) -> None:
|
||||
"""
|
||||
Extract models from prompt and record their usage.
|
||||
|
||||
This is the main entry point called from the PromptExecutor hook.
|
||||
"""
|
||||
models = self.extract_models_from_prompt(prompt)
|
||||
self.record_usage(models)
|
||||
|
||||
def get_usage_data(
|
||||
self,
|
||||
timeframe: str = "all",
|
||||
sort_by: str = "last_used",
|
||||
category: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get filtered and sorted usage data.
|
||||
|
||||
Args:
|
||||
timeframe: "all", "month", or "week"
|
||||
sort_by: "last_used", "usage_count", or "name"
|
||||
category: Optional category filter
|
||||
|
||||
Returns:
|
||||
Dict with 'models' list and 'metadata'
|
||||
"""
|
||||
models_list = []
|
||||
now = datetime.utcnow()
|
||||
|
||||
# Calculate timeframe boundaries
|
||||
if timeframe == "week":
|
||||
cutoff = now - timedelta(days=7)
|
||||
elif timeframe == "month":
|
||||
cutoff = now - timedelta(days=30)
|
||||
else:
|
||||
cutoff = None
|
||||
|
||||
cutoff_date = cutoff.date().isoformat() if cutoff else None
|
||||
|
||||
for model_id, model_data in self.data["models"].items():
|
||||
# Filter by category if specified
|
||||
if category and model_data["category"] != category:
|
||||
continue
|
||||
|
||||
# Calculate timeframe-specific usage count
|
||||
if cutoff_date:
|
||||
timeframe_count = sum(
|
||||
entry["count"]
|
||||
for entry in model_data.get("usage_log", [])
|
||||
if entry["date"] >= cutoff_date
|
||||
)
|
||||
else:
|
||||
timeframe_count = model_data["usage_count"]
|
||||
|
||||
models_list.append({
|
||||
"model_id": model_id,
|
||||
"category": model_data["category"],
|
||||
"name": model_data["name"],
|
||||
"first_used": model_data["first_used"],
|
||||
"last_used": model_data["last_used"],
|
||||
"usage_count": model_data["usage_count"],
|
||||
"timeframe_count": timeframe_count,
|
||||
})
|
||||
|
||||
# Sort the results
|
||||
if sort_by == "usage_count":
|
||||
models_list.sort(key=lambda x: x["timeframe_count"], reverse=True)
|
||||
elif sort_by == "name":
|
||||
models_list.sort(key=lambda x: x["name"].lower())
|
||||
else: # last_used (default)
|
||||
models_list.sort(key=lambda x: x["last_used"], reverse=True)
|
||||
|
||||
return {
|
||||
"models": models_list,
|
||||
"metadata": self.data["metadata"],
|
||||
"timeframe": timeframe,
|
||||
"sort_by": sort_by,
|
||||
}
|
||||
|
||||
def get_model_detail(self, model_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get detailed usage data for a specific model.
|
||||
|
||||
Args:
|
||||
model_id: The model identifier (category/filename)
|
||||
|
||||
Returns:
|
||||
Model data dict or None if not found
|
||||
"""
|
||||
model_data = self.data["models"].get(model_id)
|
||||
if not model_data:
|
||||
return None
|
||||
|
||||
return {
|
||||
"model_id": model_id,
|
||||
**model_data,
|
||||
}
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset all tracking data."""
|
||||
from .storage import create_empty_data
|
||||
|
||||
self._data = create_empty_data()
|
||||
self._save()
|
||||
|
||||
def cleanup(self, max_days: int = 365) -> None:
|
||||
"""Clean up old usage logs."""
|
||||
self._data = cleanup_old_usage_logs(self.data, max_days)
|
||||
self._save()
|
||||
Reference in New Issue
Block a user