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.
124pub fn default_backend() -> SharedBackend {
125    use std::sync::Arc;
126    Arc::new(super::llm_backend::LlmBackend::new(
127        super::llm_backend::LlmExtractorConfig::default(),
128    ))
129}
130
131/// Factory that builds a backend from a CLI flag.
132pub fn backend_from_kind(kind: BackendKind) -> SharedBackend {
133    use std::sync::Arc;
134    match kind {
135        BackendKind::Llm => default_backend(),
136        BackendKind::Embedding => Arc::new(super::embedding_backend::EmbeddingBackend::new()),
137        BackendKind::None => Arc::new(super::none_backend::NoneBackend::new()),
138        BackendKind::Composite => {
139            let llm: SharedBackend = default_backend();
140            let embedding: SharedBackend =
141                Arc::new(super::embedding_backend::EmbeddingBackend::new());
142            Arc::new(CompositeBackend::new(vec![llm, embedding]))
143        }
144    }
145}
sqlite_graphrag/extract/composite_backend.rs

sqlite_graphrag/extract/
composite_backend.rs
