Coverage for src/local_deep_research/embeddings/providers/implementations/ollama.py: 97%
99 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""Ollama embedding provider."""
3import weakref
4from typing import Any, Dict, List, Optional
6from httpx import Timeout
7from langchain_ollama import OllamaEmbeddings
8from langchain_core.embeddings import Embeddings
10from ....config.thread_settings import get_setting_from_snapshot
11from ....security.secure_logging import logger
12from ....security.ssrf_validator import assert_base_url_safe
13from ....utilities.llm_utils import (
14 _close_inner_ollama_clients,
15 get_ollama_base_url,
16)
17from ..base import BaseEmbeddingProvider, Exposure
18from ....security import redact_url_for_log, safe_get, safe_post
19from ....security.log_sanitizer import scrub_error
22class OllamaEmbeddingsProvider(BaseEmbeddingProvider):
23 """
24 Ollama embedding provider.
26 Uses Ollama API for local embedding models.
27 No API key required, runs locally.
28 """
30 provider_name = "Ollama"
31 provider_key = "OLLAMA"
32 requires_api_key = False
33 supports_local = True
34 egress_exposure = Exposure.CONTAINED
35 default_model = "nomic-embed-text" # type: ignore[assignment]
37 @classmethod
38 def create_embeddings(
39 cls,
40 model: Optional[str] = None,
41 settings_snapshot: Optional[Dict[str, Any]] = None,
42 **kwargs,
43 ) -> Embeddings:
44 """
45 Create Ollama embeddings instance.
47 Args:
48 model: Model name (defaults to nomic-embed-text)
49 settings_snapshot: Optional settings snapshot
50 **kwargs: Additional parameters (base_url, etc.)
52 Returns:
53 OllamaEmbeddings instance
54 """
55 # Get model from settings if not specified
56 if model is None:
57 model = get_setting_from_snapshot(
58 "embeddings.ollama.model",
59 default=cls.default_model,
60 settings_snapshot=settings_snapshot,
61 )
63 # Get Ollama URL
64 base_url = kwargs.get("base_url")
65 if base_url is None:
66 base_url = get_ollama_base_url(settings_snapshot)
68 # SSRF guard: the availability/discovery probes above route through
69 # safe_get/safe_post, but the embedding traffic itself flows through
70 # the OllamaEmbeddings -> ollama.Client -> httpx path, which never
71 # consults the SSRF policy. Without this, a user-configured
72 # ``embeddings.ollama.url`` pointing at a cloud-metadata endpoint
73 # (169.254.169.254 etc.) would exfiltrate to it on every embed call.
74 # Mirrors the LLM-side Ollama provider, which validates the same way.
75 # allow_localhost/allow_private_ips (inside assert_base_url_safe) keep
76 # the default localhost:11434 and RFC1918 Ollama hosts working; the
77 # ALWAYS_BLOCKED metadata IPs are refused, and so is a hostname the
78 # local resolver cannot resolve at construction time (for example a
79 # Docker service that is not up yet), with a distinct "could not
80 # resolve host" message.
81 base_url = assert_base_url_safe(
82 base_url, setting_key="embeddings.ollama.url"
83 )
85 # Without an explicit num_ctx, Ollama uses the model's modelfile
86 # default (often 2048). Inputs longer than that return HTTP 500
87 # ("input length exceeds the context length") rather than being
88 # truncated, which aborts indexing mid-batch.
89 num_ctx = get_setting_from_snapshot(
90 "embeddings.ollama.num_ctx",
91 default=8192,
92 settings_snapshot=settings_snapshot,
93 )
95 logger.info(
96 f"Creating OllamaEmbeddings with model={model}, "
97 f"base_url={redact_url_for_log(base_url)}, num_ctx={num_ctx}"
98 )
100 ollama_kwargs: Dict[str, Any] = {
101 "model": model,
102 "base_url": base_url,
103 }
104 if num_ctx:
105 ollama_kwargs["num_ctx"] = int(num_ctx)
107 # Same bound as the LLM side: langchain_ollama otherwise builds
108 # its httpx clients with ``timeout=None``, so a stalled embedding
109 # request would hold an indexing worker forever. client_kwargs
110 # reaches both the sync and the async ``ollama.Client``. Imported
111 # lazily — the ``llm.providers`` package runs provider
112 # auto-discovery at import time.
113 from ....llm.providers._helpers import (
114 build_httpx_timeout,
115 resolve_request_timeout,
116 )
117 from ....llm.providers.implementations.ollama import OllamaProvider
119 # follow_redirects=False: ollama.Client defaults it to True, so an
120 # allowed host could otherwise bounce every embed request (and the
121 # bearer header) to a metadata endpoint that the
122 # assert_base_url_safe() check above refuses when named directly.
123 client_kwargs: Dict[str, Any] = {
124 "timeout": build_httpx_timeout(
125 resolve_request_timeout(settings_snapshot), Timeout
126 ),
127 "follow_redirects": False,
128 }
130 # An authenticated Ollama instance sits behind the same proxy for
131 # embeddings as for chat, and there is a single ``llm.ollama.api_key``
132 # for both. OllamaEmbeddings has no ``headers`` field either, so the
133 # header has to travel via client_kwargs exactly as it does on the
134 # chat path — without this, chat authenticates and indexing 401s.
135 auth_headers = OllamaProvider.build_bearer_header(
136 settings_snapshot=settings_snapshot
137 )
138 if auth_headers:
139 client_kwargs["headers"] = auth_headers
141 ollama_kwargs["client_kwargs"] = client_kwargs
143 instance = OllamaEmbeddings(**ollama_kwargs)
145 # Safety net for callers that bypass LocalEmbeddingManager (e.g.,
146 # the programmatic-API examples in examples/api_usage, direct
147 # constructions in test fixtures). The manager-driven explicit
148 # close remains the load-bearing primary path; this finalizer
149 # only fires when the instance is GC'd without an explicit
150 # close. We pass the inner sync/async ``ollama.Client`` objects
151 # rather than ``instance`` itself — a strong reference back to
152 # the wrapping instance would defeat the finalizer's purpose by
153 # keeping the instance alive forever.
154 try:
155 weakref.finalize(
156 instance,
157 _close_inner_ollama_clients,
158 instance._client,
159 instance._async_client,
160 )
161 except AttributeError:
162 # Future langchain_ollama versions may reshape the private
163 # attrs; don't crash the factory if the introspection misses.
164 logger.debug(
165 "OllamaEmbeddings shape changed — finalizer not registered"
166 )
168 return instance
170 @classmethod
171 def is_available(
172 cls, settings_snapshot: Optional[Dict[str, Any]] = None
173 ) -> bool:
174 """Check if Ollama is available."""
175 try:
176 import requests
178 # Get Ollama URL
179 base_url = get_ollama_base_url(settings_snapshot)
181 # Check if Ollama is running
182 try:
183 response = safe_get(
184 f"{base_url}/api/tags",
185 timeout=3,
186 allow_localhost=True,
187 allow_private_ips=True,
188 )
189 return response.status_code == 200
190 except requests.exceptions.RequestException:
191 return False
193 except Exception as e:
194 safe_msg = scrub_error(e)
195 logger.exception(
196 f"Error checking Ollama availability ({type(e).__name__}): {safe_msg}"
197 )
198 return False
200 @classmethod
201 def _get_model_capabilities(
202 cls, base_url: str, model_name: str
203 ) -> Optional[List[str]]:
204 """Query Ollama /api/show for a model's capabilities.
206 Returns the capabilities list (e.g. ["embedding"]) or None on failure.
207 """
208 try:
209 response = safe_post(
210 f"{base_url}/api/show",
211 json={"model": model_name},
212 timeout=5,
213 allow_localhost=True,
214 allow_private_ips=True,
215 )
216 if response.status_code == 200:
217 return response.json().get("capabilities") # type: ignore[no-any-return]
218 except Exception:
219 logger.debug(f"Could not fetch capabilities for {model_name}")
220 return None
222 @classmethod
223 def is_embedding_model(
224 cls,
225 model: str,
226 settings_snapshot: Optional[Dict[str, Any]] = None,
227 ) -> Optional[bool]:
228 """Check whether an Ollama model supports embeddings.
230 Uses the /api/show capabilities field. Returns ``None`` when the
231 capability list isn't available (older Ollama servers) — the
232 provider doesn't guess from the model name. Callers must treat
233 ``None`` as "unknown", not as "no", so models stay listed even
234 when their capability can't be confirmed.
235 """
236 base_url = get_ollama_base_url(settings_snapshot)
237 caps = cls._get_model_capabilities(base_url, model)
239 # No name-based fallback on purpose — see method docstring.
240 if caps is None:
241 return None
242 return "embedding" in caps
244 @classmethod
245 def get_available_models(
246 cls, settings_snapshot: Optional[Dict[str, Any]] = None
247 ) -> List[Dict[str, Any]]:
248 """Get all Ollama models, tagged when /api/show reports support.
250 No filtering on the model list itself — every model the Ollama
251 server reports is returned. We only *tag* entries with
252 ``is_embedding`` when ``/api/show`` exposes a real capabilities
253 list (so the UI can sort them); we don't guess from the model
254 name. Older Ollama servers without capabilities → models are
255 returned untagged and the user decides.
256 """
257 from ....utilities.llm_utils import (
258 OLLAMA_MODEL_LIST_TIMEOUT,
259 fetch_ollama_models,
260 )
262 base_url = get_ollama_base_url(settings_snapshot)
263 # fetch_ollama_models returns every installed model. We pass it
264 # through unfiltered — no name heuristic, no exclusions. Same generous
265 # model-list timeout as the LLM provider (see OLLAMA_MODEL_LIST_TIMEOUT).
266 all_models = fetch_ollama_models(
267 base_url, timeout=OLLAMA_MODEL_LIST_TIMEOUT
268 )
270 if not all_models:
271 return []
273 embedding_models: List[Dict[str, Any]] = []
274 untagged_models: List[Dict[str, Any]] = []
275 other_models: List[Dict[str, Any]] = []
277 for model in all_models:
278 model_name = model["value"]
279 caps = cls._get_model_capabilities(base_url, model_name)
281 entry: Dict[str, Any] = dict(model)
282 if caps is None:
283 # No capability signal from the server → don't guess.
284 # Keep the model in the list so the user can still
285 # select it.
286 untagged_models.append(entry)
287 continue
289 # /api/show capabilities is an API-driven signal (not a
290 # name match), so it's safe to use for the flag.
291 is_embed = "embedding" in caps
292 entry["is_embedding"] = is_embed
293 if is_embed:
294 embedding_models.append(entry)
295 else:
296 other_models.append(entry)
298 logger.info(
299 "Found {} embedding-capable, {} non-embedding, and {} "
300 "untagged models from Ollama",
301 len(embedding_models),
302 len(other_models),
303 len(untagged_models),
304 )
306 # Embedding-tagged first so they're the default pick; untagged
307 # next (capability unknown — user decides); then explicit
308 # non-embedding. Nothing is dropped.
309 return embedding_models + untagged_models + other_models