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(
123 files.len(),
124 args.max_files,
125 ),
126 ));
127 }
128
129 files.sort();
130
131 emit_json(&serde_json::json!({
132 "phase": "scan",
133 "dir": args.dir.display().to_string(),
134 "files_total": files.len(),
135 "files_new": files.len(),
136 "files_existing": 0,
137 }));
138
139 if args.dry_run {
140 for (idx, file) in files.iter().enumerate() {
141 let (name, truncated, orig) =
142 super::ingest::derive_kebab_name(file, args.max_name_length);
143 emit_json(&serde_json::json!({
144 "file": file.display().to_string(),
145 "name": name,
146 "status": "preview",
147 "index": idx + 1,
148 "total": files.len(),
149 "truncated": truncated,
150 "original_name": orig,
151 }));
152 }
153 emit_json(&serde_json::json!({
154 "summary": true,
155 "files_total": files.len(),
156 "completed": 0,
157 "failed": 0,
158 "skipped": 0,
159 "entities_total": 0,
160 "rels_total": 0,
161 "cost_usd": 0.0,
162 "elapsed_ms": started.elapsed().as_millis() as u64,
163 }));
164 return Ok(());
165 }
166
167 let ns = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
172 let app_paths = crate::paths::AppPaths::resolve(args.db.as_deref())?;
173
174 let mut completed = 0usize;
175 let mut failed = 0usize;
176 let mut skipped = 0usize;
177 let mut entities_total = 0usize;
178 let mut rels_total = 0usize;
179 let mut cost_total: f64 = 0.0;
180
181 for (idx, file) in files.iter().enumerate() {
182 let (name, truncated, orig) = super::ingest::derive_kebab_name(file, args.max_name_length);
183
184 let body = match std::fs::read_to_string(file) {
185 Ok(b) => b,
186 Err(e) => {
187 emit_json(&serde_json::json!({
188 "file": file.display().to_string(),
189 "name": name,
190 "status": "failed",
191 "error": format!("read error: {e}"),
192 "index": idx + 1,
193 "total": files.len(),
194 }));
195 failed += 1;
196 if args.fail_fast {
197 break;
198 }
199 continue;
200 }
201 };
202
203 if body.len() > 512_000 {
204 emit_json(&serde_json::json!({
205 "file": file.display().to_string(),
206 "name": name,
207 "status": "skipped",
208 "error": format!("file exceeds 512KB limit ({} bytes)", body.len()),
209 "index": idx + 1,
210 "total": files.len(),
211 }));
212 skipped += 1;
213 continue;
214 }
215
216 let file_started = std::time::Instant::now();
217
218 let fut = extract_with_opencode(&binary, &model, &body, &name, timeout);
220 let extraction = match tokio::runtime::Handle::try_current() {
221 Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fut)),
222 Err(_) => crate::embedder::shared_runtime()?.block_on(fut),
223 };
224
225 match extraction {
226 Ok((result, cost, _tokens)) => {
227 let ent_count = result.entities.len();
228 let rel_count = result.relationships.len();
229
230 let graph_payload = serde_json::json!({
231 "body": body,
232 "entities": result.entities.iter().map(|e| {
233 serde_json::json!({"name": e.name, "entity_type": e.entity_type})
234 }).collect::<Vec<_>>(),
235 "relationships": result.relationships.iter().map(|r| {
236 serde_json::json!({
237 "source": r.source,
238 "target": r.target,
239 "relation": r.relation,
240 "strength": r.strength
241 })
242 }).collect::<Vec<_>>(),
243 });
244
245 let remember_result = persist_memory_with_graph(
246 &app_paths.db,
247 &ns,
248 &name,
249 &format!("{:?}", args.r#type).to_lowercase(),
250 &format!("ingested from {} via opencode", file.display()),
251 &graph_payload,
252 );
253
254 match remember_result {
255 Ok(memory_id) => {
256 entities_total += ent_count;
257 rels_total += rel_count;
258 cost_total += cost;
259 completed += 1;
260
261 emit_json(&serde_json::json!({
262 "file": file.display().to_string(),
263 "name": name,
264 "status": "done",
265 "memory_id": memory_id,
266 "entities": ent_count,
267 "rels": rel_count,
268 "cost_usd": cost,
269 "elapsed_ms": file_started.elapsed().as_millis() as u64,
270 "index": idx + 1,
271 "total": files.len(),
272 "truncated": truncated,
273 "original_name": orig,
274 }));
275 }
276 Err(e) => {
277 failed += 1;
278 emit_json(&serde_json::json!({
279 "file": file.display().to_string(),
280 "name": name,
281 "status": "failed",
282 "error": format!("persist error: {e}"),
283 "elapsed_ms": file_started.elapsed().as_millis() as u64,
284 "index": idx + 1,
285 "total": files.len(),
286 }));
287 if args.fail_fast {
288 break;
289 }
290 }
291 }
292 }
293 Err(e) => {
294 failed += 1;
295 emit_json(&serde_json::json!({
296 "file": file.display().to_string(),
297 "name": name,
298 "status": "failed",
299 "error": format!("extraction error: {e}"),
300 "elapsed_ms": file_started.elapsed().as_millis() as u64,
301 "index": idx + 1,
302 "total": files.len(),
303 }));
304 if args.fail_fast {
305 break;
306 }
307 }
308 }
309 }
310
311 emit_json(&serde_json::json!({
312 "summary": true,
313 "files_total": files.len(),
314 "completed": completed,
315 "failed": failed,
316 "skipped": skipped,
317 "entities_total": entities_total,
318 "rels_total": rels_total,
319 "cost_usd": cost_total,
320 "elapsed_ms": started.elapsed().as_millis() as u64,
321 }));
322
323 Ok(())
324}
325
326fn persist_memory_with_graph(
327 db_path: &Path,
328 namespace: &str,
329 name: &str,
330 memory_type: &str,
331 description: &str,
332 graph_payload: &serde_json::Value,
333) -> Result<i64, AppError> {
334 let conn = crate::storage::connection::open_rw(db_path)?;
335
336 let existing = conn
337 .query_row(
338 "SELECT id FROM memories WHERE name = ?1 AND namespace = ?2",
339 rusqlite::params![name, namespace],
340 |row| row.get::<_, i64>(0),
341 )
342 .ok();
343
344 let body = graph_payload
345 .get("body")
346 .and_then(|b| b.as_str())
347 .unwrap_or("");
348 let body_hash = blake3::hash(body.as_bytes()).to_hex().to_string();
349
350 let memory_id = if let Some(id) = existing {
351 conn.execute(
352 "UPDATE memories SET body = ?1, description = ?2, type = ?3, body_hash = ?4, updated_at = strftime('%s','now') WHERE id = ?5",
353 rusqlite::params![body, description, memory_type, body_hash, id],
354 )
355 .map_err(AppError::Database)?;
356 id
357 } else {
358 conn.execute(
359 "INSERT INTO memories (name, namespace, type, description, body, body_hash, created_at, updated_at) \
360 VALUES (?1, ?2, ?3, ?4, ?5, ?6, strftime('%s','now'), strftime('%s','now'))",
361 rusqlite::params![name, namespace, memory_type, description, body, body_hash],
362 )
363 .map_err(AppError::Database)?;
364 conn.last_insert_rowid()
365 };
366
367 if let Some(entities) = graph_payload.get("entities").and_then(|e| e.as_array()) {
368 for ent in entities {
369 let ent_name = ent.get("name").and_then(|n| n.as_str()).unwrap_or("");
370 let ent_type = ent
371 .get("entity_type")
372 .and_then(|t| t.as_str())
373 .unwrap_or("concept");
374 if ent_name.len() < 2 {
375 continue;
376 }
377 let normalized = normalize_entity_name(ent_name);
378 conn.execute(
379 "INSERT OR IGNORE INTO entities (name, type, namespace) VALUES (?1, ?2, ?3)",
380 rusqlite::params![normalized, ent_type, namespace],
381 )
382 .map_err(AppError::Database)?;
383
384 let entity_id: i64 = conn
385 .query_row(
386 "SELECT id FROM entities WHERE name = ?1 AND namespace = ?2",
387 rusqlite::params![normalized, namespace],
388 |row| row.get(0),
389 )
390 .map_err(AppError::Database)?;
391
392 conn.execute(
393 "INSERT OR IGNORE INTO memory_entities (memory_id, entity_id) VALUES (?1, ?2)",
394 rusqlite::params![memory_id, entity_id],
395 )
396 .map_err(AppError::Database)?;
397 }
398 }
399
400 if let Some(rels) = graph_payload
401 .get("relationships")
402 .and_then(|r| r.as_array())
403 {
404 for rel in rels {
405 let source = rel.get("source").and_then(|s| s.as_str()).unwrap_or("");
406 let target = rel.get("target").and_then(|t| t.as_str()).unwrap_or("");
407 let relation = rel
408 .get("relation")
409 .and_then(|r| r.as_str())
410 .unwrap_or("related");
411 let strength = rel.get("strength").and_then(|s| s.as_f64()).unwrap_or(0.5);
412
413 if source.len() < 2 || target.len() < 2 {
414 continue;
415 }
416
417 let src_norm = normalize_entity_name(source);
418 let tgt_norm = normalize_entity_name(target);
419
420 for name_val in [&src_norm, &tgt_norm] {
421 conn.execute(
422 "INSERT OR IGNORE INTO entities (name, type, namespace) VALUES (?1, 'concept', ?2)",
423 rusqlite::params![name_val, namespace],
424 )
425 .map_err(AppError::Database)?;
426 }
427
428 let src_id: i64 = conn
429 .query_row(
430 "SELECT id FROM entities WHERE name = ?1 AND namespace = ?2",
431 rusqlite::params![src_norm, namespace],
432 |row| row.get(0),
433 )
434 .map_err(AppError::Database)?;
435
436 let tgt_id: i64 = conn
437 .query_row(
438 "SELECT id FROM entities WHERE name = ?1 AND namespace = ?2",
439 rusqlite::params![tgt_norm, namespace],
440 |row| row.get(0),
441 )
442 .map_err(AppError::Database)?;
443
444 let rel_normalized = relation.replace('-', "_");
445 conn.execute(
446 "INSERT OR IGNORE INTO relationships (source_id, target_id, relation, weight, namespace) \
447 VALUES (?1, ?2, ?3, ?4, ?5)",
448 rusqlite::params![src_id, tgt_id, rel_normalized, strength, namespace],
449 )
450 .map_err(AppError::Database)?;
451 }
452 }
453
454 Ok(memory_id)
455}
456
457#[cfg(test)]
458mod tests {
459 use super::*;
460
461 #[test]
462 fn extraction_result_deserializes_empty() {
463 let json = r#"{"entities":[],"relationships":[]}"#;
464 let result: ExtractionResult = serde_json::from_str(json).unwrap();
465 assert!(result.entities.is_empty());
466 assert!(result.relationships.is_empty());
467 }
468
469 #[test]
470 fn extraction_result_deserializes_with_data() {
471 let json = r#"{
472 "entities": [
473 {"name": "sqlite-graphrag", "entity_type": "project"},
474 {"name": "opencode", "entity_type": "tool"}
475 ],
476 "relationships": [
477 {"source": "sqlite-graphrag", "target": "opencode", "relation": "uses", "strength": 0.8}
478 ]
479 }"#;
480 let result: ExtractionResult = serde_json::from_str(json).unwrap();
481 assert_eq!(result.entities.len(), 2);
482 assert_eq!(result.relationships.len(), 1);
483 assert_eq!(result.relationships[0].strength, 0.8);
484 }
485
486 #[test]
487 fn extraction_result_default_strength() {
488 let json = r#"{
489 "entities": [],
490 "relationships": [
491 {"source": "a", "target": "b", "relation": "related"}
492 ]
493 }"#;
494 let result: ExtractionResult = serde_json::from_str(json).unwrap();
495 assert_eq!(result.relationships[0].strength, 0.5);
496 }
497}