Skip to main content

lunaris_extract/
ollama.rs

1//! [`OllamaExtractor`] — thin wrapper over [`crate::LlmExtractor`] +
2//! `lunaris_llm::OllamaBackend`.
3//!
4//! Phase 12a duplication-delete: the load / HTTP / parse code that previously
5//! lived here has been collapsed into the shared [`lunaris_llm::OllamaBackend`]
6//! + [`crate::llm_extractor::LlmExtractor`] stack. The public API surface
7//!   (`OllamaExtractor`, `OllamaExtractorOpts`, the `Extractor` impl) is
8//!   byte-identical to v0.2 so downstream callers are unaffected.
9//!
10//! ## Constraint handling (preserved from v0.2)
11//!
12//! Ollama receives the JSON-schema translation of the extraction grammar
13//! as the `format` field via [`lunaris_llm::SchemaConstraint::JsonSchema`].
14//! GBNF is NOT passed through — Ollama (≤ 0.5) does not consume GBNF; the
15//! grammar is embedded in the prompt by the `OllamaBackend` only if
16//! `SchemaConstraint::Gbnf` is used, but for extraction we use the
17//! JSON-schema mode (existing behaviour). The `LlmExtractor` `opts.gbnf`
18//! field is left `None` here; the `OllamaBackend` will use
19//! [`lunaris_llm::SchemaConstraint::None`] which is the lighter-prompt
20//! path and matches the v0.2 behaviour (the JSON-schema was the
21//! enforcement mechanism, not GBNF).
22//!
23//! ## Failure modes (unchanged)
24//!
25//! Errors propagate from `OllamaBackend::generate` as
26//! `LunarisError::Storage(StorageError::Backend("ollama: ..."))`, same
27//! shape as v0.2.
28
29use std::sync::Arc;
30
31use async_trait::async_trait;
32use lunaris_core::LunarisError;
33use lunaris_llm::{OllamaBackend, OllamaBackendOpts};
34use std::time::Duration;
35use ulid::Ulid;
36
37use crate::Extractor;
38use crate::llm_extractor::{LlmExtractor, LlmExtractorOpts};
39use crate::types::{ChunkInput, RawExtractionBatch};
40
41/// Default Ollama endpoint.
42const DEFAULT_ENDPOINT: &str = "http://localhost:11434";
43
44/// Default Ollama model identifier (matches `ollama pull gemma3:4b`).
45const DEFAULT_MODEL: &str = "gemma3:4b";
46
47/// Default per-batch timeout per D-02 (matches candle backend).
48const DEFAULT_BATCH_TIMEOUT_MS: u64 = 150;
49
50/// Historical implicit per-chunk / HTTP-transport timeout (was hardcoded via
51/// `LlmExtractorOpts::default().per_chunk_timeout_ms`, 450ms — tuned for a
52/// local candle/Ollama instance). Kept as the default so existing callers see
53/// no behavior change; cloud-routed backends should override `timeout_ms`.
54const DEFAULT_TIMEOUT_MS: u64 = 450;
55
56/// Construction options for [`OllamaExtractor`].
57///
58/// `Default` resolves `endpoint` from `OLLAMA_URL` env (falls back to
59/// `http://localhost:11434`) and `model` from `OLLAMA_EXTRACT_MODEL` env
60/// (falls back to `gemma3:4b`).
61#[derive(Clone, Debug)]
62pub struct OllamaExtractorOpts {
63    pub endpoint: String,
64    pub model: String,
65    pub batch_timeout_ms: u64,
66    /// Per-chunk `generate()` timeout AND the backend's HTTP transport
67    /// timeout — both set to this same value. The tighter of the two always
68    /// won anyway (today that's the 450ms per-chunk wrap), so one dial is
69    /// simpler and behaviorally identical to the old hardcoded pairing.
70    /// Cloud-routed models need this in the tens of seconds, not milliseconds.
71    pub timeout_ms: u64,
72}
73
74impl Default for OllamaExtractorOpts {
75    fn default() -> Self {
76        Self {
77            endpoint: std::env::var("OLLAMA_URL").unwrap_or_else(|_| DEFAULT_ENDPOINT.to_string()),
78            model: std::env::var("OLLAMA_EXTRACT_MODEL")
79                .unwrap_or_else(|_| DEFAULT_MODEL.to_string()),
80            batch_timeout_ms: DEFAULT_BATCH_TIMEOUT_MS,
81            timeout_ms: DEFAULT_TIMEOUT_MS,
82        }
83    }
84}
85
86/// Ollama-backed extractor — thin wrapper over [`LlmExtractor`].
87///
88/// Holds a `LlmExtractor` built from `OllamaBackend` + extraction opts.
89/// All `Extractor::extract` calls delegate directly to the inner extractor.
90#[derive(Clone, Debug)]
91pub struct OllamaExtractor {
92    inner: LlmExtractor,
93}
94
95impl OllamaExtractor {
96    /// Construct a new Ollama-backed extractor. Builds an `OllamaBackend`
97    /// with `opts.timeout_ms` as both the HTTP transport timeout and the
98    /// per-chunk generation timeout, and wraps it in a `LlmExtractor`.
99    pub fn new(opts: OllamaExtractorOpts) -> Result<Self, LunarisError> {
100        let backend_opts = OllamaBackendOpts {
101            endpoint: opts.endpoint,
102            model: opts.model,
103            http_timeout: Duration::from_millis(opts.timeout_ms),
104        };
105        let backend =
106            Arc::new(OllamaBackend::new(backend_opts)?) as Arc<dyn lunaris_llm::LlmBackend>;
107        let extractor_opts = LlmExtractorOpts {
108            batch_timeout_ms: opts.batch_timeout_ms,
109            per_chunk_timeout_ms: opts.timeout_ms,
110            // Ollama uses JSON-schema structured output, not inline GBNF.
111            // The OllamaBackend receives SchemaConstraint::None here and
112            // enforces the extraction schema via the JSON-schema `format`
113            // field (OllamaBackend handles the format field translation).
114            // Leave gbnf = None so we don't double-embed the grammar.
115            gbnf: None,
116            ..LlmExtractorOpts::default()
117        };
118        Ok(Self { inner: LlmExtractor::with_opts(backend, extractor_opts) })
119    }
120}
121
122#[async_trait]
123impl Extractor for OllamaExtractor {
124    async fn extract(
125        &self,
126        episode_id: Ulid,
127        chunks: &[ChunkInput],
128    ) -> Result<RawExtractionBatch, LunarisError> {
129        self.inner.extract(episode_id, chunks).await
130    }
131
132    fn applies(&self) -> bool {
133        self.inner.applies()
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn opts_default_resolves_to_localhost_ollama() {
143        // Avoid env-var contamination from the test runner — read the actual
144        // resolved values and just check they're non-empty.
145        let opts = OllamaExtractorOpts::default();
146        assert!(!opts.endpoint.is_empty());
147        assert!(!opts.model.is_empty());
148        assert_eq!(opts.batch_timeout_ms, DEFAULT_BATCH_TIMEOUT_MS);
149        // Historical implicit default (LlmExtractorOpts::default().per_chunk_timeout_ms)
150        // preserved exactly so existing callers see no behavior change.
151        assert_eq!(opts.timeout_ms, DEFAULT_TIMEOUT_MS);
152    }
153
154    #[test]
155    fn extractor_construction_succeeds_with_defaults() {
156        let _e = OllamaExtractor::new(OllamaExtractorOpts::default()).expect("client builds");
157    }
158
159    #[test]
160    fn extractor_construction_succeeds_with_generous_cloud_timeout() {
161        // A cloud-routed model behind a shim (e.g. the LongMemEval eval
162        // harness's graph-pipeline prototype) needs far more than the
163        // local-model 450ms default on both the transport and per-chunk
164        // layers -- this just proves a large override builds cleanly.
165        let opts = OllamaExtractorOpts {
166            timeout_ms: 60_000,
167            batch_timeout_ms: 120_000,
168            ..OllamaExtractorOpts::default()
169        };
170        let _e = OllamaExtractor::new(opts).expect("client builds with a generous timeout");
171    }
172}