recall_echo/graph/
store.rs1#[cfg(all(feature = "embedded", feature = "server"))]
4compile_error!("Features `embedded` and `server` are mutually exclusive. Choose one.");
5
6use surrealdb::Surreal;
7
8use super::error::GraphError;
9
10#[cfg(feature = "embedded")]
11use std::path::Path;
12
13#[cfg(feature = "embedded")]
14use surrealdb::engine::local::SurrealKv;
15
16#[cfg(feature = "embedded")]
17pub type Db = surrealdb::engine::local::Db;
18
19#[cfg(feature = "server")]
20pub type Db = surrealdb::engine::remote::ws::Client;
21
22#[cfg(feature = "server")]
24#[derive(Clone)]
25pub struct ServerConfig {
26 pub url: String,
27 pub username: String,
28 pub password: String,
29 pub namespace: String,
30 pub database: String,
31}
32
33#[cfg(feature = "server")]
34impl std::fmt::Debug for ServerConfig {
35 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36 f.debug_struct("ServerConfig")
37 .field("url", &self.url)
38 .field("username", &self.username)
39 .field("password", &"[REDACTED]")
40 .field("namespace", &self.namespace)
41 .field("database", &self.database)
42 .finish()
43 }
44}
45
46#[cfg(feature = "embedded")]
48pub async fn open(path: &Path) -> Result<Surreal<Db>, GraphError> {
49 let surreal_path = path.join("surreal");
50 std::fs::create_dir_all(&surreal_path)?;
51
52 let path_str = surreal_path.to_str().ok_or_else(|| {
53 GraphError::Io(std::io::Error::new(
54 std::io::ErrorKind::InvalidData,
55 "graph store path contains non-UTF8 characters",
56 ))
57 })?;
58 let db: Surreal<Db> = Surreal::new::<SurrealKv>(path_str).await?;
59 db.use_ns("recall").use_db("graph").await?;
60
61 Ok(db)
62}
63
64#[cfg(feature = "server")]
66pub async fn connect(config: &ServerConfig) -> Result<Surreal<Db>, GraphError> {
67 let db = Surreal::new::<surrealdb::engine::remote::ws::Ws>(&config.url).await?;
68 db.signin(surrealdb::opt::auth::Database {
69 namespace: config.namespace.clone(),
70 database: config.database.clone(),
71 username: config.username.clone(),
72 password: config.password.clone(),
73 })
74 .await?;
75 db.use_ns(&config.namespace)
76 .use_db(&config.database)
77 .await?;
78
79 Ok(db)
80}
81
82pub async fn init_schema(db: &Surreal<Db>) -> Result<(), GraphError> {
84 db.query(
85 r#"
86 DEFINE TABLE IF NOT EXISTS entity SCHEMAFULL;
87 DEFINE FIELD IF NOT EXISTS name ON entity TYPE string;
88 DEFINE FIELD IF NOT EXISTS entity_type ON entity TYPE string;
89 DEFINE FIELD IF NOT EXISTS abstract ON entity TYPE string;
90 DEFINE FIELD IF NOT EXISTS overview ON entity TYPE string;
91 DEFINE FIELD IF NOT EXISTS content ON entity TYPE option<string>;
92 DEFINE FIELD IF NOT EXISTS attributes ON entity TYPE option<object> FLEXIBLE;
93 DEFINE FIELD IF NOT EXISTS embedding ON entity TYPE option<array<float>>;
94 DEFINE FIELD IF NOT EXISTS mutable ON entity TYPE bool DEFAULT true;
95 DEFINE FIELD IF NOT EXISTS access_count ON entity TYPE int DEFAULT 0;
96 DEFINE FIELD IF NOT EXISTS utility_score ON entity TYPE float DEFAULT 0.5;
97 DEFINE FIELD IF NOT EXISTS utility_updates ON entity TYPE int DEFAULT 0;
98 DEFINE FIELD IF NOT EXISTS created_at ON entity TYPE datetime DEFAULT time::now();
99 DEFINE FIELD IF NOT EXISTS updated_at ON entity TYPE datetime DEFAULT time::now();
100 DEFINE FIELD IF NOT EXISTS source ON entity TYPE option<string>;
101
102 DEFINE INDEX IF NOT EXISTS entity_name ON entity FIELDS name;
103 DEFINE INDEX IF NOT EXISTS entity_type ON entity FIELDS entity_type;
104 DEFINE INDEX IF NOT EXISTS entity_vector ON entity FIELDS embedding HNSW DIMENSION 384 DIST COSINE;
105
106 -- Pipeline attribute indexes
107 DEFINE INDEX IF NOT EXISTS entity_pipeline_stage ON entity FIELDS attributes.pipeline_stage;
108 DEFINE INDEX IF NOT EXISTS entity_pipeline_status ON entity FIELDS attributes.pipeline_status;
109
110 DEFINE TABLE IF NOT EXISTS relates_to SCHEMAFULL TYPE RELATION;
111 DEFINE FIELD IF NOT EXISTS rel_type ON relates_to TYPE string;
112 DEFINE FIELD IF NOT EXISTS description ON relates_to TYPE option<string>;
113 DEFINE FIELD IF NOT EXISTS valid_from ON relates_to TYPE datetime DEFAULT time::now();
114 DEFINE FIELD IF NOT EXISTS valid_until ON relates_to TYPE option<datetime>;
115 DEFINE FIELD IF NOT EXISTS confidence ON relates_to TYPE float DEFAULT 1.0;
116 DEFINE FIELD IF NOT EXISTS last_reinforced ON relates_to TYPE option<datetime>;
117 DEFINE FIELD IF NOT EXISTS source ON relates_to TYPE option<string>;
118
119 DEFINE INDEX IF NOT EXISTS rel_type_idx ON relates_to FIELDS rel_type;
120
121 DEFINE TABLE IF NOT EXISTS episode SCHEMAFULL;
122 DEFINE FIELD IF NOT EXISTS session_id ON episode TYPE string;
123 DEFINE FIELD IF NOT EXISTS timestamp ON episode TYPE datetime DEFAULT time::now();
124 DEFINE FIELD IF NOT EXISTS abstract ON episode TYPE string;
125 DEFINE FIELD IF NOT EXISTS overview ON episode TYPE option<string>;
126 DEFINE FIELD IF NOT EXISTS content ON episode TYPE option<string>;
127 DEFINE FIELD IF NOT EXISTS embedding ON episode TYPE option<array<float>>;
128 DEFINE FIELD IF NOT EXISTS log_number ON episode TYPE option<int>;
129 DEFINE FIELD IF NOT EXISTS extracted ON episode TYPE bool DEFAULT false;
130
131 DEFINE INDEX IF NOT EXISTS episode_session ON episode FIELDS session_id;
132 DEFINE INDEX IF NOT EXISTS episode_time ON episode FIELDS timestamp;
133 DEFINE INDEX IF NOT EXISTS episode_vector ON episode FIELDS embedding HNSW DIMENSION 384 DIST COSINE;
134
135 DEFINE TABLE IF NOT EXISTS contributed_to SCHEMAFULL TYPE RELATION;
136 DEFINE FIELD IF NOT EXISTS outcome_result ON contributed_to TYPE string;
137 DEFINE FIELD IF NOT EXISTS was_used ON contributed_to TYPE bool DEFAULT true;
138 DEFINE FIELD IF NOT EXISTS session_id ON contributed_to TYPE string;
139 DEFINE FIELD IF NOT EXISTS timestamp ON contributed_to TYPE datetime DEFAULT time::now();
140
141 DEFINE INDEX IF NOT EXISTS ct_session ON contributed_to FIELDS session_id;
142 "#,
143 )
144 .await?
145 .check()?;
146
147 Ok(())
148}