forked from akai/readest
feat(assistant): Add embedding model option to AIPanel (#3090)
* Add embedding model option to AIPanel * Minor code reformatting * Edit tests to account for embedding model * Minor reformatting
This commit is contained in:
@@ -65,9 +65,15 @@ describe('OllamaProvider', () => {
|
||||
test('healthCheck should verify model exists', async () => {
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () => Promise.resolve({ models: [{ name: 'llama3.2:latest' }] }),
|
||||
json: () =>
|
||||
Promise.resolve({ models: [{ name: 'llama3.2:latest' }, { name: 'nomic-embed:latest' }] }),
|
||||
});
|
||||
const settings: AISettings = { ...DEFAULT_AI_SETTINGS, enabled: true, ollamaModel: 'llama3.2' };
|
||||
const settings: AISettings = {
|
||||
...DEFAULT_AI_SETTINGS,
|
||||
enabled: true,
|
||||
ollamaModel: 'llama3.2',
|
||||
ollamaEmbeddingModel: 'nomic-embed',
|
||||
};
|
||||
const provider = new OllamaProvider(settings);
|
||||
|
||||
const result = await provider.healthCheck();
|
||||
@@ -77,9 +83,15 @@ describe('OllamaProvider', () => {
|
||||
test('healthCheck should return false if model not found', async () => {
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () => Promise.resolve({ models: [{ name: 'other-model' }] }),
|
||||
json: () =>
|
||||
Promise.resolve({ models: [{ name: 'other-model' }, { name: 'nomic-embed:latest' }] }),
|
||||
});
|
||||
const settings: AISettings = { ...DEFAULT_AI_SETTINGS, enabled: true, ollamaModel: 'llama3.2' };
|
||||
const settings: AISettings = {
|
||||
...DEFAULT_AI_SETTINGS,
|
||||
enabled: true,
|
||||
ollamaModel: 'llama3.2',
|
||||
ollamaEmbeddingModel: 'nomic-embed',
|
||||
};
|
||||
const provider = new OllamaProvider(settings);
|
||||
|
||||
const result = await provider.healthCheck();
|
||||
|
||||
@@ -71,6 +71,7 @@ const AIPanel: React.FC = () => {
|
||||
const [provider, setProvider] = useState<AIProviderName>(aiSettings.provider);
|
||||
const [ollamaUrl, setOllamaUrl] = useState(aiSettings.ollamaBaseUrl);
|
||||
const [ollamaModel, setOllamaModel] = useState(aiSettings.ollamaModel);
|
||||
const [ollamaEmbeddingModel, setOllamaEmbeddingModel] = useState(aiSettings.ollamaEmbeddingModel);
|
||||
const [ollamaModels, setOllamaModels] = useState<string[]>([]);
|
||||
const [fetchingModels, setFetchingModels] = useState(false);
|
||||
const [gatewayKey, setGatewayKey] = useState(aiSettings.aiGatewayApiKey ?? '');
|
||||
@@ -181,6 +182,14 @@ const AIPanel: React.FC = () => {
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [ollamaModel]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isMounted.current) return;
|
||||
if (ollamaEmbeddingModel !== aiSettings.ollamaEmbeddingModel) {
|
||||
saveAiSetting('ollamaEmbeddingModel', ollamaEmbeddingModel);
|
||||
}
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [ollamaEmbeddingModel]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isMounted.current) return;
|
||||
if (gatewayKey !== (aiSettings.aiGatewayApiKey ?? '')) {
|
||||
@@ -276,6 +285,7 @@ const AIPanel: React.FC = () => {
|
||||
provider,
|
||||
ollamaBaseUrl: ollamaUrl,
|
||||
ollamaModel,
|
||||
ollamaEmbeddingModel,
|
||||
aiGatewayApiKey: gatewayKey,
|
||||
aiGatewayModel: effectiveModel,
|
||||
};
|
||||
@@ -375,21 +385,38 @@ const AIPanel: React.FC = () => {
|
||||
/>
|
||||
</div>
|
||||
{ollamaModels.length > 0 ? (
|
||||
<div className='config-item !h-auto flex-col !items-start gap-2 py-3'>
|
||||
<span>{_('AI Model')}</span>
|
||||
<select
|
||||
className='select select-bordered select-sm bg-base-100 text-base-content w-full'
|
||||
value={ollamaModel}
|
||||
onChange={(e) => setOllamaModel(e.target.value)}
|
||||
disabled={!enabled}
|
||||
>
|
||||
{ollamaModels.map((model) => (
|
||||
<option key={model} value={model}>
|
||||
{model}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</div>
|
||||
<>
|
||||
<div className='config-item !h-auto flex-col !items-start gap-2 py-3'>
|
||||
<span>{_('AI Model')}</span>
|
||||
<select
|
||||
className='select select-bordered select-sm bg-base-100 text-base-content w-full'
|
||||
value={ollamaModel}
|
||||
onChange={(e) => setOllamaModel(e.target.value)}
|
||||
disabled={!enabled}
|
||||
>
|
||||
{ollamaModels.map((model) => (
|
||||
<option key={model} value={model}>
|
||||
{model}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</div>
|
||||
<div className='config-item !h-auto flex-col !items-start gap-2 py-3'>
|
||||
<span>{_('Embedding Model')}</span>
|
||||
<select
|
||||
className='select select-bordered select-sm bg-base-100 text-base-content w-full'
|
||||
value={ollamaEmbeddingModel}
|
||||
onChange={(e) => setOllamaEmbeddingModel(e.target.value)}
|
||||
disabled={!enabled}
|
||||
>
|
||||
{ollamaModels.map((model) => (
|
||||
<option key={model} value={model}>
|
||||
{model}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</div>
|
||||
</>
|
||||
) : !fetchingModels ? (
|
||||
<div className='config-item'>
|
||||
<span className='text-warning text-sm'>{_('No models detected')}</span>
|
||||
|
||||
@@ -53,7 +53,11 @@ export class OllamaProvider implements AIProvider {
|
||||
if (!response.ok) return false;
|
||||
const data = await response.json();
|
||||
const modelName = this.settings.ollamaModel?.split(':')[0] ?? '';
|
||||
return data.models?.some((m: { name: string }) => m.name.includes(modelName));
|
||||
const embeddingModelName = this.settings.ollamaEmbeddingModel?.split(':')[0] ?? '';
|
||||
return (
|
||||
data.models?.some((m: { name: string }) => m.name.includes(modelName)) &&
|
||||
data.models?.some((m: { name: string }) => m.name.includes(embeddingModelName))
|
||||
);
|
||||
} catch (e) {
|
||||
aiLogger.provider.error('ollama', (e as Error).message);
|
||||
return false;
|
||||
|
||||
Reference in New Issue
Block a user