Add OpenRouter custom model scheme

This commit is contained in:
Saoud Rizwan
2024-10-03 19:46:09 -04:00
parent d5b3bd7788
commit 7cb0c524e5
12 changed files with 262 additions and 54 deletions

View File

@@ -23,7 +23,7 @@ import {
openAiNativeDefaultModelId,
openAiNativeModels,
openRouterDefaultModelId,
openRouterModels,
openRouterDefaultModelInfo,
vertexDefaultModelId,
vertexModels,
} from "../../../../src/shared/api"
@@ -66,8 +66,8 @@ const ApiOptions = ({ showModelOptions, apiErrorMessage }: ApiOptionsProps) => {
const handleMessage = useCallback((event: MessageEvent) => {
const message: ExtensionMessage = event.data
if (message.type === "ollamaModels" && message.models) {
setOllamaModels(message.models)
if (message.type === "ollamaModels" && message.ollamaModels) {
setOllamaModels(message.ollamaModels)
}
}, [])
useEvent("message", handleMessage)
@@ -529,23 +529,25 @@ const ApiOptions = ({ showModelOptions, apiErrorMessage }: ApiOptionsProps) => {
</p>
)}
{selectedProvider !== "openai" && selectedProvider !== "ollama" && showModelOptions && (
<>
<div className="dropdown-container">
<label htmlFor="model-id">
<span style={{ fontWeight: 500 }}>Model</span>
</label>
{selectedProvider === "anthropic" && createDropdown(anthropicModels)}
{selectedProvider === "openrouter" && createDropdown(openRouterModels)}
{selectedProvider === "bedrock" && createDropdown(bedrockModels)}
{selectedProvider === "vertex" && createDropdown(vertexModels)}
{selectedProvider === "gemini" && createDropdown(geminiModels)}
{selectedProvider === "openai-native" && createDropdown(openAiNativeModels)}
</div>
{selectedProvider !== "openrouter" &&
selectedProvider !== "openai" &&
selectedProvider !== "ollama" &&
showModelOptions && (
<>
<div className="dropdown-container">
<label htmlFor="model-id">
<span style={{ fontWeight: 500 }}>Model</span>
</label>
{selectedProvider === "anthropic" && createDropdown(anthropicModels)}
{selectedProvider === "bedrock" && createDropdown(bedrockModels)}
{selectedProvider === "vertex" && createDropdown(vertexModels)}
{selectedProvider === "gemini" && createDropdown(geminiModels)}
{selectedProvider === "openai-native" && createDropdown(openAiNativeModels)}
</div>
<ModelInfoView selectedModelId={selectedModelId} modelInfo={selectedModelInfo} />
</>
)}
<ModelInfoView selectedModelId={selectedModelId} modelInfo={selectedModelInfo} />
</>
)}
</div>
)
}
@@ -563,7 +565,7 @@ export const formatPrice = (price: number) => {
}).format(price)
}
const ModelInfoView = ({ selectedModelId, modelInfo }: { selectedModelId: string; modelInfo: ModelInfo }) => {
export const ModelInfoView = ({ selectedModelId, modelInfo }: { selectedModelId: string; modelInfo: ModelInfo }) => {
const isGemini = Object.keys(geminiModels).includes(selectedModelId)
const isO1 = selectedModelId && selectedModelId.includes("o1")
return (
@@ -690,8 +692,6 @@ export function normalizeApiConfiguration(apiConfiguration?: ApiConfiguration) {
switch (provider) {
case "anthropic":
return getProviderData(anthropicModels, anthropicDefaultModelId)
case "openrouter":
return getProviderData(openRouterModels, openRouterDefaultModelId)
case "bedrock":
return getProviderData(bedrockModels, bedrockDefaultModelId)
case "vertex":
@@ -700,6 +700,12 @@ export function normalizeApiConfiguration(apiConfiguration?: ApiConfiguration) {
return getProviderData(geminiModels, geminiDefaultModelId)
case "openai-native":
return getProviderData(openAiNativeModels, openAiNativeDefaultModelId)
case "openrouter":
return {
selectedProvider: provider,
selectedModelId: apiConfiguration?.openRouterModelId ?? openRouterDefaultModelId,
selectedModelInfo: apiConfiguration?.openRouterModelInfo ?? openRouterDefaultModelInfo,
}
case "openai":
return {
selectedProvider: provider,

View File

@@ -0,0 +1,59 @@
import { VSCodeDropdown, VSCodeOption } from "@vscode/webview-ui-toolkit/react"
import React, { useMemo } from "react"
import { useExtensionState } from "../../context/ExtensionStateContext"
import { ModelInfoView, normalizeApiConfiguration } from "./ApiOptions"
interface OpenRouterModelPickerProps {}
const OpenRouterModelPicker: React.FC<OpenRouterModelPickerProps> = () => {
const { apiConfiguration, setApiConfiguration, openRouterModels } = useExtensionState()
const handleModelChange = (event: any) => {
const newModelId = event.target.value
// get info
setApiConfiguration({ ...apiConfiguration, openRouterModelId: newModelId })
}
const { selectedModelId, selectedModelInfo } = useMemo(() => {
return normalizeApiConfiguration(apiConfiguration)
}, [apiConfiguration])
return (
<div style={{ display: "flex", flexDirection: "column", gap: 5 }}>
<div className="dropdown-container">
<label htmlFor="model-id">
<span style={{ fontWeight: 500 }}>Model</span>
</label>
<VSCodeDropdown
id="model-id"
value={selectedModelId}
onChange={handleModelChange}
style={{ width: "100%" }}>
<VSCodeOption value="">Select a model...</VSCodeOption>
{Object.keys(openRouterModels).map((modelId) => (
<VSCodeOption
key={modelId}
value={modelId}
style={{
whiteSpace: "normal",
wordWrap: "break-word",
maxWidth: "100%",
}}>
{modelId}
</VSCodeOption>
))}
</VSCodeDropdown>
</div>
{selectedModelInfo.description && (
<p style={{ fontSize: "12px", marginTop: "2px", color: "var(--vscode-descriptionForeground)" }}>
{selectedModelInfo.description}
</p>
)}
<ModelInfoView selectedModelId={selectedModelId} modelInfo={selectedModelInfo} />
</div>
)
}
export default OpenRouterModelPicker

View File

@@ -1,7 +1,7 @@
import React, { createContext, useCallback, useContext, useEffect, useState } from "react"
import { useEvent } from "react-use"
import { ExtensionMessage, ExtensionState } from "../../../src/shared/ExtensionMessage"
import { ApiConfiguration } from "../../../src/shared/api"
import { ApiConfiguration, ModelInfo } from "../../../src/shared/api"
import { vscode } from "../utils/vscode"
import { convertTextMateToHljs } from "../utils/textMateToHljs"
import { findLastIndex } from "../../../src/shared/array"
@@ -10,6 +10,7 @@ interface ExtensionStateContextType extends ExtensionState {
didHydrateState: boolean
showWelcome: boolean
theme: any
openRouterModels: Record<string, ModelInfo>
filePaths: string[]
setApiConfiguration: (config: ApiConfiguration) => void
setCustomInstructions: (value?: string) => void
@@ -30,6 +31,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
const [showWelcome, setShowWelcome] = useState(false)
const [theme, setTheme] = useState<any>(undefined)
const [filePaths, setFilePaths] = useState<string[]>([])
const [openRouterModels, setOpenRouterModels] = useState<Record<string, ModelInfo>>({})
const handleMessage = useCallback((event: MessageEvent) => {
const message: ExtensionMessage = event.data
@@ -75,6 +77,11 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
}
return prevState
})
break
}
case "openRouterModels": {
setOpenRouterModels(message.openRouterModels ?? {})
break
}
}
}, [])
@@ -90,6 +97,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
didHydrateState,
showWelcome,
theme,
openRouterModels,
filePaths,
setApiConfiguration: (value) => setState((prevState) => ({ ...prevState, apiConfiguration: value })),
setCustomInstructions: (value) => setState((prevState) => ({ ...prevState, customInstructions: value })),

View File

@@ -14,8 +14,8 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s
}
break
case "openrouter":
if (!apiConfiguration.openRouterApiKey) {
return "You must provide a valid API key or choose a different provider."
if (!apiConfiguration.openRouterApiKey || !apiConfiguration.openRouterModelId) {
return "You must provide a valid API key and model ID or choose a different provider."
}
break
case "vertex":