sqlite_graphrag/extract/
composite_backend.rs1use super::{
8 BackendHealth, BackendKind, ExtractionBackend, ExtractionHints, ExtractionOutput, SharedBackend,
9};
10use crate::errors::AppError;
11use async_trait::async_trait;
12use std::time::Instant;
13
14pub struct CompositeBackend {
16 backends: Vec<SharedBackend>,
17}
18
19impl CompositeBackend {
20 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
124pub 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
132pub 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}