oxirag 0.1.1

A four-layer RAG engine with SMT-based logic verification and knowledge graph support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! Traits for the Graph (knowledge graph) layer.

use async_trait::async_trait;

use crate::error::GraphError;
use crate::types::Document;

use super::types::{
    Direction, EntityId, EntityType, GraphEntity, GraphPath, GraphQuery, GraphRelationship,
    RelationshipType,
};

/// Extracts entities from text.
#[async_trait]
pub trait EntityExtractor: Send + Sync {
    /// Extract entities from the given text.
    async fn extract_entities(&self, text: &str) -> Result<Vec<GraphEntity>, GraphError>;

    /// Get the entity types supported by this extractor.
    fn supported_entity_types(&self) -> Vec<EntityType>;
}

/// Extracts relationships between entities from text.
#[async_trait]
pub trait RelationshipExtractor: Send + Sync {
    /// Extract relationships between the given entities from text.
    async fn extract_relationships(
        &self,
        text: &str,
        entities: &[GraphEntity],
    ) -> Result<Vec<GraphRelationship>, GraphError>;

    /// Get the relationship types supported by this extractor.
    fn supported_relationship_types(&self) -> Vec<RelationshipType>;
}

/// Storage for knowledge graph data with traversal capabilities.
#[async_trait]
pub trait GraphStore: Send + Sync {
    /// Add an entity to the graph.
    async fn add_entity(&mut self, entity: GraphEntity) -> Result<EntityId, GraphError>;

    /// Add multiple entities to the graph.
    async fn add_entities(
        &mut self,
        entities: Vec<GraphEntity>,
    ) -> Result<Vec<EntityId>, GraphError>;

    /// Add a relationship to the graph.
    async fn add_relationship(
        &mut self,
        relationship: GraphRelationship,
    ) -> Result<String, GraphError>;

    /// Add multiple relationships to the graph.
    async fn add_relationships(
        &mut self,
        relationships: Vec<GraphRelationship>,
    ) -> Result<Vec<String>, GraphError>;

    /// Get an entity by its ID.
    async fn get_entity(&self, id: &EntityId) -> Result<Option<GraphEntity>, GraphError>;

    /// Get neighboring entities connected by relationships.
    async fn get_neighbors(
        &self,
        id: &EntityId,
        direction: Direction,
    ) -> Result<Vec<(GraphRelationship, GraphEntity)>, GraphError>;

    /// Find entities by their type.
    async fn find_entities_by_type(
        &self,
        entity_type: &EntityType,
    ) -> Result<Vec<GraphEntity>, GraphError>;

    /// Find entities by name (exact or partial match).
    async fn find_entities_by_name(&self, name: &str) -> Result<Vec<GraphEntity>, GraphError>;

    /// Traverse the graph from starting entities.
    async fn traverse(&self, query: &GraphQuery) -> Result<Vec<GraphPath>, GraphError>;

    /// Get the total number of entities in the graph.
    async fn entity_count(&self) -> usize;

    /// Get the total number of relationships in the graph.
    async fn relationship_count(&self) -> usize;

    /// Clear all entities and relationships from the graph.
    async fn clear(&mut self) -> Result<(), GraphError>;
}

/// The Graph layer: combines entity extraction, relationship extraction, and graph storage.
#[async_trait]
pub trait Graph: Send + Sync {
    /// Index a document by extracting entities and relationships.
    async fn index_document(&mut self, document: &Document) -> Result<(), GraphError>;

    /// Index multiple documents.
    async fn index_documents(&mut self, documents: &[Document]) -> Result<(), GraphError>;

    /// Query the graph for related entities and paths.
    async fn query(&self, query: &GraphQuery) -> Result<Vec<GraphPath>, GraphError>;

    /// Find entities related to a given entity name.
    async fn find_related(
        &self,
        entity_name: &str,
        max_hops: usize,
    ) -> Result<Vec<GraphPath>, GraphError>;

    /// Get entity count.
    async fn entity_count(&self) -> usize;

    /// Get relationship count.
    async fn relationship_count(&self) -> usize;

    /// Clear the graph.
    async fn clear(&mut self) -> Result<(), GraphError>;
}

#[cfg(test)]
mod tests {
    use super::*;

    // Test that traits compile correctly
    #[test]
    fn test_trait_object_safety() {
        fn _assert_entity_extractor_object_safe(_: &dyn EntityExtractor) {}
        fn _assert_relationship_extractor_object_safe(_: &dyn RelationshipExtractor) {}
        fn _assert_graph_store_object_safe(_: &dyn GraphStore) {}
        fn _assert_graph_object_safe(_: &dyn Graph) {}
    }
}