Skip to content

Commit 2572284

Browse files
dmikushinclaude
andcommitted
feat: add Ollama embedder and interactive embedder selection dialog
OllamaEmbedder (embedding/ollama.py): - Free local embeddings via ollama's OpenAI-compatible /v1/embeddings API - Default model: nomic-embed-text (768 dims) - Supports mxbai-embed-large, all-minilm, snowflake-arctic-embed, etc. - Graceful errors: tells user to run 'ollama serve' or 'ollama pull <model>' Interactive embedder dialog (ui.py → interactive_embedder_select): - Shows table: ollama / openai / gemini / mock with availability status - Smart default: preselects ollama when free-code or ollama is the LLM provider - Shown automatically in init_cmd after provider selection when --embedder flag is not explicitly set init_cmd.py: - Calls interactive_embedder_select after interactive_provider_select - Adds 'ollama' branch to the embedder switch - Updates --embedder help text to include ollama Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 086a3e8 commit 2572284

3 files changed

Lines changed: 283 additions & 1 deletion

File tree

packages/cli/src/repowise/cli/commands/init_cmd.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ async def _persist_result(
188188
"embedder_name",
189189
default=None,
190190
type=click.Choice(["gemini", "openai", "mock"]),
191-
help="Embedder for RAG: gemini | openai | mock (default: auto-detect).",
191+
help="Embedder for RAG: ollama | gemini | openai | mock (default: interactive).",
192192
)
193193
@click.option("--skip-tests", is_flag=True, default=False, help="Skip test files.")
194194
@click.option("--skip-infra", is_flag=True, default=False, help="Skip infrastructure files.")
@@ -377,9 +377,13 @@ def init_command(
377377
)
378378
else:
379379
if not is_interactive and provider_name is None and sys.stdin.isatty():
380+
from repowise.cli.ui import interactive_embedder_select as _ies
380381
from repowise.cli.ui import interactive_provider_select as _ips
381382

382383
provider_name, model = _ips(console, model)
384+
# Show embedder dialog only when not explicitly set via --embedder flag
385+
if embedder_name is None:
386+
embedder_name_resolved = _ies(console, llm_provider=provider_name)
383387

384388
provider = resolve_provider(provider_name, model, repo_path)
385389
if not is_interactive:
@@ -500,6 +504,13 @@ def init_command(
500504
embedder_impl = OpenAIEmbedder()
501505
except Exception:
502506
embedder_impl = MockEmbedder()
507+
elif embedder_name_resolved == "ollama":
508+
try:
509+
from repowise.core.providers.embedding.ollama import OllamaEmbedder
510+
511+
embedder_impl = OllamaEmbedder()
512+
except Exception:
513+
embedder_impl = MockEmbedder()
503514
else:
504515
embedder_impl = MockEmbedder()
505516

packages/cli/src/repowise/cli/ui.py

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,120 @@ def interactive_provider_select(
335335
return chosen, model
336336

337337

338+
# ---------------------------------------------------------------------------
339+
# Interactive embedder selection
340+
# ---------------------------------------------------------------------------
341+
342+
_EMBEDDER_DEFS: list[dict] = [
343+
{
344+
"id": "ollama",
345+
"label": "Ollama (local, free)",
346+
"hint": "nomic-embed-text · requires: ollama pull nomic-embed-text",
347+
"requires_key": False,
348+
"env_keys": [],
349+
"default_model": "nomic-embed-text",
350+
},
351+
{
352+
"id": "openai",
353+
"label": "OpenAI",
354+
"hint": "text-embedding-3-small · requires OPENAI_API_KEY",
355+
"requires_key": True,
356+
"env_keys": ["OPENAI_API_KEY"],
357+
"default_model": "text-embedding-3-small",
358+
},
359+
{
360+
"id": "gemini",
361+
"label": "Google Gemini",
362+
"hint": "gemini-embedding-exp-03-07 · requires GEMINI_API_KEY",
363+
"requires_key": True,
364+
"env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"],
365+
"default_model": "gemini-embedding-exp-03-07",
366+
},
367+
{
368+
"id": "mock",
369+
"label": "Mock (no cost, no quality)",
370+
"hint": "deterministic 8-dim vectors — OK for testing / index-only",
371+
"requires_key": False,
372+
"env_keys": [],
373+
"default_model": "",
374+
},
375+
]
376+
377+
_EMBEDDER_BY_ID = {e["id"]: e for e in _EMBEDDER_DEFS}
378+
379+
380+
def _detect_embedder_key(emb: dict) -> bool:
381+
"""Return True if the embedder's key is available in the environment."""
382+
if not emb["requires_key"]:
383+
return True
384+
return any(os.environ.get(k) for k in emb["env_keys"])
385+
386+
387+
def interactive_embedder_select(
388+
console: Console,
389+
*,
390+
llm_provider: str | None = None,
391+
) -> str:
392+
"""Show embedder table and let the user choose.
393+
394+
Returns the embedder id string (e.g. 'ollama', 'openai', 'gemini', 'mock').
395+
"""
396+
table = Table(
397+
show_header=True,
398+
box=None,
399+
padding=(0, 2),
400+
title="[bold]Embedder Setup[/bold]",
401+
title_style="",
402+
)
403+
table.add_column("#", style=BRAND_STYLE, width=4)
404+
table.add_column("Embedder", style="bold", min_width=24)
405+
table.add_column("Info", style="dim")
406+
407+
available: list[str] = []
408+
default_idx = "1"
409+
410+
for idx, emb in enumerate(_EMBEDDER_DEFS, 1):
411+
has_key = _detect_embedder_key(emb)
412+
if has_key:
413+
status_char = f"[{OK}]✓[/]"
414+
else:
415+
status_char = "[dim]✗[/dim]"
416+
table.add_row(f"[{idx}]", f"{status_char} {emb['label']}", emb["hint"])
417+
available.append(emb["id"])
418+
# Smart default: prefer ollama when llm_provider is free-code or ollama,
419+
# otherwise prefer the first embedder whose key is set
420+
if default_idx == "1":
421+
if llm_provider in ("free-code", "ollama") and emb["id"] == "ollama":
422+
default_idx = str(idx)
423+
elif has_key and emb["id"] not in ("mock",):
424+
default_idx = str(idx)
425+
426+
console.print()
427+
console.print(table)
428+
console.print()
429+
430+
valid_choices = [str(i) for i in range(1, len(available) + 1)]
431+
chosen_idx = Prompt.ask(
432+
" Select embedder",
433+
choices=valid_choices,
434+
default=default_idx,
435+
console=console,
436+
)
437+
chosen = available[int(chosen_idx) - 1]
438+
439+
if chosen == "ollama":
440+
base_url = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
441+
console.print()
442+
console.print(f" [dim]Ollama endpoint: [cyan]{base_url}[/cyan][/dim]")
443+
console.print(
444+
" [dim]Make sure the model is pulled: "
445+
"[bold]ollama pull nomic-embed-text[/bold][/dim]"
446+
)
447+
console.print()
448+
449+
return chosen
450+
451+
338452
def _prompt_api_key(
339453
console: Console,
340454
provider: str,
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
"""Ollama embedding support for repowise semantic search.
2+
3+
Uses Ollama's OpenAI-compatible /v1/embeddings endpoint so no extra
4+
dependencies are needed beyond the openai package already required.
5+
6+
Prerequisites:
7+
1. Install Ollama: https://ollama.com/download
8+
2. Pull an embedding model: ollama pull nomic-embed-text
9+
10+
Popular embedding models:
11+
nomic-embed-text 768 dims — fast, good quality (recommended default)
12+
mxbai-embed-large 1024 dims — higher quality, slower
13+
all-minilm 384 dims — very fast, lower quality
14+
15+
Usage:
16+
from repowise.core.providers.embedding.ollama import OllamaEmbedder
17+
embedder = OllamaEmbedder(model="nomic-embed-text")
18+
vecs = await embedder.embed(["hello world"])
19+
"""
20+
21+
from __future__ import annotations
22+
23+
import asyncio
24+
import math
25+
import os
26+
27+
_DEFAULT_BASE_URL = "http://localhost:11434"
28+
_DEFAULT_MODEL = "nomic-embed-text"
29+
30+
# Known dimensions for common ollama embedding models.
31+
# For unknown models we run a probe request at construction time.
32+
_KNOWN_DIMS: dict[str, int] = {
33+
"nomic-embed-text": 768,
34+
"mxbai-embed-large": 1024,
35+
"all-minilm": 384,
36+
"snowflake-arctic-embed": 1024,
37+
"bge-large": 1024,
38+
"bge-m3": 1024,
39+
}
40+
41+
42+
class OllamaEmbedder:
43+
"""Ollama embedding adapter implementing the repowise Embedder protocol.
44+
45+
Calls Ollama's OpenAI-compatible /v1/embeddings endpoint. No API key
46+
required — Ollama runs locally.
47+
48+
Args:
49+
model: Ollama embedding model name. Default: nomic-embed-text.
50+
base_url: Ollama server URL. Default: http://localhost:11434.
51+
Override via OLLAMA_BASE_URL env var.
52+
"""
53+
54+
def __init__(
55+
self,
56+
model: str = _DEFAULT_MODEL,
57+
base_url: str | None = None,
58+
) -> None:
59+
resolved_url = (
60+
base_url
61+
or os.environ.get("OLLAMA_BASE_URL")
62+
or _DEFAULT_BASE_URL
63+
)
64+
# Normalise: strip trailing slash, ensure no /v1 suffix
65+
resolved_url = resolved_url.rstrip("/")
66+
if resolved_url.endswith("/v1"):
67+
resolved_url = resolved_url[:-3]
68+
69+
self._model = model
70+
self._base_url = resolved_url
71+
self._openai_base_url = f"{resolved_url}/v1"
72+
self._dims: int | None = _KNOWN_DIMS.get(model)
73+
self._client: object | None = None
74+
75+
@property
76+
def dimensions(self) -> int:
77+
if self._dims is None:
78+
# Probe synchronously on first access if unknown
79+
import asyncio as _asyncio
80+
try:
81+
loop = _asyncio.get_event_loop()
82+
if loop.is_running():
83+
# Can't block — return a reasonable default
84+
return 768
85+
self._dims = loop.run_until_complete(self._probe_dimensions())
86+
except Exception:
87+
self._dims = 768
88+
return self._dims
89+
90+
async def _probe_dimensions(self) -> int:
91+
"""Embed a single token to discover the model's output dimension."""
92+
vecs = await self.embed(["probe"])
93+
return len(vecs[0]) if vecs else 768
94+
95+
def _get_client(self) -> object:
96+
if self._client is None:
97+
import openai # type: ignore[import-untyped]
98+
self._client = openai.OpenAI(
99+
api_key="ollama", # Ollama ignores the key
100+
base_url=self._openai_base_url,
101+
timeout=30.0,
102+
)
103+
return self._client
104+
105+
async def embed(self, texts: list[str]) -> list[list[float]]:
106+
"""Embed a batch of texts using the local Ollama embedding model.
107+
108+
Args:
109+
texts: Non-empty list of strings to embed.
110+
111+
Returns:
112+
List of L2-normalized float vectors.
113+
114+
Raises:
115+
RuntimeError: If Ollama is not reachable or the model is not pulled.
116+
"""
117+
if not texts:
118+
return []
119+
120+
model = self._model
121+
122+
def _embed_sync() -> list[list[float]]:
123+
import openai # type: ignore[import-untyped]
124+
125+
client = self._get_client()
126+
try:
127+
response = client.embeddings.create(model=model, input=texts) # type: ignore[union-attr]
128+
except openai.APIConnectionError as exc:
129+
raise RuntimeError(
130+
f"Cannot reach Ollama at {self._base_url}. "
131+
"Is it running? Start with: ollama serve"
132+
) from exc
133+
except openai.APIStatusError as exc:
134+
if exc.status_code == 404:
135+
raise RuntimeError(
136+
f"Ollama model '{model}' not found. "
137+
f"Pull it first: ollama pull {model}"
138+
) from exc
139+
raise RuntimeError(f"Ollama embedding error: {exc}") from exc
140+
141+
raw_vectors = [list(item.embedding) for item in response.data]
142+
result = [_l2_normalize(v) for v in raw_vectors]
143+
144+
# Cache discovered dimensions
145+
if self._dims is None and result:
146+
self._dims = len(result[0])
147+
148+
return result
149+
150+
return await asyncio.to_thread(_embed_sync)
151+
152+
153+
def _l2_normalize(vec: list[float]) -> list[float]:
154+
norm = math.sqrt(sum(x * x for x in vec))
155+
if norm == 0.0:
156+
norm = 1.0
157+
return [x / norm for x in vec]

0 commit comments

Comments
 (0)