1use std::str::FromStr;
37use std::sync::Arc;
38use std::time::Duration;
39
40use async_trait::async_trait;
41use lunaris_core::{LunarisError, StorageError};
42use lunaris_llm::{CloudBackend, CloudBackendOpts, GenOpts, LlmBackend, SchemaConstraint};
43use ulid::Ulid;
44
45use crate::Extractor;
46use crate::types::{ChunkInput, Entity, EntityId, RawExtraction, RawExtractionBatch};
47use crate::validator::{TRANSIENT_SENTINEL_NAME, TRANSIENT_SENTINEL_TYPE};
48
49const DEFAULT_BATCH_TIMEOUT_MS: u64 = 150;
51
52const DEFAULT_MAX_RETRIES: u8 = 1;
54
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
58pub enum CloudProvider {
59 Anthropic,
60 OpenAI,
61 Gemini,
62 MiniMax,
63 OpenAiCompat,
68}
69
70impl FromStr for CloudProvider {
71 type Err = LunarisError;
72 fn from_str(s: &str) -> Result<Self, Self::Err> {
73 match s.trim().to_ascii_lowercase().as_str() {
74 "anthropic" | "claude" => Ok(Self::Anthropic),
75 "openai" | "gpt" => Ok(Self::OpenAI),
76 "gemini" | "google" => Ok(Self::Gemini),
77 "minimax" => Ok(Self::MiniMax),
78 "openai-compat" | "openai-compatible" => Ok(Self::OpenAiCompat),
79 other => Err(LunarisError::Storage(StorageError::Backend(format!(
80 "cloud-api: unknown provider {other:?} \
81 (expected anthropic|openai|gemini|minimax|openai-compat)"
82 )))),
83 }
84 }
85}
86
87impl CloudProvider {
88 fn default_model(self) -> &'static str {
89 match self {
90 Self::Anthropic => "claude-3-5-haiku-latest",
91 Self::OpenAI => "gpt-4o-mini",
92 Self::Gemini => "gemini-2.5-flash",
93 Self::MiniMax => "MiniMax-M3",
94 Self::OpenAiCompat => "",
98 }
99 }
100 fn api_key_env(self) -> &'static str {
101 match self {
102 Self::Anthropic => "ANTHROPIC_API_KEY",
103 Self::OpenAI => "OPENAI_API_KEY",
104 Self::Gemini => "GEMINI_API_KEY",
105 Self::MiniMax => "MINIMAX_API_KEY",
106 Self::OpenAiCompat => "LUNARIS_OPENAI_COMPAT_API_KEY",
107 }
108 }
109 fn model_env(self) -> &'static str {
110 match self {
111 Self::Anthropic => "ANTHROPIC_EXTRACT_MODEL",
112 Self::OpenAI => "OPENAI_EXTRACT_MODEL",
113 Self::Gemini => "GEMINI_EXTRACT_MODEL",
114 Self::MiniMax => "MINIMAX_EXTRACT_MODEL",
115 Self::OpenAiCompat => "OPENAI_COMPAT_EXTRACT_MODEL",
116 }
117 }
118}
119
120impl From<CloudProvider> for lunaris_llm::CloudProvider {
123 fn from(p: CloudProvider) -> Self {
124 match p {
125 CloudProvider::Anthropic => lunaris_llm::CloudProvider::Anthropic,
126 CloudProvider::OpenAI => lunaris_llm::CloudProvider::OpenAI,
127 CloudProvider::Gemini => lunaris_llm::CloudProvider::Gemini,
128 CloudProvider::MiniMax => lunaris_llm::CloudProvider::MiniMax,
129 CloudProvider::OpenAiCompat => lunaris_llm::CloudProvider::OpenAiCompat,
130 }
131 }
132}
133
134#[derive(Clone, Debug)]
141pub struct CloudApiExtractorOpts {
142 pub provider: CloudProvider,
143 pub model: String,
144 pub api_key: String,
145 pub batch_timeout_ms: u64,
146 pub max_retries: u8,
147 pub max_tokens: u32,
155 pub concurrency: usize,
159 pub base_url: Option<String>,
163}
164
165const DEFAULT_MAX_TOKENS: u32 = 512;
167
168const DEFAULT_EXTRACT_CONCURRENCY: usize = 4;
179
180async fn extract_chunks_buffered<Fut>(futs: Vec<Fut>, concurrency: usize) -> Vec<RawExtraction>
186where
187 Fut: std::future::Future<Output = RawExtraction>,
188{
189 use futures::StreamExt;
190 futures::stream::iter(futs).buffered(concurrency.max(1)).collect().await
191}
192
193impl Default for CloudApiExtractorOpts {
194 fn default() -> Self {
195 let provider = std::env::var("LUNARIS_EXTRACT_PROVIDER")
196 .ok()
197 .and_then(|s| CloudProvider::from_str(&s).ok())
198 .unwrap_or(CloudProvider::Anthropic);
199 let model = std::env::var(provider.model_env())
200 .unwrap_or_else(|_| provider.default_model().to_string());
201 let api_key = std::env::var(provider.api_key_env()).unwrap_or_default();
202 let base_url = match provider {
203 CloudProvider::OpenAiCompat => {
204 std::env::var("LUNARIS_OPENAI_COMPAT_BASE_URL").ok().filter(|s| !s.is_empty())
205 }
206 _ => None,
207 };
208 Self {
209 provider,
210 model,
211 api_key,
212 batch_timeout_ms: DEFAULT_BATCH_TIMEOUT_MS,
213 max_retries: DEFAULT_MAX_RETRIES,
214 max_tokens: DEFAULT_MAX_TOKENS,
215 concurrency: DEFAULT_EXTRACT_CONCURRENCY,
216 base_url,
217 }
218 }
219}
220
221#[derive(Clone)]
224pub struct CloudApiExtractor {
225 backend: Arc<CloudBackend>,
226 provider: CloudProvider,
227 batch_timeout_ms: u64,
228 gen_opts: GenOpts,
232 concurrency: usize,
234}
235
236impl std::fmt::Debug for CloudApiExtractor {
237 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
238 f.debug_struct("CloudApiExtractor")
241 .field("provider", &self.provider)
242 .field("model_id", &self.backend.model_id())
243 .field("batch_timeout_ms", &self.batch_timeout_ms)
244 .finish()
245 }
246}
247
248impl CloudApiExtractor {
249 pub fn new(opts: CloudApiExtractorOpts) -> Result<Self, LunarisError> {
251 if opts.provider == CloudProvider::OpenAiCompat && opts.model.trim().is_empty() {
256 return Err(LunarisError::Storage(StorageError::Backend(
257 "cloud-api: openai-compat model is empty — set OPENAI_COMPAT_EXTRACT_MODEL \
258 (or CloudApiExtractorOpts.model) to the model your server hosts"
259 .to_string(),
260 )));
261 }
262 let llm_provider = lunaris_llm::CloudProvider::from(opts.provider);
263 let backend_opts = CloudBackendOpts {
264 provider: llm_provider,
265 model: opts.model,
266 api_key: opts.api_key,
267 max_retries: opts.max_retries,
268 base_url: opts.base_url,
269 };
270 let backend = Arc::new(CloudBackend::new(backend_opts)?);
271 let gen_opts = GenOpts {
274 max_tokens: opts.max_tokens,
275 temperature: 0.0,
276 timeout: Duration::from_millis(opts.batch_timeout_ms),
277 };
278 Ok(Self {
279 backend,
280 provider: opts.provider,
281 batch_timeout_ms: opts.batch_timeout_ms,
282 gen_opts,
283 concurrency: opts.concurrency.max(1),
284 })
285 }
286
287 async fn extract_one_with_sentinel(&self, chunk: &ChunkInput) -> RawExtraction {
292 let prompt = crate::llm_extractor::build_prompt(chunk);
297 match self.backend.generate(&prompt, SchemaConstraint::None, self.gen_opts).await {
298 Ok(text) => crate::llm_extractor::parse_extraction_json_pub(&text, chunk.chunk_id),
299 Err(e) => {
300 tracing::warn!(
301 err = %e,
302 chunk_id = %chunk.chunk_id,
303 model_id = self.backend.model_id(),
304 "cloud-api retry exhausted; emitting transient-after-retry sentinel"
305 );
306 let err_text = e.to_string();
307 let sentinel = Entity {
308 id: EntityId::from_name_and_type(
309 TRANSIENT_SENTINEL_NAME,
310 TRANSIENT_SENTINEL_TYPE,
311 ),
312 name: TRANSIENT_SENTINEL_NAME.into(),
313 aliases: Vec::new(),
314 entity_type: TRANSIENT_SENTINEL_TYPE.into(),
315 confidence: 0.0,
316 valid_from_iso: format!("transient: {err_text}"),
317 valid_to_iso: None,
318 };
319 RawExtraction {
320 source_chunk_id: chunk.chunk_id,
321 entities: vec![sentinel],
322 relations: Vec::new(),
323 facts: Vec::new(),
324 }
325 }
326 }
327 }
328}
329
330#[async_trait]
331impl Extractor for CloudApiExtractor {
332 async fn extract(
333 &self,
334 _episode_id: Ulid,
335 chunks: &[ChunkInput],
336 ) -> Result<RawExtractionBatch, LunarisError> {
337 if chunks.is_empty() {
338 return Ok(RawExtractionBatch::default());
339 }
340
341 let batch_timeout = Duration::from_millis(self.batch_timeout_ms);
344 let chunks_owned: Vec<ChunkInput> = chunks.to_vec();
345 let this = self.clone();
346 let concurrency = self.concurrency;
347 let batch_fut = async move {
348 let futs: Vec<_> =
349 chunks_owned.iter().map(|c| this.extract_one_with_sentinel(c)).collect();
350 let by_chunk = extract_chunks_buffered(futs, concurrency).await;
351 RawExtractionBatch { by_chunk }
352 };
353
354 match tokio::time::timeout(batch_timeout, batch_fut).await {
355 Ok(b) => Ok(b),
356 Err(_elapsed) => {
357 tracing::warn!(
358 batch_size = chunks.len(),
359 timeout_ms = self.batch_timeout_ms,
360 "cloud-api batch timeout; falling back to per-chunk"
361 );
362 let futs: Vec<_> =
363 chunks.iter().map(|c| self.extract_one_with_sentinel(c)).collect();
364 let by_chunk = extract_chunks_buffered(futs, self.concurrency).await;
365 Ok(RawExtractionBatch { by_chunk })
366 }
367 }
368 }
369
370 fn applies(&self) -> bool {
371 self.backend.applies()
372 }
373}
374
375#[cfg(test)]
376mod tests {
377 use super::*;
378
379 #[tokio::test]
380 async fn extract_chunks_buffered_overlaps_up_to_concurrency_and_preserves_order() {
381 use std::sync::Arc;
390 use std::sync::atomic::{AtomicUsize, Ordering};
391 let chunks: Vec<ChunkInput> = (0..8)
392 .map(|i| ChunkInput {
393 chunk_id: Ulid::new(),
394 text: format!("chunk {i}"),
395 heading_path: vec![],
396 reference_time_iso: None,
397 })
398 .collect();
399 let in_flight = Arc::new(AtomicUsize::new(0));
400 let max_seen = Arc::new(AtomicUsize::new(0));
401 let futs: Vec<_> = chunks
402 .iter()
403 .map(|c| {
404 let id = c.chunk_id;
405 let in_flight = Arc::clone(&in_flight);
406 let max_seen = Arc::clone(&max_seen);
407 async move {
408 let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
409 max_seen.fetch_max(now, Ordering::SeqCst);
410 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
411 in_flight.fetch_sub(1, Ordering::SeqCst);
412 RawExtraction { source_chunk_id: id, ..Default::default() }
413 }
414 })
415 .collect();
416 let out = extract_chunks_buffered(futs, 4).await;
417 assert_eq!(out.len(), chunks.len());
418 for (c, r) in chunks.iter().zip(&out) {
419 assert_eq!(c.chunk_id, r.source_chunk_id, "buffered output must preserve input order");
420 }
421 let peak = max_seen.load(Ordering::SeqCst);
422 assert!(peak >= 3, "expected >=3 overlapping extractions at concurrency 4, saw {peak}");
423 }
424
425 #[tokio::test]
426 async fn extract_chunks_buffered_concurrency_1_stays_strictly_serial() {
427 use std::sync::Arc;
432 use std::sync::atomic::{AtomicUsize, Ordering};
433 let chunks: Vec<ChunkInput> = (0..4)
434 .map(|i| ChunkInput {
435 chunk_id: Ulid::new(),
436 text: format!("chunk {i}"),
437 heading_path: vec![],
438 reference_time_iso: None,
439 })
440 .collect();
441 let in_flight = Arc::new(AtomicUsize::new(0));
442 let max_seen = Arc::new(AtomicUsize::new(0));
443 let futs: Vec<_> = chunks
444 .iter()
445 .map(|c| {
446 let id = c.chunk_id;
447 let in_flight = Arc::clone(&in_flight);
448 let max_seen = Arc::clone(&max_seen);
449 async move {
450 let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
451 max_seen.fetch_max(now, Ordering::SeqCst);
452 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
453 in_flight.fetch_sub(1, Ordering::SeqCst);
454 RawExtraction { source_chunk_id: id, ..Default::default() }
455 }
456 })
457 .collect();
458 let out = extract_chunks_buffered(futs, 1).await;
459 assert_eq!(out.len(), chunks.len());
460 assert_eq!(max_seen.load(Ordering::SeqCst), 1, "concurrency=1 must never overlap");
461 }
462
463 #[test]
464 fn cloud_provider_from_str_parses_all_three() {
465 assert_eq!(CloudProvider::from_str("anthropic").unwrap(), CloudProvider::Anthropic);
466 assert_eq!(CloudProvider::from_str("Claude").unwrap(), CloudProvider::Anthropic);
467 assert_eq!(CloudProvider::from_str("openai").unwrap(), CloudProvider::OpenAI);
468 assert_eq!(CloudProvider::from_str("GPT").unwrap(), CloudProvider::OpenAI);
469 assert_eq!(CloudProvider::from_str("gemini").unwrap(), CloudProvider::Gemini);
470 assert_eq!(CloudProvider::from_str("Google").unwrap(), CloudProvider::Gemini);
471 }
472
473 #[test]
474 fn cloud_provider_from_str_rejects_unknown() {
475 let err = CloudProvider::from_str("deepseek").unwrap_err();
476 assert!(err.to_string().contains("unknown provider"));
477 }
478
479 #[test]
480 fn empty_api_key_rejected() {
481 let opts = CloudApiExtractorOpts {
482 provider: CloudProvider::Anthropic,
483 model: "claude-3-5-haiku-latest".into(),
484 api_key: "".into(),
485 batch_timeout_ms: 150,
486 max_retries: 1,
487 max_tokens: 512,
488 concurrency: 1,
489 base_url: None,
490 };
491 let err = CloudApiExtractor::new(opts).expect_err("empty key must error");
492 let msg = err.to_string();
493 assert!(msg.contains("api_key is empty"), "got: {msg}");
494 assert!(msg.contains("ANTHROPIC_API_KEY"), "got: {msg}");
495 }
496
497 #[test]
498 fn debug_impl_redacts_api_key() {
499 let opts = CloudApiExtractorOpts {
501 provider: CloudProvider::Anthropic,
502 model: "claude-3-5-haiku-latest".into(),
503 api_key: "sk-ant-SECRET-KEY-ABCDEF".into(),
504 batch_timeout_ms: 150,
505 max_retries: 1,
506 max_tokens: 512,
507 concurrency: 1,
508 base_url: None,
509 };
510 let extractor = CloudApiExtractor::new(opts).unwrap();
511 let dbg = format!("{extractor:?}");
512 assert!(!dbg.contains("SECRET"), "Debug must redact api_key, got: {dbg}");
513 assert!(!dbg.contains("sk-ant"), "Debug must redact api_key, got: {dbg}");
514 }
515
516 #[test]
517 fn cloud_provider_bridge_maps_all_three() {
518 assert!(matches!(
519 lunaris_llm::CloudProvider::from(CloudProvider::Anthropic),
520 lunaris_llm::CloudProvider::Anthropic
521 ));
522 assert!(matches!(
523 lunaris_llm::CloudProvider::from(CloudProvider::OpenAI),
524 lunaris_llm::CloudProvider::OpenAI
525 ));
526 assert!(matches!(
527 lunaris_llm::CloudProvider::from(CloudProvider::Gemini),
528 lunaris_llm::CloudProvider::Gemini
529 ));
530 }
531
532 #[test]
533 fn cloud_provider_from_str_parses_minimax() {
534 assert_eq!(CloudProvider::from_str("minimax").unwrap(), CloudProvider::MiniMax);
535 assert_eq!(CloudProvider::from_str("MiniMax").unwrap(), CloudProvider::MiniMax);
536 }
537
538 #[test]
539 fn minimax_default_model_and_envs() {
540 assert_eq!(CloudProvider::MiniMax.default_model(), "MiniMax-M3");
541 assert_eq!(CloudProvider::MiniMax.api_key_env(), "MINIMAX_API_KEY");
542 assert_eq!(CloudProvider::MiniMax.model_env(), "MINIMAX_EXTRACT_MODEL");
543 }
544
545 #[test]
546 fn cloud_provider_bridge_maps_minimax() {
547 assert!(matches!(
548 lunaris_llm::CloudProvider::from(CloudProvider::MiniMax),
549 lunaris_llm::CloudProvider::MiniMax
550 ));
551 }
552
553 #[test]
554 fn openai_compat_parses_and_constructs_keyless() {
555 assert_eq!(CloudProvider::from_str("openai-compat").unwrap(), CloudProvider::OpenAiCompat);
558 assert_eq!(
559 CloudProvider::from_str("openai-compatible").unwrap(),
560 CloudProvider::OpenAiCompat
561 );
562 assert!(matches!(
563 lunaris_llm::CloudProvider::from(CloudProvider::OpenAiCompat),
564 lunaris_llm::CloudProvider::OpenAiCompat
565 ));
566 let e = CloudApiExtractor::new(CloudApiExtractorOpts {
567 provider: CloudProvider::OpenAiCompat,
568 model: "qwen3:4b".into(),
569 api_key: String::new(),
570 base_url: Some("http://localhost:11434/v1".into()),
571 ..CloudApiExtractorOpts::default()
572 })
573 .expect("keyless openai-compat extractor must construct");
574 let dbg = format!("{e:?}");
575 assert!(dbg.contains("OpenAiCompat"), "got: {dbg}");
576 }
577
578 #[test]
579 fn openai_compat_requires_model_and_base_url() {
580 let err = CloudApiExtractor::new(CloudApiExtractorOpts {
581 provider: CloudProvider::OpenAiCompat,
582 model: String::new(),
583 api_key: String::new(),
584 base_url: Some("http://localhost:11434/v1".into()),
585 ..CloudApiExtractorOpts::default()
586 })
587 .expect_err("empty model must fail fast");
588 assert!(err.to_string().contains("OPENAI_COMPAT_EXTRACT_MODEL"), "got: {err}");
589
590 let err = CloudApiExtractor::new(CloudApiExtractorOpts {
591 provider: CloudProvider::OpenAiCompat,
592 model: "qwen3:4b".into(),
593 api_key: String::new(),
594 base_url: None,
595 ..CloudApiExtractorOpts::default()
596 })
597 .expect_err("missing base_url must fail fast");
598 assert!(err.to_string().contains("LUNARIS_OPENAI_COMPAT_BASE_URL"), "got: {err}");
599 }
600
601 #[test]
602 fn max_tokens_defaults_to_512_and_is_configurable() {
603 let default_opts = CloudApiExtractorOpts::default();
610 assert_eq!(default_opts.max_tokens, 512);
611 let opts = CloudApiExtractorOpts {
612 provider: CloudProvider::MiniMax,
613 model: "MiniMax-M3".into(),
614 api_key: "dummy".into(),
615 batch_timeout_ms: 150,
616 max_retries: 1,
617 max_tokens: 2048,
618 concurrency: 4,
619 base_url: None,
620 };
621 let _extractor =
622 CloudApiExtractor::new(opts).expect("client builds with custom max_tokens");
623 }
624}