From 4818741a77292d44bcbf12d3aebccc43f1ad9fc1 Mon Sep 17 00:00:00 2001 From: Andrew Beal Date: Sun, 5 Oct 2025 18:17:47 +0100 Subject: [PATCH] Refactor Gemini and StreamingService to enhance streamRequest method, improve error handling in parseStreamChunk, and update service registration for StreamingService. --- AIClasses/Gemini/Gemini.ts | 37 +++++++++++++++++++++++++-- Services/ServiceRegistration.ts | 2 ++ Services/Services.ts | 2 -- Services/StreamingService.ts | 45 +++++---------------------------- 4 files changed, 43 insertions(+), 43 deletions(-) diff --git a/AIClasses/Gemini/Gemini.ts b/AIClasses/Gemini/Gemini.ts index 09634fe..8c29d46 100644 --- a/AIClasses/Gemini/Gemini.ts +++ b/AIClasses/Gemini/Gemini.ts @@ -5,6 +5,7 @@ import type { IPrompt } from "AIClasses/IPrompt"; import { StreamingService, type StreamChunk } from "Services/StreamingService"; import type { Conversation } from "Conversations/Conversation"; import { Role } from "Enums/Role"; +import { AIProviderURL } from "Enums/ApiProvider"; export class Gemini implements IAIClass { private readonly apiKey: string; @@ -14,7 +15,7 @@ export class Gemini implements IAIClass { public constructor(apiKey: string) { this.apiKey = apiKey; this.aiPrompt = Resolve(Services.IPrompt); - this.streamingService = new StreamingService(); + this.streamingService = Resolve(Services.StreamingService); } public async* streamRequest(conversation: Conversation): AsyncGenerator { @@ -48,6 +49,38 @@ export class Gemini implements IAIClass { ] }; - yield* this.streamingService.streamGeminiRequest(this.apiKey, requestBody); + yield* this.streamingService.streamRequest( + AIProviderURL.Gemini.replace("API_KEY", this.apiKey), + requestBody, + this.parseStreamChunk + ); + } + + private parseStreamChunk(chunk: string): StreamChunk { + try { + const data = JSON.parse(chunk); + + let text = ""; + const candidate = data.candidates?.[0]; + + if (candidate) { + if (candidate.content?.parts?.[0]?.text) { + text = candidate.content.parts[0].text; + } else if (candidate.text) { + text = candidate.text; + } + } + + const isComplete = !!candidate?.finishReason; + + return { + content: text, + isComplete: isComplete, + }; + } catch (error) { + const message = error instanceof Error ? error.message : "Unknown parsing error"; + console.error("Failed to parse stream chunk:", message, "Chunk:", chunk); + return { content: "", isComplete: false, error: `Failed to parse chunk: ${message}` }; + } } } \ No newline at end of file diff --git a/Services/ServiceRegistration.ts b/Services/ServiceRegistration.ts index 0443756..42f3112 100644 --- a/Services/ServiceRegistration.ts +++ b/Services/ServiceRegistration.ts @@ -13,6 +13,7 @@ import { ConversationFileSystemService } from "./ConversationFileSystemService"; import { ConversationHistoryModal } from "Modals/ConversationHistoryModal"; import type { App } from "obsidian"; import { AIFunctionService } from "./AIFunctionService"; +import { StreamingService } from "./StreamingService"; export function RegisterDependencies(plugin: AIAgentPlugin) { RegisterSingleton(Services.MessageService, new MessageService()); @@ -23,6 +24,7 @@ export function RegisterDependencies(plugin: AIAgentPlugin) { RegisterSingleton(Services.IPrompt, new AIPrompt()); RegisterSingleton(Services.AIFunctionService, new AIFunctionService()); + RegisterSingleton(Services.StreamingService, new StreamingService()); RegisterTransient(Services.StreamingMarkdownService, () => new StreamingMarkdownService()); diff --git a/Services/Services.ts b/Services/Services.ts index 917b8f7..abd578e 100644 --- a/Services/Services.ts +++ b/Services/Services.ts @@ -12,8 +12,6 @@ export class Services { // interfaces static IAIClass = Symbol("IAIClass"); static IPrompt = Symbol("IPrompt"); - static IActioner = Symbol("IActioner"); - static IActionDefinitions = Symbol("IActionDefinitions"); // modals static ConversationHistoryModal = Symbol("ConversationHistoryModal"); diff --git a/Services/StreamingService.ts b/Services/StreamingService.ts index 27795cd..6d09133 100644 --- a/Services/StreamingService.ts +++ b/Services/StreamingService.ts @@ -1,5 +1,3 @@ -import { AIProviderURL } from "Enums/ApiProvider"; - export interface StreamChunk { content: string; isComplete: boolean; @@ -7,17 +5,14 @@ export interface StreamChunk { } export class StreamingService { - /** - * Fetches data from Gemini API with streaming support - * Since Obsidian's request() doesn't support streaming, we use native fetch - */ - public async* streamGeminiRequest( - apiKey: string, - requestBody: unknown + public async* streamRequest( + url: string, + requestBody: unknown, + parseStreamChunk: (chunk: string) => StreamChunk ): AsyncGenerator { try { const response = await fetch( - AIProviderURL.Gemini.replace("API_KEY", apiKey), + url, { method: "POST", headers: { @@ -51,7 +46,7 @@ export interface StreamChunk { for (const line of lines) { if (line.trim().startsWith("data:")) { const jsonStr = line.trim().substring(5); - const chunk = this.parseStreamChunk(jsonStr); + const chunk = parseStreamChunk(jsonStr); lastChunkWasComplete = chunk.isComplete; yield chunk; } @@ -74,32 +69,4 @@ export interface StreamChunk { }; } } - - private parseStreamChunk(chunk: string): StreamChunk { - try { - const data = JSON.parse(chunk); - - let text = ""; - const candidate = data.candidates?.[0]; - - if (candidate) { - if (candidate.content?.parts?.[0]?.text) { - text = candidate.content.parts[0].text; - } else if (candidate.text) { - text = candidate.text; - } - } - - const isComplete = !!candidate?.finishReason; - - return { - content: text, - isComplete: isComplete, - }; - } catch (error) { - const message = error instanceof Error ? error.message : "Unknown parsing error"; - console.error("Failed to parse stream chunk:", message, "Chunk:", chunk); - return { content: "", isComplete: false, error: `Failed to parse chunk: ${message}` }; - } - } } \ No newline at end of file