obsidian_ollama/src/ollama-client.ts

// src/ollama-client.ts

import type { OllamaMessage, OllamaTool } from './types';
import { ApiError } from './types';
import { Logger } from './utils';

interface OllamaChatResponse {
  message?: Partial<OllamaMessage>;
}

export class OllamaClient {
  private baseURL: string;
  private model: string;
  private fetchFn: typeof fetch;
  private readonly maxRetries: number = 3;
  private currentStreamController: AbortController | null = null;

  constructor(baseURL: string, model: string, fetchFn?: typeof fetch) {
    this.baseURL = baseURL;
    this.model = model;
    this.fetchFn = fetchFn ?? fetch;
  }

  cancelStream(): void {
    if (this.currentStreamController) {
      this.currentStreamController.abort();
      this.currentStreamController = null;
    }
  }

  async *streamChat(
    messages: OllamaMessage[],
    tools: OllamaTool[] = []
  ): AsyncGenerator<OllamaMessage, void, unknown> {
    yield* this.streamChatWithRetry(messages, tools, 0);
  }

  async streamChatAsPromise(
    messages: OllamaMessage[],
    tools: OllamaTool[] = []
  ): Promise<OllamaMessage[]> {
    const chunks: OllamaMessage[] = [];
    for await (const chunk of this.streamChat(messages, tools)) {
      chunks.push(chunk);
    }
    return chunks;
  }

  private async *streamChatWithRetry(
    messages: OllamaMessage[],
    tools: OllamaTool[] = [],
    attempt: number = 0
  ): AsyncGenerator<OllamaMessage, void, unknown> {
    // Create a local controller for this request instead of using the instance variable
    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,
          tools,
          stream: true,
        }),
        signal: controller.signal,
      });

      if (!response.ok) {
        if (response.status >= 500 && attempt < this.maxRetries) {
          const retryDelay = Math.pow(2, attempt) * 100;
          Logger.warn(
            `Network error (status ${response.status}), retrying in ${retryDelay}ms (attempt ${attempt + 1}/${this.maxRetries})`,
            'ollama-client'
          );
          if (attempt < this.maxRetries - 1) {
            const retryTimeout = new Promise((resolve) => setTimeout(resolve, retryDelay));
            const abortListener = () => {
              Logger.info('Retry aborted by user', 'ollama-client');
            };
            const signal = controller.signal;
            if (signal) {
              signal.addEventListener('abort', abortListener);
              try {
                await Promise.race([
                  retryTimeout,
                  new Promise<void>((resolve) => {
                    signal.addEventListener('abort', () => resolve(), {
                      once: true,
                    });
                  }),
                ]);
              } finally {
                signal.removeEventListener('abort', abortListener);
              }
            } else {
              await retryTimeout;
            }
          }
          yield* this.streamChatWithRetry(messages, tools, attempt + 1);
          return;
        }
        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('ndjson') && !contentType.includes('json'))) {
        throw new Error('Invalid response format');
      }

      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = '';
      let malformedCount = 0;
      const maxMalformed = 50;

      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;

          buffer += decoder.decode(value, { stream: true });

          const lines = buffer.split('\n');
          buffer = lines.pop() ?? '';

          for (const line of lines) {
            if (!line.trim()) continue;

            try {
              const parsed = JSON.parse(line) as Record<string, unknown>;
              this.throwIfOllamaError(parsed);

              const message = this.toOllamaMessage(parsed.message);
              if (!message) {
                continue;
              }

              malformedCount = 0;
              yield message;
            } catch (error) {
              if (error instanceof Error && error.message.startsWith('Ollama error:')) {
                throw error;
              }

              malformedCount++;
              if (malformedCount > maxMalformed) {
                throw new Error('Too many malformed chunks in stream');
              }

              Logger.warn(
                `Skipped malformed chunk: ${line.substring(0, 80)}... - ${(error as Error).message}`,
                'ollama-client'
              );
            }
          }
        }

        if (buffer.trim()) {
          try {
            const parsed = JSON.parse(buffer) as Record<string, unknown>;
            this.throwIfOllamaError(parsed);

            const message = this.toOllamaMessage(parsed.message);
            if (message) {
              yield message;
            }
          } catch (error) {
            if (error instanceof Error && error.message.startsWith('Ollama error:')) {
              throw error;
            }
            Logger.warn(
              `Failed to parse final chunk: ${buffer.substring(0, 80)}...`,
              'ollama-client'
            );
          }
        }
      } finally {
        reader.releaseLock();
      }
    } finally {
      // Abort the local controller
      controller.abort();
    }
    // Clean up the reference
    this.currentStreamController = null;
  }

  async chat(messages: OllamaMessage[], tools: OllamaTool[] = []): Promise<OllamaMessage> {
    return this.chatWithRetry(messages, tools, 0);
  }

  private async chatWithRetry(
    messages: OllamaMessage[],
    tools: OllamaTool[] = [],
    attempt: number = 0
  ): Promise<OllamaMessage> {
    // Create a local controller for this request instead of using the instance variable
    const controller = new AbortController();
    try {
      const response = await this.fetchFn(`${this.baseURL}/api/chat`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          model: this.model,
          messages,
          tools,
          stream: false,
        }),
        signal: controller.signal,
      });

      if (!response.ok) {
        if (response.status >= 500 && attempt < this.maxRetries) {
          const retryDelay = Math.pow(2, attempt) * 100;
          Logger.warn(
            `Network error (status ${response.status}), retrying in ${retryDelay}ms (attempt ${attempt + 1}/${this.maxRetries})`,
            'ollama-client'
          );
          if (attempt < this.maxRetries - 1) {
            const retryTimeout = new Promise((resolve) => setTimeout(resolve, retryDelay));
            const abortListener = () => {
              Logger.info('Retry aborted by user', 'ollama-client');
            };
            const signal = controller.signal;
            if (signal) {
              signal.addEventListener('abort', abortListener);
              try {
                await Promise.race([
                  retryTimeout,
                  new Promise<void>((resolve) => {
                    signal.addEventListener('abort', () => resolve(), {
                      once: true,
                    });
                  }),
                ]);
              } finally {
                signal.removeEventListener('abort', abortListener);
              }
            } else {
              await retryTimeout;
            }
          }
          return this.chatWithRetry(messages, tools, attempt + 1);
        }
        throw new ApiError(`Ollama API error: ${response.status}`, response.status);
      }

      const data = (await response.json()) as OllamaChatResponse;
      return (
        this.toOllamaMessage(data.message) ?? { role: 'assistant', content: '', tool_calls: [] }
      );
    } finally {
      // Abort the local controller
      controller.abort();
    }
  }

  private throwIfOllamaError(parsed: Record<string, unknown>): void {
    if (parsed.error) {
      throw new Error(`Ollama error: ${String(parsed.error)}`);
    }
  }

  private toOllamaMessage(value: unknown): OllamaMessage | null {
    if (!value || typeof value !== 'object') {
      return null;
    }

    const record = value as Partial<OllamaMessage>;
    return {
      role: record.role ?? 'assistant',
      content: typeof record.content === 'string' ? record.content : '',
      tool_calls: record.tool_calls ?? [],
    };
  }
}