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//! Was selected by an extraction-backend switch that no longer exists: the flag
5//! was removed in v1.0.79 together with the fastembed pipeline.
6
7use super::{
8    BackendHealth, BackendKind, ExtractionBackend, ExtractionHints, ExtractionOutput, SharedBackend,
9};
10use crate::errors::AppError;
11use async_trait::async_trait;
12use std::time::Instant;
13
14/// Composite backend.
15pub struct CompositeBackend {
16    backends: Vec<SharedBackend>,
17}
18
19impl CompositeBackend {
20    /// Create a new instance.
21    pub fn new(backends: Vec<SharedBackend>) -> Self {
22        Self { backends }
23    }
24}
25
26#[async_trait]
27impl ExtractionBackend for CompositeBackend {
28    fn kind(&self) -> BackendKind {
29        BackendKind::Composite
30    }
31
32    fn model_name(&self) -> String {
33        self.backends
34            .iter()
35            .map(|b| b.model_name())
36            .collect::<Vec<_>>()
37            .join("+")
38    }
39
40    async fn extract(
41        &self,
42        content: &str,
43        hints: &ExtractionHints,
44    ) -> Result<ExtractionOutput, AppError> {
45        let start = Instant::now();
46        let mut merged = ExtractionOutput {
47            backend: self.kind().as_str().to_string(),
48            ..Default::default()
49        };
50        let mut first_embedding: Option<Vec<f32>> = None;
51        let mut any_error: Option<AppError> = None;
52
53        for backend in &self.backends {
54            match backend.extract(content, hints).await {
55                Ok(out) => {
56                    for entity in out.entities {
57                        if !merged.entities.iter().any(|e| e.name == entity.name) {
58                            merged.entities.push(entity);
59                        }
60                    }
61                    for rel in out.relationships {
62                        let exists = merged.relationships.iter().any(|r| {
63                            r.source == rel.source
64                                && r.target == rel.target
65                                && r.relation == rel.relation
66                        });
67                        if !exists {
68                            merged.relationships.push(rel);
69                        }
70                    }
71                    if first_embedding.is_none() && out.embedding.is_some() {
72                        first_embedding = out.embedding;
73                    }
74                }
75                Err(err) => {
76                    if any_error.is_none() {
77                        any_error = Some(err);
78                    }
79                }
80            }
81        }
82
83        merged.embedding = first_embedding;
84        merged.elapsed_ms = start.elapsed().as_millis() as u64;
85
86        if merged.entities.is_empty() && merged.relationships.is_empty() {
87            if let Some(err) = any_error {
88                return Err(err);
89            }
90        }
91        Ok(merged)
92    }
93
94    async fn health(&self) -> Result<BackendHealth, AppError> {
95        let mut healthy = true;
96        let mut messages = Vec::new();
97        for backend in &self.backends {
98            match backend.health().await {
99                Ok(h) => {
100                    if !h.healthy {
101                        healthy = false;
102                    }
103                    messages.push(format!(
104                        "{}:{}",
105                        h.kind.as_str(),
106                        if h.healthy { "ok" } else { "degraded" }
107                    ));
108                }
109                Err(err) => {
110                    healthy = false;
111                    messages.push(format!("err:{err}"));
112                }
113            }
114        }
115        Ok(BackendHealth {
116            kind: self.kind(),
117            healthy,
118            model_name: self.model_name(),
119            message: messages.join(" "),
120        })
121    }
122}
123
124/// Factory that builds the default backend for the current build configuration.
125pub fn default_backend() -> SharedBackend {
126    use std::sync::Arc;
127    Arc::new(super::llm_backend::LlmBackend::new(
128        super::llm_backend::LlmExtractorConfig::default(),
129    ))
130}
131
132/// Factory that builds a backend from a CLI flag.
133pub fn backend_from_kind(kind: BackendKind) -> SharedBackend {
134    use std::sync::Arc;
135    match kind {
136        BackendKind::Llm => default_backend(),
137        BackendKind::Embedding => Arc::new(super::embedding_backend::EmbeddingBackend::new()),
138        BackendKind::None => Arc::new(super::none_backend::NoneBackend::new()),
139        BackendKind::Composite => {
140            let llm: SharedBackend = default_backend();
141            let embedding: SharedBackend =
142                Arc::new(super::embedding_backend::EmbeddingBackend::new());
143            Arc::new(CompositeBackend::new(vec![llm, embedding]))
144        }
145    }
146}