Coverage for src/local_deep_research/embeddings/providers/implementations/ollama.py: 96%
89 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""Ollama embedding provider."""
3import weakref
4from typing import Any, Dict, List, Optional
6from langchain_ollama import OllamaEmbeddings
7from langchain_core.embeddings import Embeddings
9from ....config.thread_settings import get_setting_from_snapshot
10from ....security.secure_logging import logger
11from ....utilities.llm_utils import (
12 _close_inner_ollama_clients,
13 get_ollama_base_url,
14)
15from ..base import BaseEmbeddingProvider, Exposure
16from ....security import redact_url_for_log, safe_get, safe_post
17from ....security.log_sanitizer import scrub_error
20class OllamaEmbeddingsProvider(BaseEmbeddingProvider):
21 """
22 Ollama embedding provider.
24 Uses Ollama API for local embedding models.
25 No API key required, runs locally.
26 """
28 provider_name = "Ollama"
29 provider_key = "OLLAMA"
30 requires_api_key = False
31 supports_local = True
32 egress_exposure = Exposure.CONTAINED
33 default_model = "nomic-embed-text" # type: ignore[assignment]
35 @classmethod
36 def create_embeddings(
37 cls,
38 model: Optional[str] = None,
39 settings_snapshot: Optional[Dict[str, Any]] = None,
40 **kwargs,
41 ) -> Embeddings:
42 """
43 Create Ollama embeddings instance.
45 Args:
46 model: Model name (defaults to nomic-embed-text)
47 settings_snapshot: Optional settings snapshot
48 **kwargs: Additional parameters (base_url, etc.)
50 Returns:
51 OllamaEmbeddings instance
52 """
53 # Get model from settings if not specified
54 if model is None:
55 model = get_setting_from_snapshot(
56 "embeddings.ollama.model",
57 default=cls.default_model,
58 settings_snapshot=settings_snapshot,
59 )
61 # Get Ollama URL
62 base_url = kwargs.get("base_url")
63 if base_url is None:
64 base_url = get_ollama_base_url(settings_snapshot)
66 # Without an explicit num_ctx, Ollama uses the model's modelfile
67 # default (often 2048). Inputs longer than that return HTTP 500
68 # ("input length exceeds the context length") rather than being
69 # truncated, which aborts indexing mid-batch.
70 num_ctx = get_setting_from_snapshot(
71 "embeddings.ollama.num_ctx",
72 default=8192,
73 settings_snapshot=settings_snapshot,
74 )
76 logger.info(
77 f"Creating OllamaEmbeddings with model={model}, "
78 f"base_url={redact_url_for_log(base_url)}, num_ctx={num_ctx}"
79 )
81 ollama_kwargs: Dict[str, Any] = {
82 "model": model,
83 "base_url": base_url,
84 }
85 if num_ctx:
86 ollama_kwargs["num_ctx"] = int(num_ctx)
88 instance = OllamaEmbeddings(**ollama_kwargs)
90 # Safety net for callers that bypass LocalEmbeddingManager (e.g.,
91 # the programmatic-API examples in examples/api_usage, direct
92 # constructions in test fixtures). The manager-driven explicit
93 # close remains the load-bearing primary path; this finalizer
94 # only fires when the instance is GC'd without an explicit
95 # close. We pass the inner sync/async ``ollama.Client`` objects
96 # rather than ``instance`` itself — a strong reference back to
97 # the wrapping instance would defeat the finalizer's purpose by
98 # keeping the instance alive forever.
99 try:
100 weakref.finalize(
101 instance,
102 _close_inner_ollama_clients,
103 instance._client,
104 instance._async_client,
105 )
106 except AttributeError:
107 # Future langchain_ollama versions may reshape the private
108 # attrs; don't crash the factory if the introspection misses.
109 logger.debug(
110 "OllamaEmbeddings shape changed — finalizer not registered"
111 )
113 return instance
115 @classmethod
116 def is_available(
117 cls, settings_snapshot: Optional[Dict[str, Any]] = None
118 ) -> bool:
119 """Check if Ollama is available."""
120 try:
121 import requests
123 # Get Ollama URL
124 base_url = get_ollama_base_url(settings_snapshot)
126 # Check if Ollama is running
127 try:
128 response = safe_get(
129 f"{base_url}/api/tags",
130 timeout=3,
131 allow_localhost=True,
132 allow_private_ips=True,
133 )
134 return response.status_code == 200
135 except requests.exceptions.RequestException:
136 return False
138 except Exception as e:
139 safe_msg = scrub_error(e)
140 logger.exception(
141 f"Error checking Ollama availability ({type(e).__name__}): {safe_msg}"
142 )
143 return False
145 @classmethod
146 def _get_model_capabilities(
147 cls, base_url: str, model_name: str
148 ) -> Optional[List[str]]:
149 """Query Ollama /api/show for a model's capabilities.
151 Returns the capabilities list (e.g. ["embedding"]) or None on failure.
152 """
153 try:
154 response = safe_post(
155 f"{base_url}/api/show",
156 json={"model": model_name},
157 timeout=5,
158 allow_localhost=True,
159 allow_private_ips=True,
160 )
161 if response.status_code == 200:
162 return response.json().get("capabilities") # type: ignore[no-any-return]
163 except Exception:
164 logger.debug(f"Could not fetch capabilities for {model_name}")
165 return None
167 @classmethod
168 def is_embedding_model(
169 cls,
170 model: str,
171 settings_snapshot: Optional[Dict[str, Any]] = None,
172 ) -> Optional[bool]:
173 """Check whether an Ollama model supports embeddings.
175 Uses the /api/show capabilities field. Returns ``None`` when the
176 capability list isn't available (older Ollama servers) — the
177 provider doesn't guess from the model name. Callers must treat
178 ``None`` as "unknown", not as "no", so models stay listed even
179 when their capability can't be confirmed.
180 """
181 base_url = get_ollama_base_url(settings_snapshot)
182 caps = cls._get_model_capabilities(base_url, model)
184 # No name-based fallback on purpose — see method docstring.
185 if caps is None:
186 return None
187 return "embedding" in caps
189 @classmethod
190 def get_available_models(
191 cls, settings_snapshot: Optional[Dict[str, Any]] = None
192 ) -> List[Dict[str, Any]]:
193 """Get all Ollama models, tagged when /api/show reports support.
195 No filtering on the model list itself — every model the Ollama
196 server reports is returned. We only *tag* entries with
197 ``is_embedding`` when ``/api/show`` exposes a real capabilities
198 list (so the UI can sort them); we don't guess from the model
199 name. Older Ollama servers without capabilities → models are
200 returned untagged and the user decides.
201 """
202 from ....utilities.llm_utils import fetch_ollama_models
204 base_url = get_ollama_base_url(settings_snapshot)
205 # fetch_ollama_models returns every installed model. We pass it
206 # through unfiltered — no name heuristic, no exclusions.
207 all_models = fetch_ollama_models(base_url, timeout=3.0)
209 if not all_models:
210 return []
212 embedding_models: List[Dict[str, Any]] = []
213 untagged_models: List[Dict[str, Any]] = []
214 other_models: List[Dict[str, Any]] = []
216 for model in all_models:
217 model_name = model["value"]
218 caps = cls._get_model_capabilities(base_url, model_name)
220 entry: Dict[str, Any] = dict(model)
221 if caps is None:
222 # No capability signal from the server → don't guess.
223 # Keep the model in the list so the user can still
224 # select it.
225 untagged_models.append(entry)
226 continue
228 # /api/show capabilities is an API-driven signal (not a
229 # name match), so it's safe to use for the flag.
230 is_embed = "embedding" in caps
231 entry["is_embedding"] = is_embed
232 if is_embed:
233 embedding_models.append(entry)
234 else:
235 other_models.append(entry)
237 logger.info(
238 "Found {} embedding-capable, {} non-embedding, and {} "
239 "untagged models from Ollama",
240 len(embedding_models),
241 len(other_models),
242 len(untagged_models),
243 )
245 # Embedding-tagged first so they're the default pick; untagged
246 # next (capability unknown — user decides); then explicit
247 # non-embedding. Nothing is dropped.
248 return embedding_models + untagged_models + other_models