fegger
138890b9d2
- src/ollama-client.ts: Detect HTTP 404 on /api/chat and throw a descriptive ApiError with the model name and the exact ollama pull command needed. - src/error-handler.ts: For API_ERROR type, return the error message directly instead of prefixing with 'API error: ', so the user-friendly 404 message is shown cleanly in the Obsidian notice. - tests/ollama-client.test.ts: Update 404 assertions to match the new descriptive error message.
394 lines
12 KiB
TypeScript
394 lines
12 KiB
TypeScript
// src/ollama-client.ts
|
|
|
|
import type { CacheConfig, OllamaMessage, OllamaTool } from './types';
|
|
import { ApiError } from './types';
|
|
import { Logger } from './utils';
|
|
import { SemanticCacheService } from './semantic-cache';
|
|
|
|
interface OllamaChatResponse {
|
|
message?: Partial<OllamaMessage>;
|
|
error?: string;
|
|
}
|
|
|
|
export class OllamaClient {
|
|
private baseURL: string;
|
|
private model: string;
|
|
private fetchFn: typeof fetch;
|
|
private readonly maxRetries: number = 3;
|
|
private readonly maxMalformedChunks: number = 50;
|
|
private currentStreamController: AbortController | null = null;
|
|
private cacheService?: SemanticCacheService;
|
|
|
|
constructor(baseURL: string, model: string, fetchFn?: typeof fetch, cacheConfig?: CacheConfig) {
|
|
this.baseURL = baseURL;
|
|
this.model = model;
|
|
this.fetchFn = fetchFn ?? ((url, init) => fetch(url, init));
|
|
|
|
if (cacheConfig?.enabled) {
|
|
this.cacheService = new SemanticCacheService(baseURL, cacheConfig);
|
|
}
|
|
}
|
|
|
|
async initializeCache(): Promise<void> {
|
|
if (this.cacheService) {
|
|
await this.cacheService.initialize();
|
|
}
|
|
}
|
|
|
|
async clearCache(): Promise<void> {
|
|
if (this.cacheService) {
|
|
await this.cacheService.clearCache();
|
|
}
|
|
}
|
|
|
|
cancelStream(): void {
|
|
if (this.currentStreamController) {
|
|
this.currentStreamController.abort();
|
|
this.currentStreamController = null;
|
|
}
|
|
}
|
|
|
|
async *streamChat(
|
|
messages: OllamaMessage[],
|
|
tools: OllamaTool[] = []
|
|
): AsyncGenerator<OllamaMessage, void, unknown> {
|
|
// Bypass cache if tools are involved to prevent state corruption
|
|
if (tools.length > 0) {
|
|
yield* this.streamChatWithRetry(messages, tools, 0);
|
|
return;
|
|
}
|
|
|
|
// Find the last user message
|
|
const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
|
|
if (lastUserMsg && this.cacheService) {
|
|
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
|
if (cached) {
|
|
yield { role: 'assistant', content: cached, tool_calls: [] };
|
|
return;
|
|
}
|
|
}
|
|
|
|
const chunks: OllamaMessage[] = [];
|
|
for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) {
|
|
chunks.push(chunk);
|
|
yield chunk;
|
|
}
|
|
|
|
const fullContent = chunks.map((c) => c.content).join('');
|
|
if (this.cacheService && lastUserMsg) {
|
|
void this.cacheService.setCache(lastUserMsg.content, fullContent);
|
|
}
|
|
}
|
|
|
|
async chat(messages: OllamaMessage[], tools: OllamaTool[] = []): Promise<OllamaMessage> {
|
|
// Bypass cache if tools are involved to prevent state corruption
|
|
if (tools.length > 0) {
|
|
return this.chatWithRetry(messages, tools, 0);
|
|
}
|
|
|
|
// Find the last user message
|
|
const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
|
|
if (lastUserMsg && this.cacheService) {
|
|
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
|
if (cached) {
|
|
return { role: 'assistant', content: cached };
|
|
}
|
|
}
|
|
|
|
const response = await this.chatWithRetry(messages, tools, 0);
|
|
if (this.cacheService && lastUserMsg) {
|
|
void this.cacheService.setCache(lastUserMsg.content, response.content);
|
|
}
|
|
return response;
|
|
}
|
|
|
|
async streamChatAsPromise(
|
|
messages: OllamaMessage[],
|
|
tools: OllamaTool[] = []
|
|
): Promise<OllamaMessage> {
|
|
let content = '';
|
|
let role: OllamaMessage['role'] = 'assistant';
|
|
let toolCalls: OllamaMessage['tool_calls'];
|
|
|
|
for await (const chunk of this.streamChat(messages, tools)) {
|
|
role = chunk.role ?? role;
|
|
content += chunk.content ?? '';
|
|
if (chunk.tool_calls) {
|
|
toolCalls = [...(toolCalls ?? []), ...chunk.tool_calls];
|
|
}
|
|
}
|
|
|
|
return { role, content, tool_calls: toolCalls };
|
|
}
|
|
|
|
async *streamChatWithRetry(
|
|
messages: OllamaMessage[],
|
|
tools: OllamaTool[] = [],
|
|
retryCount: number
|
|
): AsyncGenerator<OllamaMessage, void, unknown> {
|
|
const controller = new AbortController();
|
|
this.currentStreamController = controller;
|
|
let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
|
|
|
|
try {
|
|
const response = await this.fetchFn(`${this.baseURL}/api/chat`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
},
|
|
body: JSON.stringify({
|
|
model: this.model,
|
|
messages: messages,
|
|
tools: tools,
|
|
stream: true,
|
|
}),
|
|
signal: controller.signal,
|
|
});
|
|
|
|
if (!response.ok) {
|
|
if (response.status === 404) {
|
|
throw new ApiError(
|
|
`Model "${this.model}" not found. Run \`ollama pull ${this.model}\` first.`,
|
|
404
|
|
);
|
|
}
|
|
throw new ApiError(`Ollama API error: ${response.status}`, response.status);
|
|
}
|
|
|
|
if (!response.body) {
|
|
throw new Error('No response body');
|
|
}
|
|
|
|
const contentType = response.headers?.get?.('content-type');
|
|
if (contentType && !contentType.includes('application/x-ndjson')) {
|
|
throw new Error('Invalid response format');
|
|
}
|
|
|
|
reader = response.body.getReader();
|
|
const decoder = new TextDecoder();
|
|
let buffer = '';
|
|
let malformedChunks = 0;
|
|
|
|
while (true) {
|
|
const { done, value } = await reader.read();
|
|
if (done) {
|
|
break;
|
|
}
|
|
|
|
buffer += decoder.decode(value);
|
|
const lines = buffer.split('\n');
|
|
buffer = lines.pop() || '';
|
|
|
|
for (const line of lines) {
|
|
if (line.trim() === '') {
|
|
continue;
|
|
}
|
|
|
|
let parsed: OllamaChatResponse;
|
|
try {
|
|
parsed = this.parseChatResponse(line);
|
|
} catch (error) {
|
|
malformedChunks++;
|
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
Logger.warn(
|
|
`Skipped malformed chunk: ${line.slice(0, 50)}... - ${errorMessage}`,
|
|
'ollama-client'
|
|
);
|
|
if (malformedChunks > this.maxMalformedChunks) {
|
|
throw new Error('Too many malformed chunks in Ollama response');
|
|
}
|
|
continue;
|
|
}
|
|
|
|
if (parsed.error) {
|
|
const errorMsg =
|
|
typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error);
|
|
throw new Error(`Ollama error: ${errorMsg}`);
|
|
}
|
|
|
|
yield this.normalizeMessage(parsed.message);
|
|
}
|
|
}
|
|
|
|
if (buffer.trim() !== '') {
|
|
let parsed: OllamaChatResponse | null = null;
|
|
try {
|
|
parsed = this.parseChatResponse(buffer);
|
|
} catch (error) {
|
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
Logger.warn(
|
|
`Skipped malformed chunk: ${buffer.slice(0, 50)}... - ${errorMessage}`,
|
|
'ollama-client'
|
|
);
|
|
}
|
|
|
|
if (parsed?.error) {
|
|
const errorMsg =
|
|
typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error);
|
|
throw new Error(`Ollama error: ${errorMsg}`);
|
|
}
|
|
|
|
if (parsed?.message) {
|
|
yield this.normalizeMessage(parsed.message);
|
|
}
|
|
}
|
|
} catch (error) {
|
|
if (retryCount < this.maxRetries && this.isRetryableError(error, controller)) {
|
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
Logger.warn(
|
|
`Retrying after error (attempt ${retryCount + 1}): ${errorMessage}`,
|
|
'ollama-client'
|
|
);
|
|
await new Promise((resolve) => setTimeout(resolve, 10 * Math.pow(2, retryCount)));
|
|
yield* this.streamChatWithRetry(messages, tools, retryCount + 1);
|
|
} else {
|
|
throw error;
|
|
}
|
|
} finally {
|
|
reader?.releaseLock();
|
|
if (this.currentStreamController === controller) {
|
|
this.currentStreamController = null;
|
|
}
|
|
}
|
|
}
|
|
|
|
async chatWithRetry(
|
|
messages: OllamaMessage[],
|
|
tools: OllamaTool[] = [],
|
|
retryCount: number
|
|
): Promise<OllamaMessage> {
|
|
const controller = new AbortController();
|
|
this.currentStreamController = controller;
|
|
|
|
try {
|
|
const response = await this.fetchFn(`${this.baseURL}/api/chat`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
},
|
|
body: JSON.stringify({
|
|
model: this.model,
|
|
messages: messages,
|
|
tools: tools,
|
|
stream: false,
|
|
}),
|
|
signal: controller.signal,
|
|
});
|
|
|
|
if (!response.ok) {
|
|
if (response.status === 404) {
|
|
throw new ApiError(
|
|
`Model "${this.model}" not found. Run \`ollama pull ${this.model}\` first.`,
|
|
404
|
|
);
|
|
}
|
|
throw new ApiError(`Ollama API error: ${response.status}`, response.status);
|
|
}
|
|
|
|
const data = (await response.json()) as unknown;
|
|
if (!this.isChatResponse(data)) {
|
|
return this.normalizeMessage();
|
|
}
|
|
return this.normalizeMessage(data.message);
|
|
} catch (error) {
|
|
if (retryCount < this.maxRetries && this.isRetryableError(error, controller)) {
|
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
Logger.warn(
|
|
`Retrying after error (attempt ${retryCount + 1}): ${errorMessage}`,
|
|
'ollama-client'
|
|
);
|
|
await new Promise((resolve) => setTimeout(resolve, 10 * Math.pow(2, retryCount)));
|
|
return this.chatWithRetry(messages, tools, retryCount + 1);
|
|
} else {
|
|
throw error;
|
|
}
|
|
} finally {
|
|
if (this.currentStreamController === controller) {
|
|
this.currentStreamController = null;
|
|
}
|
|
}
|
|
}
|
|
|
|
private normalizeMessage(message?: Partial<OllamaMessage>): OllamaMessage {
|
|
return {
|
|
role: message?.role ?? 'assistant',
|
|
content: message?.content ?? '',
|
|
tool_calls: message?.tool_calls ?? [],
|
|
tool_call_id: message?.tool_call_id,
|
|
};
|
|
}
|
|
|
|
private parseChatResponse(raw: string): OllamaChatResponse {
|
|
const parsed: unknown = JSON.parse(raw);
|
|
if (!this.isChatResponse(parsed)) {
|
|
throw new Error('Invalid chat response');
|
|
}
|
|
return parsed;
|
|
}
|
|
|
|
private isChatResponse(data: unknown): data is OllamaChatResponse {
|
|
if (typeof data !== 'object' || data === null) {
|
|
return false;
|
|
}
|
|
|
|
const response = data as { message?: unknown; error?: unknown };
|
|
return (
|
|
(response.error === undefined || typeof response.error === 'string') &&
|
|
(response.message === undefined || this.isPartialMessage(response.message))
|
|
);
|
|
}
|
|
|
|
private isPartialMessage(data: unknown): data is Partial<OllamaMessage> {
|
|
if (typeof data !== 'object' || data === null) {
|
|
return false;
|
|
}
|
|
|
|
const message = data as {
|
|
role?: unknown;
|
|
content?: unknown;
|
|
tool_calls?: unknown;
|
|
tool_call_id?: unknown;
|
|
};
|
|
const validRole =
|
|
message.role === undefined ||
|
|
message.role === 'system' ||
|
|
message.role === 'user' ||
|
|
message.role === 'assistant' ||
|
|
message.role === 'tool';
|
|
|
|
return (
|
|
validRole &&
|
|
(message.content === undefined || typeof message.content === 'string') &&
|
|
(message.tool_calls === undefined || Array.isArray(message.tool_calls)) &&
|
|
(message.tool_call_id === undefined || typeof message.tool_call_id === 'string')
|
|
);
|
|
}
|
|
|
|
private isRetryableError(error: unknown, controller: AbortController): boolean {
|
|
if (controller.signal.aborted) {
|
|
return false;
|
|
}
|
|
|
|
if (error instanceof ApiError && error.statusCode >= 400 && error.statusCode < 500) {
|
|
return false;
|
|
}
|
|
|
|
if (error instanceof Error) {
|
|
if (error.name === 'AbortError') {
|
|
return false;
|
|
}
|
|
|
|
if (
|
|
error.message.startsWith('Ollama error:') ||
|
|
error.message.includes('Too many malformed chunks') ||
|
|
error.message === 'No response body' ||
|
|
error.message === 'Invalid response format'
|
|
) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
}
|