diff --git a/README.md b/README.md index f60fefa..26175f3 100644 --- a/README.md +++ b/README.md @@ -268,11 +268,11 @@ Run `onboard_project` to index a project's history. Here's what happens: 2. **Parses events** — extracts the 8 event types from each session file (streams files >10MB) 3. **Extracts contracts** — scans source for types, interfaces, enums, routes, Prisma models, OpenAPI schemas 4. **Loads manual contracts** — merges any `.preflight/contracts/*.yml` definitions (manual wins on name conflicts) -5. **Generates embeddings** — local [Xenova/all-MiniLM-L6-v2](https://huggingface.co/Xenova/all-MiniLM-L6-v2) by default (~90MB model download on first run, ~50 events/sec) or OpenAI if `OPENAI_API_KEY` is set (~200 events/sec) +5. **Generates embeddings** — local [Xenova/all-MiniLM-L6-v2](https://huggingface.co/Xenova/all-MiniLM-L6-v2) by default (~90MB model download on first run, ~50 events/sec), OpenAI if `OPENAI_API_KEY` is set (~200 events/sec), or [Ollama](https://ollama.com) for local GPU-accelerated embeddings with no external API 6. **Stores in LanceDB** — per-project database at `~/.preflight/projects//timeline.lance/` 7. **Updates registry** — records the project in `~/.preflight/projects/index.json` -No data leaves your machine unless you opt into OpenAI embeddings. +No data leaves your machine unless you opt into OpenAI embeddings. Ollama embeddings also stay fully local. After onboarding, you get: - 🔎 **Semantic search** — "How did I set up auth middleware last month?" actually works @@ -430,8 +430,11 @@ thresholds: # Embedding configuration embeddings: - provider: local # type: "local" | "openai" + provider: local # type: "local" | "openai" | "ollama" openai_api_key: sk-... # type: string — only needed if provider is "openai" + ollama_base_url: http://localhost:11434 # type: string — Ollama API URL (default shown) + ollama_model: nomic-embed-text # type: string — Ollama embedding model (default shown) + ollama_dimensions: 768 # type: number — vector dimensions for your model (default: 768) ``` ### `.preflight/triage.yml` @@ -495,7 +498,9 @@ Manual contract definitions that supplement auto-extraction: | `CLAUDE_PROJECT_DIR` | Project root to monitor | **Required** | | `OPENAI_API_KEY` | OpenAI key for embeddings | Uses local Xenova | | `PREFLIGHT_RELATED` | Comma-separated related project paths | None | -| `EMBEDDING_PROVIDER` | `local` or `openai` | `local` | +| `EMBEDDING_PROVIDER` | `local`, `openai`, or `ollama` | `local` | +| `OLLAMA_BASE_URL` | Ollama API base URL | `http://localhost:11434` | +| `OLLAMA_EMBED_MODEL` | Ollama embedding model name | `nomic-embed-text` | | `PROMPT_DISCIPLINE_PROFILE` | `minimal`, `standard`, or `full` | `standard` | Environment variables are **fallbacks** — `.preflight/` config takes precedence when present. diff --git a/src/lib/config.ts b/src/lib/config.ts index fc9d8f2..ad807f4 100644 --- a/src/lib/config.ts +++ b/src/lib/config.ts @@ -11,7 +11,7 @@ import { load as yamlLoad } from "js-yaml"; import { PROJECT_DIR } from "./files.js"; export type Profile = "minimal" | "standard" | "full"; -export type EmbeddingProvider = "local" | "openai"; +export type EmbeddingProvider = "local" | "openai" | "ollama"; export type TriageStrictness = "relaxed" | "standard" | "strict"; export interface RelatedProject { @@ -30,6 +30,9 @@ export interface PreflightConfig { embeddings: { provider: EmbeddingProvider; openai_api_key?: string; + ollama_base_url?: string; + ollama_model?: string; + ollama_dimensions?: number; }; triage: { rules: { @@ -125,7 +128,7 @@ function loadConfig(): PreflightConfig { // Embedding provider const envProvider = process.env.EMBEDDING_PROVIDER?.toLowerCase(); - if (envProvider === "local" || envProvider === "openai") { + if (envProvider === "local" || envProvider === "openai" || envProvider === "ollama") { config.embeddings.provider = envProvider; } @@ -133,6 +136,14 @@ function loadConfig(): PreflightConfig { if (process.env.OPENAI_API_KEY) { config.embeddings.openai_api_key = process.env.OPENAI_API_KEY; } + + // Ollama settings + if (process.env.OLLAMA_BASE_URL) { + config.embeddings.ollama_base_url = process.env.OLLAMA_BASE_URL; + } + if (process.env.OLLAMA_EMBED_MODEL) { + config.embeddings.ollama_model = process.env.OLLAMA_EMBED_MODEL; + } } return config; diff --git a/src/lib/embeddings.ts b/src/lib/embeddings.ts index 69b5883..0283e89 100644 --- a/src/lib/embeddings.ts +++ b/src/lib/embeddings.ts @@ -102,11 +102,71 @@ class OpenAIEmbeddingProvider implements EmbeddingProvider { } } +// --- Ollama Provider --- + +class OllamaEmbeddingProvider implements EmbeddingProvider { + dimensions: number; + private baseUrl: string; + private model: string; + + constructor(options?: { baseUrl?: string; model?: string; dimensions?: number }) { + this.baseUrl = options?.baseUrl ?? "http://localhost:11434"; + this.model = options?.model ?? "nomic-embed-text"; + // nomic-embed-text produces 768-dim vectors; allow override for other models + this.dimensions = options?.dimensions ?? 768; + } + + async embed(text: string): Promise { + const processed = preprocessText(text); + const resp = await fetch(`${this.baseUrl}/api/embed`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model: this.model, input: processed }), + }); + + if (!resp.ok) { + const err = await resp.text(); + throw new Error(`Ollama embeddings API error ${resp.status}: ${err}`); + } + + const data = await resp.json(); + // /api/embed returns { embeddings: number[][] } + if (!data.embeddings?.[0]) { + throw new Error("Ollama returned no embeddings"); + } + return data.embeddings[0]; + } + + async embedBatch(texts: string[]): Promise { + // Ollama /api/embed supports multiple inputs natively + const processed = texts.map(preprocessText); + const resp = await fetch(`${this.baseUrl}/api/embed`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model: this.model, input: processed }), + }); + + if (!resp.ok) { + const err = await resp.text(); + throw new Error(`Ollama embeddings API error ${resp.status}: ${err}`); + } + + const data = await resp.json(); + if (!data.embeddings || data.embeddings.length !== texts.length) { + throw new Error(`Ollama returned ${data.embeddings?.length ?? 0} embeddings, expected ${texts.length}`); + } + return data.embeddings; + } +} + // --- Factory --- export interface EmbeddingConfig { - provider: "local" | "openai"; + provider: "local" | "openai" | "ollama"; apiKey?: string; + ollamaBaseUrl?: string; + ollamaModel?: string; + ollamaDimensions?: number; } export function createEmbeddingProvider(config: EmbeddingConfig): EmbeddingProvider { @@ -114,5 +174,12 @@ export function createEmbeddingProvider(config: EmbeddingConfig): EmbeddingProvi if (!config.apiKey) throw new Error("OpenAI API key required for openai embedding provider"); return new OpenAIEmbeddingProvider(config.apiKey); } + if (config.provider === "ollama") { + return new OllamaEmbeddingProvider({ + baseUrl: config.ollamaBaseUrl, + model: config.ollamaModel, + dimensions: config.ollamaDimensions, + }); + } return new LocalEmbeddingProvider(); } diff --git a/src/lib/timeline-db.ts b/src/lib/timeline-db.ts index 49b4f78..d984bad 100644 --- a/src/lib/timeline-db.ts +++ b/src/lib/timeline-db.ts @@ -55,9 +55,12 @@ export interface ProjectInfo { } export interface TimelineConfig { - embedding_provider: "local" | "openai"; + embedding_provider: "local" | "openai" | "ollama"; embedding_model: string; openai_api_key?: string; + ollama_base_url?: string; + ollama_model?: string; + ollama_dimensions?: number; indexed_projects: Record { _embedder = createEmbeddingProvider({ provider: config.embedding_provider, apiKey: config.openai_api_key, + ollamaBaseUrl: config.ollama_base_url, + ollamaModel: config.ollama_model, + ollamaDimensions: config.ollama_dimensions, }); } return _embedder; diff --git a/src/tools/onboard-project.ts b/src/tools/onboard-project.ts index bca91a0..3566af5 100644 --- a/src/tools/onboard-project.ts +++ b/src/tools/onboard-project.ts @@ -30,15 +30,17 @@ export function registerOnboardProject(server: McpServer) { "Index a project's Claude Code sessions and git history into the timeline database for semantic search and chronological viewing.", { project_dir: z.string().describe("Absolute path to the project directory"), - embedding_provider: z.enum(["local", "openai"]).default("local"), + embedding_provider: z.enum(["local", "openai", "ollama"]).default("local"), openai_api_key: z.string().optional(), + ollama_base_url: z.string().optional().describe("Ollama API base URL (default: http://localhost:11434)"), + ollama_model: z.string().optional().describe("Ollama embedding model (default: nomic-embed-text)"), git_depth: z.enum(["all", "6months", "1year", "3months"]).default("all"), git_since: z.string().optional().describe("Override git_depth with exact start date (ISO: '2025-08-01')"), git_authors: z.array(z.string()).optional().describe("Filter git commits to these authors. If omitted, auto-detects the primary author (most commits)."), reindex: z.boolean().default(false).describe("If true, drop existing data and rebuild from scratch"), }, async (params) => { - const { project_dir, embedding_provider, openai_api_key, git_depth, git_since, git_authors, reindex } = params; + const { project_dir, embedding_provider, openai_api_key, ollama_base_url, ollama_model, git_depth, git_since, git_authors, reindex } = params; // 1. Validate project_dir if (!fs.existsSync(project_dir)) { @@ -174,6 +176,8 @@ export function registerOnboardProject(server: McpServer) { const embedder = createEmbeddingProvider({ provider: embedding_provider, apiKey: openai_api_key, + ollamaBaseUrl: ollama_base_url, + ollamaModel: ollama_model, }); const BATCH_SIZE = 50; diff --git a/tests/lib/embeddings.test.ts b/tests/lib/embeddings.test.ts index 68b0ccf..4168c75 100644 --- a/tests/lib/embeddings.test.ts +++ b/tests/lib/embeddings.test.ts @@ -1,66 +1,48 @@ -import { describe, it, expect } from "vitest"; -import { preprocessText, createEmbeddingProvider } from "../../src/lib/embeddings.js"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { createEmbeddingProvider, preprocessText } from "../../src/lib/embeddings.js"; describe("preprocessText", () => { - it("strips code blocks", () => { - const result = preprocessText("before ```const x = 1;``` after"); - expect(result).not.toContain("const x"); - expect(result).toContain("before"); - expect(result).toContain("after"); - }); - - it("strips inline code", () => { - const result = preprocessText("use `foo()` here"); - expect(result).not.toContain("foo()"); - expect(result).toContain("use"); - expect(result).toContain("here"); - }); - - it("removes markdown characters", () => { - const result = preprocessText("# Heading **bold** _italic_"); + it("strips markdown formatting", () => { + const result = preprocessText("# Hello **world** `code` [link](http://x.com)"); expect(result).not.toContain("#"); expect(result).not.toContain("**"); + expect(result).not.toContain("`"); + expect(result).toContain("link"); }); - it("strips markdown link brackets and parens", () => { - const result = preprocessText("[Click here](https://example.com)"); - // Brackets and parens are removed by markdown char stripping - expect(result).not.toContain("["); - expect(result).not.toContain("]"); - expect(result).not.toContain("("); - expect(result).not.toContain(")"); - expect(result).toContain("Click here"); - }); - - it("normalizes whitespace", () => { - const result = preprocessText("hello \n\n world"); - expect(result).toBe("hello world"); - }); - - it("truncates to 2048 chars", () => { - const long = "a".repeat(3000); - const result = preprocessText(long); - expect(result.length).toBeLessThanOrEqual(2048); + it("truncates long text to 2048 chars", () => { + const long = "a".repeat(5000); + expect(preprocessText(long).length).toBe(2048); }); }); describe("createEmbeddingProvider", () => { - it("returns local provider with 384 dimensions", () => { + it("creates local provider by default", () => { const provider = createEmbeddingProvider({ provider: "local" }); expect(provider.dimensions).toBe(384); }); - it("throws when openai provider has no API key", () => { - expect(() => - createEmbeddingProvider({ provider: "openai" }), - ).toThrow("API key required"); + it("creates openai provider with api key", () => { + const provider = createEmbeddingProvider({ provider: "openai", apiKey: "sk-test" }); + expect(provider.dimensions).toBe(1536); + }); + + it("throws when openai provider missing api key", () => { + expect(() => createEmbeddingProvider({ provider: "openai" })).toThrow("OpenAI API key required"); }); - it("returns openai provider with 1536 dimensions when key provided", () => { + it("creates ollama provider with defaults", () => { + const provider = createEmbeddingProvider({ provider: "ollama" }); + expect(provider.dimensions).toBe(768); + }); + + it("creates ollama provider with custom config", () => { const provider = createEmbeddingProvider({ - provider: "openai", - apiKey: "sk-test-key", + provider: "ollama", + ollamaBaseUrl: "http://myhost:11434", + ollamaModel: "mxbai-embed-large", + ollamaDimensions: 1024, }); - expect(provider.dimensions).toBe(1536); + expect(provider.dimensions).toBe(1024); }); });