atheneum 0.4.0

Agent coordination graph database - episodic and semantic memory for multi-agent workflows
Documentation
use anyhow::Result;
use chrono::Utc;
use serde_json::{json, Value};
use sqlitegraph::GraphEntity;

use super::{AtheneumGraph, EdgeType, EntityType};

#[derive(Debug, Clone)]
struct MagellanSymbol {
    name: String,
    file_path: Option<String>,
    data: Value,
}

fn read_magellan_symbols(
    db_path: &std::path::Path,
    name_filter: Option<&str>,
    limit: Option<usize>,
) -> Result<Vec<MagellanSymbol>> {
    let conn =
        rusqlite::Connection::open_with_flags(db_path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)
            .map_err(|e| {
                anyhow::anyhow!("Failed to open magellan DB {}: {}", db_path.display(), e)
            })?;

    let mut sql =
        String::from("SELECT name, file_path, data FROM graph_entities WHERE kind = 'Symbol'");
    if name_filter.is_some() {
        sql.push_str(" AND name = ?1");
    }
    if let Some(n) = limit {
        sql.push_str(&format!(" LIMIT {}", n));
    }

    let mut stmt = conn.prepare(&sql)?;
    let mapper = |row: &rusqlite::Row<'_>| -> rusqlite::Result<MagellanSymbol> {
        let data_str: String = row.get(2)?;
        let data: Value = serde_json::from_str(&data_str).unwrap_or(Value::Null);
        Ok(MagellanSymbol {
            name: row.get::<_, String>(0)?,
            file_path: row.get::<_, Option<String>>(1)?,
            data,
        })
    };

    let rows: Vec<MagellanSymbol> = if let Some(filter) = name_filter {
        stmt.query_map(rusqlite::params![filter], mapper)?
            .collect::<Result<Vec<_>, _>>()?
    } else {
        stmt.query_map([], mapper)?.collect::<Result<Vec<_>, _>>()?
    };
    Ok(rows)
}

fn symbol_to_metadata(sym: &MagellanSymbol) -> Value {
    let start_line = sym.data.get("start_line").cloned().unwrap_or(Value::Null);
    let end_line = sym.data.get("end_line").cloned().unwrap_or(Value::Null);
    let fqn = sym.data.get("fqn").cloned().unwrap_or(Value::Null);
    let kind = sym.data.get("kind").cloned().unwrap_or(Value::Null);
    json!({
        "file": sym.file_path,
        "start_line": start_line,
        "end_line": end_line,
        "fqn": fqn,
        "kind": kind,
        "magellan_raw": sym.data,
    })
}

impl AtheneumGraph {
    pub fn import_symbol_from_magellan(
        &self,
        magellan_db_path: &std::path::Path,
        symbol_name: &str,
        agent_name: &str,
        project_id: Option<&str>,
    ) -> Result<Option<i64>> {
        let symbols = read_magellan_symbols(magellan_db_path, Some(symbol_name), None)?;
        if symbols.is_empty() {
            return Ok(None);
        }
        let sym = &symbols[0];
        let metadata = symbol_to_metadata(sym);
        self.upsert_symbol_discovery(agent_name, symbol_name, project_id, metadata)
            .map(Some)
    }

    pub fn import_all_symbols_from_magellan(
        &self,
        magellan_db_path: &std::path::Path,
        agent_name: &str,
        project_id: Option<&str>,
        limit: Option<usize>,
    ) -> Result<usize> {
        let symbols = read_magellan_symbols(magellan_db_path, None, limit)?;
        let mut count = 0;
        for sym in &symbols {
            let metadata = symbol_to_metadata(sym);
            let target = sym.name.clone();
            self.upsert_symbol_discovery(agent_name, &target, project_id, metadata)?;
            count += 1;
        }
        Ok(count)
    }

    pub fn link_wiki_to_symbols(
        &self,
        magellan_db_path: &std::path::Path,
        agent_name: &str,
        project_id: Option<&str>,
    ) -> Result<usize> {
        let pages = self.list_wiki_pages(project_id)?;
        let mut linked = 0;

        for page in &pages {
            for wikilink in &page.wikilinks {
                let matches = read_magellan_symbols(magellan_db_path, Some(wikilink), None)?;
                if matches.is_empty() {
                    continue;
                }

                let sym = &matches[0];
                let metadata = symbol_to_metadata(sym);
                let symbol_entity_id =
                    self.upsert_symbol_discovery(agent_name, &sym.name, project_id, metadata)?;

                let existing_edge = self.with_raw_connection(|conn| {
                    Ok(conn
                        .query_row(
                            "SELECT id FROM graph_edges WHERE from_id = ?1 AND to_id = ?2 AND edge_type = ?3 LIMIT 1",
                            rusqlite::params![page.id, symbol_entity_id, EdgeType::Explains.as_str()],
                            |r| r.get::<_, i64>(0),
                        )
                        .ok())
                })?;

                if existing_edge.is_none() {
                    self.insert_edge(
                        page.id,
                        symbol_entity_id,
                        EdgeType::Explains,
                        json!({"wikilink": wikilink, "symbol": sym.name}),
                    )?;
                    linked += 1;
                }
            }
        }

        Ok(linked)
    }

    fn upsert_symbol_discovery(
        &self,
        agent_name: &str,
        target: &str,
        project_id: Option<&str>,
        mut metadata: Value,
    ) -> Result<i64> {
        let agent_s = agent_name.to_string();
        let target_s = target.to_string();
        let project_s = project_id.map(|s| s.to_string());
        let existing_sql_id: Option<i64> = self.with_raw_connection(|conn| {
            let row = conn
                .query_row(
                    "SELECT id FROM discoveries
                     WHERE agent_name = ?1 AND target = ?2 AND discovery_type = 'Symbol'
                       AND ((project_id IS NULL AND ?3 IS NULL) OR project_id = ?3)
                     LIMIT 1",
                    rusqlite::params![agent_s, target_s, project_s],
                    |r| r.get::<_, i64>(0),
                )
                .ok();
            Ok(row)
        })?;

        if let Some(sql_id) = existing_sql_id {
            let agent_s = agent_name.to_string();
            let target_s = target.to_string();
            let project_s = project_id.map(|s| s.to_string());
            if let Some(obj) = metadata.as_object_mut() {
                obj.insert("agent".to_string(), Value::String(agent_s.clone()));
                obj.insert(
                    "discovery_type".to_string(),
                    Value::String("Symbol".to_string()),
                );
                obj.insert("target".to_string(), Value::String(target_s.clone()));
                obj.insert(
                    "timestamp".to_string(),
                    Value::String(Utc::now().to_rfc3339()),
                );
                obj.insert("sql_id".to_string(), Value::Number(sql_id.into()));
                if let Some(ref pid) = project_s {
                    obj.insert("project_id".to_string(), Value::String(pid.clone()));
                }
            }
            let metadata_str = super::json_to_string(&metadata)?;
            let metadata_for_entity = metadata.clone();
            self.with_raw_connection(|conn| {
                conn.execute(
                    "UPDATE discoveries SET metadata = ?1, project_id = ?2 WHERE id = ?3",
                    rusqlite::params![metadata_str, project_s, sql_id],
                )?;
                Ok(())
            })?;

            let entity_id: Option<i64> = self.with_raw_connection(|conn| {
                let row = conn
                    .query_row(
                        "SELECT id FROM graph_entities
                         WHERE kind = 'Discovery'
                           AND json_extract(data, '$.sql_id') = ?1
                         LIMIT 1",
                        rusqlite::params![sql_id],
                        |r| r.get::<_, i64>(0),
                    )
                    .ok();
                Ok(row)
            })?;

            if let Some(entity_id) = entity_id {
                self.update_entity_data(entity_id, &metadata_for_entity)?;
                Ok(entity_id)
            } else {
                let name = format!("{}: {}", agent_name, target);
                let entity = GraphEntity {
                    id: 0,
                    kind: EntityType::Discovery.as_str().to_string(),
                    name,
                    file_path: None,
                    data: metadata_for_entity,
                };
                self.inner.insert_entity(&entity).map_err(Into::into)
            }
        } else {
            self.store_discovery_in_project(agent_name, "Symbol", target, project_id, metadata)
        }
    }
}