diff --git a/.env.example b/.env.example index dceefec4..962ecfb0 100644 --- a/.env.example +++ b/.env.example @@ -2,7 +2,14 @@ MNEMON_DATA_DIR=~/.mnemon MNEMON_STORE=default -# Optional embeddings through Ollama. -# Enable only when an Ollama service is available. +# Optional embeddings. The defaults below use Ollama. MNEMON_EMBED_ENDPOINT=http://localhost:11434 MNEMON_EMBED_MODEL=nomic-embed-text + +# For an OpenAI-compatible server, replace the endpoint and model above. +# The remaining settings are optional: +# MNEMON_EMBED_ENDPOINT=http://127.0.0.1:18000/v1 +# MNEMON_EMBED_MODEL=bge-m3-mlx-8bit +# MNEMON_EMBED_PROTOCOL=openai +# MNEMON_EMBED_API_KEY=sk-... +# MNEMON_EMBED_DIMENSIONS=256 diff --git a/README.md b/README.md index 0aec3e29..b812d5e5 100644 --- a/README.md +++ b/README.md @@ -352,7 +352,7 @@ memory is useful. - **Built-in deduplication** — `remember` auto-detects duplicates and conflicts; skips or auto-replaces - **Retention lifecycle** — importance decay, access-count boosting, and garbage collection - **Privacy-safe receipts** — export hashed operation receipts for memory-boundary audits without raw memory contents or queries -- **Optional embeddings** — works fully without Ollama; add local [Ollama](https://ollama.ai) for enhanced vector+keyword hybrid search +- **Optional embeddings** — works fully without an embedding provider; add local [Ollama](https://ollama.ai) or an OpenAI-compatible server for enhanced vector+keyword hybrid search ## Vision @@ -446,12 +446,27 @@ Mnemon architecture. | `MNEMON_DATA_DIR` | `~/.mnemon` | Base data directory | | `MNEMON_STORE` | *(active file or `default`)* | Named memory store for data isolation | -**Ollama-specific** (only relevant if using embeddings): +**Embedding** (only relevant if using embeddings): | Environment Variable | Default | Description | |---|---|---| -| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | Ollama API endpoint | +| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | Embedding API endpoint | | `MNEMON_EMBED_MODEL` | `nomic-embed-text` | Embedding model name | +| `MNEMON_EMBED_PROTOCOL` | *(auto-detect)* | `ollama` or `openai`; auto-detected from an endpoint ending in `/v1` | +| `MNEMON_EMBED_API_KEY` | *(none)* | Bearer token for OpenAI-compatible servers (oMLX, vLLM, etc.) | +| `MNEMON_EMBED_DIMENSIONS` | *(native)* | Optional Matryoshka dimension truncation | + +The embedding client speaks the Ollama API by default and the +OpenAI-compatible embeddings API when the endpoint ends in `/v1` (or when +`MNEMON_EMBED_PROTOCOL=openai` is set). For example, a local server such as +[oMLX](https://omlx.dev) can be configured with: + +```bash +export MNEMON_EMBED_ENDPOINT=http://127.0.0.1:18000/v1 +export MNEMON_EMBED_MODEL=bge-m3-mlx-8bit +export MNEMON_EMBED_API_KEY=sk-... # omit for keyless local servers +mnemon embed --status +``` ## Development diff --git a/SECURITY.md b/SECURITY.md index 181c541d..0c572cc6 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -16,7 +16,7 @@ Mnemon runs locally and stores data in `~/.mnemon/`. Key security considerations - **SQLite database** — contains all stored insights; protected by filesystem permissions (`0644`). - **Hook scripts** — shell scripts executed by the LLM CLI at lifecycle events; written with `0755` permissions. -- **Ollama connection** — optional HTTP calls to a local Ollama instance; no TLS by default. If `MNEMON_EMBED_ENDPOINT` is pointed at a remote server, traffic is unencrypted unless the endpoint uses HTTPS. +- **Embedding provider connection** — optional requests send insight or query text to the configured Ollama or OpenAI-compatible server. The default local Ollama endpoint does not use TLS. If `MNEMON_EMBED_ENDPOINT` points outside a trusted local network, use HTTPS to protect content and any `MNEMON_EMBED_API_KEY` bearer token in transit. ## Supported Versions diff --git a/cmd/memory/embed.go b/cmd/memory/embed.go index 0ac9be70..10ad4988 100644 --- a/cmd/memory/embed.go +++ b/cmd/memory/embed.go @@ -16,8 +16,8 @@ var ( var embedCmd = &cobra.Command{ Use: "embed [id]", - Short: "Generate embeddings for insights via Ollama", - Long: `Generate embedding vectors for insights using a local Ollama model. + Short: "Generate embeddings for insights", + Long: `Generate embedding vectors for insights using the configured provider. Modes: mnemon embed --status Show embedding coverage statistics @@ -40,19 +40,25 @@ Modes: if err != nil { return fmt.Errorf("embedding stats: %w", err) } + available := ec.Available() output := map[string]interface{}{ - "total_insights": total, - "embedded": embedded, - "coverage": fmt.Sprintf("%.0f%%", float64(embedded)/float64(max(total, 1))*100), - "ollama_available": ec.Available(), - "model": ec.Model(), + "total_insights": total, + "embedded": embedded, + "coverage": fmt.Sprintf("%.0f%%", float64(embedded)/float64(max(total, 1))*100), + "embedding_available": available, + "ollama_available": available, // Backward-compatible alias. + "protocol": ec.Protocol(), + "model": ec.Model(), } return enc.Encode(output) } - // Check Ollama availability + // Check embedding provider availability. if !ec.Available() { - return fmt.Errorf("Ollama not available at %s — install with: brew install ollama && ollama pull %s", ec.Endpoint(), ec.Model()) + if ec.Protocol() == embed.ProtocolOllama { + return fmt.Errorf("Ollama embedding provider not available at %s — install with: brew install ollama && ollama pull %s", ec.Endpoint(), ec.Model()) + } + return fmt.Errorf("OpenAI-compatible embedding provider not available at %s", ec.Endpoint()) } // Single insight mode diff --git a/cmd/memory/root.go b/cmd/memory/root.go index 08ffc6fa..d68b3cf1 100644 --- a/cmd/memory/root.go +++ b/cmd/memory/root.go @@ -41,7 +41,7 @@ func init() { rootCmd.PersistentFlags().StringVar(&storeName, "store", "", "named memory store (overrides MNEMON_STORE and active file)") rootCmd.PersistentFlags().BoolVar(&readOnly, "readonly", false, "open database in read-only mode (no WAL files, safe for read-only mounts)") rootCmd.PersistentFlags().StringVar(&embedModel, "embed-model", "", - fmt.Sprintf("Ollama embedding model (env: MNEMON_EMBED_MODEL; default: %s)", embed.DefaultModel)) + fmt.Sprintf("embedding model (env: MNEMON_EMBED_MODEL; default: %s)", embed.DefaultModel)) } // resolveEmbedModel returns the embedding model selector that should be diff --git a/cmd/memory/root_test.go b/cmd/memory/root_test.go index eb2a79ac..a10d1f30 100644 --- a/cmd/memory/root_test.go +++ b/cmd/memory/root_test.go @@ -75,7 +75,7 @@ func TestOpenDBRejectsInvalidStoreNameFromFlag(t *testing.T) { // TestResolveEmbedModelChain exercises the full cmd → embed pipeline for the // --embed-model flag and MNEMON_EMBED_MODEL env var, mirroring how cobra // will hand the value off at runtime. The test runs against -// embed.NewClientWithModel directly so it does not require a live Ollama. +// embed.NewClientWithModel directly so it does not require a live provider. func TestResolveEmbedModelChain(t *testing.T) { oldEmbedModel := embedModel t.Cleanup(func() { embedModel = oldEmbedModel }) diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index 2c05d94e..b0419ef7 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -75,7 +75,8 @@ make compose-down ## Optional Embeddings -Mnemon works without embeddings. To use Ollama-backed vector search in the Compose environment: +Mnemon works without embeddings. The Compose embeddings profile provides the +default Ollama-backed vector search setup: ```bash docker compose --profile embeddings up -d ollama @@ -87,9 +88,18 @@ The relevant environment variables are: - `MNEMON_EMBED_ENDPOINT` - `MNEMON_EMBED_MODEL` +- `MNEMON_EMBED_PROTOCOL` +- `MNEMON_EMBED_API_KEY` +- `MNEMON_EMBED_DIMENSIONS` For host-based Ollama, set `MNEMON_EMBED_ENDPOINT=http://host.docker.internal:11434` on Docker Desktop, or use the host gateway address for Linux deployments. +An external OpenAI-compatible server can be selected with an endpoint ending +in `/v1`, for example `MNEMON_EMBED_ENDPOINT=http://host.docker.internal:18000/v1`. +Set `MNEMON_EMBED_MODEL` to a model exposed by that server and +`MNEMON_EMBED_API_KEY` when authentication is required. Use HTTPS whenever the +server is not on a trusted local network. + ## Release Deployment Tagged releases are handled by GoReleaser through `.github/workflows/release.yml`. diff --git a/docs/USAGE.md b/docs/USAGE.md index a29ca2f0..c5a57079 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -12,7 +12,7 @@ These root flags configure Memory commands: |---|---|---| | `--store ` | (auto) | Named memory store (overrides `MNEMON_STORE` and active file) | | `--data-dir ` | `~/.mnemon` | Base data directory | -| `--embed-model ` | `nomic-embed-text` | Ollama embedding model (overrides `MNEMON_EMBED_MODEL`) | +| `--embed-model ` | `nomic-embed-text` | Embedding model (overrides `MNEMON_EMBED_MODEL`) | | `--readonly` | `false` | Open the Memory database read-only, without creating WAL files | | `--version` | | Print version and exit | @@ -242,8 +242,10 @@ Nodes are colored by category (decision, fact, insight, preference, context); ed |---|---|---| | `MNEMON_DATA_DIR` | `~/.mnemon` | Base data directory | | `MNEMON_STORE` | `default` | Active named store | -| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | Ollama API endpoint | -| `MNEMON_EMBED_MODEL` | `nomic-embed-text` | Ollama embedding model | +| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | Embedding API endpoint | +| `MNEMON_EMBED_MODEL` | `nomic-embed-text` | Embedding model | +| `MNEMON_EMBED_PROTOCOL` | (auto-detect) | `ollama` or `openai`; endpoints ending in `/v1` select `openai` | +| `MNEMON_EMBED_API_KEY` | (none) | Bearer token for OpenAI-compatible servers | | `MNEMON_EMBED_DIMENSIONS` | (native) | Embedding dimensions; set to truncate (e.g., `256` for Matryoshka models) | | `MNEMON_MAX_INSIGHTS` | `1000` | Active-insight ceiling before auto-pruning starts; `0` disables auto-pruning | @@ -251,26 +253,41 @@ Nodes are colored by category (decision, fact, insight, preference, context); ed ## Embedding Support (Optional) -Mnemon works fully without Ollama — all core features (remember, recall, link, graph traversal) function out of the box. Adding Ollama enhances recall precision through vector similarity, but is never required. +Mnemon works fully without an embedding provider — all core features (remember, recall, link, graph traversal) function out of the box. Configuring Ollama or an OpenAI-compatible server enhances recall precision through vector similarity, but is never required. ### What changes with and without embeddings -| Capability | Without Ollama | With Ollama | +| Capability | Without embeddings | With embeddings | |---|---|---| | **Recall anchors** | Keyword + recency | Keyword + vector + recency (RRF hybrid) | | **Semantic edges** | Token overlap (coarser) | Cosine similarity ≥ 0.50 (precise) | | **Traversal scoring** | Pure structural | Structural + semantic | | **Rerank weights** | Keyword 45%, Entity 25%, Graph 30% | Keyword 30%, Entity 15%, Similarity 35%, Graph 20% | -When Ollama is unavailable, the reranking system automatically redistributes similarity weight to keyword and graph signals — no configuration needed, no degraded mode flag. The system detects Ollama availability at runtime with a 2-second timeout. +When the configured provider is unavailable, the reranking system automatically redistributes similarity weight to keyword and graph signals — no configuration or degraded-mode flag is needed. Mnemon checks provider availability at runtime with a 2-second timeout. ### Setup +Ollama remains the default provider: + ```bash brew install ollama # or see https://ollama.ai ollama pull nomic-embed-text # download the embedding model ``` +For an OpenAI-compatible server, point the endpoint at its `/v1` base URL and +select the server's embedding model. The API key is optional for keyless local +servers: + +```bash +export MNEMON_EMBED_ENDPOINT=http://127.0.0.1:18000/v1 +export MNEMON_EMBED_MODEL=bge-m3-mlx-8bit +export MNEMON_EMBED_API_KEY=sk-... # omit for keyless local servers +``` + +Set `MNEMON_EMBED_PROTOCOL=openai` explicitly only when the compatible endpoint +does not end in `/v1`. + Verify with: ```bash @@ -282,14 +299,19 @@ mnemon embed --status "total_insights": 87, "embedded": 87, "coverage": "100%", + "embedding_available": true, "ollama_available": true, + "protocol": "ollama", "model": "nomic-embed-text" } ``` +`ollama_available` is retained as a compatibility alias for existing scripts; +new integrations should use `embedding_available` and `protocol`. + ### Backfilling existing insights -If you install Ollama after already using mnemon, existing insights won't have embeddings. Backfill them in one command: +If you configure an embedding provider after already using mnemon, existing insights won't have embeddings. Backfill them in one command: ```bash mnemon embed --all @@ -316,8 +338,8 @@ This generates embeddings for all un-embedded insights and automatically creates retrieve. │ │ causal │ │ │ │ semantic │ │ ┌──────────────────┐ │ ├────────────┤ │ - │ Ollama │ (optional) │ │ Embeddings │ │ - │ nomic-embed-text│ ◄───────────── │ └────────────┘ │ + │ Embedding server │ (optional) │ │ Embeddings │ │ + │ configured model │ ◄───────────── │ └────────────┘ │ └──────────────────┘ └──────────────────┘ ``` diff --git a/docs/zh/README.md b/docs/zh/README.md index 250a2304..e5b64df6 100644 --- a/docs/zh/README.md +++ b/docs/zh/README.md @@ -313,7 +313,7 @@ store 可见。**Remind** 触发 recall 判断。**Nudge** 触发 writeback 判 - **意图感知召回** — 图遍历 + 可选向量搜索(RRF 融合),所有查询默认启用 - **内置去重** — `remember` 自动检测重复和冲突;跳过或自动替换 - **保留度生命周期** — 重要性衰减、访问计数提升、免疫规则、垃圾回收 -- **可选嵌入向量** — 本地 [Ollama](https://ollama.ai) 集成,支持混合向量+关键词搜索 +- **可选嵌入向量** — 可使用本地 [Ollama](https://ollama.ai) 或 OpenAI 兼容服务器,支持混合向量+关键词搜索 ## 愿景 @@ -396,8 +396,22 @@ Sub-agent 委派是可选执行策略。当 runtime 支持时,主 agent 可以 |---------|-------|------| | `MNEMON_DATA_DIR` | `~/.mnemon` | 基础数据目录 | | `MNEMON_STORE` | *(active 文件或 `default`)* | 命名记忆体,用于数据隔离 | -| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | Ollama API 端点 | +| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | 嵌入 API 端点 | | `MNEMON_EMBED_MODEL` | `nomic-embed-text` | 嵌入模型名称 | +| `MNEMON_EMBED_PROTOCOL` | *(自动探测)* | `ollama` 或 `openai`;端点以 `/v1` 结尾时自动切换 | +| `MNEMON_EMBED_API_KEY` | *(无)* | OpenAI 兼容服务器(oMLX、vLLM 等)的 Bearer 令牌 | +| `MNEMON_EMBED_DIMENSIONS` | *(原生维度)* | 可选的 Matryoshka 维度截断 | + +嵌入客户端默认使用 Ollama API;当端点以 `/v1` 结尾(或显式设置 +`MNEMON_EMBED_PROTOCOL=openai`)时改用 OpenAI 兼容的 embeddings API。例如, +可通过以下配置对接 [oMLX](https://omlx.dev) 等本地服务器: + +```bash +export MNEMON_EMBED_ENDPOINT=http://127.0.0.1:18000/v1 +export MNEMON_EMBED_MODEL=bge-m3-mlx-8bit +export MNEMON_EMBED_API_KEY=sk-... # 无需认证的本地服务器可省略 +mnemon embed --status +``` 也可在命令上使用 `--data-dir` 或 `--store` 标志覆盖。 @@ -415,7 +429,7 @@ make help # 显示所有目标 **依赖**:Go 1.24+、`modernc.org/sqlite`、`spf13/cobra`、`google/uuid` -**可选**:[Ollama](https://ollama.ai) + `nomic-embed-text` 嵌入支持 +**可选**:[Ollama](https://ollama.ai) 或 OpenAI 兼容的嵌入服务器 ## 文档 diff --git a/docs/zh/USAGE.md b/docs/zh/USAGE.md index ec2ae635..70d0ee3a 100644 --- a/docs/zh/USAGE.md +++ b/docs/zh/USAGE.md @@ -12,7 +12,7 @@ |---|---|---| | `--store ` | (自动) | 命名记忆体(覆盖 `MNEMON_STORE` 和 active 文件) | | `--data-dir ` | `~/.mnemon` | 基础数据目录 | -| `--embed-model ` | `nomic-embed-text` | Ollama 嵌入模型(覆盖 `MNEMON_EMBED_MODEL`) | +| `--embed-model ` | `nomic-embed-text` | 嵌入模型(覆盖 `MNEMON_EMBED_MODEL`) | | `--readonly` | `false` | 以只读模式打开 Memory 数据库,不创建 WAL 文件 | | `--version` | | 打印版本并退出 | @@ -246,8 +246,10 @@ open graph.html |---|---|---| | `MNEMON_DATA_DIR` | `~/.mnemon` | 基础数据目录 | | `MNEMON_STORE` | `default` | 活跃命名记忆体 | -| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | Ollama API 端点 | -| `MNEMON_EMBED_MODEL` | `nomic-embed-text` | Ollama 嵌入模型 | +| `MNEMON_EMBED_ENDPOINT` | `http://localhost:11434` | 嵌入 API 端点 | +| `MNEMON_EMBED_MODEL` | `nomic-embed-text` | 嵌入模型 | +| `MNEMON_EMBED_PROTOCOL` | (自动探测) | `ollama` 或 `openai`;以 `/v1` 结尾的端点自动选择 `openai` | +| `MNEMON_EMBED_API_KEY` | (无) | OpenAI 兼容服务器的 Bearer 令牌 | | `MNEMON_EMBED_DIMENSIONS` | (原生维度) | 嵌入向量维度;可设置截断值(例如 Matryoshka 模型使用 `256`) | | `MNEMON_MAX_INSIGHTS` | `1000` | 触发自动清理的活跃洞察数量上限;设为 `0` 可关闭自动清理 | @@ -255,26 +257,38 @@ open graph.html ## 嵌入向量支持(可选) -Mnemon 无需 Ollama 即可完整运行 — 所有核心功能(remember、recall、link、图遍历)开箱即用。添加 Ollama 可通过向量相似度增强召回精度,但从不是必需的。 +Mnemon 无需嵌入服务即可完整运行 — 所有核心功能(remember、recall、link、图遍历)开箱即用。配置 Ollama 或 OpenAI 兼容服务器可通过向量相似度增强召回精度,但从不是必需的。 ### 有无嵌入的对比 -| 能力 | 无 Ollama | 有 Ollama | +| 能力 | 无嵌入向量 | 有嵌入向量 | |---|---|---| | **召回锚点** | 关键词 + 时间 | 关键词 + 向量 + 时间(RRF 混合) | | **语义边** | Token 重叠(较粗) | 余弦相似度 ≥ 0.50(精确) | | **遍历评分** | 纯结构分 | 结构 + 语义 | | **重排序权重** | 关键词 45%、实体 25%、图 30% | 关键词 30%、实体 15%、相似度 35%、图 20% | -Ollama 不可用时,重排序系统自动将相似度权重重新分配给关键词和图信号 — 无需配置,无降级模式标志。系统在运行时以 2 秒超时检测 Ollama 可用性。 +配置的嵌入服务不可用时,重排序系统会自动将相似度权重重新分配给关键词和图信号 — 无需额外配置或降级模式标志。Mnemon 在运行时以 2 秒超时检测服务可用性。 ### 安装 +Ollama 仍是默认服务: + ```bash brew install ollama # 或参见 https://ollama.ai ollama pull nomic-embed-text # 下载嵌入模型 ``` +使用 OpenAI 兼容服务器时,将端点指向其 `/v1` 基础 URL,并选择服务器上的嵌入模型。无需认证的本地服务器可省略 API key: + +```bash +export MNEMON_EMBED_ENDPOINT=http://127.0.0.1:18000/v1 +export MNEMON_EMBED_MODEL=bge-m3-mlx-8bit +export MNEMON_EMBED_API_KEY=sk-... # 无需认证的本地服务器可省略 +``` + +仅当兼容端点不以 `/v1` 结尾时,才需要显式设置 `MNEMON_EMBED_PROTOCOL=openai`。 + 验证: ```bash @@ -286,14 +300,19 @@ mnemon embed --status "total_insights": 87, "embedded": 87, "coverage": "100%", + "embedding_available": true, "ollama_available": true, + "protocol": "ollama", "model": "nomic-embed-text" } ``` +为兼容现有脚本,`ollama_available` 字段会继续保留;新集成应使用 +`embedding_available` 和 `protocol`。 + ### 回填已有洞察 -如果在使用 mnemon 之后才安装 Ollama,已有洞察不会有嵌入向量。一条命令即可回填: +如果在使用 mnemon 之后才配置嵌入服务,已有洞察不会有嵌入向量。一条命令即可回填: ```bash mnemon embed --all @@ -320,8 +339,8 @@ mnemon embed --all retrieve. │ │ causal │ │ │ │ semantic │ │ ┌──────────────────┐ │ ├────────────┤ │ - │ Ollama │ (optional) │ │ Embeddings │ │ - │ nomic-embed-text│ ◄───────────── │ └────────────┘ │ + │ Embedding server │ (optional) │ │ Embeddings │ │ + │ configured model │ ◄───────────── │ └────────────┘ │ └──────────────────┘ └──────────────────┘ ``` diff --git a/internal/memory/embed/ollama.go b/internal/memory/embed/ollama.go index 452635a4..ff2f5a55 100644 --- a/internal/memory/embed/ollama.go +++ b/internal/memory/embed/ollama.go @@ -7,37 +7,57 @@ import ( "fmt" "net" "net/http" + "net/url" "os" "strconv" + "strings" "time" ) +// Protocol identifies the wire protocol used to reach the embedding server. +type Protocol string + +const ( + // ProtocolOllama is the Ollama /api/embed protocol (default). + ProtocolOllama Protocol = "ollama" + // ProtocolOpenAI is the OpenAI-compatible /v1/embeddings protocol + // (e.g. oMLX, llama.cpp server, vLLM, LM Studio). + ProtocolOpenAI Protocol = "openai" +) + // DefaultModel is the default Ollama embedding model. const DefaultModel = "nomic-embed-text" // DefaultEndpoint is the default Ollama API endpoint. const DefaultEndpoint = "http://localhost:11434" -// Client communicates with an Ollama instance for embedding generation. +// Client communicates with an embedding server (an Ollama instance or an +// OpenAI-compatible server) for embedding generation. type Client struct { endpoint string model string dims int // 0 means use native dimensions + apiKey string + protocol Protocol http *http.Client } -// NewClient creates an Ollama embedding client. -// It checks MNEMON_EMBED_ENDPOINT, MNEMON_EMBED_MODEL, and -// MNEMON_EMBED_DIMENSIONS env vars. +// NewClient creates an embedding client. +// It checks MNEMON_EMBED_ENDPOINT, MNEMON_EMBED_MODEL, +// MNEMON_EMBED_DIMENSIONS, MNEMON_EMBED_API_KEY, and +// MNEMON_EMBED_PROTOCOL env vars. func NewClient() *Client { return NewClientWithModel("") } -// NewClientWithModel creates an Ollama embedding client with an explicit -// model override. Resolution order for the model: explicit argument > -// MNEMON_EMBED_MODEL env var > DefaultModel. The endpoint and dimensions -// continue to be resolved from MNEMON_EMBED_ENDPOINT and -// MNEMON_EMBED_DIMENSIONS env vars. +// NewClientWithModel creates an embedding client with an explicit model +// override. Resolution order for the model: explicit argument > +// MNEMON_EMBED_MODEL env var > DefaultModel. The endpoint, dimensions, +// API key, and protocol continue to be resolved from environment vars. +// +// Protocol resolution: MNEMON_EMBED_PROTOCOL ("ollama" | "openai") wins +// when set; otherwise the protocol is auto-detected — an endpoint whose +// URL path ends in /v1 is assumed to be an OpenAI-compatible server. func NewClientWithModel(model string) *Client { endpoint := os.Getenv("MNEMON_EMBED_ENDPOINT") if endpoint == "" { @@ -55,14 +75,37 @@ func NewClientWithModel(model string) *Client { dims = v } } + protocol := ProtocolOllama + explicit := false + if p := os.Getenv("MNEMON_EMBED_PROTOCOL"); p != "" { + switch Protocol(strings.ToLower(p)) { + case ProtocolOllama, ProtocolOpenAI: + protocol = Protocol(strings.ToLower(p)) + explicit = true + default: + fmt.Fprintf(os.Stderr, "warning: invalid MNEMON_EMBED_PROTOCOL %q, falling back to auto-detect\n", p) + } + } + if !explicit { + // Auto-detect: OpenAI-compatible servers conventionally serve the + // API under a /v1 path prefix. + if u, err := url.Parse(endpoint); err == nil { + trimmed := strings.TrimRight(u.Path, "/") + if strings.HasSuffix(trimmed, "/v1") { + protocol = ProtocolOpenAI + } + } + } return &Client{ endpoint: endpoint, model: model, dims: dims, + apiKey: os.Getenv("MNEMON_EMBED_API_KEY"), + protocol: protocol, http: &http.Client{ Timeout: 30 * time.Second, Transport: &http.Transport{ - // Bypass system proxy for localhost Ollama connections. + // Bypass system proxy for localhost connections. Proxy: nil, DialContext: (&net.Dialer{ Timeout: 5 * time.Second, @@ -73,15 +116,44 @@ func NewClientWithModel(model string) *Client { } } -// Available returns true if the Ollama server is reachable and the model is loaded. -// Uses a 2s timeout to avoid blocking the CLI on unresponsive servers. +// Protocol returns the active wire protocol. +func (c *Client) Protocol() Protocol { + return c.protocol +} + +// endpointURL resolves a provider route relative to the configured endpoint. +// url.JoinPath keeps both /v1 and /v1/ endpoint forms equivalent while +// preserving any path prefix used by an OpenAI-compatible server. +func (c *Client) endpointURL(route string) (string, error) { + endpointURL, err := url.JoinPath(c.endpoint, route) + if err != nil { + return "", fmt.Errorf("join embedding endpoint: %w", err) + } + return endpointURL, nil +} + +// Available returns true if the embedding server's discovery endpoint +// responds successfully. Uses a 2s timeout to avoid blocking the CLI on +// unresponsive servers. func (c *Client) Available() bool { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.endpoint+"/api/tags", nil) + var route string + switch c.protocol { + case ProtocolOpenAI: + route = "models" + default: + route = "api/tags" + } + endpointURL, err := c.endpointURL(route) if err != nil { return false } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpointURL, nil) + if err != nil { + return false + } + c.applyAuth(req) resp, err := c.http.Do(req) if err != nil { return false @@ -95,22 +167,38 @@ func (c *Client) Model() string { return c.model } -// Endpoint returns the configured Ollama endpoint URL. +// Endpoint returns the configured embedding endpoint URL. func (c *Client) Endpoint() string { return c.endpoint } +// applyAuth attaches the Bearer token for OpenAI-compatible servers. +// Ollama requires no authentication. +func (c *Client) applyAuth(req *http.Request) { + if c.protocol == ProtocolOpenAI && c.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+c.apiKey) + } +} + type embedRequest struct { Model string `json:"model"` Input string `json:"input"` Dimensions int `json:"dimensions,omitempty"` } -type embedResponse struct { +type ollamaEmbedResponse struct { Embeddings [][]float64 `json:"embeddings"` } +type openaiEmbedResponse struct { + Data []struct { + Embedding []float64 `json:"embedding"` + } `json:"data"` +} + // Embed generates an embedding vector for the given text. +// The request body is identical for both protocols; only the endpoint +// path and the response shape differ. func (c *Client) Embed(text string) ([]float64, error) { req := embedRequest{Model: c.model, Input: text} if c.dims > 0 { @@ -121,24 +209,52 @@ func (c *Client) Embed(text string) ([]float64, error) { return nil, fmt.Errorf("marshal request: %w", err) } - resp, err := c.http.Post(c.endpoint+"/api/embed", "application/json", bytes.NewReader(body)) + var route string + switch c.protocol { + case ProtocolOpenAI: + route = "embeddings" + default: + route = "api/embed" + } + endpointURL, err := c.endpointURL(route) if err != nil { - return nil, fmt.Errorf("ollama request: %w", err) + return nil, err } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("ollama returned status %d", resp.StatusCode) + httpReq, err := http.NewRequest(http.MethodPost, endpointURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build request: %w", err) } + httpReq.Header.Set("Content-Type", "application/json") + c.applyAuth(httpReq) - var result embedResponse - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return nil, fmt.Errorf("decode response: %w", err) + resp, err := c.http.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("embed request: %w", err) } + defer resp.Body.Close() - if len(result.Embeddings) == 0 || len(result.Embeddings[0]) == 0 { - return nil, fmt.Errorf("empty embedding returned") + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("embedding provider returned status %d", resp.StatusCode) } - return result.Embeddings[0], nil + switch c.protocol { + case ProtocolOpenAI: + var result openaiEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + if len(result.Data) == 0 || len(result.Data[0].Embedding) == 0 { + return nil, fmt.Errorf("empty embedding returned") + } + return result.Data[0].Embedding, nil + default: + var result ollamaEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + if len(result.Embeddings) == 0 || len(result.Embeddings[0]) == 0 { + return nil, fmt.Errorf("empty embedding returned") + } + return result.Embeddings[0], nil + } } diff --git a/internal/memory/embed/ollama_test.go b/internal/memory/embed/ollama_test.go index aa8e8e69..006befc1 100644 --- a/internal/memory/embed/ollama_test.go +++ b/internal/memory/embed/ollama_test.go @@ -1,6 +1,8 @@ package embed import ( + "net/http" + "net/http/httptest" "testing" ) @@ -78,3 +80,42 @@ func TestNewClientWithModel_ExplicitEmptyTreatedAsUnset(t *testing.T) { t.Errorf("explicit empty + no env should fall through to default: want %q, got %q", DefaultModel, c.Model()) } } + +func TestOllamaEndpointWithTrailingSlash(t *testing.T) { + t.Setenv("MNEMON_EMBED_PROTOCOL", "ollama") + t.Setenv("MNEMON_EMBED_API_KEY", "must-not-be-sent") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("expected Ollama request without Authorization header, got %q", got) + } + switch r.URL.Path { + case "/api/tags": + w.WriteHeader(http.StatusOK) + case "/api/embed": + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + if got := r.Header.Get("Content-Type"); got != "application/json" { + t.Errorf("expected application/json, got %q", got) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"embeddings":[[0.1,0.2,0.3]]}`)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/") + c := NewClient() + if !c.Available() { + t.Fatal("expected Available() true for trailing-slash Ollama endpoint") + } + vec, err := c.Embed("hello") + if err != nil { + t.Fatalf("Embed with trailing-slash Ollama endpoint: %v", err) + } + if len(vec) != 3 { + t.Fatalf("expected 3 dims, got %d", len(vec)) + } +} diff --git a/internal/memory/embed/openai_test.go b/internal/memory/embed/openai_test.go new file mode 100644 index 00000000..e3fdff86 --- /dev/null +++ b/internal/memory/embed/openai_test.go @@ -0,0 +1,168 @@ +package embed + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestProtocolAutoDetect(t *testing.T) { + t.Setenv("MNEMON_EMBED_ENDPOINT", "http://127.0.0.1:18000/v1") + c := NewClient() + if c.Protocol() != ProtocolOpenAI { + t.Fatalf("expected openai protocol for /v1 endpoint, got %q", c.Protocol()) + } + + t.Setenv("MNEMON_EMBED_ENDPOINT", "http://localhost:11434") + c = NewClient() + if c.Protocol() != ProtocolOllama { + t.Fatalf("expected ollama protocol for default endpoint, got %q", c.Protocol()) + } + + // Explicit protocol override wins over auto-detection. + t.Setenv("MNEMON_EMBED_ENDPOINT", "http://127.0.0.1:18000/v1") + t.Setenv("MNEMON_EMBED_PROTOCOL", "ollama") + c = NewClient() + if c.Protocol() != ProtocolOllama { + t.Fatalf("expected explicit protocol override to win, got %q", c.Protocol()) + } + + t.Setenv("MNEMON_EMBED_PROTOCOL", "openai") + t.Setenv("MNEMON_EMBED_ENDPOINT", "http://localhost:11434") + c = NewClient() + if c.Protocol() != ProtocolOpenAI { + t.Fatalf("expected explicit openai protocol, got %q", c.Protocol()) + } +} + +func TestOpenAIAvailable(t *testing.T) { + t.Setenv("MNEMON_EMBED_API_KEY", "sk-test") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("expected /v1/models, got %s", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer sk-test" { + t.Errorf("expected Bearer sk-test, got %q", got) + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + if !c.Available() { + t.Fatal("expected Available() true for 200 /v1/models") + } +} + +func TestOpenAIEndpointWithTrailingSlash(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/models": + w.WriteHeader(http.StatusOK) + case "/v1/embeddings": + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"embedding":[1.0,2.0]}]}`)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1/") + c := NewClient() + if c.Protocol() != ProtocolOpenAI { + t.Fatalf("expected openai protocol for /v1/ endpoint, got %q", c.Protocol()) + } + if !c.Available() { + t.Fatal("expected Available() true for trailing-slash endpoint") + } + vec, err := c.Embed("hello") + if err != nil { + t.Fatalf("Embed with trailing-slash endpoint: %v", err) + } + if len(vec) != 2 { + t.Fatalf("expected 2 dims, got %d", len(vec)) + } +} + +func TestOpenAIEmbed(t *testing.T) { + t.Setenv("MNEMON_EMBED_MODEL", "bge-m3-mlx-8bit") + t.Setenv("MNEMON_EMBED_API_KEY", "sk-test") + var gotAuth string + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/embeddings" { + t.Errorf("expected /v1/embeddings, got %s", r.URL.Path) + } + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + gotAuth = r.Header.Get("Authorization") + if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil { + t.Errorf("decode request body: %v", err) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"object":"list","data":[{"object":"embedding","embedding":[0.1,0.2,0.3],"index":0}]}`)) + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + vec, err := c.Embed("跨会话记忆测试") + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(vec) != 3 { + t.Fatalf("expected 3 dims, got %d", len(vec)) + } + if gotAuth != "Bearer sk-test" { + t.Errorf("expected Bearer sk-test, got %q", gotAuth) + } + if gotBody["model"] != "bge-m3-mlx-8bit" { + t.Errorf("expected model in body, got %v", gotBody["model"]) + } + if input, _ := gotBody["input"].(string); input != "跨会话记忆测试" { + t.Errorf("expected input text, got %v", gotBody["input"]) + } +} + +func TestOpenAIEmbedWithoutKey(t *testing.T) { + // Keyless OpenAI-compatible servers must still work: no Authorization + // header should be sent when MNEMON_EMBED_API_KEY is unset. + var gotAuth string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"data":[{"embedding":[1.0]}]}`)) + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + vec, err := c.Embed("hello") + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(vec) != 1 { + t.Fatalf("expected 1 dim, got %d", len(vec)) + } + if gotAuth != "" { + t.Errorf("expected no Authorization header without API key, got %q", gotAuth) + } +} + +func TestOpenAIEmbedEmptyResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"data":[]}`)) + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + if _, err := c.Embed("hello"); err == nil { + t.Fatal("expected error for empty embedding response") + } +}