1use crate::commands::ingest::IngestArgs;
8use crate::commands::opencode_runner;
9use crate::errors::AppError;
10use crate::parsers::normalize_entity_name;
11use serde::{Deserialize, Serialize};
12use std::io::Write;
13use std::path::{Path, PathBuf};
14
15const EXTRACTION_SCHEMA: &str = r#"Return ONLY a valid JSON object with this exact structure (no markdown, no explanation):
16{
17 "entities": [
18 {"name": "entity-name-in-kebab-case", "entity_type": "concept|project|tool|person|file|incident|decision|organization|location|date"}
19 ],
20 "relationships": [
21 {"source": "entity-a", "target": "entity-b", "relation": "applies-to|uses|depends-on|causes|fixes|contradicts|supports|follows|related|replaces|tracked-in", "strength": 0.7}
22 ]
23}"#;
24
25#[derive(Debug, Deserialize, Serialize)]
27pub struct ExtractionResult {
28 #[serde(default)]
30 pub entities: Vec<ExtractedEntity>,
31 #[serde(default)]
33 pub relationships: Vec<ExtractedRelationship>,
34}
35
36#[derive(Debug, Deserialize, Serialize, Clone)]
38pub struct ExtractedEntity {
39 pub name: String,
41 pub entity_type: String,
43}
44
45#[derive(Debug, Deserialize, Serialize, Clone)]
47pub struct ExtractedRelationship {
48 pub source: String,
50 pub target: String,
52 pub relation: String,
54 #[serde(default = "default_strength")]
56 pub strength: f64,
57}
58
59fn default_strength() -> f64 {
60 0.5
61}
62
63pub async fn extract_with_opencode(
65 binary: &Path,
66 model: &str,
67 body: &str,
68 memory_name: &str,
69 timeout_secs: u64,
70) -> Result<(ExtractionResult, f64, u64), AppError> {
71 let prompt = format!(
72 "Analyze the following document and extract domain-specific entities and their relationships.\n\
73 Memory name: {memory_name}\n\n\
74 {EXTRACTION_SCHEMA}\n\n\
75 Document content:\n{body}"
76 );
77
78 opencode_runner::call_opencode::<ExtractionResult>(binary, model, &prompt, timeout_secs).await
79}
80
81fn emit_json(value: &serde_json::Value) {
82 let _ = writeln!(
83 std::io::stdout(),
84 "{}",
85 serde_json::to_string(value).unwrap_or_default()
86 );
87 let _ = std::io::stdout().flush();
88}
89
90pub fn run_opencode_ingest(args: &IngestArgs) -> Result<(), AppError> {
92 let started = std::time::Instant::now();
93
94 if !args.dir.exists() {
95 return Err(AppError::Validation(
96 crate::i18n::validation::directory_not_found(&args.dir.display().to_string()),
97 ));
98 }
99
100 let binary =
101 opencode_runner::find_opencode_binary_with_override(args.opencode_binary.as_deref())?;
102 let version = opencode_runner::validate_opencode_version(&binary)?;
103 let model = opencode_runner::resolve_opencode_model(args.opencode_model.as_deref());
104 let timeout = opencode_runner::resolve_opencode_timeout(if args.opencode_timeout != 300 {
105 Some(args.opencode_timeout)
106 } else {
107 None
108 });
109
110 emit_json(&serde_json::json!({
111 "phase": "validate",
112 "opencode_path": binary.display().to_string(),
113 "version": format!("{}.{}.{}", version.0, version.1, version.2),
114 "model": &model,
115 }));
116
117 let mut files: Vec<PathBuf> = Vec::new();
118 super::ingest::collect_files(&args.dir, &args.pattern, args.recursive, &mut files)?;
119
120 if files.len() > args.max_files {
121 return Err(AppError::Validation(
122 crate::i18n::validation::max_files_exceeded_all_or_nothing(files.len(), args.max_files),
123 ));
124 }
125
126 files.sort();
127
128 emit_json(&serde_json::json!({
129 "phase": "scan",
130 "dir": args.dir.display().to_string(),
131 "files_total": files.len(),
132 "files_new": files.len(),
133 "files_existing": 0,
134 }));
135
136 if args.dry_run {
137 for (idx, file) in files.iter().enumerate() {
138 let (name, truncated, orig) =
139 super::ingest::derive_kebab_name(file, args.max_name_length);
140 emit_json(&serde_json::json!({
141 "file": file.display().to_string(),
142 "name": name,
143 "status": "preview",
144 "index": idx + 1,
145 "total": files.len(),
146 "truncated": truncated,
147 "original_name": orig,
148 }));
149 }
150 emit_json(&serde_json::json!({
151 "summary": true,
152 "files_total": files.len(),
153 "completed": 0,
154 "failed": 0,
155 "skipped": 0,
156 "entities_total": 0,
157 "rels_total": 0,
158 "cost_usd": 0.0,
159 "elapsed_ms": started.elapsed().as_millis() as u64,
160 }));
161 return Ok(());
162 }
163
164 let ns = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
169 let app_paths = crate::paths::AppPaths::resolve(args.db.as_deref())?;
170
171 let mut completed = 0usize;
172 let mut failed = 0usize;
173 let mut skipped = 0usize;
174 let mut entities_total = 0usize;
175 let mut rels_total = 0usize;
176 let mut cost_total: f64 = 0.0;
177
178 for (idx, file) in files.iter().enumerate() {
179 let (name, truncated, orig) = super::ingest::derive_kebab_name(file, args.max_name_length);
180
181 let body = match std::fs::read_to_string(file) {
182 Ok(b) => b,
183 Err(e) => {
184 emit_json(&serde_json::json!({
185 "file": file.display().to_string(),
186 "name": name,
187 "status": "failed",
188 "error": format!("read error: {e}"),
189 "index": idx + 1,
190 "total": files.len(),
191 }));
192 failed += 1;
193 if args.fail_fast {
194 break;
195 }
196 continue;
197 }
198 };
199
200 if body.len() > 512_000 {
201 emit_json(&serde_json::json!({
202 "file": file.display().to_string(),
203 "name": name,
204 "status": "skipped",
205 "error": format!("file exceeds 512KB limit ({} bytes)", body.len()),
206 "index": idx + 1,
207 "total": files.len(),
208 }));
209 skipped += 1;
210 continue;
211 }
212
213 let file_started = std::time::Instant::now();
214
215 let fut = extract_with_opencode(&binary, &model, &body, &name, timeout);
217 let extraction = match tokio::runtime::Handle::try_current() {
218 Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fut)),
219 Err(_) => crate::embedder::shared_runtime()?.block_on(fut),
220 };
221
222 match extraction {
223 Ok((result, cost, _tokens)) => {
224 let ent_count = result.entities.len();
225 let rel_count = result.relationships.len();
226
227 let graph_payload = serde_json::json!({
228 "body": body,
229 "entities": result.entities.iter().map(|e| {
230 serde_json::json!({"name": e.name, "entity_type": e.entity_type})
231 }).collect::<Vec<_>>(),
232 "relationships": result.relationships.iter().map(|r| {
233 serde_json::json!({
234 "source": r.source,
235 "target": r.target,
236 "relation": r.relation,
237 "strength": r.strength
238 })
239 }).collect::<Vec<_>>(),
240 });
241
242 let remember_result = persist_memory_with_graph(
243 &app_paths.db,
244 &ns,
245 &name,
246 &format!("{:?}", args.r#type).to_lowercase(),
247 &format!("ingested from {} via opencode", file.display()),
248 &graph_payload,
249 );
250
251 match remember_result {
252 Ok(memory_id) => {
253 entities_total += ent_count;
254 rels_total += rel_count;
255 cost_total += cost;
256 completed += 1;
257
258 emit_json(&serde_json::json!({
259 "file": file.display().to_string(),
260 "name": name,
261 "status": "done",
262 "memory_id": memory_id,
263 "entities": ent_count,
264 "rels": rel_count,
265 "cost_usd": cost,
266 "elapsed_ms": file_started.elapsed().as_millis() as u64,
267 "index": idx + 1,
268 "total": files.len(),
269 "truncated": truncated,
270 "original_name": orig,
271 }));
272 }
273 Err(e) => {
274 failed += 1;
275 emit_json(&serde_json::json!({
276 "file": file.display().to_string(),
277 "name": name,
278 "status": "failed",
279 "error": format!("persist error: {e}"),
280 "elapsed_ms": file_started.elapsed().as_millis() as u64,
281 "index": idx + 1,
282 "total": files.len(),
283 }));
284 if args.fail_fast {
285 break;
286 }
287 }
288 }
289 }
290 Err(e) => {
291 failed += 1;
292 emit_json(&serde_json::json!({
293 "file": file.display().to_string(),
294 "name": name,
295 "status": "failed",
296 "error": format!("extraction error: {e}"),
297 "elapsed_ms": file_started.elapsed().as_millis() as u64,
298 "index": idx + 1,
299 "total": files.len(),
300 }));
301 if args.fail_fast {
302 break;
303 }
304 }
305 }
306 }
307
308 emit_json(&serde_json::json!({
309 "summary": true,
310 "files_total": files.len(),
311 "completed": completed,
312 "failed": failed,
313 "skipped": skipped,
314 "entities_total": entities_total,
315 "rels_total": rels_total,
316 "cost_usd": cost_total,
317 "elapsed_ms": started.elapsed().as_millis() as u64,
318 }));
319
320 Ok(())
321}
322
323fn persist_memory_with_graph(
324 db_path: &Path,
325 namespace: &str,
326 name: &str,
327 memory_type: &str,
328 description: &str,
329 graph_payload: &serde_json::Value,
330) -> Result<i64, AppError> {
331 let conn = crate::storage::connection::open_rw(db_path)?;
332
333 let existing = conn
334 .query_row(
335 "SELECT id FROM memories WHERE name = ?1 AND namespace = ?2",
336 rusqlite::params![name, namespace],
337 |row| row.get::<_, i64>(0),
338 )
339 .ok();
340
341 let body = graph_payload
342 .get("body")
343 .and_then(|b| b.as_str())
344 .unwrap_or("");
345 let body_hash = blake3::hash(body.as_bytes()).to_hex().to_string();
346
347 let memory_id = if let Some(id) = existing {
348 conn.execute(
349 "UPDATE memories SET body = ?1, description = ?2, type = ?3, body_hash = ?4, updated_at = strftime('%s','now') WHERE id = ?5",
350 rusqlite::params![body, description, memory_type, body_hash, id],
351 )
352 .map_err(AppError::Database)?;
353 id
354 } else {
355 conn.execute(
356 "INSERT INTO memories (name, namespace, type, description, body, body_hash, created_at, updated_at) \
357 VALUES (?1, ?2, ?3, ?4, ?5, ?6, strftime('%s','now'), strftime('%s','now'))",
358 rusqlite::params![name, namespace, memory_type, description, body, body_hash],
359 )
360 .map_err(AppError::Database)?;
361 conn.last_insert_rowid()
362 };
363
364 if let Some(entities) = graph_payload.get("entities").and_then(|e| e.as_array()) {
365 for ent in entities {
366 let ent_name = ent.get("name").and_then(|n| n.as_str()).unwrap_or("");
367 let ent_type = ent
368 .get("entity_type")
369 .and_then(|t| t.as_str())
370 .unwrap_or("concept");
371 if ent_name.len() < 2 {
372 continue;
373 }
374 let normalized = normalize_entity_name(ent_name);
375 conn.execute(
376 "INSERT OR IGNORE INTO entities (name, type, namespace) VALUES (?1, ?2, ?3)",
377 rusqlite::params![normalized, ent_type, namespace],
378 )
379 .map_err(AppError::Database)?;
380
381 let entity_id: i64 = conn
382 .query_row(
383 "SELECT id FROM entities WHERE name = ?1 AND namespace = ?2",
384 rusqlite::params![normalized, namespace],
385 |row| row.get(0),
386 )
387 .map_err(AppError::Database)?;
388
389 conn.execute(
390 "INSERT OR IGNORE INTO memory_entities (memory_id, entity_id) VALUES (?1, ?2)",
391 rusqlite::params![memory_id, entity_id],
392 )
393 .map_err(AppError::Database)?;
394 }
395 }
396
397 if let Some(rels) = graph_payload
398 .get("relationships")
399 .and_then(|r| r.as_array())
400 {
401 for rel in rels {
402 let source = rel.get("source").and_then(|s| s.as_str()).unwrap_or("");
403 let target = rel.get("target").and_then(|t| t.as_str()).unwrap_or("");
404 let relation = rel
405 .get("relation")
406 .and_then(|r| r.as_str())
407 .unwrap_or("related");
408 let strength = rel.get("strength").and_then(|s| s.as_f64()).unwrap_or(0.5);
409
410 if source.len() < 2 || target.len() < 2 {
411 continue;
412 }
413
414 let src_norm = normalize_entity_name(source);
415 let tgt_norm = normalize_entity_name(target);
416
417 for name_val in [&src_norm, &tgt_norm] {
418 conn.execute(
419 "INSERT OR IGNORE INTO entities (name, type, namespace) VALUES (?1, 'concept', ?2)",
420 rusqlite::params![name_val, namespace],
421 )
422 .map_err(AppError::Database)?;
423 }
424
425 let src_id: i64 = conn
426 .query_row(
427 "SELECT id FROM entities WHERE name = ?1 AND namespace = ?2",
428 rusqlite::params![src_norm, namespace],
429 |row| row.get(0),
430 )
431 .map_err(AppError::Database)?;
432
433 let tgt_id: i64 = conn
434 .query_row(
435 "SELECT id FROM entities WHERE name = ?1 AND namespace = ?2",
436 rusqlite::params![tgt_norm, namespace],
437 |row| row.get(0),
438 )
439 .map_err(AppError::Database)?;
440
441 let rel_normalized = relation.replace('-', "_");
442 conn.execute(
443 "INSERT OR IGNORE INTO relationships (source_id, target_id, relation, weight, namespace) \
444 VALUES (?1, ?2, ?3, ?4, ?5)",
445 rusqlite::params![src_id, tgt_id, rel_normalized, strength, namespace],
446 )
447 .map_err(AppError::Database)?;
448 }
449 }
450
451 Ok(memory_id)
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457
458 #[test]
459 fn extraction_result_deserializes_empty() {
460 let json = r#"{"entities":[],"relationships":[]}"#;
461 let result: ExtractionResult = serde_json::from_str(json).unwrap();
462 assert!(result.entities.is_empty());
463 assert!(result.relationships.is_empty());
464 }
465
466 #[test]
467 fn extraction_result_deserializes_with_data() {
468 let json = r#"{
469 "entities": [
470 {"name": "sqlite-graphrag", "entity_type": "project"},
471 {"name": "opencode", "entity_type": "tool"}
472 ],
473 "relationships": [
474 {"source": "sqlite-graphrag", "target": "opencode", "relation": "uses", "strength": 0.8}
475 ]
476 }"#;
477 let result: ExtractionResult = serde_json::from_str(json).unwrap();
478 assert_eq!(result.entities.len(), 2);
479 assert_eq!(result.relationships.len(), 1);
480 assert_eq!(result.relationships[0].strength, 0.8);
481 }
482
483 #[test]
484 fn extraction_result_default_strength() {
485 let json = r#"{
486 "entities": [],
487 "relationships": [
488 {"source": "a", "target": "b", "relation": "related"}
489 ]
490 }"#;
491 let result: ExtractionResult = serde_json::from_str(json).unwrap();
492 assert_eq!(result.relationships[0].strength, 0.5);
493 }
494}