sqlite_graphrag/extract/
mod.rs

1//! Extraction backend abstraction (v1.0.75 — G21 solution)
2//!
3//! Provides the `ExtractionBackend` trait with concrete implementations for
4//! LLM-only (default in v1.0.75), Embedding (legacy), None (no extraction),
5//! and Composite (orchestrates multiple backends in parallel).
6//!
7//! The trait enables backend-agnostic ingest/enrich/remember pipelines.
8
9use crate::errors::AppError;
10use async_trait::async_trait;
11use serde::{Deserialize, Serialize};
12use std::sync::Arc;
13
14/// Hint configuration forwarded to the extraction backend.
15#[derive(Debug, Clone, Default, Serialize, Deserialize)]
16pub struct ExtractionHints {
17    /// Memory name to be remembered (kebab-case)
18    pub memory_name: Option<String>,
19    /// Memory type to be remembered
20    pub memory_type: Option<String>,
21    /// Existing entity names to avoid duplicates
22    pub existing_entities: Vec<String>,
23    /// Whether to skip relation extraction
24    pub skip_relations: bool,
25    /// Backend-specific seed for determinism
26    pub seed: Option<u64>,
27}
28
29/// Entity extracted from content.
30#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
31pub struct ExtractedEntity {
32    /// Name of this item.
33    pub name: String,
34    /// Entity type label.
35    pub entity_type: String,
36    /// Human-readable description.
37    pub description: Option<String>,
38    /// Confidence score in `[0, 1]`.
39    pub confidence: Option<f32>,
40}
41
42/// Relationship extracted from content.
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44pub struct ExtractedRelationship {
45    /// Source side of the relationship.
46    pub source: String,
47    /// Target side of the relationship.
48    pub target: String,
49    /// Relationship type.
50    pub relation: String,
51    /// Strength.
52    pub strength: f32,
53}
54
55/// Output of extraction backend.
56#[derive(Debug, Clone, Default, Serialize, Deserialize)]
57pub struct ExtractionOutput {
58    /// Extracted entities.
59    pub entities: Vec<ExtractedEntity>,
60    /// Relationships.
61    pub relationships: Vec<ExtractedRelationship>,
62    /// Optional embedding vector (only populated by EmbeddingBackend)
63    pub embedding: Option<Vec<f32>>,
64    /// Backend that produced this output
65    pub backend: String,
66    /// Latency in milliseconds
67    pub elapsed_ms: u64,
68}
69
70/// Backend kind enumeration used for selection and telemetry.
71#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
72#[serde(rename_all = "kebab-case")]
73pub enum BackendKind {
74    /// LLM variant.
75    Llm,
76    /// Embedding variant.
77    Embedding,
78    /// None variant.
79    None,
80    /// Composite variant.
81    Composite,
82}
83
84impl BackendKind {
85    /// Return the canonical string representation.
86    pub fn as_str(self) -> &'static str {
87        match self {
88            BackendKind::Llm => "llm",
89            BackendKind::Embedding => "embedding",
90            BackendKind::None => "none",
91            BackendKind::Composite => "composite",
92        }
93    }
94
95    /// Parse from a string.
96    pub fn parse(s: &str) -> Option<Self> {
97        match s.to_ascii_lowercase().as_str() {
98            "llm" => Some(BackendKind::Llm),
99            "embedding" => Some(BackendKind::Embedding),
100            "none" => Some(BackendKind::None),
101            "both" | "composite" => Some(BackendKind::Composite),
102            _ => None,
103        }
104    }
105}
106
107/// Trait abstraction for any extraction backend (LLM, Embedding, None, Composite).
108///
109/// G21 HIGH solution: the trait allows the rest of the codebase to remain
110/// agnostic of the underlying extraction mechanism. New backends can be added
111/// without touching call sites.
112#[async_trait]
113pub trait ExtractionBackend: Send + Sync {
114    /// Identify this backend (used in metrics, logs and ExtractionOutput)
115    fn kind(&self) -> BackendKind;
116
117    /// Identify the underlying model/CLI being used (e.g. "codex-0.137.0")
118    fn model_name(&self) -> String;
119
120    /// Extract entities and relationships from `content`.
121    ///
122    /// `hints` provides optional context (memory name, type, etc.).
123    /// Returns `ExtractionOutput` with entities, relationships, and optional embedding.
124    async fn extract(
125        &self,
126        content: &str,
127        hints: &ExtractionHints,
128    ) -> Result<ExtractionOutput, AppError>;
129
130    /// Health check: whether this backend is ready to operate.
131    async fn health(&self) -> Result<BackendHealth, AppError>;
132}
133
134/// Health status of a backend.
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct BackendHealth {
137    /// Kind discriminator.
138    pub kind: BackendKind,
139    /// Healthy.
140    pub healthy: bool,
141    /// Model name.
142    pub model_name: String,
143    /// Message text.
144    pub message: String,
145}
146
147/// Type alias for shared backend references.
148pub type SharedBackend = Arc<dyn ExtractionBackend>;
149
150pub mod composite_backend;
151pub mod embedding_backend;
152pub mod llm_backend;
153pub mod none_backend;
154
155pub use composite_backend::{backend_from_kind, default_backend, CompositeBackend};
156pub use embedding_backend::EmbeddingBackend;
157pub use llm_backend::{LlmBackend, LlmExtractorConfig};
158pub use none_backend::NoneBackend;
sqlite_graphrag/extract/mod.rs

sqlite_graphrag/extract/
mod.rs
