diff --git a/mem0-ts/src/oss/src/llms/google.ts b/mem0-ts/src/oss/src/llms/google.ts index c852baf37..d9a461b98 100644 --- a/mem0-ts/src/oss/src/llms/google.ts +++ b/mem0-ts/src/oss/src/llms/google.ts @@ -11,12 +11,8 @@ export class GoogleLLM implements LLM { this.model = config.model || "gemini-2.0-flash"; } - async generateResponse( - messages: Message[], - responseFormat?: { type: string }, - tools?: any[], - ): Promise { - const contents = messages.map((msg) => ({ + private formatContents(messages: Message[]) { + return messages.map((msg) => ({ parts: [ { text: @@ -27,6 +23,14 @@ export class GoogleLLM implements LLM { ], role: msg.role === "system" ? "model" : "user", })); + } + + async generateResponse( + messages: Message[], + responseFormat?: { type: string }, + tools?: any[], + ): Promise { + const contents = this.formatContents(messages); // Build config with tools if provided const config: Record = {}; @@ -69,13 +73,18 @@ export class GoogleLLM implements LLM { async generateChat(messages: Message[]): Promise { const completion = await this.google.models.generateContent({ - contents: messages, + contents: this.formatContents(messages), model: this.model, }); - const response = completion.candidates![0].content; + const response = completion.candidates?.[0]?.content; + const content = + response?.parts?.map((part) => part.text || "").join("") || + completion.text || + ""; + return { - content: response!.parts![0].text || "", - role: response!.role!, + content, + role: response?.role || "assistant", }; } } diff --git a/mem0-ts/src/oss/tests/google-llm.test.ts b/mem0-ts/src/oss/tests/google-llm.test.ts index 40a3c4e42..373229665 100644 --- a/mem0-ts/src/oss/tests/google-llm.test.ts +++ b/mem0-ts/src/oss/tests/google-llm.test.ts @@ -184,4 +184,45 @@ describe("GoogleLLM (unit)", () => { expect(response.toolCalls[0].name).toBe("add_graph_memory"); expect(response.toolCalls[1].name).toBe("add_graph_memory"); }); + + it("formats generateChat messages and joins Gemini response parts", async () => { + mockGenerateContent.mockResolvedValueOnce({ + candidates: [ + { + content: { + role: "model", + parts: [{ text: "Hello" }, { text: ", world" }], + }, + }, + ], + }); + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateChat([ + { role: "system", content: "Be concise" }, + { role: "user", content: "Say hello" }, + ]); + + expect(mockGenerateContent).toHaveBeenCalledWith( + expect.objectContaining({ + contents: [ + { role: "model", parts: [{ text: "Be concise" }] }, + { role: "user", parts: [{ text: "Say hello" }] }, + ], + }), + ); + expect(result).toEqual({ content: "Hello, world", role: "model" }); + }); + + it("returns an empty assistant response when generateChat has no candidates", async () => { + mockGenerateContent.mockResolvedValueOnce({ + candidates: [], + text: "", + }); + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateChat([{ role: "user", content: "Hi" }]); + + expect(result).toEqual({ content: "", role: "assistant" }); + }); });