From 02a0788dfa4e1fea78aac9e9aaf58226e8fd7d9c Mon Sep 17 00:00:00 2001 From: mrlonely <116348059+mameikagou@users.noreply.github.com> Date: Tue, 25 Aug 2026 01:07:53 +0800 Subject: [PATCH] Rebase OpenRouter provider on latest main --- .env.example | 10 +- README.md | 2 +- dashboard/app/api/settings/route.ts | 2 + dashboard/app/settings/page.tsx | 23 +++- docker-compose.yml | 5 + docs/mcp.md | 4 +- models.yml | 6 +- packages/openmemory-js/README.md | 7 +- packages/openmemory-js/src/core/cfg.ts | 10 ++ packages/openmemory-js/src/core/models.ts | 11 ++ packages/openmemory-js/src/memory/embed.ts | 83 +++++++++++- .../openmemory-js/tests/openrouter.test.ts | 118 ++++++++++++++++++ .../src/openmemory/ai/__init__.py | 3 +- .../src/openmemory/ai/openrouter.py | 52 ++++++++ .../src/openmemory/core/config.py | 5 + .../src/openmemory/memory/embed.py | 3 + .../openmemory-py/tests/test_openrouter.py | 69 ++++++++++ 17 files changed, 402 insertions(+), 11 deletions(-) create mode 100644 packages/openmemory-js/tests/openrouter.test.ts create mode 100644 packages/openmemory-py/src/openmemory/ai/openrouter.py create mode 100644 packages/openmemory-py/tests/test_openrouter.py diff --git a/.env.example b/.env.example index 169b9a1c..cb9f9cc4 100644 --- a/.env.example +++ b/.env.example @@ -68,7 +68,7 @@ OM_WEAVIATE_CLASS=OpenMemory # -------------------------------------------- # Embeddings Configuration # -------------------------------------------- -# Available providers: openai, gemini, aws, ollama, local, synthetic, orcarouter +# Available providers: openai, openrouter, gemini, aws, ollama, local, synthetic, orcarouter # Embedding models per sector can be configured in models.yaml # # NOTE: Your selected TIER (fast/smart/deep) affects how embeddings work: @@ -105,6 +105,11 @@ OM_EMBED_DELAY_MS=200 # Model override for all sector embeddings (leave empty to use defaults) # OM_OPENAI_MODEL=text-embedding-qwen3-embedding-4b +# OpenRouter Embeddings (OpenAI-compatible gateway) +# OM_OPENROUTER_BASE_URL=https://openrouter.ai/api/v1 +# Model override for all sector embeddings (leave empty to use defaults) +# OM_OPENROUTER_EMBEDDING_MODEL=openai/text-embedding-3-small + # Alibaba Cloud Model Studio (Bailian / DashScope) uses the same # OpenAI-compatible provider. Set OPENAI_API_KEY and OM_OPENAI_API_KEY empty # when using DASHSCOPE_API_KEY so the explicit OpenAI keys keep precedence. @@ -127,6 +132,9 @@ OM_MAX_PAYLOAD_SIZE=1000000 # OpenAI Embeddings OPENAI_API_KEY=your-openai-api-key-here +# OpenRouter Embeddings +OPENROUTER_API_KEY=sk-or-v1-your-key-here + # Google Gemini Embeddings GEMINI_API_KEY=your-gemini-api-key-here diff --git a/README.md b/README.md index 325b517a..e848cf24 100644 --- a/README.md +++ b/README.md @@ -308,7 +308,7 @@ results = mem.search("allergies", user_id="user123") “Waypoint” traces that show exactly which nodes were used in context. - **Embeddings** - OpenAI, Gemini, Ollama, AWS, [OrcaRouter](https://www.orcarouter.ai), synthetic fallback. + OpenAI, [OpenRouter](https://openrouter.ai), Gemini, Ollama, AWS, [OrcaRouter](https://www.orcarouter.ai), synthetic fallback. - **Integrations** LangChain, CrewAI, AutoGen, Streamlit, MCP, VS Code, IDEs. diff --git a/dashboard/app/api/settings/route.ts b/dashboard/app/api/settings/route.ts index 7b815ad4..5611b150 100644 --- a/dashboard/app/api/settings/route.ts +++ b/dashboard/app/api/settings/route.ts @@ -47,6 +47,8 @@ export async function GET() { const masked = { ...settings } if (masked.OPENAI_API_KEY) masked.OPENAI_API_KEY = '***' + if (masked.OPENROUTER_API_KEY) masked.OPENROUTER_API_KEY = '***' + if (masked.OM_OPENROUTER_API_KEY) masked.OM_OPENROUTER_API_KEY = '***' if (masked.DASHSCOPE_API_KEY) masked.DASHSCOPE_API_KEY = '***' if (masked.GEMINI_API_KEY) masked.GEMINI_API_KEY = '***' if (masked.AWS_SECRET_ACCESS_KEY) masked.AWS_SECRET_ACCESS_KEY = "***" diff --git a/dashboard/app/settings/page.tsx b/dashboard/app/settings/page.tsx index f2ae894b..92084e2d 100644 --- a/dashboard/app/settings/page.tsx +++ b/dashboard/app/settings/page.tsx @@ -188,7 +188,7 @@ const SETTING_METADATA: Record = { label: 'Embedding Provider', description: 'AI provider for generating embeddings (used in SMART/DEEP tiers)', type: 'select', - options: ['openai', 'gemini', 'aws', 'ollama', 'local', 'synthetic', 'orcarouter'] + options: ['openai', 'openrouter', 'gemini', 'aws', 'ollama', 'local', 'synthetic', 'orcarouter'] }, OM_VEC_DIM: { category: 'Embeddings', @@ -231,6 +231,20 @@ const SETTING_METADATA: Record = { type: 'text', placeholder: 'text-embedding-3-small' }, + OM_OPENROUTER_BASE_URL: { + category: 'Embeddings', + label: 'OpenRouter Base URL', + description: 'OpenRouter OpenAI-compatible API endpoint', + type: 'text', + placeholder: 'https://openrouter.ai/api/v1' + }, + OM_OPENROUTER_EMBEDDING_MODEL: { + category: 'Embeddings', + label: 'OpenRouter Model Override', + description: 'Override the default OpenRouter embedding model', + type: 'text', + placeholder: 'openai/text-embedding-3-small' + }, OM_ORCAROUTER_BASE_URL: { category: 'Embeddings', label: 'OrcaRouter Base URL', @@ -259,6 +273,13 @@ const SETTING_METADATA: Record = { type: 'password', placeholder: 'sk-...' }, + OPENROUTER_API_KEY: { + category: 'API Keys', + label: 'OpenRouter API Key', + description: 'API key for OpenRouter embeddings', + type: 'password', + placeholder: 'sk-or-v1-...' + }, DASHSCOPE_API_KEY: { category: 'API Keys', label: 'DashScope API Key', diff --git a/docker-compose.yml b/docker-compose.yml index 33ab8e26..cd16bab7 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -39,6 +39,11 @@ services: - DASHSCOPE_API_KEY=${DASHSCOPE_API_KEY:-} - OM_OPENAI_BASE_URL=${OM_OPENAI_BASE_URL:-https://api.openai.com/v1} - OM_OPENAI_MODEL=${OM_OPENAI_MODEL:-} + - OPENROUTER_API_KEY=${OPENROUTER_API_KEY:-} + - OM_OPENROUTER_API_KEY=${OM_OPENROUTER_API_KEY:-} + - OM_OPENROUTER_BASE_URL=${OM_OPENROUTER_BASE_URL:-https://openrouter.ai/api/v1} + - OM_OPENROUTER_MODEL=${OM_OPENROUTER_MODEL:-} + - OM_OPENROUTER_EMBEDDING_MODEL=${OM_OPENROUTER_EMBEDDING_MODEL:-} # Gemini Provider - GEMINI_API_KEY=${GEMINI_API_KEY:-} diff --git a/docs/mcp.md b/docs/mcp.md index 2ac933c4..a0b96bef 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -38,7 +38,7 @@ Configure your MCP client (`.mcp.json` for Cursor/Windsurf, or Claude settings): #### Environment Variables - `OM_METADATA_BACKEND`: Database type (`sqlite` or `postgres`) -- `OM_EMBEDDINGS`: Embedding provider (`openai`, `aws`, `ollama`, `synthetic`, `orcarouter`) +- `OM_EMBEDDINGS`: Embedding provider (`openai`, `openrouter`, `aws`, `ollama`, `synthetic`, `orcarouter`) - Database credentials (if using Postgres): `OM_PG_HOST`, `OM_PG_DB`, `OM_PG_USER`, `OM_PG_PASSWORD` ## Memory Systems @@ -612,7 +612,7 @@ If `type="factual"` returns no results: Ensure environment variables are set: - Postgres: `OM_PG_HOST`, `OM_PG_DB`, `OM_PG_USER`, `OM_PG_PASSWORD` -- Embeddings: `OPENAI_API_KEY` (or a compatible provider key such as `DASHSCOPE_API_KEY` or `ORCAROUTER_API_KEY`) +- Embeddings: `OPENAI_API_KEY` (or a compatible provider key such as `OPENROUTER_API_KEY`, `DASHSCOPE_API_KEY`, or `ORCAROUTER_API_KEY`) ## Performance diff --git a/models.yml b/models.yml index f7a8841a..8fbacda3 100644 --- a/models.yml +++ b/models.yml @@ -8,6 +8,7 @@ episodic: ollama: nomic-embed-text openai: text-embedding-3-small + openrouter: openai/text-embedding-3-small gemini: models/gemini-embedding-001 aws: amazon.titan-embed-text-v2:0 local: all-MiniLM-L6-v2 @@ -16,6 +17,7 @@ episodic: semantic: ollama: nomic-embed-text openai: text-embedding-3-small + openrouter: openai/text-embedding-3-small gemini: models/gemini-embedding-001 aws: amazon.titan-embed-text-v2:0 local: all-MiniLM-L6-v2 @@ -24,6 +26,7 @@ semantic: procedural: ollama: nomic-embed-text openai: text-embedding-3-small + openrouter: openai/text-embedding-3-small gemini: models/gemini-embedding-001 aws: amazon.titan-embed-text-v2:0 local: all-MiniLM-L6-v2 @@ -32,6 +35,7 @@ procedural: emotional: ollama: nomic-embed-text openai: text-embedding-3-small + openrouter: openai/text-embedding-3-small gemini: models/gemini-embedding-001 aws: amazon.titan-embed-text-v2:0 local: all-MiniLM-L6-v2 @@ -40,6 +44,7 @@ emotional: reflective: ollama: nomic-embed-text openai: text-embedding-3-large + openrouter: openai/text-embedding-3-small gemini: models/gemini-embedding-001 aws: amazon.titan-embed-text-v2:0 local: all-mpnet-base-v2 @@ -61,4 +66,3 @@ reflective: #AWS models: # - amazon.titan-embed-text-v2:0 (1024d, 512d, 256) - diff --git a/packages/openmemory-js/README.md b/packages/openmemory-js/README.md index fc8ae6ea..19b2b54c 100644 --- a/packages/openmemory-js/README.md +++ b/packages/openmemory-js/README.md @@ -93,7 +93,7 @@ OM_DB_PATH=./data/om.db # sqlite file path (default: ./data/openmem OM_DB_URL=sqlite://:memory: # or use in-memory db # embeddings -OM_EMBEDDINGS=ollama # synthetic | openai | gemini | ollama +OM_EMBEDDINGS=ollama # synthetic | openai | openrouter | gemini | ollama OM_OLLAMA_URL=http://localhost:11434 OM_OLLAMA_MODEL=embeddinggemma # or nomic-embed-text, mxbai-embed-large @@ -102,6 +102,11 @@ OM_OLLAMA_MODEL=embeddinggemma # or nomic-embed-text, mxbai-embed-large # OM_OPENAI_BASE_URL=https://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/compatible-mode/v1 # OM_OPENAI_MODEL=text-embedding-v4 +# OpenRouter through its OpenAI-compatible endpoint +# OPENROUTER_API_KEY=sk-or-v1-... +# OM_OPENROUTER_BASE_URL=https://openrouter.ai/api/v1 +# OM_OPENROUTER_EMBEDDING_MODEL=openai/text-embedding-3-small + # openai OPENAI_API_KEY=sk-... OM_OPENAI_MODEL=text-embedding-3-small diff --git a/packages/openmemory-js/src/core/cfg.ts b/packages/openmemory-js/src/core/cfg.ts index d5247500..821c93be 100644 --- a/packages/openmemory-js/src/core/cfg.ts +++ b/packages/openmemory-js/src/core/cfg.ts @@ -52,6 +52,16 @@ export const env = { "https://api.openai.com/v1", ), openai_model: process.env.OM_OPENAI_MODEL, + openrouter_key: + process.env.OPENROUTER_API_KEY || + process.env.OM_OPENROUTER_API_KEY || + "", + openrouter_base_url: str( + process.env.OM_OPENROUTER_BASE_URL, + "https://openrouter.ai/api/v1", + ), + openrouter_model: process.env.OM_OPENROUTER_MODEL, + openrouter_embedding_model: process.env.OM_OPENROUTER_EMBEDDING_MODEL, gemini_key: process.env.GEMINI_API_KEY || process.env.OM_GEMINI_API_KEY || "", AWS_REGION: process.env.AWS_REGION || "", diff --git a/packages/openmemory-js/src/core/models.ts b/packages/openmemory-js/src/core/models.ts index ba6ea6b0..049c4018 100644 --- a/packages/openmemory-js/src/core/models.ts +++ b/packages/openmemory-js/src/core/models.ts @@ -51,6 +51,7 @@ const get_defaults = (): model_cfg => ({ episodic: { ollama: "nomic-embed-text", openai: "text-embedding-3-small", + openrouter: "openai/text-embedding-3-small", gemini: "models/gemini-embedding-001", aws: "amazon.titan-embed-text-v2:0", siray: "text-embedding-3-small", @@ -60,6 +61,7 @@ const get_defaults = (): model_cfg => ({ semantic: { ollama: "nomic-embed-text", openai: "text-embedding-3-small", + openrouter: "openai/text-embedding-3-small", gemini: "models/gemini-embedding-001", aws: "amazon.titan-embed-text-v2:0", siray: "text-embedding-3-small", @@ -71,6 +73,7 @@ const get_defaults = (): model_cfg => ({ openai: "text-embedding-3-small", gemini: "models/gemini-embedding-001", aws: "amazon.titan-embed-text-v2:0", + openrouter: "openai/text-embedding-3-small", orcarouter: "orcarouter/auto", local: "all-MiniLM-L6-v2", }, @@ -79,6 +82,7 @@ const get_defaults = (): model_cfg => ({ openai: "text-embedding-3-small", gemini: "models/gemini-embedding-001", aws: "amazon.titan-embed-text-v2:0", + openrouter: "openai/text-embedding-3-small", orcarouter: "orcarouter/auto", local: "all-MiniLM-L6-v2", }, @@ -87,6 +91,7 @@ const get_defaults = (): model_cfg => ({ openai: "text-embedding-3-large", gemini: "models/gemini-embedding-001", aws: "amazon.titan-embed-text-v2:0", + openrouter: "openai/text-embedding-3-small", orcarouter: "orcarouter/auto", local: "all-mpnet-base-v2", }, @@ -99,6 +104,12 @@ export const get_model = (sector: string, provider: string): string => { if (provider === "openai" && process.env.OM_OPENAI_MODEL) { return process.env.OM_OPENAI_MODEL; } + if ( + provider === "openrouter" && + process.env.OM_OPENROUTER_EMBEDDING_MODEL + ) { + return process.env.OM_OPENROUTER_EMBEDDING_MODEL; + } if (provider === "orcarouter" && process.env.OM_ORCAROUTER_EMBEDDING_MODEL) { return process.env.OM_ORCAROUTER_EMBEDDING_MODEL; } diff --git a/packages/openmemory-js/src/memory/embed.ts b/packages/openmemory-js/src/memory/embed.ts index 0be4840c..9eb88057 100644 --- a/packages/openmemory-js/src/memory/embed.ts +++ b/packages/openmemory-js/src/memory/embed.ts @@ -109,14 +109,20 @@ export async function embedQueryForAllSectors( return result; } - if (env.emb_kind === "gemini" && env.gemini_key) { + if ( + env.embed_mode === "simple" && + (env.emb_kind === "gemini" || + env.emb_kind === "openai" || + env.emb_kind === "openrouter" || + env.emb_kind === "orcarouter") + ) { try { const txts: Record = {}; for (const s of sectors) txts[s] = query; - return await emb_gemini(txts); + return await emb_batch_with_fallback(txts); } catch (e) { console.error( - `[EMBED] Gemini batch failed, falling back to sequential: ${e}`, + `[EMBED] Batch query embedding failed, falling back to sequential: ${e}`, ); } } @@ -134,6 +140,8 @@ async function embed_with_provider( switch (provider) { case "openai": return await emb_openai(t, s); + case "openrouter": + return await emb_openrouter(t, s); case "gemini": return (await emb_gemini({ [s]: t }))[s]; case "ollama": @@ -205,6 +213,9 @@ async function emb_batch_with_fallback( case "openai": result = await emb_batch_openai(txts); break; + case "openrouter": + result = await emb_batch_openrouter(txts); + break; case "orcarouter": result = await emb_batch_orcarouter(txts); break; @@ -291,6 +302,58 @@ async function emb_batch_openai( return out; } +async function emb_openrouter(t: string, s: string): Promise { + if (!env.openrouter_key) throw new Error("OpenRouter key missing"); + const m = get_model(s, "openrouter"); + const r = await fetchWithTimeout( + `${env.openrouter_base_url.replace(/\/$/, "")}/embeddings`, + { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${env.openrouter_key}`, + }, + body: JSON.stringify({ + input: t, + model: env.openrouter_embedding_model || m, + ...(env.vec_dim ? { dimensions: env.vec_dim } : {}), + }), + }, + ); + if (!r.ok) throw new Error(`OpenRouter: ${r.status}`); + return ((await r.json()) as any).data[0].embedding; +} + +async function emb_batch_openrouter( + txts: Record, +): Promise> { + if (!env.openrouter_key) throw new Error("OpenRouter key missing"); + const secs = Object.keys(txts), + m = get_model("semantic", "openrouter"); + const r = await fetchWithTimeout( + `${env.openrouter_base_url.replace(/\/$/, "")}/embeddings`, + { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${env.openrouter_key}`, + }, + body: JSON.stringify({ + input: Object.values(txts), + model: env.openrouter_embedding_model || m, + ...(env.vec_dim ? { dimensions: env.vec_dim } : {}), + }), + }, + ); + if (!r.ok) throw new Error(`OpenRouter batch: ${r.status}`); + const d = (await r.json()) as any, + out: Record = {}; + secs.forEach((sector, index) => { + out[sector] = d.data[index].embedding; + }); + return out; +} + const task_map: Record = { episodic: "RETRIEVAL_DOCUMENT", semantic: "SEMANTIC_SIMILARITY", @@ -631,6 +694,7 @@ export async function embedMultiSector( simp && (env.emb_kind === "gemini" || env.emb_kind === "openai" || + env.emb_kind === "openrouter" || env.emb_kind === "orcarouter") ) { console.error( @@ -738,6 +802,7 @@ export const getEmbeddingInfo = () => { env.embed_mode === "simple" && (env.emb_kind === "gemini" || env.emb_kind === "openai" || + env.emb_kind === "openrouter" || env.emb_kind === "orcarouter"), advanced_parallel: env.adv_embed_parallel, embed_delay_ms: env.embed_delay_ms, @@ -754,6 +819,18 @@ export const getEmbeddingInfo = () => { emotional: get_model("emotional", "openai"), reflective: get_model("reflective", "openai"), }; + } else if (env.emb_kind === "openrouter") { + i.configured = !!env.openrouter_key; + i.base_url = env.openrouter_base_url; + i.model_override = env.openrouter_embedding_model || null; + i.batch_api = env.embed_mode === "simple"; + i.models = { + episodic: get_model("episodic", "openrouter"), + semantic: get_model("semantic", "openrouter"), + procedural: get_model("procedural", "openrouter"), + emotional: get_model("emotional", "openrouter"), + reflective: get_model("reflective", "openrouter"), + }; } else if (env.emb_kind === "gemini") { i.configured = !!env.gemini_key; i.batch_api = env.embed_mode === "simple"; diff --git a/packages/openmemory-js/tests/openrouter.test.ts b/packages/openmemory-js/tests/openrouter.test.ts new file mode 100644 index 00000000..0396709e --- /dev/null +++ b/packages/openmemory-js/tests/openrouter.test.ts @@ -0,0 +1,118 @@ +import { + afterAll, + afterEach, + beforeAll, + describe, + expect, + it, + vi, +} from "vitest"; +import { get_model } from "../src/core/models"; + +describe("openrouter provider wiring", () => { + let env: any; + let embedQueryForAllSectors: any; + let getEmbeddingInfo: any; + let getEmbeddingProvider: any; + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + beforeAll(async () => { + vi.resetModules(); + process.env.OM_EMBEDDINGS = "openrouter"; + process.env.OM_TIER = "deep"; + process.env.OPENROUTER_API_KEY = "sk-or-v1-test"; + process.env.OM_OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"; + process.env.OM_OPENROUTER_EMBEDDING_MODEL = + "openai/text-embedding-3-small"; + process.env.OM_EMBED_MODE = "simple"; + const cfg = await import("../src/core/cfg"); + const embed = await import("../src/memory/embed"); + env = cfg.env; + getEmbeddingInfo = embed.getEmbeddingInfo; + getEmbeddingProvider = embed.getEmbeddingProvider; + embedQueryForAllSectors = embed.embedQueryForAllSectors; + }); + + afterAll(() => { + delete process.env.OM_EMBEDDINGS; + delete process.env.OM_TIER; + delete process.env.OPENROUTER_API_KEY; + delete process.env.OM_OPENROUTER_BASE_URL; + delete process.env.OM_OPENROUTER_EMBEDDING_MODEL; + delete process.env.OM_EMBED_MODE; + }); + + it("exposes OpenRouter configuration", () => { + expect(env.openrouter_key).toBe("sk-or-v1-test"); + expect(env.openrouter_base_url).toBe("https://openrouter.ai/api/v1"); + expect(env.emb_kind).toBe("openrouter"); + }); + + it("resolves the OpenRouter default for every sector", () => { + for (const sector of [ + "episodic", + "semantic", + "procedural", + "emotional", + "reflective", + ]) { + expect(get_model(sector, "openrouter")).toBe( + "openai/text-embedding-3-small", + ); + } + }); + + it("honors the model override", () => { + const previous = process.env.OM_OPENROUTER_EMBEDDING_MODEL; + process.env.OM_OPENROUTER_EMBEDDING_MODEL = "vendor/custom-embedding"; + try { + expect(get_model("semantic", "openrouter")).toBe( + "vendor/custom-embedding", + ); + } finally { + process.env.OM_OPENROUTER_EMBEDDING_MODEL = previous; + } + }); + + it("reports OpenRouter as the active batch-capable provider", () => { + expect(getEmbeddingProvider()).toBe("openrouter"); + const info = getEmbeddingInfo(); + expect(info.provider).toBe("openrouter"); + expect(info.configured).toBe(true); + expect(info.base_url).toBe("https://openrouter.ai/api/v1"); + expect(info.models.semantic).toBe("openai/text-embedding-3-small"); + expect(info.batch_support).toBe(true); + }); + + it("batches query embeddings through the OpenRouter endpoint", async () => { + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + data: [{ embedding: [1, 0] }, { embedding: [0, 1] }], + }), + }); + vi.stubGlobal("fetch", fetchMock); + + const result = await embedQueryForAllSectors("hello", [ + "semantic", + "episodic", + ]); + + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(fetchMock.mock.calls[0][0]).toBe( + "https://openrouter.ai/api/v1/embeddings", + ); + const request = fetchMock.mock.calls[0][1]; + expect(JSON.parse(request.body)).toMatchObject({ + input: ["hello", "hello"], + model: "openai/text-embedding-3-small", + }); + expect(result).toEqual({ + semantic: [1, 0], + episodic: [0, 1], + }); + }); +}); diff --git a/packages/openmemory-py/src/openmemory/ai/__init__.py b/packages/openmemory-py/src/openmemory/ai/__init__.py index 1e477078..d427243d 100644 --- a/packages/openmemory-py/src/openmemory/ai/__init__.py +++ b/packages/openmemory-py/src/openmemory/ai/__init__.py @@ -7,5 +7,6 @@ from .siray import SirayAdapter from .minimax import MiniMaxAdapter from .orcarouter import OrcaRouterAdapter +from .openrouter import OpenRouterAdapter -__all__ = ["AIAdapter", "OpenAIAdapter", "OllamaAdapter", "GeminiAdapter", "AwsAdapter", "SyntheticAdapter", "SirayAdapter", "MiniMaxAdapter", "OrcaRouterAdapter"] +__all__ = ["AIAdapter", "OpenAIAdapter", "OllamaAdapter", "GeminiAdapter", "AwsAdapter", "SyntheticAdapter", "SirayAdapter", "MiniMaxAdapter", "OrcaRouterAdapter", "OpenRouterAdapter"] diff --git a/packages/openmemory-py/src/openmemory/ai/openrouter.py b/packages/openmemory-py/src/openmemory/ai/openrouter.py new file mode 100644 index 00000000..67b0ff5b --- /dev/null +++ b/packages/openmemory-py/src/openmemory/ai/openrouter.py @@ -0,0 +1,52 @@ +from typing import List, Dict + +from openai import AsyncOpenAI + +from ..core.config import env +from .adapter import AIAdapter + + +class OpenRouterAdapter(AIAdapter): + """OpenRouter adapter for chat completions and embeddings.""" + + def __init__(self, api_key: str = None, base_url: str = None): + self.api_key = api_key or env.openrouter_key + self.base_url = base_url or env.openrouter_base_url + self.client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url) + + async def chat( + self, messages: List[Dict[str, str]], model: str = None, **kwargs + ) -> str: + selected_model = model or env.openrouter_model or "openrouter/auto" + response = await self.client.chat.completions.create( + model=selected_model, + messages=messages, + **kwargs, + ) + return response.choices[0].message.content or "" + + async def embed(self, text: str, model: str = None) -> List[float]: + selected_model = ( + model + or env.openrouter_embedding_model + or "openai/text-embedding-3-small" + ) + response = await self.client.embeddings.create( + input=text, + model=selected_model, + ) + return response.data[0].embedding + + async def embed_batch( + self, texts: List[str], model: str = None + ) -> List[List[float]]: + selected_model = ( + model + or env.openrouter_embedding_model + or "openai/text-embedding-3-small" + ) + response = await self.client.embeddings.create( + input=texts, + model=selected_model, + ) + return [item.embedding for item in response.data] diff --git a/packages/openmemory-py/src/openmemory/core/config.py b/packages/openmemory-py/src/openmemory/core/config.py index 1c5dc729..481451ac 100644 --- a/packages/openmemory-py/src/openmemory/core/config.py +++ b/packages/openmemory-py/src/openmemory/core/config.py @@ -55,6 +55,11 @@ def get(sec: str, key: str, env_var: str, default: Any) -> Any: self.openai_base_url = get("ai", "openai_base", "OM_OPENAI_BASE_URL", "https://api.openai.com/v1") self.openai_model = get("ai", "openai_model", "OM_OPENAI_MODEL", None) + self.openrouter_key = get("ai", "openrouter_key", "OPENROUTER_API_KEY", "") or os.getenv("OM_OPENROUTER_API_KEY") + self.openrouter_base_url = get("ai", "openrouter_base", "OM_OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1") + self.openrouter_model = get("ai", "openrouter_model", "OM_OPENROUTER_MODEL", None) + self.openrouter_embedding_model = os.getenv("OM_OPENROUTER_EMBEDDING_MODEL") + self.ollama_url = get("ai", "ollama_url", "OLLAMA_URL", "http://localhost:11434") self.emb_kind = get("ai", "embedding_provider", "OM_EMBED_KIND", "synthetic") diff --git a/packages/openmemory-py/src/openmemory/memory/embed.py b/packages/openmemory-py/src/openmemory/memory/embed.py index 723e32fd..1815db29 100644 --- a/packages/openmemory-py/src/openmemory/memory/embed.py +++ b/packages/openmemory-py/src/openmemory/memory/embed.py @@ -21,12 +21,15 @@ from ..ai.synthetic import SyntheticAdapter from ..ai.minimax import MiniMaxAdapter from ..ai.orcarouter import OrcaRouterAdapter +from ..ai.openrouter import OpenRouterAdapter async def emb_dispatch(provider: str, t: str, s: str) -> List[float]: if provider == "synthetic": return await SyntheticAdapter(env.vec_dim or 768).embed(t, model=s) if provider == "openai": return await OpenAIAdapter().embed(t, model=env.openai_model) + if provider == "openrouter": + return await OpenRouterAdapter().embed(t, model=env.openrouter_embedding_model) if provider == "ollama": return await OllamaAdapter().embed(t, model=env.ollama_embedding_model) if provider == "gemini": diff --git a/packages/openmemory-py/tests/test_openrouter.py b/packages/openmemory-py/tests/test_openrouter.py new file mode 100644 index 00000000..54270caa --- /dev/null +++ b/packages/openmemory-py/tests/test_openrouter.py @@ -0,0 +1,69 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from openmemory.ai.adapter import AIAdapter +from openmemory.ai.openrouter import OpenRouterAdapter + + +def test_openrouter_adapter_implements_interface(): + assert issubclass(OpenRouterAdapter, AIAdapter) + + +def test_openrouter_adapter_uses_configured_endpoint(): + with patch("openmemory.ai.openrouter.env") as mock_env: + mock_env.openrouter_key = "sk-or-v1-test" + mock_env.openrouter_base_url = "https://openrouter.ai/api/v1" + adapter = OpenRouterAdapter() + + assert adapter.api_key == "sk-or-v1-test" + assert adapter.base_url == "https://openrouter.ai/api/v1" + + +@pytest.mark.asyncio +async def test_openrouter_embed_uses_default_model(): + with patch("openmemory.ai.openrouter.env") as mock_env: + mock_env.openrouter_key = "sk-or-v1-test" + mock_env.openrouter_base_url = "https://openrouter.ai/api/v1" + mock_env.openrouter_embedding_model = None + adapter = OpenRouterAdapter() + + item = MagicMock() + item.embedding = [0.1, 0.2] + response = MagicMock() + response.data = [item] + adapter.client.embeddings.create = AsyncMock(return_value=response) + + result = await adapter.embed("hello") + + assert result == [0.1, 0.2] + adapter.client.embeddings.create.assert_awaited_once_with( + input="hello", + model="openai/text-embedding-3-small", + ) + + +@pytest.mark.asyncio +async def test_openrouter_embed_batch_preserves_order(): + adapter = OpenRouterAdapter( + api_key="sk-or-v1-test", + base_url="https://openrouter.ai/api/v1", + ) + first = MagicMock() + first.embedding = [0.1] + second = MagicMock() + second.embedding = [0.2] + response = MagicMock() + response.data = [first, second] + adapter.client.embeddings.create = AsyncMock(return_value=response) + + result = await adapter.embed_batch( + ["first", "second"], + model="vendor/custom-embedding", + ) + + assert result == [[0.1], [0.2]] + adapter.client.embeddings.create.assert_awaited_once_with( + input=["first", "second"], + model="vendor/custom-embedding", + )