mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
refactor: simplify featured models and fix HuggingFace search types
This commit is contained in:
Douwe Osinga
@@ -1510,6 +1510,7 @@ async fn handle_local_models_command(command: LocalModelsCommand) -> Result<()>
|
||||
source_url: file.download_url.clone(),
|
||||
settings: Default::default(),
|
||||
size_bytes: file.size_bytes,
|
||||
context_limit: 8192, // Default context limit
|
||||
};
|
||||
|
||||
{
|
||||
|
||||
@@ -604,10 +604,11 @@ derive_utoipa!(Icon as IconSchema);
|
||||
super::routes::local_inference::LocalModelResponse,
|
||||
super::routes::local_inference::ModelDownloadStatus,
|
||||
super::routes::local_inference::DownloadProgressResponse,
|
||||
|
||||
super::routes::local_inference::HfSearchResult,
|
||||
super::routes::local_inference::HfQuantVariant,
|
||||
super::routes::local_inference::DownloadModelRequest,
|
||||
goose::providers::local_inference::hf_models::HfModelInfo,
|
||||
goose::providers::local_inference::hf_models::HfGgufFile,
|
||||
goose::providers::local_inference::hf_models::HfQuantVariant,
|
||||
goose::providers::local_inference::local_model_registry::ModelSettings,
|
||||
goose::providers::local_inference::local_model_registry::SamplingConfig,
|
||||
DownloadProgress,
|
||||
|
||||
@@ -46,6 +46,7 @@ pub struct LocalModelResponse {
|
||||
pub size_bytes: u64,
|
||||
pub status: ModelDownloadStatus,
|
||||
pub recommended: bool,
|
||||
pub context_limit: u32,
|
||||
pub settings: ModelSettings,
|
||||
}
|
||||
|
||||
@@ -53,6 +54,20 @@ fn handle_internal_error<T: std::fmt::Display>(msg: T) -> ErrorResponse {
|
||||
ErrorResponse::internal(format!("{}", msg))
|
||||
}
|
||||
|
||||
/// Get context limit for a model based on its repo_id
|
||||
fn get_context_limit_for_model(repo_id: &str) -> u32 {
|
||||
// Known context limits for featured models
|
||||
if repo_id.contains("Llama-3.2") {
|
||||
131072 // 128K context
|
||||
} else if repo_id.contains("Hermes-2-Pro") {
|
||||
8192 // 8K context
|
||||
} else if repo_id.contains("Mistral-Small") {
|
||||
32768 // 32K context
|
||||
} else {
|
||||
8192 // Default 8K context
|
||||
}
|
||||
}
|
||||
|
||||
/// Recommend the best model based on available system memory
|
||||
fn recommend_model_id() -> Option<String> {
|
||||
let mut sys = System::new_all();
|
||||
@@ -141,6 +156,7 @@ async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
|
||||
source_url: hf_file.download_url,
|
||||
settings: ModelSettings::default(),
|
||||
size_bytes: hf_file.size_bytes,
|
||||
context_limit: get_context_limit_for_model(repo_id),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -213,6 +229,7 @@ pub async fn list_local_models() -> Result<Json<Vec<LocalModelResponse>>, ErrorR
|
||||
size_bytes,
|
||||
status,
|
||||
recommended: recommended_id.as_deref() == Some(&entry.id),
|
||||
context_limit: entry.context_limit,
|
||||
settings: entry.settings.clone(),
|
||||
});
|
||||
}
|
||||
@@ -376,6 +393,7 @@ pub async fn download_hf_model(
|
||||
source_url: download_url.clone(),
|
||||
settings: ModelSettings::default(),
|
||||
size_bytes: hf_file.size_bytes,
|
||||
context_limit: get_context_limit_for_model(&repo_id),
|
||||
};
|
||||
|
||||
{
|
||||
|
||||
@@ -201,6 +201,12 @@ pub struct LocalModelEntry {
|
||||
pub settings: ModelSettings,
|
||||
#[serde(default)]
|
||||
pub size_bytes: u64,
|
||||
#[serde(default = "default_context_limit")]
|
||||
pub context_limit: u32,
|
||||
}
|
||||
|
||||
fn default_context_limit() -> u32 {
|
||||
8192 // Default context limit for most models
|
||||
}
|
||||
|
||||
impl LocalModelEntry {
|
||||
@@ -307,6 +313,7 @@ impl LocalModelRegistry {
|
||||
source_url: legacy.source_url.to_string(),
|
||||
settings: ModelSettings::default(),
|
||||
size_bytes: 0,
|
||||
context_limit: 8192, // Default for legacy models
|
||||
});
|
||||
changed = true;
|
||||
}
|
||||
|
||||
+73
-171
@@ -1837,7 +1837,7 @@
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HfDownloadRequest"
|
||||
"$ref": "#/components/schemas/DownloadModelRequest"
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -1847,18 +1847,15 @@
|
||||
"202": {
|
||||
"description": "Download started",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"text/plain": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HfDownloadResponse"
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"400": {
|
||||
"description": "Invalid request"
|
||||
},
|
||||
"500": {
|
||||
"description": "Internal server error"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1940,7 +1937,7 @@
|
||||
}
|
||||
},
|
||||
"404": {
|
||||
"description": "Download not found"
|
||||
"description": "No active download"
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -1964,7 +1961,7 @@
|
||||
"description": "Download cancelled"
|
||||
},
|
||||
"404": {
|
||||
"description": "Download not found"
|
||||
"description": "No active download"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2072,17 +2069,17 @@
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "GGUF quantization variants in repo",
|
||||
"description": "GGUF files in the repo",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/RepoVariantsResponse"
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/HfQuantVariant"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"500": {
|
||||
"description": "Failed to fetch repo files"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2102,17 +2099,6 @@
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "limit",
|
||||
"in": "query",
|
||||
"description": "Max results (default 20)",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"type": "integer",
|
||||
"nullable": true,
|
||||
"minimum": 0
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
@@ -2123,7 +2109,7 @@
|
||||
"schema": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/HfModelInfo"
|
||||
"$ref": "#/components/schemas/HfSearchResult"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4259,6 +4245,25 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"DownloadModelRequest": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filename": {
|
||||
"type": "string",
|
||||
"nullable": true
|
||||
},
|
||||
"repo_id": {
|
||||
"type": "string",
|
||||
"description": "Alternative: provide repo_id and filename separately",
|
||||
"nullable": true
|
||||
},
|
||||
"spec": {
|
||||
"type": "string",
|
||||
"description": "Model spec like \"bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M\"",
|
||||
"nullable": true
|
||||
}
|
||||
}
|
||||
},
|
||||
"DownloadProgress": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
@@ -4320,7 +4325,8 @@
|
||||
"model_id",
|
||||
"status",
|
||||
"bytes_downloaded",
|
||||
"total_bytes"
|
||||
"total_bytes",
|
||||
"progress_percent"
|
||||
],
|
||||
"properties": {
|
||||
"bytes_downloaded": {
|
||||
@@ -4337,6 +4343,10 @@
|
||||
"model_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"progress_percent": {
|
||||
"type": "number",
|
||||
"format": "float"
|
||||
},
|
||||
"speed_bps": {
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
@@ -4932,34 +4942,6 @@
|
||||
}
|
||||
]
|
||||
},
|
||||
"HfDownloadRequest": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filename": {
|
||||
"type": "string",
|
||||
"nullable": true
|
||||
},
|
||||
"repo_id": {
|
||||
"type": "string",
|
||||
"nullable": true
|
||||
},
|
||||
"spec": {
|
||||
"type": "string",
|
||||
"nullable": true
|
||||
}
|
||||
}
|
||||
},
|
||||
"HfDownloadResponse": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"model_id"
|
||||
],
|
||||
"properties": {
|
||||
"model_id": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"HfGgufFile": {
|
||||
"type": "object",
|
||||
"description": "A single downloadable GGUF file (used internally and for downloads).",
|
||||
@@ -5020,19 +5002,14 @@
|
||||
},
|
||||
"HfQuantVariant": {
|
||||
"type": "object",
|
||||
"description": "A quantization variant — groups sharded files into one logical entry.",
|
||||
"required": [
|
||||
"quantization",
|
||||
"size_bytes",
|
||||
"filename",
|
||||
"size_bytes",
|
||||
"quantization",
|
||||
"download_url",
|
||||
"description",
|
||||
"quality_rank"
|
||||
],
|
||||
"properties": {
|
||||
"description": {
|
||||
"type": "string"
|
||||
},
|
||||
"download_url": {
|
||||
"type": "string"
|
||||
},
|
||||
@@ -5054,6 +5031,37 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"HfSearchResult": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"id",
|
||||
"author",
|
||||
"name",
|
||||
"downloads",
|
||||
"likes"
|
||||
],
|
||||
"properties": {
|
||||
"author": {
|
||||
"type": "string"
|
||||
},
|
||||
"downloads": {
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
"minimum": 0
|
||||
},
|
||||
"id": {
|
||||
"type": "string"
|
||||
},
|
||||
"likes": {
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
"minimum": 0
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"Icon": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
@@ -5241,7 +5249,7 @@
|
||||
},
|
||||
"LocalModelResponse": {
|
||||
"type": "object",
|
||||
"description": "Response for a single local model",
|
||||
"description": "Response for a local model",
|
||||
"required": [
|
||||
"id",
|
||||
"display_name",
|
||||
@@ -5251,13 +5259,13 @@
|
||||
"size_bytes",
|
||||
"status",
|
||||
"recommended",
|
||||
"context_limit",
|
||||
"settings"
|
||||
],
|
||||
"properties": {
|
||||
"context_limit": {
|
||||
"type": "integer",
|
||||
"format": "int32",
|
||||
"nullable": true,
|
||||
"minimum": 0
|
||||
},
|
||||
"display_name": {
|
||||
@@ -5287,85 +5295,7 @@
|
||||
"minimum": 0
|
||||
},
|
||||
"status": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"required": [
|
||||
"state"
|
||||
],
|
||||
"properties": {
|
||||
"state": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"NotDownloaded"
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"required": [
|
||||
"progress_percent",
|
||||
"bytes_downloaded",
|
||||
"total_bytes",
|
||||
"state"
|
||||
],
|
||||
"properties": {
|
||||
"bytes_downloaded": {
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
"minimum": 0
|
||||
},
|
||||
"progress_percent": {
|
||||
"type": "number",
|
||||
"format": "float"
|
||||
},
|
||||
"speed_bps": {
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
"nullable": true,
|
||||
"minimum": 0
|
||||
},
|
||||
"state": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"Downloading"
|
||||
]
|
||||
},
|
||||
"total_bytes": {
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
"minimum": 0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"required": [
|
||||
"state"
|
||||
],
|
||||
"properties": {
|
||||
"state": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"Downloaded"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"description": "Download status for a local model (mirrors goose::providers::local_inference::local_model_registry::ModelDownloadStatus)",
|
||||
"discriminator": {
|
||||
"propertyName": "state"
|
||||
}
|
||||
},
|
||||
"tier": {
|
||||
"allOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/ModelTier"
|
||||
}
|
||||
],
|
||||
"nullable": true
|
||||
"$ref": "#/components/schemas/ModelDownloadStatus"
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -5957,7 +5887,7 @@
|
||||
}
|
||||
}
|
||||
],
|
||||
"description": "Download status for a local model (mirrors goose::providers::local_inference::local_model_registry::ModelDownloadStatus)",
|
||||
"description": "Download status for local models",
|
||||
"discriminator": {
|
||||
"propertyName": "state"
|
||||
}
|
||||
@@ -6151,15 +6081,6 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"ModelTier": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"Tiny",
|
||||
"Small",
|
||||
"Medium",
|
||||
"Large"
|
||||
]
|
||||
},
|
||||
"ParseRecipeRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
@@ -6748,25 +6669,6 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"RepoVariantsResponse": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"variants"
|
||||
],
|
||||
"properties": {
|
||||
"recommended_index": {
|
||||
"type": "integer",
|
||||
"nullable": true,
|
||||
"minimum": 0
|
||||
},
|
||||
"variants": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/HfQuantVariant"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"ResourceContents": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -232,6 +232,18 @@ export type DictationProviderStatus = {
|
||||
uses_provider_config: boolean;
|
||||
};
|
||||
|
||||
export type DownloadModelRequest = {
|
||||
filename?: string | null;
|
||||
/**
|
||||
* Alternative: provide repo_id and filename separately
|
||||
*/
|
||||
repo_id?: string | null;
|
||||
/**
|
||||
* Model spec like "bartowski/Llama-3.2-3B-Instruct-GGUF:Q4_K_M"
|
||||
*/
|
||||
spec?: string | null;
|
||||
};
|
||||
|
||||
export type DownloadProgress = {
|
||||
/**
|
||||
* Bytes downloaded so far
|
||||
@@ -268,6 +280,7 @@ export type DownloadProgressResponse = {
|
||||
bytes_downloaded: number;
|
||||
eta_seconds?: number | null;
|
||||
model_id: string;
|
||||
progress_percent: number;
|
||||
speed_bps?: number | null;
|
||||
status: string;
|
||||
total_bytes: number;
|
||||
@@ -455,16 +468,6 @@ export type GooseApp = McpAppResource & (WindowProps | null) & {
|
||||
prd?: string | null;
|
||||
};
|
||||
|
||||
export type HfDownloadRequest = {
|
||||
filename?: string | null;
|
||||
repo_id?: string | null;
|
||||
spec?: string | null;
|
||||
};
|
||||
|
||||
export type HfDownloadResponse = {
|
||||
model_id: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* A single downloadable GGUF file (used internally and for downloads).
|
||||
*/
|
||||
@@ -483,11 +486,7 @@ export type HfModelInfo = {
|
||||
repo_id: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* A quantization variant — groups sharded files into one logical entry.
|
||||
*/
|
||||
export type HfQuantVariant = {
|
||||
description: string;
|
||||
download_url: string;
|
||||
filename: string;
|
||||
quality_rank: number;
|
||||
@@ -495,6 +494,14 @@ export type HfQuantVariant = {
|
||||
size_bytes: number;
|
||||
};
|
||||
|
||||
export type HfSearchResult = {
|
||||
author: string;
|
||||
downloads: number;
|
||||
id: string;
|
||||
likes: number;
|
||||
name: string;
|
||||
};
|
||||
|
||||
export type Icon = {
|
||||
mimeType?: string;
|
||||
sizes?: Array<string>;
|
||||
@@ -561,10 +568,10 @@ export type LoadedProvider = {
|
||||
};
|
||||
|
||||
/**
|
||||
* Response for a single local model
|
||||
* Response for a local model
|
||||
*/
|
||||
export type LocalModelResponse = {
|
||||
context_limit?: number | null;
|
||||
context_limit: number;
|
||||
display_name: string;
|
||||
filename: string;
|
||||
id: string;
|
||||
@@ -573,21 +580,7 @@ export type LocalModelResponse = {
|
||||
repo_id: string;
|
||||
settings: ModelSettings;
|
||||
size_bytes: number;
|
||||
/**
|
||||
* Download status for a local model (mirrors goose::providers::local_inference::local_model_registry::ModelDownloadStatus)
|
||||
*/
|
||||
status: {
|
||||
state: 'NotDownloaded';
|
||||
} | {
|
||||
bytes_downloaded: number;
|
||||
progress_percent: number;
|
||||
speed_bps?: number | null;
|
||||
state: 'Downloading';
|
||||
total_bytes: number;
|
||||
} | {
|
||||
state: 'Downloaded';
|
||||
};
|
||||
tier?: ModelTier | null;
|
||||
status: ModelDownloadStatus;
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -718,7 +711,7 @@ export type ModelConfig = {
|
||||
};
|
||||
|
||||
/**
|
||||
* Download status for a local model (mirrors goose::providers::local_inference::local_model_registry::ModelDownloadStatus)
|
||||
* Download status for local models
|
||||
*/
|
||||
export type ModelDownloadStatus = {
|
||||
state: 'NotDownloaded';
|
||||
@@ -801,8 +794,6 @@ export type ModelSettings = {
|
||||
use_mlock?: boolean;
|
||||
};
|
||||
|
||||
export type ModelTier = 'Tiny' | 'Small' | 'Medium' | 'Large';
|
||||
|
||||
export type ParseRecipeRequest = {
|
||||
content: string;
|
||||
};
|
||||
@@ -1022,11 +1013,6 @@ export type RemoveExtensionRequest = {
|
||||
session_id: string;
|
||||
};
|
||||
|
||||
export type RepoVariantsResponse = {
|
||||
recommended_index?: number | null;
|
||||
variants: Array<HfQuantVariant>;
|
||||
};
|
||||
|
||||
export type ResourceContents = {
|
||||
_meta?: {
|
||||
[key: string]: unknown;
|
||||
@@ -2971,7 +2957,7 @@ export type StartTetrateSetupResponses = {
|
||||
export type StartTetrateSetupResponse = StartTetrateSetupResponses[keyof StartTetrateSetupResponses];
|
||||
|
||||
export type DownloadHfModelData = {
|
||||
body: HfDownloadRequest;
|
||||
body: DownloadModelRequest;
|
||||
path?: never;
|
||||
query?: never;
|
||||
url: '/local-inference/download';
|
||||
@@ -2982,17 +2968,13 @@ export type DownloadHfModelErrors = {
|
||||
* Invalid request
|
||||
*/
|
||||
400: unknown;
|
||||
/**
|
||||
* Internal server error
|
||||
*/
|
||||
500: unknown;
|
||||
};
|
||||
|
||||
export type DownloadHfModelResponses = {
|
||||
/**
|
||||
* Download started
|
||||
*/
|
||||
202: HfDownloadResponse;
|
||||
202: string;
|
||||
};
|
||||
|
||||
export type DownloadHfModelResponse = DownloadHfModelResponses[keyof DownloadHfModelResponses];
|
||||
@@ -3047,7 +3029,7 @@ export type CancelLocalModelDownloadData = {
|
||||
|
||||
export type CancelLocalModelDownloadErrors = {
|
||||
/**
|
||||
* Download not found
|
||||
* No active download
|
||||
*/
|
||||
404: unknown;
|
||||
};
|
||||
@@ -3070,7 +3052,7 @@ export type GetLocalModelDownloadProgressData = {
|
||||
|
||||
export type GetLocalModelDownloadProgressErrors = {
|
||||
/**
|
||||
* Download not found
|
||||
* No active download
|
||||
*/
|
||||
404: unknown;
|
||||
};
|
||||
@@ -3148,18 +3130,11 @@ export type GetRepoFilesData = {
|
||||
url: '/local-inference/repo/{author}/{repo}/files';
|
||||
};
|
||||
|
||||
export type GetRepoFilesErrors = {
|
||||
/**
|
||||
* Failed to fetch repo files
|
||||
*/
|
||||
500: unknown;
|
||||
};
|
||||
|
||||
export type GetRepoFilesResponses = {
|
||||
/**
|
||||
* GGUF quantization variants in repo
|
||||
* GGUF files in the repo
|
||||
*/
|
||||
200: RepoVariantsResponse;
|
||||
200: Array<HfQuantVariant>;
|
||||
};
|
||||
|
||||
export type GetRepoFilesResponse = GetRepoFilesResponses[keyof GetRepoFilesResponses];
|
||||
@@ -3172,10 +3147,6 @@ export type SearchHfModelsData = {
|
||||
* Search query
|
||||
*/
|
||||
q: string;
|
||||
/**
|
||||
* Max results (default 20)
|
||||
*/
|
||||
limit?: number | null;
|
||||
};
|
||||
url: '/local-inference/search';
|
||||
};
|
||||
@@ -3184,7 +3155,7 @@ export type SearchHfModelsResponses = {
|
||||
/**
|
||||
* Search results
|
||||
*/
|
||||
200: Array<HfModelInfo>;
|
||||
200: Array<HfSearchResult>;
|
||||
};
|
||||
|
||||
export type SearchHfModelsResponse = SearchHfModelsResponses[keyof SearchHfModelsResponses];
|
||||
|
||||
@@ -5,7 +5,7 @@ import {
|
||||
searchHfModels,
|
||||
getRepoFiles,
|
||||
downloadHfModel,
|
||||
type HfModelInfo,
|
||||
type HfSearchResult,
|
||||
type HfQuantVariant,
|
||||
} from '../../../api';
|
||||
|
||||
@@ -34,7 +34,7 @@ interface Props {
|
||||
|
||||
export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
const [query, setQuery] = useState('');
|
||||
const [results, setResults] = useState<HfModelInfo[]>([]);
|
||||
const [results, setResults] = useState<HfSearchResult[]>([]);
|
||||
const [expandedRepo, setExpandedRepo] = useState<string | null>(null);
|
||||
const [repoData, setRepoData] = useState<Record<string, RepoData>>({});
|
||||
const [searching, setSearching] = useState(false);
|
||||
@@ -44,6 +44,11 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const debounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
// Build repo_id from search result
|
||||
const getRepoId = (result: HfSearchResult): string => {
|
||||
return `${result.author}/${result.name}`;
|
||||
};
|
||||
|
||||
const doSearch = useCallback(async (q: string) => {
|
||||
if (!q.trim()) {
|
||||
setResults([]);
|
||||
@@ -54,17 +59,18 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
setError(null);
|
||||
try {
|
||||
const response = await searchHfModels({
|
||||
query: { q, limit: 20 },
|
||||
query: { q },
|
||||
});
|
||||
if (response.data) {
|
||||
// Pre-fetch variants for all results and filter out repos with no suitable quantizations
|
||||
const modelsWithVariants = await Promise.all(
|
||||
response.data.map(async (model) => {
|
||||
try {
|
||||
const [author, repo] = model.repo_id.split('/');
|
||||
const repoId = getRepoId(model);
|
||||
const [author, repo] = repoId.split('/');
|
||||
const filesResponse = await getRepoFiles({ path: { author, repo } });
|
||||
if (filesResponse.data && filesResponse.data.variants.length > 0) {
|
||||
return { model, data: filesResponse.data };
|
||||
if (filesResponse.data && filesResponse.data.length > 0) {
|
||||
return { model, variants: filesResponse.data };
|
||||
}
|
||||
} catch {
|
||||
// Skip repos we can't fetch
|
||||
@@ -74,17 +80,30 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
);
|
||||
|
||||
const validResults = modelsWithVariants.filter(Boolean) as {
|
||||
model: HfModelInfo;
|
||||
data: { variants: HfQuantVariant[]; recommended_index?: number | null };
|
||||
model: HfSearchResult;
|
||||
variants: HfQuantVariant[];
|
||||
}[];
|
||||
|
||||
setResults(validResults.map((r) => r.model));
|
||||
setRepoData((prev) => {
|
||||
const next = { ...prev };
|
||||
for (const r of validResults) {
|
||||
next[r.model.repo_id] = {
|
||||
variants: r.data.variants,
|
||||
recommendedIndex: r.data.recommended_index ?? null,
|
||||
const repoId = getRepoId(r.model);
|
||||
// Find the best variant (highest quality_rank that's not too big)
|
||||
let recommendedIdx: number | null = null;
|
||||
const sorted = [...r.variants].sort((a, b) => b.quality_rank - a.quality_rank);
|
||||
// Recommend the highest quality variant under 8GB, or the smallest if all are big
|
||||
const under8GB = sorted.filter((v) => v.size_bytes < 8 * 1024 * 1024 * 1024);
|
||||
if (under8GB.length > 0) {
|
||||
recommendedIdx = r.variants.findIndex((v) => v.filename === under8GB[0].filename);
|
||||
} else {
|
||||
// Pick smallest
|
||||
const smallest = [...r.variants].sort((a, b) => a.size_bytes - b.size_bytes)[0];
|
||||
recommendedIdx = r.variants.findIndex((v) => v.filename === smallest.filename);
|
||||
}
|
||||
next[repoId] = {
|
||||
variants: r.variants,
|
||||
recommendedIndex: recommendedIdx,
|
||||
};
|
||||
}
|
||||
return next;
|
||||
@@ -102,40 +121,49 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Search failed:', e);
|
||||
setError('Search failed. Please try again.');
|
||||
setError('Search failed. Check console for details.');
|
||||
} finally {
|
||||
setSearching(false);
|
||||
}
|
||||
}, []);
|
||||
|
||||
const handleQueryChange = (value: string) => {
|
||||
setQuery(value);
|
||||
if (debounceRef.current) clearTimeout(debounceRef.current);
|
||||
debounceRef.current = setTimeout(() => doSearch(value), 300);
|
||||
};
|
||||
const handleSearch = useCallback(
|
||||
(q: string) => {
|
||||
setQuery(q);
|
||||
if (debounceRef.current) clearTimeout(debounceRef.current);
|
||||
debounceRef.current = setTimeout(() => doSearch(q), 400);
|
||||
},
|
||||
[doSearch]
|
||||
);
|
||||
|
||||
const toggleRepo = async (repoId: string) => {
|
||||
const expandRepo = async (repoId: string) => {
|
||||
if (expandedRepo === repoId) {
|
||||
setExpandedRepo(null);
|
||||
return;
|
||||
}
|
||||
setExpandedRepo(repoId);
|
||||
|
||||
if (!repoData[repoId]?.variants.length) {
|
||||
// If we don't have variants yet, fetch them
|
||||
if (!repoData[repoId]) {
|
||||
setLoadingFiles((prev) => new Set(prev).add(repoId));
|
||||
try {
|
||||
const [author, repo] = repoId.split('/');
|
||||
const response = await getRepoFiles({
|
||||
path: { author, repo },
|
||||
});
|
||||
const response = await getRepoFiles({ path: { author, repo } });
|
||||
if (response.data) {
|
||||
const variants = response.data.variants;
|
||||
const variants = response.data;
|
||||
// Find recommended
|
||||
let recommendedIdx: number | null = null;
|
||||
const sorted = [...variants].sort((a, b) => b.quality_rank - a.quality_rank);
|
||||
const under8GB = sorted.filter((v) => v.size_bytes < 8 * 1024 * 1024 * 1024);
|
||||
if (under8GB.length > 0) {
|
||||
recommendedIdx = variants.findIndex((v) => v.filename === under8GB[0].filename);
|
||||
} else {
|
||||
const smallest = [...variants].sort((a, b) => a.size_bytes - b.size_bytes)[0];
|
||||
recommendedIdx = variants.findIndex((v) => v.filename === smallest.filename);
|
||||
}
|
||||
setRepoData((prev) => ({
|
||||
...prev,
|
||||
[repoId]: {
|
||||
variants,
|
||||
recommendedIndex: response.data!.recommended_index ?? null,
|
||||
},
|
||||
[repoId]: { variants, recommendedIndex: recommendedIdx },
|
||||
}));
|
||||
}
|
||||
} catch (e) {
|
||||
@@ -151,38 +179,38 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
};
|
||||
|
||||
const startDownload = async (repoId: string, filename: string) => {
|
||||
const key = `${repoId}/${filename}`;
|
||||
setDownloading((prev) => new Set(prev).add(key));
|
||||
const dlKey = `${repoId}/${filename}`;
|
||||
setDownloading((prev) => new Set(prev).add(dlKey));
|
||||
try {
|
||||
const response = await downloadHfModel({
|
||||
body: { repo_id: repoId, filename },
|
||||
});
|
||||
if (response.data) {
|
||||
onDownloadStarted(response.data.model_id);
|
||||
} else {
|
||||
console.error('Download error:', response.error);
|
||||
onDownloadStarted(response.data);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Download failed:', e);
|
||||
} finally {
|
||||
setDownloading((prev) => {
|
||||
const next = new Set(prev);
|
||||
next.delete(key);
|
||||
next.delete(dlKey);
|
||||
return next;
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
const startDirectDownload = async () => {
|
||||
if (!directSpec.trim()) return;
|
||||
const key = `direct:${directSpec}`;
|
||||
setDownloading((prev) => new Set(prev).add(key));
|
||||
const spec = directSpec.trim();
|
||||
if (!spec) return;
|
||||
|
||||
const dlKey = `direct:${spec}`;
|
||||
setDownloading((prev) => new Set(prev).add(dlKey));
|
||||
try {
|
||||
const response = await downloadHfModel({
|
||||
body: { spec: directSpec.trim() },
|
||||
body: { spec },
|
||||
});
|
||||
if (response.data) {
|
||||
onDownloadStarted(response.data.model_id);
|
||||
onDownloadStarted(response.data);
|
||||
setDirectSpec('');
|
||||
}
|
||||
} catch (e) {
|
||||
@@ -190,7 +218,7 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
} finally {
|
||||
setDownloading((prev) => {
|
||||
const next = new Set(prev);
|
||||
next.delete(key);
|
||||
next.delete(dlKey);
|
||||
return next;
|
||||
});
|
||||
}
|
||||
@@ -201,46 +229,46 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
<div>
|
||||
<h4 className="text-sm font-medium text-text-default mb-2">Search HuggingFace</h4>
|
||||
<div className="relative">
|
||||
<Search className="absolute left-3 top-1/2 -translate-y-1/2 w-4 h-4 text-text-muted" />
|
||||
<input
|
||||
type="text"
|
||||
value={query}
|
||||
onChange={(e) => handleQueryChange(e.target.value)}
|
||||
onChange={(e) => handleSearch(e.target.value)}
|
||||
placeholder="Search for GGUF models..."
|
||||
className="w-full pl-9 pr-4 py-2 text-sm border border-border-subtle rounded-lg bg-background-default text-text-default placeholder:text-text-muted focus:outline-none focus:border-accent-primary"
|
||||
className="w-full px-3 py-2 pl-9 text-sm border border-border-subtle rounded-lg bg-background-default text-text-default placeholder:text-text-muted focus:outline-none focus:border-accent-primary"
|
||||
/>
|
||||
<Search className="w-4 h-4 absolute left-3 top-1/2 -translate-y-1/2 text-text-muted" />
|
||||
{searching && (
|
||||
<Loader2 className="absolute right-3 top-1/2 -translate-y-1/2 w-4 h-4 text-text-muted animate-spin" />
|
||||
<Loader2 className="w-4 h-4 absolute right-3 top-1/2 -translate-y-1/2 text-text-muted animate-spin" />
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{error && !searching && <p className="text-xs text-text-muted">{error}</p>}
|
||||
{error && <p className="text-sm text-red-500">{error}</p>}
|
||||
|
||||
{results.length > 0 && (
|
||||
<div className="space-y-1 max-h-96 overflow-y-auto">
|
||||
<div className="space-y-2 max-h-[400px] overflow-y-auto">
|
||||
{results.map((model) => {
|
||||
const isExpanded = expandedRepo === model.repo_id;
|
||||
const data = repoData[model.repo_id];
|
||||
const repoId = getRepoId(model);
|
||||
const isExpanded = expandedRepo === repoId;
|
||||
const data = repoData[repoId];
|
||||
const variants = data?.variants || [];
|
||||
const recommendedIndex = data?.recommendedIndex ?? null;
|
||||
const isLoading = loadingFiles.has(repoId);
|
||||
|
||||
return (
|
||||
<div key={model.repo_id} className="border border-border-subtle rounded-lg">
|
||||
<div
|
||||
key={model.id}
|
||||
className="border border-border-subtle rounded-lg overflow-hidden"
|
||||
>
|
||||
<button
|
||||
onClick={() => toggleRepo(model.repo_id)}
|
||||
className="w-full flex items-center justify-between p-3 text-left hover:bg-background-subtle rounded-lg"
|
||||
className="w-full flex items-center justify-between p-3 text-left hover:bg-background-subtle transition-colors"
|
||||
onClick={() => expandRepo(repoId)}
|
||||
>
|
||||
<div className="flex-1 min-w-0">
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="text-sm font-medium text-text-default truncate">
|
||||
{model.repo_id}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center gap-3 mt-0.5">
|
||||
<span className="text-xs text-text-muted">
|
||||
↓ {formatDownloads(model.downloads)}
|
||||
</span>
|
||||
<div className="flex flex-col gap-1 min-w-0">
|
||||
<span className="text-sm font-medium text-text-default truncate">{repoId}</span>
|
||||
<div className="flex items-center gap-3 text-xs text-text-muted">
|
||||
<span>{formatDownloads(model.downloads)} downloads</span>
|
||||
{model.likes > 0 && <span>❤️ {formatDownloads(model.likes)}</span>}
|
||||
</div>
|
||||
</div>
|
||||
{isExpanded ? (
|
||||
@@ -251,15 +279,15 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
</button>
|
||||
|
||||
{isExpanded && (
|
||||
<div className="border-t border-border-subtle px-3 pb-3 space-y-1">
|
||||
{loadingFiles.has(model.repo_id) && (
|
||||
<div className="flex items-center gap-2 py-2 text-xs text-text-muted">
|
||||
<div className="border-t border-border-subtle p-2 space-y-1 bg-background-subtle/50">
|
||||
{isLoading && (
|
||||
<div className="flex items-center gap-2 py-2 px-2 text-sm text-text-muted">
|
||||
<Loader2 className="w-3 h-3 animate-spin" />
|
||||
Loading variants...
|
||||
</div>
|
||||
)}
|
||||
{variants.map((variant, idx) => {
|
||||
const dlKey = `${model.repo_id}/${variant.filename}`;
|
||||
const dlKey = `${repoId}/${variant.filename}`;
|
||||
const isStarting = downloading.has(dlKey);
|
||||
const isRecommended = idx === recommendedIndex;
|
||||
|
||||
@@ -287,15 +315,12 @@ export const HuggingFaceModelSearch = ({ onDownloadStarted }: Props) => {
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
{variant.description && (
|
||||
<span className="text-xs text-text-muted">{variant.description}</span>
|
||||
)}
|
||||
</div>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
disabled={isStarting}
|
||||
onClick={() => startDownload(model.repo_id, variant.filename)}
|
||||
onClick={() => startDownload(repoId, variant.filename)}
|
||||
>
|
||||
{isStarting ? (
|
||||
<Loader2 className="w-3 h-3 animate-spin" />
|
||||
|
||||
Reference in New Issue
Block a user