Skip to main content

sqlite_graphrag/extract/
composite_backend.rs

1//! Composite extraction backend (v1.0.75 — G21 orchestration)
2//!
3//! Runs multiple backends in parallel and merges their outputs.
4//! Used when the user requests `--extraction-backend both`.
5
6use super::{
7    BackendHealth, BackendKind, ExtractionBackend, ExtractionHints, ExtractionOutput, SharedBackend,
8};
9use crate::errors::AppError;
10use async_trait::async_trait;
11use std::time::Instant;
12
13/// Composite backend.
14pub struct CompositeBackend {
15    backends: Vec<SharedBackend>,
16}
17
18impl CompositeBackend {
19    /// Create a new instance.
20    pub fn new(backends: Vec<SharedBackend>) -> Self {
21        Self { backends }
22    }
23}
24
25#[async_trait]
26impl ExtractionBackend for CompositeBackend {
27    fn kind(&self) -> BackendKind {
28        BackendKind::Composite
29    }
30
31    fn model_name(&self) -> String {
32        self.backends
33            .iter()
34            .map(|b| b.model_name())
35            .collect::<Vec<_>>()
36            .join("+")
37    }
38
39    async fn extract(
40        &self,
41        content: &str,
42        hints: &ExtractionHints,
43    ) -> Result<ExtractionOutput, AppError> {
44        let start = Instant::now();
45        let mut merged = ExtractionOutput {
46            backend: self.kind().as_str().to_string(),
47            ..Default::default()
48        };
49        let mut first_embedding: Option<Vec<f32>> = None;
50        let mut any_error: Option<AppError> = None;
51
52        for backend in &self.backends {
53            match backend.extract(content, hints).await {
54                Ok(out) => {
55                    for entity in out.entities {
56                        if !merged.entities.iter().any(|e| e.name == entity.name) {
57                            merged.entities.push(entity);
58                        }
59                    }
60                    for rel in out.relationships {
61                        let exists = merged.relationships.iter().any(|r| {
62                            r.source == rel.source
63                                && r.target == rel.target
64                                && r.relation == rel.relation
65                        });
66                        if !exists {
67                            merged.relationships.push(rel);
68                        }
69                    }
70                    if first_embedding.is_none() && out.embedding.is_some() {
71                        first_embedding = out.embedding;
72                    }
73                }
74                Err(err) => {
75                    if any_error.is_none() {
76                        any_error = Some(err);
77                    }
78                }
79            }
80        }
81
82        merged.embedding = first_embedding;
83        merged.elapsed_ms = start.elapsed().as_millis() as u64;
84
85        if merged.entities.is_empty() && merged.relationships.is_empty() {
86            if let Some(err) = any_error {
87                return Err(err);
88            }
89        }
90        Ok(merged)
91    }
92
93    async fn health(&self) -> Result<BackendHealth, AppError> {
94        let mut healthy = true;
95        let mut messages = Vec::new();
96        for backend in &self.backends {
97            match backend.health().await {
98                Ok(h) => {
99                    if !h.healthy {
100                        healthy = false;
101                    }
102                    messages.push(format!(
103                        "{}:{}",
104                        h.kind.as_str(),
105                        if h.healthy { "ok" } else { "degraded" }
106                    ));
107                }
108                Err(err) => {
109                    healthy = false;
110                    messages.push(format!("err:{err}"));
111                }
112            }
113        }
114        Ok(BackendHealth {
115            kind: self.kind(),
116            healthy,
117            model_name: self.model_name(),
118            message: messages.join(" "),
119        })
120    }
121}
122
123/// Factory that builds the default backend for the current build configuration.
124///
125/// v1.0.89 (GAP-META-006): uses `detect_available_backend()` to resolve
126/// the LLM CLI at runtime instead of hardcoding codex. The user's
127/// `--llm-backend` flag (propagated via env var) is honoured by the
128/// factory infrastructure in `llm_backend.rs`.
129pub fn default_backend() -> SharedBackend {
130    use std::sync::Arc;
131    let kind = super::llm_backend::detect_available_backend()
132        .unwrap_or(super::llm_backend::LlmBackendKindFactory::None);
133    let config = super::llm_backend::LlmExtractorConfig::default();
134    match kind {
135        super::llm_backend::LlmBackendKindFactory::Codex
136        | super::llm_backend::LlmBackendKindFactory::Auto => Arc::new(
137            super::llm_backend::LlmBackend::new(super::llm_backend::LlmExtractorConfig {
138                backend: "codex".to_string(),
139                ..config
140            }),
141        ),
142        super::llm_backend::LlmBackendKindFactory::Claude => Arc::new(
143            super::llm_backend::LlmBackend::new(super::llm_backend::LlmExtractorConfig {
144                backend: "claude".to_string(),
145                ..config
146            }),
147        ),
148        super::llm_backend::LlmBackendKindFactory::Opencode => Arc::new(
149            super::llm_backend::LlmBackend::new(super::llm_backend::LlmExtractorConfig {
150                backend: "opencode".to_string(),
151                ..config
152            }),
153        ),
154        super::llm_backend::LlmBackendKindFactory::None => {
155            Arc::new(super::none_backend::NoneBackend::new())
156        }
157    }
158}
159
160/// Factory that builds a backend from a CLI flag.
161///
162/// v1.0.89 (GAP-META-006): the `Llm` variant now resolves via
163/// `detect_available_backend()` instead of hardcoding codex.
164pub fn backend_from_kind(kind: BackendKind) -> SharedBackend {
165    use std::sync::Arc;
166    match kind {
167        BackendKind::Llm => default_backend(),
168        BackendKind::Embedding => Arc::new(super::embedding_backend::EmbeddingBackend::new()),
169        BackendKind::None => Arc::new(super::none_backend::NoneBackend::new()),
170        BackendKind::Composite => {
171            let llm: SharedBackend = default_backend();
172            let embedding: SharedBackend =
173                Arc::new(super::embedding_backend::EmbeddingBackend::new());
174            Arc::new(CompositeBackend::new(vec![llm, embedding]))
175        }
176    }
177}