feat(ts-sdk): add LiteLLM as LLM provider (#5830)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
@@ -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.
|
[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
|
## Usage
|
||||||
|
|
||||||
```python
|
<CodeGroup>
|
||||||
|
```python Python
|
||||||
import os
|
import os
|
||||||
from mem0 import Memory
|
from mem0 import Memory
|
||||||
|
|
||||||
@@ -33,6 +36,33 @@ messages = [
|
|||||||
m.add(messages, user_id="alice", metadata={"category": "movies"})
|
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' } });
|
||||||
|
```
|
||||||
|
</CodeGroup>
|
||||||
|
|
||||||
## Config
|
## Config
|
||||||
|
|
||||||
All available parameters for the `litellm` config are present in [Master List of All Params in Config](../config).
|
All available parameters for the `litellm` config are present in [Master List of All Params in Config](../config).
|
||||||
@@ -18,6 +18,7 @@ export * from "./llms/ollama";
|
|||||||
export * from "./llms/lmstudio";
|
export * from "./llms/lmstudio";
|
||||||
export * from "./llms/mistral";
|
export * from "./llms/mistral";
|
||||||
export * from "./llms/langchain";
|
export * from "./llms/langchain";
|
||||||
|
export * from "./llms/litellm";
|
||||||
export * from "./vector_stores/base";
|
export * from "./vector_stores/base";
|
||||||
export * from "./vector_stores/memory";
|
export * from "./vector_stores/memory";
|
||||||
export * from "./vector_stores/qdrant";
|
export * from "./vector_stores/qdrant";
|
||||||
|
|||||||
@@ -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<string | LLMResponse> {
|
||||||
|
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<LLMResponse> {
|
||||||
|
try {
|
||||||
|
return await super.generateChat(messages);
|
||||||
|
} catch (err) {
|
||||||
|
const message = err instanceof Error ? err.message : String(err);
|
||||||
|
throw new Error(`LiteLLM failed: ${message}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -22,6 +22,7 @@ import { RedisDB } from "../vector_stores/redis";
|
|||||||
import { OllamaLLM } from "../llms/ollama";
|
import { OllamaLLM } from "../llms/ollama";
|
||||||
import { LMStudioLLM } from "../llms/lmstudio";
|
import { LMStudioLLM } from "../llms/lmstudio";
|
||||||
import { DeepSeekLLM } from "../llms/deepseek";
|
import { DeepSeekLLM } from "../llms/deepseek";
|
||||||
|
import { LiteLLM } from "../llms/litellm";
|
||||||
import { SupabaseDB } from "../vector_stores/supabase";
|
import { SupabaseDB } from "../vector_stores/supabase";
|
||||||
import { SQLiteManager } from "../storage/SQLiteManager";
|
import { SQLiteManager } from "../storage/SQLiteManager";
|
||||||
import { MemoryHistoryManager } from "../storage/MemoryHistoryManager";
|
import { MemoryHistoryManager } from "../storage/MemoryHistoryManager";
|
||||||
@@ -85,6 +86,8 @@ export class LLMFactory {
|
|||||||
return new LangchainLLM(config);
|
return new LangchainLLM(config);
|
||||||
case "deepseek":
|
case "deepseek":
|
||||||
return new DeepSeekLLM(config);
|
return new DeepSeekLLM(config);
|
||||||
|
case "litellm":
|
||||||
|
return new LiteLLM(config);
|
||||||
default:
|
default:
|
||||||
throw new Error(`Unsupported LLM provider: ${provider}`);
|
throw new Error(`Unsupported LLM provider: ${provider}`);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -92,6 +92,11 @@ jest.mock("../src/llms/deepseek", () => ({
|
|||||||
.fn()
|
.fn()
|
||||||
.mockImplementation((config) => ({ type: "deepseek-llm", config })),
|
.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", () => ({
|
jest.mock("../src/vector_stores/qdrant", () => ({
|
||||||
Qdrant: jest
|
Qdrant: jest
|
||||||
@@ -206,6 +211,7 @@ describe("LLMFactory", () => {
|
|||||||
["langchain"],
|
["langchain"],
|
||||||
["lmstudio"],
|
["lmstudio"],
|
||||||
["deepseek"],
|
["deepseek"],
|
||||||
|
["litellm"],
|
||||||
])("creates LLM for provider '%s'", (provider) => {
|
])("creates LLM for provider '%s'", (provider) => {
|
||||||
expect(() => LLMFactory.create(provider, dummyLLMConfig)).not.toThrow();
|
expect(() => LLMFactory.create(provider, dummyLLMConfig)).not.toThrow();
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -0,0 +1,131 @@
|
|||||||
|
/// <reference types="jest" />
|
||||||
|
/**
|
||||||
|
* 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;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user