diff --git a/LLM.md b/LLM.md index d8b07b34ec..d6179276f7 100644 --- a/LLM.md +++ b/LLM.md @@ -273,7 +273,8 @@ config = MemoryConfig( ### Supported Providers -#### LLM Providers (20 supported) +#### LLM Providers (21 supported) +- **aimlapi** - aimlapi.com gateway (350+ models) - **openai** - OpenAI GPT models (default) - **anthropic** - Claude models - **gemini** - Google Gemini @@ -293,7 +294,8 @@ config = MemoryConfig( - **openai_structured** - OpenAI with structured output - **azure_openai_structured** - Azure OpenAI with structured output -#### Embedding Providers (10 supported) +#### Embedding Providers (11 supported) +- **aimlapi** - aimlapi.com embeddings - **openai** - OpenAI embeddings (default) - **ollama** - Ollama embeddings - **huggingface** - HuggingFace models @@ -419,6 +421,7 @@ config = MemoryConfig( ``` #### LLM Providers +- **aimlapi.com** - OpenAI-compatible gateway to 350+ models - **OpenAI** - GPT-4, GPT-3.5-turbo, and structured outputs - **Anthropic** - Claude models with advanced reasoning - **Google AI** - Gemini models for multimodal applications diff --git a/docs/components/embedders/config.mdx b/docs/components/embedders/config.mdx index 8bde10ba75..01e22534ba 100644 --- a/docs/components/embedders/config.mdx +++ b/docs/components/embedders/config.mdx @@ -85,6 +85,7 @@ Here's a comprehensive list of all parameters that can be used across different | `memory_update_embedding_type` | The type of embedding to use for the update memory action | VertexAI | | `memory_search_embedding_type` | The type of embedding to use for the search memory action | VertexAI | | `lmstudio_base_url` | Base URL for LM Studio API | LM Studio | +| `aimlapi_base_url` | Base URL for the aimlapi.com API | aimlapi.com | | Parameter | Description | Provider | diff --git a/docs/components/embedders/models/aimlapi.mdx b/docs/components/embedders/models/aimlapi.mdx new file mode 100644 index 0000000000..f4218ec43f --- /dev/null +++ b/docs/components/embedders/models/aimlapi.mdx @@ -0,0 +1,91 @@ +--- +title: aimlapi.com +description: "Configure aimlapi.com as an embedding provider in Mem0, with OpenAI, Voyage, Google and Alibaba embedding models behind one OpenAI-compatible endpoint." +--- + +[aimlapi.com](https://aimlapi.com) serves embedding models from several vendors over the OpenAI-compatible `/v1/embeddings` route. Set the `AIMLAPI_API_KEY` environment variable with a key from the [aimlapi.com dashboard](https://aimlapi.com/app/keys). + +Model ids are namespaced, for example `openai/text-embedding-3-small` (1536 dims) or `openai/text-embedding-3-large` (3072 dims). The live catalog is at `https://api.aimlapi.com/v1/models`. + +### Usage + + The `embedding_model_dims` parameter for `vector_store` must match the model you pick — `1536` for `openai/text-embedding-3-small`, `3072` for `openai/text-embedding-3-large`. + + +```python Python +import os +from mem0 import Memory + +os.environ["AIMLAPI_API_KEY"] = "your-api-key" + +config = { + "llm": { + "provider": "aimlapi", + "config": { + "model": "openai/gpt-5-mini" + } + }, + "embedder": { + "provider": "aimlapi", + "config": { + "model": "openai/text-embedding-3-small" + } + } +} + +m = Memory.from_config(config) +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."} +] +m.add(messages, user_id="john") +``` + +```typescript TypeScript +import { Memory } from 'mem0ai/oss'; + +const config = { + embedder: { + provider: 'aimlapi', + config: { + apiKey: process.env.AIMLAPI_API_KEY || '', + model: 'openai/text-embedding-3-small', + embeddingDims: 1536, + }, + }, +}; + +const memory = new Memory(config); +await memory.add("I'm visiting Paris", { userId: "john" }); +``` + + + + +The endpoint accepts a string or a list of strings as `input` and rejects the pre-tokenised integer-array form with a 400. If you route aimlapi.com through the `langchain` embedder instead of this one, configure the LangChain embedding class to send text rather than token ids. + + +### Config + +Here are the parameters available for configuring the aimlapi.com embedder: + + + +| Parameter | Description | Default Value | +| --- | --- | --- | +| `model` | The name of the embedding model to use | `openai/text-embedding-3-small` | +| `embedding_dims` | Dimensions of the embedding model; sent as `dimensions` only when set | `1536` | +| `api_key` | The aimlapi.com API key | `AIMLAPI_API_KEY` | +| `aimlapi_base_url` | Base URL for the API | `https://api.aimlapi.com/v1` | + + +| Parameter | Description | Default Value | +| --- | --- | --- | +| `model` | The name of the embedding model to use | `openai/text-embedding-3-small` | +| `embeddingDims` | Dimensions of the embedding model for vector store configuration | `1536` | +| `apiKey` | The aimlapi.com API key | `AIMLAPI_API_KEY` | +| `baseURL` | Base URL for the API | `https://api.aimlapi.com/v1` | + + diff --git a/docs/components/embedders/overview.mdx b/docs/components/embedders/overview.mdx index 0a1b66961f..80ec180747 100644 --- a/docs/components/embedders/overview.mdx +++ b/docs/components/embedders/overview.mdx @@ -10,10 +10,11 @@ Mem0 offers support for various embedding models, allowing users to choose the o See the list of supported embedders below. - All embedders listed below are supported in the Python implementation. The TypeScript implementation supports: **OpenAI**, **Azure OpenAI**, **AWS Bedrock**, **FastEmbed**, **Google AI**, **Hugging Face**, **Langchain**, **LM Studio**, **Ollama**, **Together**, and **Vertex AI**. + All embedders listed below are supported in the Python implementation. The TypeScript implementation supports: **OpenAI**, **Azure OpenAI**, **AWS Bedrock**, **FastEmbed**, **Google AI**, **Hugging Face**, **Langchain**, **LM Studio**, **Ollama**, **Together**, **Vertex AI**, and **aimlapi.com**. + diff --git a/docs/components/llms/config.mdx b/docs/components/llms/config.mdx index 1f183a6bdf..399be66ad5 100644 --- a/docs/components/llms/config.mdx +++ b/docs/components/llms/config.mdx @@ -110,6 +110,7 @@ Here's a comprehensive list of all parameters that can be used across different | `deepseek_base_url` | Base URL for DeepSeek API | DeepSeek | | `xai_base_url` | Base URL for XAI API | XAI | | `sarvam_base_url` | Base URL for Sarvam API | Sarvam | + | `aimlapi_base_url` | Base URL for the aimlapi.com API | aimlapi.com | | `reasoning_effort` | Reasoning level (low, medium, high) | All | | `frequency_penalty` | Penalize frequent tokens (-2.0 to 2.0) | Sarvam | | `presence_penalty` | Penalize existing tokens (-2.0 to 2.0) | Sarvam | diff --git a/docs/components/llms/models/aimlapi.mdx b/docs/components/llms/models/aimlapi.mdx new file mode 100644 index 0000000000..a852af94b2 --- /dev/null +++ b/docs/components/llms/models/aimlapi.mdx @@ -0,0 +1,103 @@ +--- +title: aimlapi.com +description: "Configure aimlapi.com as an LLM provider in Mem0, an OpenAI-compatible gateway to 350+ chat models from OpenAI, Anthropic, Google, DeepSeek and others." +--- + +[aimlapi.com](https://aimlapi.com) exposes 350+ chat models behind a single OpenAI-compatible endpoint, so one key and one base URL reach OpenAI, Anthropic, Google, DeepSeek, Alibaba and others. Mem0 can use it for the LLM, the [embedder](/components/embedders/models/aimlapi), or both. + +Set the `AIMLAPI_API_KEY` environment variable with a key from the [aimlapi.com dashboard](https://aimlapi.com/app/keys). + +Model ids are namespaced, for example `openai/gpt-5-mini`, `anthropic/claude-sonnet-4.6` or `google/gemini-2.5-flash`. The live catalog is at `https://api.aimlapi.com/v1/models`. + +## Usage + + +```python Python +import os +from mem0 import Memory + +os.environ["AIMLAPI_API_KEY"] = "your-api-key" + +config = { + "llm": { + "provider": "aimlapi", + "config": { + "model": "openai/gpt-5-mini", + "temperature": 0.1, + } + }, + "embedder": { + "provider": "aimlapi", + "config": { + "model": "openai/text-embedding-3-small" + } + } +} + +m = Memory.from_config(config) +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."} +] +m.add(messages, user_id="alex") +``` + +```typescript TypeScript +import { Memory } from 'mem0ai/oss'; + +const config = { + llm: { + provider: 'aimlapi', + config: { + apiKey: process.env.AIMLAPI_API_KEY || '', + model: 'openai/gpt-5-mini', + temperature: 0.1, + }, + }, + embedder: { + provider: 'aimlapi', + config: { + apiKey: process.env.AIMLAPI_API_KEY || '', + model: 'openai/text-embedding-3-small', + embeddingDims: 1536, + }, + }, +}; + +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: 'alex' }); +``` + + + +## Pointing at a different endpoint + +`aimlapi_base_url` (Python) or `baseURL` (TypeScript) overrides the endpoint, and `AIMLAPI_API_BASE` does the same from the environment. Mem0's attribution headers are only attached when the resolved host is `api.aimlapi.com`, so they are not forwarded to a gateway of your own. + +```python +config = { + "llm": { + "provider": "aimlapi", + "config": { + "model": "anthropic/claude-sonnet-4.6", + "aimlapi_base_url": "https://my-gateway.internal/v1", + } + } +} +``` + + +`max_tokens` does not bound reasoning tokens on every model behind the gateway, and `completion_tokens` in the usage block excludes reasoning tokens on some of them. Treat it as a length hint, not a cost ceiling. + + +## Config + +All available parameters for the `aimlapi` config are present in [Master List of All Params in Config](../config). diff --git a/docs/components/llms/overview.mdx b/docs/components/llms/overview.mdx index 12b280f47d..3ff73e8e81 100644 --- a/docs/components/llms/overview.mdx +++ b/docs/components/llms/overview.mdx @@ -16,10 +16,11 @@ For a comprehensive list of available parameters for llm configuration, please r See the list of supported LLMs below. - All LLMs are supported in Python. The following LLMs are also supported in TypeScript: **OpenAI**, **Anthropic**, **AWS Bedrock**, **Groq**, **Azure OpenAI**, **DeepSeek**, **Google AI**, **Langchain**, **LM Studio**, **Mistral AI**, and **Ollama**. + All LLMs are supported in Python. The following LLMs are also supported in TypeScript: **OpenAI**, **Anthropic**, **AWS Bedrock**, **Groq**, **Azure OpenAI**, **DeepSeek**, **Google AI**, **Langchain**, **LM Studio**, **Mistral AI**, **Ollama**, and **aimlapi.com**. + diff --git a/docs/docs.json b/docs/docs.json index 54a53eca65..fa69306ba8 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -171,6 +171,7 @@ "group": "Supported LLMs", "icon": "list", "pages": [ + "components/llms/models/aimlapi", "components/llms/models/openai", "components/llms/models/anthropic", "components/llms/models/azure_openai", @@ -243,6 +244,7 @@ "group": "Supported Embedding Models", "icon": "list", "pages": [ + "components/embedders/models/aimlapi", "components/embedders/models/openai", "components/embedders/models/azure_openai", "components/embedders/models/ollama", diff --git a/docs/images/provider-icons/aimlapi.svg b/docs/images/provider-icons/aimlapi.svg new file mode 100644 index 0000000000..b0129c980d --- /dev/null +++ b/docs/images/provider-icons/aimlapi.svg @@ -0,0 +1,10 @@ + + + + + diff --git a/docs/llms.txt b/docs/llms.txt index 1ea94c4d87..c97de4e304 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -450,6 +450,7 @@ Everything below is OSS-only provider configuration. Skip this entire section wh ### LLM Providers [OSS] - [LLM Overview](https://docs.mem0.ai/components/llms/overview) [OSS]: Use when the user is choosing an LLM for memory extraction. - [LLM Configuration](https://docs.mem0.ai/components/llms/config) [OSS]: Use for the `llm` config schema. +- [aimlapi.com](https://docs.mem0.ai/components/llms/models/aimlapi) [OSS]: Use when the LLM is reached through the aimlapi.com gateway. - [OpenAI](https://docs.mem0.ai/components/llms/models/openai) [OSS]: Use when the extraction LLM is OpenAI. - [Anthropic](https://docs.mem0.ai/components/llms/models/anthropic) [OSS]: Use when the extraction LLM is Claude. - [Azure OpenAI](https://docs.mem0.ai/components/llms/models/azure_openai) [OSS]: Use when the user is on Azure-hosted OpenAI. @@ -471,6 +472,7 @@ Everything below is OSS-only provider configuration. Skip this entire section wh ### Embedding Providers [OSS] - [Embeddings Overview](https://docs.mem0.ai/components/embedders/overview) [OSS]: Use when choosing an embedding model. - [Embeddings Configuration](https://docs.mem0.ai/components/embedders/config) [OSS]: Use for the `embedder` config schema. +- [aimlapi.com Embeddings](https://docs.mem0.ai/components/embedders/models/aimlapi) [OSS]: Use when embeddings are reached through the aimlapi.com gateway. - [OpenAI Embeddings](https://docs.mem0.ai/components/embedders/models/openai) [OSS]: Use when embeddings come from OpenAI. - [Azure OpenAI Embeddings](https://docs.mem0.ai/components/embedders/models/azure_openai) [OSS]: Use for Azure-hosted OpenAI embeddings. - [AWS Bedrock Embeddings](https://docs.mem0.ai/components/embedders/models/aws_bedrock) [OSS]: Use for Bedrock-hosted embeddings. diff --git a/mem0-ts/src/oss/src/embeddings/aimlapi.ts b/mem0-ts/src/oss/src/embeddings/aimlapi.ts new file mode 100644 index 0000000000..e50ce9883e --- /dev/null +++ b/mem0-ts/src/oss/src/embeddings/aimlapi.ts @@ -0,0 +1,34 @@ +import { OpenAIEmbedder } from "./openai"; +import { EmbeddingConfig } from "../types"; +import { buildAimlapiHeaders, resolveAimlapiBaseURL } from "../utils/aimlapi"; + +const DEFAULT_MODEL = "openai/text-embedding-3-small"; + +/** + * aimlapi.com embedder — the OpenAI-compatible `/v1/embeddings` route. + * + * Mirrors `mem0/embeddings/aimlapi.py` in the Python SDK. Model ids are the + * gateway's namespaced ids, e.g. `openai/text-embedding-3-small` (1536 dims) or + * `openai/text-embedding-3-large` (3072 dims). + */ +export class AimlapiEmbedder extends OpenAIEmbedder { + constructor(config: EmbeddingConfig) { + const apiKey = config.apiKey || process.env.AIMLAPI_API_KEY; + if (!apiKey) { + throw new Error( + "aimlapi.com API key is required. Set AIMLAPI_API_KEY or pass apiKey in the config.", + ); + } + const baseURL = resolveAimlapiBaseURL( + config.baseURL || config.url, + process.env.AIMLAPI_API_BASE, + ); + super({ + ...config, + apiKey, + baseURL, + model: config.model || DEFAULT_MODEL, + defaultHeaders: buildAimlapiHeaders(baseURL, config.defaultHeaders), + }); + } +} diff --git a/mem0-ts/src/oss/src/embeddings/openai.ts b/mem0-ts/src/oss/src/embeddings/openai.ts index 4fe90beeef..f62618d7c0 100644 --- a/mem0-ts/src/oss/src/embeddings/openai.ts +++ b/mem0-ts/src/oss/src/embeddings/openai.ts @@ -11,6 +11,7 @@ export class OpenAIEmbedder implements Embedder { this.openai = new OpenAI({ apiKey: config.apiKey, baseURL: config.baseURL || config.url, + ...(config.defaultHeaders && { defaultHeaders: config.defaultHeaders }), }); this.model = config.model || "text-embedding-3-small"; this.embeddingDims = config.embeddingDims; diff --git a/mem0-ts/src/oss/src/llms/aimlapi.ts b/mem0-ts/src/oss/src/llms/aimlapi.ts new file mode 100644 index 0000000000..2b32d52b36 --- /dev/null +++ b/mem0-ts/src/oss/src/llms/aimlapi.ts @@ -0,0 +1,60 @@ +import { OpenAILLM } from "./openai"; +import { LLMConfig, Message } from "../types"; +import { LLMResponse } from "./base"; +import { buildAimlapiHeaders, resolveAimlapiBaseURL } from "../utils/aimlapi"; + +/** + * aimlapi.com LLM provider — an OpenAI-compatible gateway over 350+ chat models. + * + * The API is OpenAI-compatible, so this reuses {@link OpenAILLM} and overrides the + * connection defaults — mirroring `mem0/llms/aimlapi.py` in the Python SDK. The API + * key resolves from `config.apiKey` or `AIMLAPI_API_KEY`, and the base URL from + * `config.baseURL`, `AIMLAPI_API_BASE`, else `https://api.aimlapi.com/v1`. + * + * Model ids are the gateway's namespaced ids, e.g. `openai/gpt-5-mini`, + * `anthropic/claude-sonnet-4.6`. The catalog is at + * `https://api.aimlapi.com/v1/models`. + */ +export class AimlapiLLM extends OpenAILLM { + constructor(config: LLMConfig) { + const apiKey = config.apiKey || process.env.AIMLAPI_API_KEY; + if (!apiKey) { + throw new Error( + "aimlapi.com API key is required. Set AIMLAPI_API_KEY or pass apiKey in the config.", + ); + } + const baseURL = resolveAimlapiBaseURL( + config.baseURL || config.url, + process.env.AIMLAPI_API_BASE, + ); + super({ + ...config, + apiKey, + baseURL, + model: config.model || "openai/gpt-5-mini", + defaultHeaders: buildAimlapiHeaders(baseURL, config.defaultHeaders), + }); + } + + 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(`aimlapi.com LLM 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(`aimlapi.com LLM failed: ${message}`); + } + } +} diff --git a/mem0-ts/src/oss/src/llms/openai.ts b/mem0-ts/src/oss/src/llms/openai.ts index 41644ca351..8402276041 100644 --- a/mem0-ts/src/oss/src/llms/openai.ts +++ b/mem0-ts/src/oss/src/llms/openai.ts @@ -10,6 +10,7 @@ export class OpenAILLM implements LLM { this.openai = new OpenAI({ apiKey: config.apiKey, baseURL: config.baseURL, + ...(config.defaultHeaders && { defaultHeaders: config.defaultHeaders }), ...(config.timeout != null && { timeout: config.timeout }), }); this.model = config.model || "gpt-5-mini"; diff --git a/mem0-ts/src/oss/src/types/index.ts b/mem0-ts/src/oss/src/types/index.ts index fb3e5b391e..bd243744da 100644 --- a/mem0-ts/src/oss/src/types/index.ts +++ b/mem0-ts/src/oss/src/types/index.ts @@ -17,6 +17,8 @@ export interface EmbeddingConfig { model?: string | any; baseURL?: string; url?: string; + // Extra headers sent on every request by OpenAI-compatible embedders. + defaultHeaders?: Record; embeddingDims?: number; modelProperties?: Record; // HuggingFace TEI / OpenAI-compatible inference endpoint base URL. @@ -61,6 +63,8 @@ export interface HistoryStoreConfig { export interface LLMConfig { provider?: string; baseURL?: string; + // Extra headers sent on every request by OpenAI-compatible providers. + defaultHeaders?: Record; vllmBaseURL?: string; vllm_base_url?: string; url?: string; diff --git a/mem0-ts/src/oss/src/utils/aimlapi.ts b/mem0-ts/src/oss/src/utils/aimlapi.ts new file mode 100644 index 0000000000..632f86d1cd --- /dev/null +++ b/mem0-ts/src/oss/src/utils/aimlapi.ts @@ -0,0 +1,58 @@ +/** + * Shared connection details for the aimlapi.com LLM and embedding providers. + * + * Mirrors `mem0/utils/aimlapi.py` in the Python SDK: both providers speak the + * OpenAI-compatible surface of https://api.aimlapi.com/v1, so base-URL resolution + * and the attribution headers live here rather than being duplicated. + */ + +export const AIMLAPI_DEFAULT_BASE_URL = "https://api.aimlapi.com/v1"; + +/** + * Host the attribution headers may be sent to. Users can point the providers at a + * gateway of their own, and in that case the headers must not ride along to a + * third party. + */ +export const AIMLAPI_ATTRIBUTION_HOST = "api.aimlapi.com"; + +/** + * Identifies Mem0 as the calling application. HTTP-Referer / X-Title are the + * OpenRouter convention and name the *host* project. + */ +const ATTRIBUTION_HEADERS: Record = { + "HTTP-Referer": "https://github.com/mem0ai/mem0", + "X-Title": "Mem0", + "X-AIMLAPI-Partner-ID": "part_JNAROikm3sdRqpewzcZxLgrK", + "X-AIMLAPI-Source": "agent/mem0", +}; + +export function resolveAimlapiBaseURL( + configured?: string, + envValue?: string, +): string { + return configured || envValue || AIMLAPI_DEFAULT_BASE_URL; +} + +/** + * Build the per-client `defaultHeaders` for a request to `baseURL`. + * + * Returns a new object every call, so the module constant is never mutated and two + * clients cannot share one header map. Caller-supplied headers win on a key clash. + * Returns undefined when `baseURL` does not point at aimlapi.com, so attribution + * never travels to another provider or to a proxy fronting this API. + */ +export function buildAimlapiHeaders( + baseURL: string, + extra?: Record, +): Record | undefined { + let host: string; + try { + host = new URL(baseURL).hostname; + } catch { + return extra ? { ...extra } : undefined; + } + if (host !== AIMLAPI_ATTRIBUTION_HOST) { + return extra ? { ...extra } : undefined; + } + return { ...ATTRIBUTION_HEADERS, ...(extra || {}) }; +} diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 16de0301f3..47f8c7fdcd 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -35,6 +35,8 @@ import { LMStudioLLM } from "../llms/lmstudio"; import { DeepSeekLLM } from "../llms/deepseek"; import { XAILLM } from "../llms/xai"; import { SarvamLLM } from "../llms/sarvam"; +import { AimlapiLLM } from "../llms/aimlapi"; +import { AimlapiEmbedder } from "../embeddings/aimlapi"; import { AWSBedrockLLM } from "../llms/aws_bedrock"; import { LiteLLM } from "../llms/litellm"; import { MiniMaxLLM } from "../llms/minimax"; @@ -99,6 +101,8 @@ export class EmbedderFactory { return new VertexAIEmbedder(config); case "huggingface": return new HuggingFaceEmbedder(config); + case "aimlapi": + return new AimlapiEmbedder(config); default: throw new Error(`Unsupported embedder provider: ${provider}`); } @@ -145,6 +149,8 @@ export class LLMFactory { return new TogetherLLM(config); case "vllm": return new VllmLLM(config); + case "aimlapi": + return new AimlapiLLM(config); default: throw new Error(`Unsupported LLM provider: ${provider}`); } diff --git a/mem0-ts/src/oss/tests/aimlapi.test.ts b/mem0-ts/src/oss/tests/aimlapi.test.ts new file mode 100644 index 0000000000..8dff43e7bf --- /dev/null +++ b/mem0-ts/src/oss/tests/aimlapi.test.ts @@ -0,0 +1,239 @@ +/// +/** + * aimlapi.com LLM + embedder — unit tests (mocked OpenAI). + */ + +import { AimlapiLLM } from "../src/llms/aimlapi"; +import { AimlapiEmbedder } from "../src/embeddings/aimlapi"; +import { buildAimlapiHeaders } from "../src/utils/aimlapi"; + +const mockCreate = jest.fn(); +const mockEmbeddingsCreate = jest.fn(); +const mockOpenAICtor = jest.fn(); + +jest.mock("openai", () => { + return jest.fn().mockImplementation((config) => { + mockOpenAICtor(config); + return { + chat: { completions: { create: mockCreate } }, + embeddings: { create: mockEmbeddingsCreate }, + }; + }); +}); + +describe("aimlapi.com providers (unit)", () => { + const ORIGINAL_ENV = process.env; + + beforeEach(() => { + jest.clearAllMocks(); + process.env = { ...ORIGINAL_ENV }; + delete process.env.AIMLAPI_API_KEY; + delete process.env.AIMLAPI_API_BASE; + mockCreate.mockResolvedValue({ + choices: [{ message: { content: "hi", role: "assistant" } }], + }); + mockEmbeddingsCreate.mockResolvedValue({ + data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }], + }); + }); + + afterAll(() => { + process.env = ORIGINAL_ENV; + }); + + describe("attribution headers", () => { + it("uses a partner id the gateway will accept", () => { + // A malformed partner id is dropped silently and earns nothing, so the + // shape is asserted here rather than discovered in production. + const headers = buildAimlapiHeaders("https://api.aimlapi.com/v1")!; + expect(headers["X-AIMLAPI-Partner-ID"]).toMatch( + /^part_[A-Za-z0-9]{1,64}$/, + ); + expect(headers["X-AIMLAPI-Source"]).toMatch( + /^(web|agent|mcp)\/[a-z0-9-]{1,32}$/, + ); + }); + + it("sends attribution to aimlapi.com and nowhere else", () => { + new AimlapiLLM({ apiKey: "test-key" }); + expect(mockOpenAICtor).toHaveBeenCalledWith( + expect.objectContaining({ + defaultHeaders: expect.objectContaining({ + "X-AIMLAPI-Source": "agent/mem0", + "HTTP-Referer": "https://github.com/mem0ai/mem0", + "X-Title": "Mem0", + }), + }), + ); + + // A user pointing the provider at their own gateway must not have our + // attribution ride along to a third party. + new AimlapiLLM({ + apiKey: "test-key", + baseURL: "https://proxy.example.com/v1", + }); + expect(mockOpenAICtor).toHaveBeenLastCalledWith( + expect.not.objectContaining({ defaultHeaders: expect.anything() }), + ); + }); + + it("does not mutate the shared header constant", () => { + const first = buildAimlapiHeaders("https://api.aimlapi.com/v1", { + "X-Title": "Other", + })!; + const second = buildAimlapiHeaders("https://api.aimlapi.com/v1")!; + expect(first["X-Title"]).toBe("Other"); + expect(second["X-Title"]).toBe("Mem0"); + expect(first).not.toBe(second); + }); + }); + + describe("AimlapiLLM", () => { + it("defaults to openai/gpt-5-mini and the aimlapi.com base URL (matching the Python provider)", async () => { + const llm = new AimlapiLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse([ + { role: "user", content: "hello" }, + ]); + + expect(mockOpenAICtor).toHaveBeenCalledWith( + expect.objectContaining({ + apiKey: "test-key", + baseURL: "https://api.aimlapi.com/v1", + }), + ); + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: "openai/gpt-5-mini" }), + ); + expect(result).toBe("hi"); + }); + + it("resolves AIMLAPI_API_KEY / AIMLAPI_API_BASE from the environment", () => { + process.env.AIMLAPI_API_KEY = "env-key"; + process.env.AIMLAPI_API_BASE = "https://gateway.example.com/v1"; + + new AimlapiLLM({}); + + expect(mockOpenAICtor).toHaveBeenCalledWith( + expect.objectContaining({ + apiKey: "env-key", + baseURL: "https://gateway.example.com/v1", + }), + ); + }); + + it("prefers explicit config over defaults and the environment", async () => { + process.env.AIMLAPI_API_KEY = "env-key"; + + const llm = new AimlapiLLM({ + apiKey: "explicit-key", + model: "anthropic/claude-sonnet-4.6", + }); + await llm.generateResponse([{ role: "user", content: "hello" }]); + + expect(mockOpenAICtor).toHaveBeenCalledWith( + expect.objectContaining({ apiKey: "explicit-key" }), + ); + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: "anthropic/claude-sonnet-4.6" }), + ); + }); + + it("throws when no API key is provided", () => { + expect(() => new AimlapiLLM({})).toThrow("API key is required"); + }); + + it("generateResponse() handles tool calls", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [ + { + message: { + content: "", + role: "assistant", + tool_calls: [ + { + function: { + name: "add_memory", + arguments: '{"data": "likes pizza"}', + }, + }, + ], + }, + }, + ], + }); + + const llm = new AimlapiLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "remember this" }], + undefined, + [{ type: "function", function: { name: "add_memory" } }], + ); + + expect(result).toEqual({ + content: "", + role: "assistant", + toolCalls: [ + { name: "add_memory", arguments: '{"data": "likes pizza"}' }, + ], + }); + }); + + it("wraps downstream errors with a provider-specific message", async () => { + mockCreate.mockRejectedValueOnce(new Error("Connection refused")); + const llm = new AimlapiLLM({ apiKey: "test-key" }); + + await expect( + llm.generateResponse([{ role: "user", content: "hi" }]), + ).rejects.toThrow("aimlapi.com LLM failed: Connection refused"); + }); + + it("generateChat() returns the LLMResponse shape", async () => { + const llm = new AimlapiLLM({ apiKey: "test-key" }); + const result = await llm.generateChat([ + { role: "user", content: "help me" }, + ]); + expect(result).toEqual({ content: "hi", role: "assistant" }); + }); + }); + + describe("AimlapiEmbedder", () => { + it("defaults to openai/text-embedding-3-small and sends string input", async () => { + const embedder = new AimlapiEmbedder({ apiKey: "test-key" }); + const result = await embedder.embed("hello world"); + + expect(mockOpenAICtor).toHaveBeenCalledWith( + expect.objectContaining({ baseURL: "https://api.aimlapi.com/v1" }), + ); + // The API rejects the pre-tokenised integer-array form of `input` with a + // 400, so the text must reach it as a string. + expect(mockEmbeddingsCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "openai/text-embedding-3-small", + input: "hello world", + encoding_format: "float", + }), + ); + expect(result).toEqual([0.1, 0.2, 0.3]); + }); + + it("throws when no API key is provided", () => { + expect(() => new AimlapiEmbedder({})).toThrow("API key is required"); + }); + + it("passes dimensions only when embeddingDims is configured", async () => { + const embedder = new AimlapiEmbedder({ + apiKey: "test-key", + model: "openai/text-embedding-3-large", + embeddingDims: 1024, + }); + await embedder.embed("hi"); + + expect(mockEmbeddingsCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "openai/text-embedding-3-large", + dimensions: 1024, + }), + ); + }); + }); +}); diff --git a/mem0-ts/src/oss/tests/factory.unit.test.ts b/mem0-ts/src/oss/tests/factory.unit.test.ts index 706c178b3b..c3e0656307 100644 --- a/mem0-ts/src/oss/tests/factory.unit.test.ts +++ b/mem0-ts/src/oss/tests/factory.unit.test.ts @@ -117,6 +117,16 @@ jest.mock("../src/llms/sarvam", () => ({ .fn() .mockImplementation((config) => ({ type: "sarvam-llm", config })), })); +jest.mock("../src/llms/aimlapi", () => ({ + AimlapiLLM: jest + .fn() + .mockImplementation((config) => ({ type: "aimlapi-llm", config })), +})); +jest.mock("../src/embeddings/aimlapi", () => ({ + AimlapiEmbedder: jest + .fn() + .mockImplementation((config) => ({ type: "aimlapi-embedder", config })), +})); jest.mock("../src/llms/litellm", () => ({ LiteLLM: jest .fn() @@ -259,6 +269,7 @@ describe("EmbedderFactory", () => { ["lmstudio"], ["vertexai"], ["together"], + ["aimlapi"], ])("creates embedder for provider '%s'", (provider) => { expect(() => EmbedderFactory.create(provider, dummyEmbedConfig), @@ -306,6 +317,7 @@ describe("LLMFactory", () => { ["minimax"], ["together"], ["vllm"], + ["aimlapi"], ])("creates LLM for provider '%s'", (provider) => { expect(() => LLMFactory.create(provider, dummyLLMConfig)).not.toThrow(); }); diff --git a/mem0/configs/embeddings/base.py b/mem0/configs/embeddings/base.py index de2f4324f4..708f0be328 100644 --- a/mem0/configs/embeddings/base.py +++ b/mem0/configs/embeddings/base.py @@ -20,6 +20,8 @@ def __init__( ollama_base_url: Optional[str] = None, # Openai specific openai_base_url: Optional[str] = None, + # aimlapi.com specific + aimlapi_base_url: Optional[str] = None, # Huggingface specific model_kwargs: Optional[dict] = None, huggingface_base_url: Optional[str] = None, @@ -58,6 +60,8 @@ def __init__( :type huggingface_base_url: Optional[str], optional :param openai_base_url: Openai base URL to be use, defaults to "https://api.openai.com/v1" :type openai_base_url: Optional[str], optional + :param aimlapi_base_url: aimlapi.com base URL to be use, defaults to "https://api.aimlapi.com/v1" + :type aimlapi_base_url: Optional[str], optional :param azure_kwargs: key-value arguments for the AzureOpenAI embedding model, defaults a dict inside init :type azure_kwargs: Optional[Dict[str, Any]], defaults a dict inside init :param http_client_proxies: The proxy server settings used to create self.http_client, defaults to None @@ -79,6 +83,9 @@ def __init__( self.openai_base_url = openai_base_url self.embedding_dims = embedding_dims + # aimlapi.com specific + self.aimlapi_base_url = aimlapi_base_url + # AzureOpenAI specific self.http_client_proxies = http_client_proxies self.http_client = build_http_client(http_client_proxies) diff --git a/mem0/configs/llms/aimlapi.py b/mem0/configs/llms/aimlapi.py new file mode 100644 index 0000000000..32c775d205 --- /dev/null +++ b/mem0/configs/llms/aimlapi.py @@ -0,0 +1,61 @@ +from typing import Optional + +from mem0.configs.llms.base import BaseLlmConfig + + +class AimlapiConfig(BaseLlmConfig): + """ + Configuration class for aimlapi.com-specific parameters. + Inherits from BaseLlmConfig and adds aimlapi.com-specific settings. + """ + + def __init__( + self, + # Base parameters + model: Optional[str] = None, + temperature: float = 0.1, + api_key: Optional[str] = None, + max_tokens: int = 2000, + top_p: float = 0.1, + top_k: int = 1, + enable_vision: bool = False, + vision_details: Optional[str] = "auto", + reasoning_effort: Optional[str] = None, + http_client_proxies: Optional[dict] = None, + is_reasoning_model: Optional[bool] = None, + # aimlapi.com-specific parameters + aimlapi_base_url: Optional[str] = None, + ): + """ + Initialize aimlapi.com configuration. + + Args: + model: aimlapi.com model id to use (e.g. "openai/gpt-5-mini"), defaults to None + temperature: Controls randomness, defaults to 0.1 + api_key: aimlapi.com API key, defaults to None + max_tokens: Maximum tokens to generate, defaults to 2000 + top_p: Nucleus sampling parameter, defaults to 0.1 + top_k: Top-k sampling parameter, defaults to 1 + enable_vision: Enable vision capabilities, defaults to False + vision_details: Vision detail level, defaults to "auto" + reasoning_effort: Effort level for reasoning models, defaults to None + http_client_proxies: HTTP client proxy settings, defaults to None + is_reasoning_model: Explicit reasoning-model override, defaults to None + aimlapi_base_url: aimlapi.com API base URL, defaults to None + """ + super().__init__( + model=model, + temperature=temperature, + api_key=api_key, + max_tokens=max_tokens, + top_p=top_p, + top_k=top_k, + enable_vision=enable_vision, + vision_details=vision_details, + reasoning_effort=reasoning_effort, + http_client_proxies=http_client_proxies, + is_reasoning_model=is_reasoning_model, + ) + + # aimlapi.com-specific parameters + self.aimlapi_base_url = aimlapi_base_url diff --git a/mem0/embeddings/aimlapi.py b/mem0/embeddings/aimlapi.py new file mode 100644 index 0000000000..2ecd03f13d --- /dev/null +++ b/mem0/embeddings/aimlapi.py @@ -0,0 +1,87 @@ +import os +from typing import Literal, Optional + +from openai import OpenAI + +from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.embeddings.base import EmbeddingBase +from mem0.utils.aimlapi import build_default_headers, resolve_base_url + + +class AimlapiEmbedding(EmbeddingBase): + """aimlapi.com embeddings — the OpenAI-compatible `/v1/embeddings` route. + + Model ids are the gateway's namespaced ids, e.g. `openai/text-embedding-3-small` + (1536 dims) or `openai/text-embedding-3-large` (3072 dims). + """ + + def __init__(self, config: Optional[BaseEmbedderConfig] = None): + super().__init__(config) + + self.config.model = self.config.model or "openai/text-embedding-3-small" + # Only pass `dimensions` to the API when the user set embedding_dims; not every + # model behind the gateway is matryoshka, and those reject the parameter. + self._pass_dimensions_to_api = self.config.embedding_dims is not None + self.config.embedding_dims = self.config.embedding_dims or 1536 + + api_key = self.config.api_key or os.getenv("AIMLAPI_API_KEY") + if not api_key: + raise ValueError( + "aimlapi.com API key is required. Set the AIMLAPI_API_KEY environment variable " + "or pass api_key in the config." + ) + + base_url = resolve_base_url(self.config.aimlapi_base_url, os.getenv("AIMLAPI_API_BASE")) + self.client = OpenAI( + api_key=api_key, + base_url=base_url, + default_headers=build_default_headers(base_url), + ) + + def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None): + """ + Get the embedding for the given text using aimlapi.com. + + Args: + text (str): The text to embed. + memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None. + Returns: + list: The embedding vector. + """ + text = text.replace("\n", " ") + # `input` must be a string or a list of strings. The API rejects the + # pre-tokenised integer-array form with a 400 naming `input`. + kwargs = { + "input": [text], + "model": self.config.model, + "encoding_format": "float", + } + if self._pass_dimensions_to_api: + kwargs["dimensions"] = self.config.embedding_dims + return self.client.embeddings.create(**kwargs).data[0].embedding + + def embed_batch(self, texts, memory_action="add"): + """Embed multiple texts in a single aimlapi.com API call. + + Automatically chunks into batches of 100 to stay within API limits. + """ + MAX_BATCH = 100 + texts = [text.replace("\n", " ") for text in texts] + all_embeddings = [] + for i in range(0, len(texts), MAX_BATCH): + chunk = texts[i : i + MAX_BATCH] + kwargs = { + "input": chunk, + "model": self.config.model, + "encoding_format": "float", + } + if self._pass_dimensions_to_api: + kwargs["dimensions"] = self.config.embedding_dims + response = self.client.embeddings.create(**kwargs) + all_embeddings.extend(item.embedding for item in sorted(response.data, key=lambda x: x.index)) + if len(all_embeddings) != len(texts): + raise ValueError( + f"aimlapi.com embed_batch() returned {len(all_embeddings)} embeddings for {len(texts)} texts" + f" using model '{self.config.model}'" + ) + return all_embeddings diff --git a/mem0/embeddings/configs.py b/mem0/embeddings/configs.py index b924263eb1..246e01be02 100644 --- a/mem0/embeddings/configs.py +++ b/mem0/embeddings/configs.py @@ -25,6 +25,7 @@ def validate_config(cls, v, values): "langchain", "aws_bedrock", "fastembed", + "aimlapi", ]: return v else: diff --git a/mem0/llms/aimlapi.py b/mem0/llms/aimlapi.py new file mode 100644 index 0000000000..d08b0e9848 --- /dev/null +++ b/mem0/llms/aimlapi.py @@ -0,0 +1,133 @@ +import json +import os +from typing import Dict, List, Optional, Union + +from openai import OpenAI + +from mem0.configs.llms.aimlapi import AimlapiConfig +from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase +from mem0.memory.utils import extract_json +from mem0.utils.aimlapi import build_default_headers, drop_none, resolve_base_url + + +class AimlapiLLM(LLMBase): + """aimlapi.com — an OpenAI-compatible gateway over 350+ chat models. + + Model ids are the gateway's own namespaced ids, e.g. `openai/gpt-5-mini`, + `anthropic/claude-sonnet-4.6`, `google/gemini-2.5-flash`. The full catalog is at + `https://api.aimlapi.com/v1/models`. + """ + + def __init__(self, config: Optional[Union[BaseLlmConfig, AimlapiConfig, Dict]] = None): + # Convert to AimlapiConfig if needed + if config is None: + config = AimlapiConfig() + elif isinstance(config, dict): + config = AimlapiConfig(**config) + elif isinstance(config, BaseLlmConfig) and not isinstance(config, AimlapiConfig): + # Convert BaseLlmConfig to AimlapiConfig so aimlapi_base_url is available + config = AimlapiConfig( + model=config.model, + temperature=config.temperature, + api_key=config.api_key, + max_tokens=config.max_tokens, + top_p=config.top_p, + top_k=config.top_k, + enable_vision=config.enable_vision, + vision_details=config.vision_details, + reasoning_effort=getattr(config, "reasoning_effort", None), + http_client_proxies=config.http_client_proxies, + is_reasoning_model=getattr(config, "is_reasoning_model", None), + ) + + super().__init__(config) + + if not self.config.model: + self.config.model = "openai/gpt-5-mini" + + api_key = self.config.api_key or os.getenv("AIMLAPI_API_KEY") + if not api_key: + raise ValueError( + "aimlapi.com API key is required. Set the AIMLAPI_API_KEY environment variable " + "or pass api_key in the config." + ) + + base_url = resolve_base_url(self.config.aimlapi_base_url, os.getenv("AIMLAPI_API_BASE")) + self.client = OpenAI( + api_key=api_key, + base_url=base_url, + default_headers=build_default_headers(base_url), + ) + + def _parse_response(self, response, tools): + """ + Process the response based on whether tools are used or not. + + Args: + response: The raw response from API. + tools: The list of tools provided in the request. + + Returns: + str or dict: The processed response. + """ + if tools: + processed_response = { + "content": response.choices[0].message.content, + "tool_calls": [], + } + + if response.choices[0].message.tool_calls: + for tool_call in response.choices[0].message.tool_calls: + processed_response["tool_calls"].append( + { + "name": tool_call.function.name, + "arguments": json.loads(extract_json(tool_call.function.arguments)), + } + ) + + return processed_response + else: + return response.choices[0].message.content + + def generate_response( + self, + messages: List[Dict[str, str]], + response_format=None, + tools: Optional[List[Dict]] = None, + tool_choice: str = "auto", + **kwargs, + ): + """ + Generate a response based on the given messages using aimlapi.com. + + Args: + messages (list): List of message dicts containing 'role' and 'content'. + response_format (str or object, optional): Format of the response. Defaults to None. + tools (list, optional): List of tools that the model can call. Defaults to None. + tool_choice (str, optional): Tool choice method. Defaults to "auto". + **kwargs: Additional provider-specific parameters. + + Returns: + str or dict: The generated response. A string when tools are not requested; + a dict ``{"content": ..., "tool_calls": [...]}`` when tools are requested. + """ + params = self._get_supported_params(messages=messages, **kwargs) + params.update( + { + "model": self.config.model, + "messages": messages, + } + ) + + if response_format: + params["response_format"] = response_format + if tools: + params["tools"] = tools + params["tool_choice"] = tool_choice + + # The API answers 400 to an explicit `null` on several fields rather than + # treating it as unset — `tools: null` is the common one, and it turns the + # second turn of every tool-clearing agent loop into an error. Omit unset keys. + response = self.client.chat.completions.create(**drop_none(params)) + return self._parse_response(response, tools) diff --git a/mem0/llms/configs.py b/mem0/llms/configs.py index 11d5880da8..214324640f 100644 --- a/mem0/llms/configs.py +++ b/mem0/llms/configs.py @@ -29,6 +29,7 @@ def validate_config(cls, v, values): "lmstudio", "vllm", "langchain", + "aimlapi", ): return v else: diff --git a/mem0/utils/aimlapi.py b/mem0/utils/aimlapi.py new file mode 100644 index 0000000000..b6da6260e2 --- /dev/null +++ b/mem0/utils/aimlapi.py @@ -0,0 +1,55 @@ +"""Shared connection details for the aimlapi.com LLM and embedding providers. + +Both providers speak the OpenAI-compatible surface of https://api.aimlapi.com/v1, +so the base-URL resolution and the attribution headers live here rather than being +duplicated in `mem0/llms/aimlapi.py` and `mem0/embeddings/aimlapi.py`. +""" + +from typing import Dict, Optional +from urllib.parse import urlsplit + +DEFAULT_BASE_URL = "https://api.aimlapi.com/v1" + +# Host that attribution headers may be sent to. Users can point the providers at a +# gateway of their own via `aimlapi_base_url` / AIMLAPI_API_BASE, and in that case the +# headers must not ride along to a third party. +ATTRIBUTION_HOST = "api.aimlapi.com" + +# Identifies Mem0 as the calling application. HTTP-Referer / X-Title are the +# OpenRouter convention and name the *host* project, matching how the OpenAI provider +# already labels OpenRouter traffic in `mem0/llms/openai.py`. +_ATTRIBUTION_HEADERS = { + "HTTP-Referer": "https://github.com/mem0ai/mem0", + "X-Title": "Mem0", + "X-AIMLAPI-Partner-ID": "part_JNAROikm3sdRqpewzcZxLgrK", + "X-AIMLAPI-Source": "agent/mem0", +} + + +def resolve_base_url(configured: Optional[str], env_value: Optional[str]) -> str: + """Resolve the base URL from config, then environment, then the public default.""" + return configured or env_value or DEFAULT_BASE_URL + + +def build_default_headers(base_url: str, extra: Optional[Dict[str, str]] = None) -> Optional[Dict[str, str]]: + """Build the per-client `default_headers` for a request to `base_url`. + + Returns a new dict every call, so the module-level constant is never mutated and + two clients cannot share (and corrupt) one header map. Caller-supplied headers win + on a key clash. Returns None when `base_url` does not point at aimlapi.com, so + attribution never travels to another provider or to a proxy fronting this API. + """ + if urlsplit(base_url).hostname != ATTRIBUTION_HOST: + return dict(extra) if extra else None + return {**_ATTRIBUTION_HEADERS, **(extra or {})} + + +def drop_none(params: Dict) -> Dict: + """Return a copy of `params` without keys whose value is None. + + The API rejects an explicit `null` on several fields — `tools` and `temperature` + among them — with a 400 rather than treating it as "unset". The OpenAI SDK + serialises a None-valued keyword straight through, so unset options must be + omitted from the payload instead of being passed as None. + """ + return {k: v for k, v in params.items() if v is not None} diff --git a/mem0/utils/factory.py b/mem0/utils/factory.py index 30a1079bb2..4046421ca0 100644 --- a/mem0/utils/factory.py +++ b/mem0/utils/factory.py @@ -3,6 +3,7 @@ from typing import Dict, Optional, Union from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.configs.llms.aimlapi import AimlapiConfig from mem0.configs.llms.anthropic import AnthropicConfig from mem0.configs.llms.aws_bedrock import AWSBedrockConfig from mem0.configs.llms.azure import AzureOpenAIConfig @@ -58,6 +59,7 @@ class LlmFactory: "lmstudio": ("mem0.llms.lmstudio.LMStudioLLM", LMStudioConfig), "vllm": ("mem0.llms.vllm.VllmLLM", VllmConfig), "langchain": ("mem0.llms.langchain.LangchainLLM", BaseLlmConfig), + "aimlapi": ("mem0.llms.aimlapi.AimlapiLLM", AimlapiConfig), } @classmethod @@ -162,6 +164,7 @@ class EmbedderFactory: "langchain": "mem0.embeddings.langchain.LangchainEmbedding", "aws_bedrock": "mem0.embeddings.aws_bedrock.AWSBedrockEmbedding", "fastembed": "mem0.embeddings.fastembed.FastEmbedEmbedding", + "aimlapi": "mem0.embeddings.aimlapi.AimlapiEmbedding", } @classmethod diff --git a/tests/embeddings/test_aimlapi_embeddings.py b/tests/embeddings/test_aimlapi_embeddings.py new file mode 100644 index 0000000000..7c697b319b --- /dev/null +++ b/tests/embeddings/test_aimlapi_embeddings.py @@ -0,0 +1,116 @@ +import os +from unittest.mock import Mock, patch + +import pytest + +from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.embeddings.aimlapi import AimlapiEmbedding +from mem0.utils.factory import EmbedderFactory + + +@pytest.fixture(autouse=True) +def clear_aimlapi_env(): + saved = {k: os.environ.pop(k, None) for k in ("AIMLAPI_API_KEY", "AIMLAPI_API_BASE")} + yield + for k, v in saved.items(): + if v is not None: + os.environ[k] = v + else: + os.environ.pop(k, None) + + +@pytest.fixture +def mock_aimlapi_client(): + with patch("mem0.embeddings.aimlapi.OpenAI") as mock_openai: + mock_client = Mock() + mock_openai.return_value = mock_client + yield mock_openai, mock_client + + +def test_default_model_and_base_url(mock_aimlapi_client): + mock_openai, _ = mock_aimlapi_client + embedder = AimlapiEmbedding(BaseEmbedderConfig(api_key="k")) + assert embedder.config.model == "openai/text-embedding-3-small" + assert embedder.config.embedding_dims == 1536 + assert mock_openai.call_args.kwargs["base_url"] == "https://api.aimlapi.com/v1" + + +def test_base_url_override(mock_aimlapi_client): + mock_openai, _ = mock_aimlapi_client + os.environ["AIMLAPI_API_BASE"] = "https://gateway.example.com/v1" + AimlapiEmbedding(BaseEmbedderConfig(api_key="k")) + assert mock_openai.call_args.kwargs["base_url"] == "https://gateway.example.com/v1" + + AimlapiEmbedding(BaseEmbedderConfig(api_key="k", aimlapi_base_url="https://config.example.com/v1")) + assert mock_openai.call_args.kwargs["base_url"] == "https://config.example.com/v1" + + +def test_requires_api_key(): + with pytest.raises(ValueError, match="API key is required"): + AimlapiEmbedding(BaseEmbedderConfig()) + + +def test_attribution_headers(mock_aimlapi_client): + mock_openai, _ = mock_aimlapi_client + AimlapiEmbedding(BaseEmbedderConfig(api_key="k")) + headers = mock_openai.call_args.kwargs["default_headers"] + assert headers["X-AIMLAPI-Partner-ID"].startswith("part_") + assert headers["X-AIMLAPI-Source"] == "agent/mem0" + + # Not on a request to somebody else's gateway. + AimlapiEmbedding(BaseEmbedderConfig(api_key="k", aimlapi_base_url="https://proxy.example.com/v1")) + assert mock_openai.call_args.kwargs["default_headers"] is None + + +def test_embed_sends_string_input(mock_aimlapi_client): + """The API rejects the pre-tokenised integer-array form of ``input`` with a 400, + so the text must reach it as a string. + """ + _, mock_client = mock_aimlapi_client + embedder = AimlapiEmbedding(BaseEmbedderConfig(api_key="k")) + mock_client.embeddings.create.return_value = Mock(data=[Mock(embedding=[0.1, 0.2, 0.3])]) + + result = embedder.embed("Hello\nworld") + + mock_client.embeddings.create.assert_called_once_with( + input=["Hello world"], + model="openai/text-embedding-3-small", + encoding_format="float", + ) + assert all(isinstance(item, str) for item in mock_client.embeddings.create.call_args.kwargs["input"]) + assert result == [0.1, 0.2, 0.3] + + +def test_embed_passes_dimensions_only_when_configured(mock_aimlapi_client): + _, mock_client = mock_aimlapi_client + embedder = AimlapiEmbedding( + BaseEmbedderConfig(api_key="k", model="openai/text-embedding-3-large", embedding_dims=1024) + ) + mock_client.embeddings.create.return_value = Mock(data=[Mock(embedding=[0.4])]) + + embedder.embed("hi") + + mock_client.embeddings.create.assert_called_once_with( + input=["hi"], + model="openai/text-embedding-3-large", + encoding_format="float", + dimensions=1024, + ) + + +def test_embed_batch(mock_aimlapi_client): + _, mock_client = mock_aimlapi_client + embedder = AimlapiEmbedding(BaseEmbedderConfig(api_key="k")) + mock_client.embeddings.create.return_value = Mock( + data=[Mock(index=1, embedding=[0.3]), Mock(index=0, embedding=[0.1])] + ) + + result = embedder.embed_batch(["a", "b"]) + + assert result == [[0.1], [0.3]] + assert mock_client.embeddings.create.call_args.kwargs["input"] == ["a", "b"] + + +def test_registered_in_factory(mock_aimlapi_client): + embedder = EmbedderFactory.create("aimlapi", {"api_key": "k"}, None) + assert isinstance(embedder, AimlapiEmbedding) diff --git a/tests/llms/test_aimlapi.py b/tests/llms/test_aimlapi.py new file mode 100644 index 0000000000..42a52cb74b --- /dev/null +++ b/tests/llms/test_aimlapi.py @@ -0,0 +1,200 @@ +import os +import re +from unittest.mock import Mock, patch + +import pytest + +from mem0.configs.llms.aimlapi import AimlapiConfig +from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.aimlapi import AimlapiLLM +from mem0.utils.aimlapi import build_default_headers +from mem0.utils.factory import LlmFactory + + +@pytest.fixture(autouse=True) +def clear_aimlapi_env(): + saved = {k: os.environ.pop(k, None) for k in ("AIMLAPI_API_KEY", "AIMLAPI_API_BASE")} + yield + for k, v in saved.items(): + if v is not None: + os.environ[k] = v + else: + os.environ.pop(k, None) + + +@pytest.fixture +def mock_aimlapi_client(): + with patch("mem0.llms.aimlapi.OpenAI") as mock_openai: + mock_client = Mock() + mock_openai.return_value = mock_client + yield mock_openai, mock_client + + +def test_aimlapi_base_url_resolution(): + # case1: default + llm = AimlapiLLM(AimlapiConfig(api_key="api_key")) + assert str(llm.client.base_url) == "https://api.aimlapi.com/v1/" + + # case2: AIMLAPI_API_BASE env var + os.environ["AIMLAPI_API_BASE"] = "https://gateway.example.com/v1" + llm = AimlapiLLM(AimlapiConfig(api_key="api_key")) + assert str(llm.client.base_url) == "https://gateway.example.com/v1/" + + # case3: config.aimlapi_base_url wins over env + llm = AimlapiLLM(AimlapiConfig(api_key="api_key", aimlapi_base_url="https://config.example.com/v1")) + assert str(llm.client.base_url) == "https://config.example.com/v1/" + + +def test_aimlapi_default_model(): + llm = AimlapiLLM(AimlapiConfig(api_key="k")) + assert llm.config.model == "openai/gpt-5-mini" + + +def test_aimlapi_reads_api_key_from_env(): + os.environ["AIMLAPI_API_KEY"] = "env_key" + llm = AimlapiLLM(AimlapiConfig()) + assert llm.client.api_key == "env_key" + + +def test_aimlapi_requires_api_key(): + with pytest.raises(ValueError, match="API key is required"): + AimlapiLLM(AimlapiConfig()) + + +def test_aimlapi_accepts_base_llm_config(): + # The factory may hand over a plain BaseLlmConfig; aimlapi_base_url must still exist. + llm = AimlapiLLM(BaseLlmConfig(model="openai/gpt-5-mini", api_key="k")) + assert isinstance(llm.config, AimlapiConfig) + assert llm.config.aimlapi_base_url is None + + +def test_aimlapi_registered_in_factory(): + llm = LlmFactory.create("aimlapi", {"model": "openai/gpt-5-mini", "api_key": "k"}) + assert isinstance(llm, AimlapiLLM) + + +# --- attribution headers ------------------------------------------------- + + +def test_partner_id_matches_gateway_pattern(): + # A malformed partner id is dropped silently by the gateway and earns nothing, + # so the shape is asserted here rather than discovered in production. + headers = build_default_headers("https://api.aimlapi.com/v1") + assert re.fullmatch(r"^part_[A-Za-z0-9]{1,64}$", headers["X-AIMLAPI-Partner-ID"]) + assert re.fullmatch(r"^(web|agent|mcp)/[a-z0-9-]{1,32}$", headers["X-AIMLAPI-Source"]) + + +def test_attribution_headers_sent_to_aimlapi(mock_aimlapi_client): + mock_openai, _ = mock_aimlapi_client + AimlapiLLM(AimlapiConfig(api_key="k")) + headers = mock_openai.call_args.kwargs["default_headers"] + assert headers["X-AIMLAPI-Source"] == "agent/mem0" + # HTTP-Referer / X-Title identify the calling application, which is Mem0. + assert headers["HTTP-Referer"] == "https://github.com/mem0ai/mem0" + assert headers["X-Title"] == "Mem0" + + +def test_attribution_headers_not_sent_to_other_hosts(mock_aimlapi_client): + # A user pointing the provider at their own gateway must not have our + # attribution ride along to a third party. + mock_openai, _ = mock_aimlapi_client + AimlapiLLM(AimlapiConfig(api_key="k", aimlapi_base_url="https://proxy.example.com/v1")) + assert mock_openai.call_args.kwargs["default_headers"] is None + + +def test_build_default_headers_does_not_mutate_shared_constant(): + first = build_default_headers("https://api.aimlapi.com/v1", {"X-Title": "Other"}) + second = build_default_headers("https://api.aimlapi.com/v1") + assert first["X-Title"] == "Other" # caller wins on a clash + assert second["X-Title"] == "Mem0" # and the constant is untouched + assert first is not second + + +# --- request payload ----------------------------------------------------- + + +def test_generate_response_without_tools(mock_aimlapi_client): + _, mock_client = mock_aimlapi_client + config = AimlapiConfig(model="openai/gpt-4o-mini", temperature=0.7, max_tokens=100, top_p=1.0, api_key="k") + llm = AimlapiLLM(config) + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello, how are you?"}, + ] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="I'm doing well!"))] + mock_client.chat.completions.create.return_value = mock_response + + response = llm.generate_response(messages) + + mock_client.chat.completions.create.assert_called_once_with( + model="openai/gpt-4o-mini", + messages=messages, + temperature=0.7, + max_tokens=100, + top_p=1.0, + ) + assert response == "I'm doing well!" + + +def test_generate_response_with_tools(mock_aimlapi_client): + _, mock_client = mock_aimlapi_client + llm = AimlapiLLM(AimlapiConfig(model="openai/gpt-4o-mini", api_key="k")) + messages = [{"role": "user", "content": "Add a new memory: Today is a sunny day."}] + tools = [ + { + "type": "function", + "function": { + "name": "add_memory", + "description": "Add a memory", + "parameters": { + "type": "object", + "properties": {"data": {"type": "string"}}, + "required": ["data"], + }, + }, + } + ] + + mock_response = Mock() + mock_message = Mock(content="I've added the memory.") + mock_tool_call = Mock() + mock_tool_call.function.name = "add_memory" + mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}' + mock_message.tool_calls = [mock_tool_call] + mock_response.choices = [Mock(message=mock_message)] + mock_client.chat.completions.create.return_value = mock_response + + response = llm.generate_response(messages, tools=tools) + + assert response["content"] == "I've added the memory." + assert response["tool_calls"][0]["name"] == "add_memory" + assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."} + + +def test_unset_parameters_are_omitted_not_nulled(mock_aimlapi_client): + """The API answers 400 to an explicit ``null`` on tools, temperature, top_p and + friends instead of treating it as unset. Every key on the wire must have a value. + """ + _, mock_client = mock_aimlapi_client + config = AimlapiConfig(model="openai/gpt-5-mini", api_key="k") + config.temperature = None + config.top_p = None + config.max_tokens = None + llm = AimlapiLLM(config) + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="ok"))] + mock_client.chat.completions.create.return_value = mock_response + + llm.generate_response( + [{"role": "user", "content": "hi"}], + response_format=None, + tools=None, + ) + + sent = mock_client.chat.completions.create.call_args.kwargs + assert None not in sent.values() + assert "tools" not in sent + assert "temperature" not in sent