oxirs_graphrag/distributed/
coordinator.rs

1//! Coordinator types: configuration, knowledge graph, entity resolution, and context building.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4use std::sync::Arc;
5use std::time::Duration;
6
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9use tokio::sync::{Mutex, RwLock, Semaphore};
10use tracing::debug;
11
12use crate::{GraphRAGError, GraphRAGResult, Triple};
13
14// ─────────────────────────────────────────────────────────────────────────────
15// Error types
16// ─────────────────────────────────────────────────────────────────────────────
17
18/// Distributed GraphRAG–specific error variants
19#[derive(Error, Debug)]
20pub enum DistributedError {
21    #[error("Endpoint {endpoint} is unreachable: {reason}")]
22    EndpointUnreachable { endpoint: String, reason: String },
23
24    #[error("Authentication failed for endpoint {endpoint}")]
25    AuthFailed { endpoint: String },
26
27    #[error("SPARQL query timeout after {timeout_ms}ms on endpoint {endpoint}")]
28    QueryTimeout { endpoint: String, timeout_ms: u64 },
29
30    #[error("Entity resolution cycle detected for URI {uri}")]
31    SameAsCycle { uri: String },
32
33    #[error("No healthy endpoints available for query")]
34    NoHealthyEndpoints,
35
36    #[error("Merge conflict: cannot reconcile {uri} across endpoints")]
37    MergeConflict { uri: String },
38
39    #[error("Configuration invalid: {0}")]
40    InvalidConfig(String),
41}
42
43impl From<DistributedError> for GraphRAGError {
44    fn from(e: DistributedError) -> Self {
45        GraphRAGError::InternalError(e.to_string())
46    }
47}
48
49// ─────────────────────────────────────────────────────────────────────────────
50// Authentication
51// ─────────────────────────────────────────────────────────────────────────────
52
53/// Authentication method for SPARQL endpoints
54#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
55#[serde(tag = "type", rename_all = "snake_case")]
56pub enum EndpointAuth {
57    /// No authentication
58    None,
59    /// HTTP Bearer token
60    Bearer { token: String },
61    /// HTTP Basic auth
62    Basic { username: String, password: String },
63    /// API key in header
64    ApiKey { header: String, key: String },
65}
66
67// ─────────────────────────────────────────────────────────────────────────────
68// Configuration
69// ─────────────────────────────────────────────────────────────────────────────
70
71/// Configuration for a single remote SPARQL endpoint
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct EndpointConfig {
74    /// Human-readable name for the endpoint
75    pub name: String,
76    /// Base URL of the SPARQL endpoint
77    pub url: String,
78    /// Authentication method
79    pub auth: EndpointAuth,
80    /// Per-endpoint query timeout in milliseconds (overrides global setting)
81    pub timeout_ms: Option<u64>,
82    /// Priority weight (higher = preferred; used when deduplicating conflicting triples)
83    pub priority: f64,
84    /// Whether this endpoint is enabled
85    pub enabled: bool,
86    /// Graph URI to restrict queries to (SPARQL FROM clause)
87    pub graph_uri: Option<String>,
88    /// Maximum triples to fetch from this endpoint per query
89    pub max_triples: usize,
90}
91
92impl Default for EndpointConfig {
93    fn default() -> Self {
94        Self {
95            name: String::new(),
96            url: String::new(),
97            auth: EndpointAuth::None,
98            timeout_ms: None,
99            priority: 1.0,
100            enabled: true,
101            graph_uri: None,
102            max_triples: 10_000,
103        }
104    }
105}
106
107/// Top-level configuration for federated GraphRAG
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct FederatedGraphRAGConfig {
110    /// List of remote endpoints to query
111    pub endpoints: Vec<EndpointConfig>,
112    /// Global query timeout in milliseconds
113    pub global_timeout_ms: u64,
114    /// Maximum concurrent endpoint requests
115    pub max_concurrency: usize,
116    /// Maximum transitive sameAs hops to follow
117    pub same_as_max_depth: usize,
118    /// Minimum endpoint priority to include in a query (0.0 = include all)
119    pub min_endpoint_priority: f64,
120    /// Whether to continue when some endpoints fail
121    pub partial_results_ok: bool,
122    /// Retry count for failed endpoint requests
123    pub retry_count: usize,
124    /// Delay between retries in milliseconds
125    pub retry_delay_ms: u64,
126}
127
128impl Default for FederatedGraphRAGConfig {
129    fn default() -> Self {
130        Self {
131            endpoints: vec![],
132            global_timeout_ms: 30_000,
133            max_concurrency: 8,
134            same_as_max_depth: 5,
135            min_endpoint_priority: 0.0,
136            partial_results_ok: true,
137            retry_count: 2,
138            retry_delay_ms: 500,
139        }
140    }
141}
142
143impl FederatedGraphRAGConfig {
144    /// Validate configuration and return an error description if invalid.
145    pub fn validate(&self) -> Result<(), DistributedError> {
146        if self.global_timeout_ms == 0 {
147            return Err(DistributedError::InvalidConfig(
148                "global_timeout_ms must be > 0".into(),
149            ));
150        }
151        if self.max_concurrency == 0 {
152            return Err(DistributedError::InvalidConfig(
153                "max_concurrency must be > 0".into(),
154            ));
155        }
156        if self.same_as_max_depth == 0 {
157            return Err(DistributedError::InvalidConfig(
158                "same_as_max_depth must be > 0".into(),
159            ));
160        }
161        for ep in &self.endpoints {
162            if ep.url.is_empty() {
163                return Err(DistributedError::InvalidConfig(format!(
164                    "Endpoint '{}' has an empty URL",
165                    ep.name
166                )));
167            }
168            if ep.max_triples == 0 {
169                return Err(DistributedError::InvalidConfig(format!(
170                    "Endpoint '{}' max_triples must be > 0",
171                    ep.name
172                )));
173            }
174        }
175        Ok(())
176    }
177
178    /// Return only enabled endpoints with priority >= `min_endpoint_priority`.
179    pub fn active_endpoints(&self) -> Vec<&EndpointConfig> {
180        self.endpoints
181            .iter()
182            .filter(|ep| ep.enabled && ep.priority >= self.min_endpoint_priority)
183            .collect()
184    }
185}
186
187// ─────────────────────────────────────────────────────────────────────────────
188// Knowledge graph result type
189// ─────────────────────────────────────────────────────────────────────────────
190
191/// A merged knowledge graph assembled from multiple endpoints
192#[derive(Debug, Clone, Default)]
193pub struct KnowledgeGraph {
194    /// All triples gathered from the federation
195    pub triples: Vec<Triple>,
196    /// Provenance: which endpoint contributed each triple index
197    pub provenance: Vec<String>,
198    /// Entity equivalence classes after sameAs resolution
199    pub equivalence_classes: Vec<HashSet<String>>,
200    /// Canonical URIs chosen for each equivalence class (representative URI)
201    pub canonical_uris: HashMap<String, String>,
202}
203
204impl KnowledgeGraph {
205    /// Create an empty knowledge graph
206    pub fn new() -> Self {
207        Self::default()
208    }
209
210    /// Return the number of distinct triples
211    pub fn triple_count(&self) -> usize {
212        self.triples.len()
213    }
214
215    /// Return true if the knowledge graph has no triples
216    pub fn is_empty(&self) -> bool {
217        self.triples.is_empty()
218    }
219
220    /// Resolve a URI to its canonical form (or return the URI unchanged)
221    pub fn canonical<'a>(&'a self, uri: &'a str) -> &'a str {
222        self.canonical_uris
223            .get(uri)
224            .map(|s| s.as_str())
225            .unwrap_or(uri)
226    }
227}
228
229// ─────────────────────────────────────────────────────────────────────────────
230// Endpoint executor trait (defined here to break the worker ↔ coordinator cycle)
231// ─────────────────────────────────────────────────────────────────────────────
232
233/// Trait for executing SPARQL queries against a remote endpoint.
234/// Defined in `coordinator` so both coordinator and worker can reference it
235/// without a circular dependency.
236#[async_trait::async_trait]
237pub trait EndpointExecutor: Send + Sync {
238    /// Execute a SPARQL CONSTRUCT query and return RDF triples
239    async fn construct(
240        &self,
241        endpoint: &EndpointConfig,
242        sparql: &str,
243        timeout: Duration,
244    ) -> GraphRAGResult<Vec<Triple>>;
245
246    /// Execute a SPARQL SELECT query and return rows (variable → value maps)
247    async fn select(
248        &self,
249        endpoint: &EndpointConfig,
250        sparql: &str,
251        timeout: Duration,
252    ) -> GraphRAGResult<Vec<HashMap<String, String>>>;
253}
254
255// ─────────────────────────────────────────────────────────────────────────────
256// DistributedEntityResolver
257// ─────────────────────────────────────────────────────────────────────────────
258
259/// Resolves entity identity across endpoints using owl:sameAs links.
260///
261/// The resolver computes the transitive sameAs closure: if A sameAs B and
262/// B sameAs C, all three are placed in the same equivalence class.
263pub struct DistributedEntityResolver<E: EndpointExecutor> {
264    pub(super) config: FederatedGraphRAGConfig,
265    pub(super) executor: Arc<E>,
266}
267
268impl<E: EndpointExecutor + 'static> DistributedEntityResolver<E> {
269    /// Create a new resolver
270    pub fn new(config: FederatedGraphRAGConfig, executor: Arc<E>) -> Self {
271        Self { config, executor }
272    }
273
274    /// Compute the transitive owl:sameAs closure for the given URIs across all
275    /// active endpoints.
276    pub async fn same_as_closure(
277        &self,
278        uris: &[String],
279    ) -> GraphRAGResult<HashMap<String, String>> {
280        if uris.is_empty() {
281            return Ok(HashMap::new());
282        }
283
284        let parent: Arc<RwLock<HashMap<String, String>>> = Arc::new(RwLock::new(HashMap::new()));
285
286        {
287            let mut p = parent.write().await;
288            for uri in uris {
289                p.insert(uri.clone(), uri.clone());
290            }
291        }
292
293        let mut frontier: VecDeque<String> = uris.iter().cloned().collect();
294        let mut visited: HashSet<String> = HashSet::from_iter(uris.iter().cloned());
295        let mut depth = 0usize;
296
297        while !frontier.is_empty() && depth < self.config.same_as_max_depth {
298            let batch: Vec<String> = frontier.drain(..).collect();
299            let batch_refs: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
300
301            let links = self.fetch_same_as_links(&batch_refs).await?;
302
303            let mut p = parent.write().await;
304            for (a, b) in links {
305                p.entry(a.clone()).or_insert_with(|| a.clone());
306                p.entry(b.clone()).or_insert_with(|| b.clone());
307
308                let root_a = find_root_path(&p, &a);
309                let root_b = find_root_path(&p, &b);
310                if root_a != root_b {
311                    let canonical = if root_a <= root_b {
312                        root_a.clone()
313                    } else {
314                        root_b.clone()
315                    };
316                    p.insert(root_a, canonical.clone());
317                    p.insert(root_b, canonical);
318                }
319
320                if !visited.contains(&b) {
321                    visited.insert(b.clone());
322                    frontier.push_back(b);
323                }
324            }
325
326            depth += 1;
327        }
328
329        let p = parent.read().await;
330        let mut result = HashMap::new();
331        for uri in p.keys() {
332            let canonical = find_root_path(&p, uri);
333            result.insert(uri.clone(), canonical);
334        }
335        Ok(result)
336    }
337
338    /// Fetch raw owl:sameAs pairs from all active endpoints for the given URIs
339    async fn fetch_same_as_links(&self, uris: &[&str]) -> GraphRAGResult<Vec<(String, String)>> {
340        let active = self.config.active_endpoints();
341        let semaphore = Arc::new(Semaphore::new(self.config.max_concurrency));
342        let pairs: Arc<Mutex<Vec<(String, String)>>> = Arc::new(Mutex::new(Vec::new()));
343        let mut handles = Vec::new();
344
345        for ep in active {
346            let ep = ep.clone();
347            let executor = Arc::clone(&self.executor);
348            let sem = Arc::clone(&semaphore);
349            let pairs = Arc::clone(&pairs);
350            let uris_owned: Vec<String> = uris.iter().map(|s| s.to_string()).collect();
351            let timeout_ms = ep.timeout_ms.unwrap_or(self.config.global_timeout_ms);
352            let timeout = Duration::from_millis(timeout_ms);
353
354            let handle = tokio::spawn(async move {
355                let _permit = match sem.acquire_owned().await {
356                    Ok(p) => p,
357                    Err(_) => return,
358                };
359
360                let sparql = build_same_as_sparql(
361                    &uris_owned.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
362                    ep.graph_uri.as_deref(),
363                );
364
365                match executor.select(&ep, &sparql, timeout).await {
366                    Ok(rows) => {
367                        let mut guard = pairs.lock().await;
368                        for row in rows {
369                            if let (Some(a), Some(b)) = (row.get("a"), row.get("b")) {
370                                guard.push((a.clone(), b.clone()));
371                            }
372                        }
373                    }
374                    Err(e) => {
375                        debug!(endpoint = %ep.name, error = %e, "sameAs fetch failed");
376                    }
377                }
378            });
379
380            handles.push(handle);
381        }
382
383        for h in handles {
384            let _ = h.await;
385        }
386
387        let guard = Arc::try_unwrap(pairs)
388            .map_err(|_| GraphRAGError::InternalError("Arc unwrap failed".into()))?
389            .into_inner();
390
391        Ok(guard)
392    }
393
394    /// Apply sameAs closure to a knowledge graph, rewriting URIs to canonical forms
395    /// and deduplicating triples that become identical after rewriting.
396    pub fn apply_to_graph(&self, kg: &mut KnowledgeGraph, canonical_map: &HashMap<String, String>) {
397        let canonicalize = |s: &str| -> String {
398            canonical_map
399                .get(s)
400                .cloned()
401                .unwrap_or_else(|| s.to_string())
402        };
403
404        let mut seen: HashSet<(String, String, String)> = HashSet::new();
405        let mut new_triples = Vec::new();
406        let mut new_provenance = Vec::new();
407
408        for (triple, prov) in kg.triples.iter().zip(kg.provenance.iter()) {
409            let s = canonicalize(&triple.subject);
410            let p = triple.predicate.clone();
411            let o = canonicalize(&triple.object);
412            let key = (s.clone(), p.clone(), o.clone());
413            if seen.insert(key) {
414                new_triples.push(Triple::new(s, p, o));
415                new_provenance.push(prov.clone());
416            }
417        }
418
419        kg.triples = new_triples;
420        kg.provenance = new_provenance;
421        kg.canonical_uris = canonical_map.clone();
422
423        let mut classes: HashMap<String, HashSet<String>> = HashMap::new();
424        for (uri, canonical) in canonical_map {
425            classes
426                .entry(canonical.clone())
427                .or_default()
428                .insert(uri.clone());
429        }
430        kg.equivalence_classes = classes.into_values().collect();
431    }
432}
433
434/// Path-compression find for the union-find structure
435pub(super) fn find_root_path(parent: &HashMap<String, String>, uri: &str) -> String {
436    let mut current = uri.to_string();
437    let mut depth = 0usize;
438    loop {
439        let next = parent
440            .get(&current)
441            .cloned()
442            .unwrap_or_else(|| current.clone());
443        if next == current || depth > 100 {
444            return current;
445        }
446        current = next;
447        depth += 1;
448    }
449}
450
451/// Build a SPARQL SELECT query for sameAs links (shared by coordinator and worker)
452pub(super) fn build_same_as_sparql(uris: &[&str], graph_uri: Option<&str>) -> String {
453    let values: Vec<String> = uris.iter().map(|s| format!("<{}>", s)).collect();
454    let values_block = values.join(" ");
455
456    let from_clause = match graph_uri {
457        Some(g) => format!("FROM <{}>", g),
458        None => String::new(),
459    };
460
461    format!(
462        r#"SELECT DISTINCT ?a ?b
463{from}
464WHERE {{
465    VALUES ?a {{ {uris} }}
466    {{
467        ?a <http://www.w3.org/2002/07/owl#sameAs> ?b .
468    }} UNION {{
469        ?b <http://www.w3.org/2002/07/owl#sameAs> ?a .
470    }}
471}}
472"#,
473        from = from_clause,
474        uris = values_block,
475    )
476}
477
478// ─────────────────────────────────────────────────────────────────────────────
479// FederatedContextBuilder
480// ─────────────────────────────────────────────────────────────────────────────
481
482/// Strategy for ordering triples in the generated context
483#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
484pub enum ContextOrderingStrategy {
485    /// Order by endpoint priority (highest first)
486    ByEndpointPriority,
487    /// Order by query latency (fastest endpoints first)
488    ByLatency,
489    /// No specific ordering (insertion order)
490    Insertion,
491}
492
493/// Configuration for the federated context builder
494#[derive(Debug, Clone, Serialize, Deserialize)]
495pub struct FederatedContextConfig {
496    /// Maximum number of triples to include in the context
497    pub max_context_triples: usize,
498    /// Maximum length (characters) of the context string
499    pub max_context_chars: usize,
500    /// Triple ordering strategy
501    pub ordering: ContextOrderingStrategy,
502    /// Whether to include provenance annotations in the context
503    pub include_provenance: bool,
504    /// Minimum endpoint priority to include triples from
505    pub min_endpoint_priority: f64,
506    /// Whether to include equivalence class annotations
507    pub include_equivalences: bool,
508}
509
510impl Default for FederatedContextConfig {
511    fn default() -> Self {
512        Self {
513            max_context_triples: 500,
514            max_context_chars: 50_000,
515            ordering: ContextOrderingStrategy::ByEndpointPriority,
516            include_provenance: false,
517            include_equivalences: false,
518            min_endpoint_priority: 0.0,
519        }
520    }
521}
522
523/// Builds RAG context strings from distributed knowledge graphs
524pub struct FederatedContextBuilder {
525    config: FederatedContextConfig,
526    /// Per-endpoint priority registry
527    endpoint_priorities: HashMap<String, f64>,
528    /// Per-endpoint latency registry (milliseconds, populated from expansion runs)
529    endpoint_latencies: Arc<RwLock<HashMap<String, u64>>>,
530}
531
532impl FederatedContextBuilder {
533    /// Create a new context builder
534    pub fn new(config: FederatedContextConfig, graphrag_config: &FederatedGraphRAGConfig) -> Self {
535        let endpoint_priorities: HashMap<String, f64> = graphrag_config
536            .endpoints
537            .iter()
538            .map(|ep| (ep.name.clone(), ep.priority))
539            .collect();
540
541        Self {
542            config,
543            endpoint_priorities,
544            endpoint_latencies: Arc::new(RwLock::new(HashMap::new())),
545        }
546    }
547
548    /// Record observed latency for an endpoint (used in ByLatency ordering)
549    pub async fn record_latency(&self, endpoint_name: &str, latency_ms: u64) {
550        let mut lats = self.endpoint_latencies.write().await;
551        lats.insert(endpoint_name.to_string(), latency_ms);
552    }
553
554    /// Build a context string from a [`KnowledgeGraph`].
555    pub async fn build_context(&self, kg: &KnowledgeGraph, query: &str) -> GraphRAGResult<String> {
556        if kg.is_empty() {
557            return Ok(String::new());
558        }
559
560        let latencies = self.endpoint_latencies.read().await;
561        let mut indexed: Vec<(usize, f64)> = kg
562            .triples
563            .iter()
564            .enumerate()
565            .map(|(i, _)| {
566                let ep = kg.provenance.get(i).map(|s| s.as_str()).unwrap_or("");
567                let sort_key = match self.config.ordering {
568                    ContextOrderingStrategy::ByEndpointPriority => {
569                        self.endpoint_priorities.get(ep).copied().unwrap_or(1.0)
570                    }
571                    ContextOrderingStrategy::ByLatency => {
572                        let lat = latencies.get(ep).copied().unwrap_or(u64::MAX);
573                        1.0 / (lat as f64 + 1.0)
574                    }
575                    ContextOrderingStrategy::Insertion => i as f64,
576                };
577                (i, sort_key)
578            })
579            .filter(|(i, _)| {
580                let ep = kg.provenance.get(*i).map(|s| s.as_str()).unwrap_or("");
581                let prio = self.endpoint_priorities.get(ep).copied().unwrap_or(1.0);
582                prio >= self.config.min_endpoint_priority
583            })
584            .collect();
585
586        match self.config.ordering {
587            ContextOrderingStrategy::ByEndpointPriority | ContextOrderingStrategy::ByLatency => {
588                indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
589            }
590            ContextOrderingStrategy::Insertion => {
591                indexed.sort_by_key(|(i, _)| *i);
592            }
593        }
594
595        let mut context = format!("## Knowledge Graph Context\n\nQuery: {}\n\n", query);
596
597        if self.config.include_equivalences && !kg.equivalence_classes.is_empty() {
598            context.push_str("### Entity Equivalences\n");
599            for class in &kg.equivalence_classes {
600                if class.len() > 1 {
601                    let mut members: Vec<&str> = class.iter().map(|s| s.as_str()).collect();
602                    members.sort();
603                    context.push_str(&format!("- {}\n", members.join(" ≡ ")));
604                }
605            }
606            context.push('\n');
607        }
608
609        context.push_str("### Facts\n\n");
610
611        for (triple_count, (idx, _)) in indexed.iter().enumerate() {
612            if triple_count >= self.config.max_context_triples {
613                break;
614            }
615            if context.len() >= self.config.max_context_chars {
616                break;
617            }
618
619            let triple = &kg.triples[*idx];
620            let line = if self.config.include_provenance {
621                let ep = kg.provenance.get(*idx).map(|s| s.as_str()).unwrap_or("?");
622                format!(
623                    "- {} → {} → {} [{}]\n",
624                    triple.subject, triple.predicate, triple.object, ep
625                )
626            } else {
627                format!(
628                    "- {} → {} → {}\n",
629                    triple.subject, triple.predicate, triple.object
630                )
631            };
632
633            context.push_str(&line);
634        }
635
636        Ok(context)
637    }
638}
oxirs_graphrag/distributed/coordinator.rs

oxirs_graphrag/distributed/
coordinator.rs
