sqlite_graphrag/extract/
mod.rs

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

sqlite_graphrag/extract/
mod.rs
