From 9269a0ad6e720ea9a7f3d2fa0b71b7527215bb82 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Thu, 25 Jun 2026 06:22:46 -0400 Subject: [PATCH] feat(ts-sdk): add LiteLLM as LLM provider (#5830) Co-authored-by: kartik-mem0 --- docs/components/llms/models/litellm.mdx | 32 ++++- mem0-ts/src/oss/src/index.ts | 1 + mem0-ts/src/oss/src/llms/litellm.ts | 39 ++++++ mem0-ts/src/oss/src/utils/factory.ts | 3 + mem0-ts/src/oss/tests/factory.unit.test.ts | 6 + mem0-ts/src/oss/tests/litellm.test.ts | 131 +++++++++++++++++++++ 6 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 mem0-ts/src/oss/src/llms/litellm.ts create mode 100644 mem0-ts/src/oss/tests/litellm.test.ts diff --git a/docs/components/llms/models/litellm.mdx b/docs/components/llms/models/litellm.mdx index 31efee473..61aecba21 100644 --- a/docs/components/llms/models/litellm.mdx +++ b/docs/components/llms/models/litellm.mdx @@ -4,9 +4,12 @@ description: "Use LiteLLM as an LLM provider in Mem0 to access over 100 language --- [Litellm](https://litellm.vercel.app/docs/) is compatible with over 100 large language models (LLMs), all using a standardized input/output format. You can explore the [available models](https://litellm.vercel.app/docs/providers) to use with Litellm. Ensure you set the `API_KEY` for the model you choose to use. +In the TypeScript SDK, run LiteLLM as a [proxy server](https://docs.litellm.ai/docs/simple_proxy) (an OpenAI-compatible endpoint) and point Mem0 at it via `LITELLM_API_BASE` (defaults to `http://localhost:4000`). + ## Usage -```python + +```python Python import os from mem0 import Memory @@ -33,6 +36,33 @@ messages = [ m.add(messages, user_id="alice", metadata={"category": "movies"}) ``` +```typescript TypeScript +import { Memory } from 'mem0ai/oss'; + +// Point Mem0 at your LiteLLM proxy. apiKey defaults to "sk-anything" +// (the proxy handles real auth); baseURL defaults to http://localhost:4000. +const config = { + llm: { + provider: 'litellm', + config: { + apiKey: process.env.LITELLM_API_KEY || 'sk-anything', + baseURL: process.env.LITELLM_API_BASE || 'http://localhost:4000', + model: 'gpt-5-mini', + }, + }, +}; + +const memory = new Memory(config); +const messages = [ + {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"}, + {"role": "assistant", "content": "How about thriller movies? They can be quite engaging."}, + {"role": "user", "content": "I’m not a big fan of thriller movies but I love sci-fi movies."}, + {"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."} +]; +await memory.add(messages, { userId: 'alice', metadata: { category: 'movies' } }); +``` + + ## Config All available parameters for the `litellm` config are present in [Master List of All Params in Config](../config). \ No newline at end of file diff --git a/mem0-ts/src/oss/src/index.ts b/mem0-ts/src/oss/src/index.ts index 68c98a90e..141edcacf 100644 --- a/mem0-ts/src/oss/src/index.ts +++ b/mem0-ts/src/oss/src/index.ts @@ -18,6 +18,7 @@ export * from "./llms/ollama"; export * from "./llms/lmstudio"; export * from "./llms/mistral"; export * from "./llms/langchain"; +export * from "./llms/litellm"; export * from "./vector_stores/base"; export * from "./vector_stores/memory"; export * from "./vector_stores/qdrant"; diff --git a/mem0-ts/src/oss/src/llms/litellm.ts b/mem0-ts/src/oss/src/llms/litellm.ts new file mode 100644 index 000000000..3642fce08 --- /dev/null +++ b/mem0-ts/src/oss/src/llms/litellm.ts @@ -0,0 +1,39 @@ +import { OpenAILLM } from "./openai"; +import { LLMConfig, Message } from "../types"; +import { LLMResponse } from "./base"; + +export class LiteLLM extends OpenAILLM { + constructor(config: LLMConfig) { + super({ + ...config, + apiKey: config.apiKey || process.env.LITELLM_API_KEY || "sk-anything", + baseURL: + config.baseURL || + process.env.LITELLM_API_BASE || + "http://localhost:4000", + model: config.model || "gpt-5-mini", + }); + } + + async generateResponse( + messages: Message[], + responseFormat?: { type: string }, + tools?: any[], + ): Promise { + try { + return await super.generateResponse(messages, responseFormat, tools); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new Error(`LiteLLM failed: ${message}`); + } + } + + async generateChat(messages: Message[]): Promise { + try { + return await super.generateChat(messages); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new Error(`LiteLLM failed: ${message}`); + } + } +} diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 35c2a21a1..f8a0708de 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -22,6 +22,7 @@ import { RedisDB } from "../vector_stores/redis"; import { OllamaLLM } from "../llms/ollama"; import { LMStudioLLM } from "../llms/lmstudio"; import { DeepSeekLLM } from "../llms/deepseek"; +import { LiteLLM } from "../llms/litellm"; import { SupabaseDB } from "../vector_stores/supabase"; import { SQLiteManager } from "../storage/SQLiteManager"; import { MemoryHistoryManager } from "../storage/MemoryHistoryManager"; @@ -85,6 +86,8 @@ export class LLMFactory { return new LangchainLLM(config); case "deepseek": return new DeepSeekLLM(config); + case "litellm": + return new LiteLLM(config); default: throw new Error(`Unsupported LLM provider: ${provider}`); } diff --git a/mem0-ts/src/oss/tests/factory.unit.test.ts b/mem0-ts/src/oss/tests/factory.unit.test.ts index c84009150..2d2bc7fc8 100644 --- a/mem0-ts/src/oss/tests/factory.unit.test.ts +++ b/mem0-ts/src/oss/tests/factory.unit.test.ts @@ -92,6 +92,11 @@ jest.mock("../src/llms/deepseek", () => ({ .fn() .mockImplementation((config) => ({ type: "deepseek-llm", config })), })); +jest.mock("../src/llms/litellm", () => ({ + LiteLLM: jest + .fn() + .mockImplementation((config) => ({ type: "litellm-llm", config })), +})); jest.mock("../src/vector_stores/qdrant", () => ({ Qdrant: jest @@ -206,6 +211,7 @@ describe("LLMFactory", () => { ["langchain"], ["lmstudio"], ["deepseek"], + ["litellm"], ])("creates LLM for provider '%s'", (provider) => { expect(() => LLMFactory.create(provider, dummyLLMConfig)).not.toThrow(); }); diff --git a/mem0-ts/src/oss/tests/litellm.test.ts b/mem0-ts/src/oss/tests/litellm.test.ts new file mode 100644 index 000000000..4e73d27d3 --- /dev/null +++ b/mem0-ts/src/oss/tests/litellm.test.ts @@ -0,0 +1,131 @@ +/// +/** + * LiteLLM — unit tests (mocked OpenAI). + */ + +import { LiteLLM } from "../src/llms/litellm"; + +const mockCreate = jest.fn(); + +jest.mock("openai", () => { + return jest.fn().mockImplementation(() => ({ + chat: { completions: { create: mockCreate } }, + })); +}); + +describe("LiteLLM (unit)", () => { + beforeEach(() => mockCreate.mockClear()); + + it("uses default baseURL when none is provided", () => { + const llm = new LiteLLM({}); + expect(llm).toBeDefined(); + }); + + it("generateResponse() returns a text response", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [ + { + message: { + content: "Hello, world!", + role: "assistant", + tool_calls: null, + }, + }, + ], + }); + + const llm = new LiteLLM({ baseURL: "http://localhost:4000" }); + const result = await llm.generateResponse([ + { role: "user", content: "Hi" }, + ]); + + expect(mockCreate).toHaveBeenCalledTimes(1); + expect(result).toBe("Hello, world!"); + }); + + it("generateResponse() handles tool calls", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [ + { + message: { + content: "", + role: "assistant", + tool_calls: [ + { + function: { + name: "get_weather", + arguments: '{"city": "London"}', + }, + }, + ], + }, + }, + ], + }); + + const llm = new LiteLLM({}); + const result = await llm.generateResponse( + [{ role: "user", content: "What is the weather?" }], + undefined, + [{ type: "function", function: { name: "get_weather" } }], + ); + + expect(result).toEqual({ + content: "", + role: "assistant", + toolCalls: [{ name: "get_weather", arguments: '{"city": "London"}' }], + }); + }); + + it("generateResponse() wraps API errors with a clear message", async () => { + mockCreate.mockRejectedValueOnce(new Error("Connection refused")); + + const llm = new LiteLLM({}); + + await expect( + llm.generateResponse([{ role: "user", content: "Hi" }]), + ).rejects.toThrow("LiteLLM failed: Connection refused"); + }); + + it("generateChat() returns LLMResponse shape", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [ + { + message: { content: "I can help with that.", role: "assistant" }, + }, + ], + }); + + const llm = new LiteLLM({}); + const result = await llm.generateChat([ + { role: "user", content: "Help me" }, + ]); + + expect(result).toEqual({ + content: "I can help with that.", + role: "assistant", + }); + }); + + it("generateChat() wraps API errors with a clear message", async () => { + mockCreate.mockRejectedValueOnce(new Error("Timeout")); + + const llm = new LiteLLM({}); + + await expect( + llm.generateChat([{ role: "user", content: "Hi" }]), + ).rejects.toThrow("LiteLLM failed: Timeout"); + }); + + it("respects LITELLM_API_BASE env var", () => { + const original = process.env.LITELLM_API_BASE; + process.env.LITELLM_API_BASE = "http://custom-proxy:8080"; + try { + const llm = new LiteLLM({}); + expect(llm).toBeDefined(); + } finally { + if (original !== undefined) process.env.LITELLM_API_BASE = original; + else delete process.env.LITELLM_API_BASE; + } + }); +});