Skip to content

B
bit-pm
Unverified 提交 993c5417 authored 作者: Will Chen's avatar Will Chen 提交者： GitHub
浏览文件
操作

Fix DB schemas (#138)

上级 f5a6a1ab
隐藏空白字符变更
内嵌 并排
正在显示 包含 272 行增加 和 127 行删除
drizzle/0005_superb_lady_mastermind.sql → drizzle/0005_clumsy_namor.sql
浏览文件 @ 993c5417
......@@ -11,11 +11,12 @@ CREATE TABLE `language_models` (
`id` integer PRIMARY KEY AUTOINCREMENT NOT NULL,
`display_name` text NOT NULL,
`api_name` text NOT NULL,
`provider_id` text NOT NULL,
`builtin_provider_id` text,
`custom_provider_id` text,
`description` text,
`max_output_tokens` integer,
`context_window` integer,
`created_at` integer DEFAULT (unixepoch()) NOT NULL,
`updated_at` integer DEFAULT (unixepoch()) NOT NULL,
FOREIGN KEY (`provider_id`) REFERENCES `language_model_providers`(`id`) ON UPDATE no action ON DELETE cascade
FOREIGN KEY (`custom_provider_id`) REFERENCES `language_model_providers`(`id`) ON UPDATE no action ON DELETE cascade
);
drizzle/meta/0005_snapshot.json
浏览文件 @ 993c5417
{
"version": "6",
"dialect": "sqlite",
"id": "424973e5-a102-4b71-8d5d-6160f1609b8c",
"id": "0a47ec41-9477-4457-b3e8-e5ecb3e3a855",
"prevId": "ceedb797-6aa3-4a50-b42f-bc85ee08b3df",
"tables": {
"apps": {
......@@ -210,11 +210,18 @@
"notNull": true,
"autoincrement": false
},
"provider_id": {
"name": "provider_id",
"builtin_provider_id": {
"name": "builtin_provider_id",
"type": "text",
"primaryKey": false,
"notNull": true,
"notNull": false,
"autoincrement": false
},
"custom_provider_id": {
"name": "custom_provider_id",
"type": "text",
"primaryKey": false,
"notNull": false,
"autoincrement": false
},
"description": {
......@@ -257,12 +264,12 @@
},
"indexes": {},
"foreignKeys": {
"language_models_provider_id_language_model_providers_id_fk": {
"name": "language_models_provider_id_language_model_providers_id_fk",
"language_models_custom_provider_id_language_model_providers_id_fk": {
"name": "language_models_custom_provider_id_language_model_providers_id_fk",
"tableFrom": "language_models",
"tableTo": "language_model_providers",
"columnsFrom": [
"provider_id"
"custom_provider_id"
],
"columnsTo": [
"id"
......
drizzle/meta/_journal.json
浏览文件 @ 993c5417
......@@ -40,8 +40,8 @@
{
"idx": 5,
"version": "6",
"when": 1747091036229,
"tag": "0005_superb_lady_mastermind",
"when": 1747095436506,
"tag": "0005_clumsy_namor",
"breakpoints": true
}
]
......
src/components/ModelPicker.tsx
浏览文件 @ 993c5417
......@@ -17,9 +17,9 @@ import {
DropdownMenuSubContent,
} from "@/components/ui/dropdown-menu";
import { useEffect, useState } from "react";
import { MODEL_OPTIONS } from "@/constants/models";
import { useLocalModels } from "@/hooks/useLocalModels";
import { useLocalLMSModels } from "@/hooks/useLMStudioModels";
import { useLanguageModelsByProviders } from "@/hooks/useLanguageModelsByProviders";
import { ChevronDown } from "lucide-react";
import { LocalModel } from "@/ipc/ipc_types";
interface ModelPickerProps {
......@@ -33,6 +33,10 @@ export function ModelPicker({
}: ModelPickerProps) {
const [open, setOpen] = useState(false);
// Cloud models from providers
const { data: modelsByProviders, isLoading: providersLoading } =
useLanguageModelsByProviders();
// Ollama Models Hook
const {
models: ollamaModels,
......@@ -74,24 +78,35 @@ export function ModelPicker({
);
}
// Fallback for cloud models
return (
MODEL_OPTIONS[selectedModel.provider]?.find(
(model) => model.name === selectedModel.name,
)?.displayName || selectedModel.name
);
// For cloud models, look up in the modelsByProviders data
if (modelsByProviders && modelsByProviders[selectedModel.provider]) {
const foundModel = modelsByProviders[selectedModel.provider].find(
(model) => model.apiName === selectedModel.name,
);
if (foundModel) {
return foundModel.displayName;
}
}
// Fallback if not found
return selectedModel.name;
};
const modelDisplayName = getModelDisplayName();
// Flatten the cloud model options
const cloudModels = Object.entries(MODEL_OPTIONS).flatMap(
([provider, models]) =>
models.map((model) => ({
...model,
provider: provider as ModelProvider,
})),
);
// Flatten the cloud models from all providers
const cloudModels =
!providersLoading && modelsByProviders
? Object.entries(modelsByProviders).flatMap(([providerId, models]) =>
models.map((model) => ({
name: model.apiName,
displayName: model.displayName,
description: model.description || "",
tag: model.tag,
provider: providerId as ModelProvider,
})),
)
: [];
// Determine availability of local models
const hasOllamaModels =
......@@ -119,43 +134,54 @@ export function ModelPicker({
{/* Increased width slightly */}
<DropdownMenuLabel>Cloud Models</DropdownMenuLabel>
<DropdownMenuSeparator />
{/* Cloud models */}
{cloudModels.map((model) => (
<Tooltip key={`${model.provider}-${model.name}`}>
<TooltipTrigger asChild>
<DropdownMenuItem
className={
selectedModel.provider === model.provider &&
selectedModel.name === model.name
? "bg-secondary"
: ""
}
onClick={() => {
onModelSelect({
name: model.name,
provider: model.provider,
});
setOpen(false);
}}
>
<div className="flex justify-between items-start w-full">
<span className="flex flex-col items-start">
<span>{model.displayName}</span>
<span className="text-xs text-muted-foreground">
{model.provider}
</span>
</span>
{model.tag && (
<span className="text-[10px] bg-primary/10 text-primary px-1.5 py-0.5 rounded-full font-medium">
{model.tag}
{/* Cloud models - loading state */}
{providersLoading ? (
<div className="text-xs text-center py-2 text-muted-foreground">
Loading models...
</div>
) : cloudModels.length === 0 ? (
<div className="text-xs text-center py-2 text-muted-foreground">
No cloud models available
</div>
) : (
/* Cloud models loaded */
cloudModels.map((model) => (
<Tooltip key={`${model.provider}-${model.name}`}>
<TooltipTrigger asChild>
<DropdownMenuItem
className={
selectedModel.provider === model.provider &&
selectedModel.name === model.name
? "bg-secondary"
: ""
}
onClick={() => {
onModelSelect({
name: model.name,
provider: model.provider,
});
setOpen(false);
}}
>
<div className="flex justify-between items-start w-full">
<span className="flex flex-col items-start">
<span>{model.displayName}</span>
<span className="text-xs text-muted-foreground">
{model.provider}
</span>
</span>
)}
</div>
</DropdownMenuItem>
</TooltipTrigger>
<TooltipContent side="right">{model.description}</TooltipContent>
</Tooltip>
))}
{model.tag && (
<span className="text-[10px] bg-primary/10 text-primary px-1.5 py-0.5 rounded-full font-medium">
{model.tag}
</span>
)}
</div>
</DropdownMenuItem>
</TooltipTrigger>
<TooltipContent side="right">{model.description}</TooltipContent>
</Tooltip>
))
)}
<DropdownMenuSeparator />
{/* Local Models Parent SubMenu */}
<DropdownMenuSub>
......
src/db/schema.ts
浏览文件 @ 993c5417
......@@ -85,9 +85,11 @@ export const language_models = sqliteTable("language_models", {
id: integer("id").primaryKey({ autoIncrement: true }),
displayName: text("display_name").notNull(),
apiName: text("api_name").notNull(),
provider_id: text("provider_id")
.notNull()
.references(() => language_model_providers.id, { onDelete: "cascade" }),
builtinProviderId: text("builtin_provider_id"),
customProviderId: text("custom_provider_id").references(
() => language_model_providers.id,
{ onDelete: "cascade" },
),
description: text("description"),
max_output_tokens: integer("max_output_tokens"),
context_window: integer("context_window"),
......@@ -111,7 +113,7 @@ export const languageModelsRelations = relations(
language_models,
({ one }) => ({
provider: one(language_model_providers, {
fields: [language_models.provider_id],
fields: [language_models.customProviderId],
references: [language_model_providers.id],
}),
}),
......
src/hooks/useLanguageModelsByProviders.ts 0 → 100644
浏览文件 @ 993c5417
import { useQuery } from "@tanstack/react-query";
import { IpcClient } from "@/ipc/ipc_client";
import type { LanguageModel } from "@/ipc/ipc_types";
/**
* Fetches all available language models grouped by their provider IDs.
*
* @returns TanStack Query result object for the language models organized by provider.
*/
export function useLanguageModelsByProviders() {
const ipcClient = IpcClient.getInstance();
return useQuery<Record<string, LanguageModel[]>, Error>({
queryKey: ["language-models-by-providers"],
queryFn: async () => {
return ipcClient.getLanguageModelsByProviders();
},
});
}
src/ipc/handlers/language_model_handlers.ts
浏览文件 @ 993c5417
......@@ -9,6 +9,7 @@ import log from "electron-log";
import {
getLanguageModelProviders,
getLanguageModels,
getLanguageModelsByProviders,
} from "../shared/language_model_helpers";
import { db } from "@/db";
import {
......@@ -64,7 +65,8 @@ export function registerLanguageModelHandlers() {
// Insert the new provider
await db.insert(languageModelProvidersSchema).values({
id,
// Make sure we will never have accidental collisions with builtin providers
id: "custom::" + id,
name,
api_base_url: apiBaseUrl,
env_var_name: envVarName || null,
......@@ -108,12 +110,8 @@ export function registerLanguageModelHandlers() {
}
// Check if provider exists
const provider = db
.select()
.from(languageModelProvidersSchema)
.where(eq(languageModelProvidersSchema.id, providerId))
.get();
const providers = await getLanguageModelProviders();
const provider = providers.find((p) => p.id === providerId);
if (!provider) {
throw new Error(`Provider with ID "${providerId}" not found`);
}
......@@ -122,7 +120,8 @@ export function registerLanguageModelHandlers() {
await db.insert(languageModelsSchema).values({
displayName,
apiName,
provider_id: providerId,
builtinProviderId: provider.type === "cloud" ? providerId : undefined,
customProviderId: provider.type === "custom" ? providerId : undefined,
description: description || null,
max_output_tokens: maxOutputTokens || null,
context_window: contextWindow || null,
......@@ -182,11 +181,22 @@ export function registerLanguageModelHandlers() {
`Attempting to delete custom model ${modelApiName} for provider ${providerId}`,
);
const providers = await getLanguageModelProviders();
const provider = providers.find((p) => p.id === providerId);
if (!provider) {
throw new Error(`Provider with ID "${providerId}" not found`);
}
if (provider.type === "local") {
throw new Error("Local models cannot be deleted");
}
const result = db
.delete(language_models)
.where(
and(
eq(language_models.provider_id, providerId),
provider.type === "cloud"
? eq(language_models.builtinProviderId, providerId)
: eq(language_models.customProviderId, providerId),
eq(language_models.apiName, modelApiName),
),
)
......@@ -243,7 +253,7 @@ export function registerLanguageModelHandlers() {
// 1. Delete associated models
const deleteModelsResult = await tx
.delete(languageModelsSchema)
.where(eq(languageModelsSchema.provider_id, providerId))
.where(eq(languageModelsSchema.customProviderId, providerId))
.run();
logger.info(
`Deleted ${deleteModelsResult.changes} model(s) associated with provider ${providerId}`,
......@@ -279,7 +289,26 @@ export function registerLanguageModelHandlers() {
if (!params || typeof params.providerId !== "string") {
throw new Error("Invalid parameters: providerId (string) is required.");
}
return getLanguageModels({ providerId: params.providerId });
const providers = await getLanguageModelProviders();
const provider = providers.find((p) => p.id === params.providerId);
if (!provider) {
throw new Error(`Provider with ID "${params.providerId}" not found`);
}
if (provider.type === "local") {
throw new Error("Local models cannot be fetched");
}
return getLanguageModels(
provider.type === "cloud"
? { builtinProviderId: params.providerId }
: { customProviderId: params.providerId },
);
},
);
handle(
"get-language-models-by-providers",
async (): Promise<Record<string, LanguageModel[]>> => {
return getLanguageModelsByProviders();
},
);
}
src/ipc/ipc_client.ts
浏览文件 @ 993c5417
......@@ -746,6 +746,12 @@ export class IpcClient {
return this.ipcRenderer.invoke("get-language-models", params);
}
public async getLanguageModelsByProviders(): Promise<
Record<string, LanguageModel[]>
> {
return this.ipcRenderer.invoke("get-language-models-by-providers");
}
public async createCustomLanguageModelProvider({
id,
name,
......
src/ipc/shared/language_model_helpers.ts
浏览文件 @ 993c5417
......@@ -139,23 +139,72 @@ export async function getLanguageModelProviders(): Promise<
* @param obj An object containing the providerId.
* @returns A promise that resolves to an array of LanguageModel objects.
*/
export async function getLanguageModels(obj: {
providerId: string;
}): Promise<LanguageModel[]> {
const { providerId } = obj;
export async function getLanguageModels(
obj:
| {
customProviderId: string;
// builtinProviderId?: undefined;
}
| {
builtinProviderId: string;
// customProviderId?: undefined;
},
): Promise<LanguageModel[]> {
const allProviders = await getLanguageModelProviders();
const provider = allProviders.find((p) => p.id === providerId);
const provider = allProviders.find(
(p) =>
p.id === (obj as { customProviderId: string }).customProviderId ||
p.id === (obj as { builtinProviderId: string }).builtinProviderId,
);
if (!provider) {
console.warn(`Provider with ID "${providerId}" not found.`);
console.warn(`Provider with ID "${JSON.stringify(obj)}" not found.`);
return [];
}
// Get custom models from DB for all provider types
let customModels: LanguageModel[] = [];
try {
const customModelsDb = await db
.select({
id: languageModelsSchema.id,
displayName: languageModelsSchema.displayName,
apiName: languageModelsSchema.apiName,
description: languageModelsSchema.description,
maxOutputTokens: languageModelsSchema.max_output_tokens,
contextWindow: languageModelsSchema.context_window,
})
.from(languageModelsSchema)
.where(
"customProviderId" in obj
? eq(languageModelsSchema.customProviderId, obj.customProviderId)
: eq(languageModelsSchema.builtinProviderId, obj.builtinProviderId),
);
customModels = customModelsDb.map((model) => ({
...model,
description: model.description ?? "",
tag: undefined,
maxOutputTokens: model.maxOutputTokens ?? undefined,
contextWindow: model.contextWindow ?? undefined,
type: "custom",
}));
} catch (error) {
console.error(
`Error fetching custom models for provider "${JSON.stringify(obj)}" from DB:`,
error,
);
// Continue with empty custom models array
}
// If it's a cloud provider, also get the hardcoded models
let hardcodedModels: LanguageModel[] = [];
const providerId = provider.id;
if (provider.type === "cloud") {
// Check if providerId is a valid key for MODEL_OPTIONS
if (providerId in MODEL_OPTIONS) {
const models = MODEL_OPTIONS[providerId as RegularModelProvider] || [];
return models.map((model) => ({
hardcodedModels = models.map((model) => ({
...model,
apiName: model.name,
type: "cloud",
......@@ -164,49 +213,54 @@ export async function getLanguageModels(obj: {
console.warn(
`Provider "${providerId}" is cloud type but not found in MODEL_OPTIONS.`,
);
return [];
}
} else if (provider.type === "custom") {
// Fetch models from the database for this custom provider
// Assuming a language_models table with necessary columns and provider_id foreign key
try {
const customModelsDb = await db
.select({
id: languageModelsSchema.id,
// Map DB columns to LanguageModel fields
displayName: languageModelsSchema.displayName,
apiName: languageModelsSchema.apiName,
// No display_name in DB, use name instead
description: languageModelsSchema.description,
// No tag in DB
maxOutputTokens: languageModelsSchema.max_output_tokens,
contextWindow: languageModelsSchema.context_window,
})
.from(languageModelsSchema)
.where(eq(languageModelsSchema.provider_id, providerId)); // Assuming eq is imported or available
return customModelsDb.map((model) => ({
...model,
// Ensure possibly null fields are handled, provide defaults or undefined if needed
description: model.description ?? "",
tag: undefined, // No tag for custom models from DB
maxOutputTokens: model.maxOutputTokens ?? undefined,
contextWindow: model.contextWindow ?? undefined,
type: "custom",
}));
} catch (error) {
console.error(
`Error fetching custom models for provider "${providerId}" from DB:`,
error,
}
// Merge the models, with custom models taking precedence over hardcoded ones
const mergedModelsMap = new Map<string, LanguageModel>();
// Add hardcoded models first
for (const model of hardcodedModels) {
mergedModelsMap.set(model.apiName, model);
}
// Then override with custom models
for (const model of customModels) {
mergedModelsMap.set(model.apiName, model);
}
return Array.from(mergedModelsMap.values());
}
/**
* Fetches all language models grouped by their provider IDs.
* @returns A promise that resolves to a Record mapping provider IDs to arrays of LanguageModel objects.
*/
export async function getLanguageModelsByProviders(): Promise<
Record<string, LanguageModel[]>
> {
const providers = await getLanguageModelProviders();
// Fetch all models concurrently
const modelPromises = providers
.filter((p) => p.type !== "local")
.map(async (provider) => {
const models = await getLanguageModels(
provider.type === "cloud"
? { builtinProviderId: provider.id }
: { customProviderId: provider.id },
);
// Depending on desired behavior, could throw, return empty, or return a specific error state
return [];
}
} else {
// Handle other types like "local" if necessary, currently ignored
console.warn(
`Provider type "${provider.type}" not handled for model fetching.`,
);
return [];
return { providerId: provider.id, models };
});
// Wait for all requests to complete
const results = await Promise.all(modelPromises);
// Convert the array of results to a record
const record: Record<string, LanguageModel[]> = {};
for (const result of results) {
record[result.providerId] = result.models;
}
return record;
}
src/preload.ts
浏览文件 @ 993c5417
......@@ -6,6 +6,7 @@ import { contextBridge, ipcRenderer } from "electron";
// Whitelist of valid channels
const validInvokeChannels = [
"get-language-models",
"get-language-models-by-providers",
"create-custom-language-model",
"get-language-model-providers",
"delete-custom-language-model-provider",
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论