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