use rusqlite::{Connection, OptionalExtension};
use serde::Deserialize;
#[derive(Debug, Deserialize)]
struct MagellanConfig {
embeddings: EmbeddingConfig,
}
#[derive(Debug, Deserialize)]
struct EmbeddingConfig {
base_url: String,
model: String,
api_key: Option<String>,
}
pub fn embed_query(
base_url: &str,
model: &str,
api_key: &str,
query: &str,
) -> anyhow::Result<Vec<f32>> {
let url = format!("{}/api/embeddings", base_url.trim_end_matches('/'));
#[derive(serde::Serialize)]
struct OllamaEmbedRequest<'a> {
model: &'a str,
prompt: &'a str,
}
#[derive(serde::Deserialize)]
struct OllamaEmbedResponse {
embedding: Vec<f32>,
}
let body = OllamaEmbedRequest {
model,
prompt: query,
};
let json_body = serde_json::to_string(&body)
.map_err(|e| anyhow::anyhow!("Failed to serialize embedding request: {e}"))?;
let mut request = ureq::post(&url).header("Content-Type", "application/json");
if !api_key.is_empty() {
request = request.header("Authorization", &format!("Bearer {api_key}"));
}
let response = request
.send(&json_body)
.map_err(|e| anyhow::anyhow!("Ollama embedding request failed: {e}"))?;
let body = response
.into_body()
.read_to_string()
.map_err(|e| anyhow::anyhow!("Failed to read Ollama response body: {e}"))?;
let resp: OllamaEmbedResponse = serde_json::from_str(&body)
.map_err(|e| anyhow::anyhow!("Failed to parse Ollama embedding response: {e}"))?;
Ok(resp.embedding)
}
pub fn read_magellan_config() -> anyhow::Result<(String, String, String)> {
let home = std::env::var("HOME")
.or_else(|_| std::env::var("USERPROFILE"))
.map_err(|_| anyhow::anyhow!("Could not determine home directory"))?;
let config_path = std::path::Path::new(&home)
.join(".config")
.join("magellan")
.join("config.toml");
let contents = std::fs::read_to_string(&config_path).map_err(|e| {
anyhow::anyhow!("Failed to read config at {:?}: {}", config_path, e)
})?;
let config: MagellanConfig =
toml::from_str(&contents).map_err(|e| anyhow::anyhow!("Failed to parse config: {e}"))?;
let api_key = config.embeddings.api_key.unwrap_or_default();
Ok((config.embeddings.base_url, config.embeddings.model, api_key))
}
pub fn resolve_semantic_query(
db_path: &std::path::Path,
query: &str,
top_k: usize,
) -> anyhow::Result<Option<Vec<(i64, f32)>>> {
let conn =
Connection::open(db_path).map_err(|e| anyhow::anyhow!("Failed to open DB: {e}"))?;
let hnsw_exists: bool = conn
.query_row(
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='hnsw_indexes'",
[],
|_| Ok(true),
)
.unwrap_or(false);
if !hnsw_exists {
return Ok(None);
}
let (base_url, model, api_key) = read_magellan_config()?;
let query_vector = embed_query(&base_url, &model, &api_key, query)?;
let index = sqlitegraph::hnsw::HnswIndex::load_with_vectors(&conn, "symbols")
.map_err(|e| anyhow::anyhow!("Failed to load HNSW index: {e}"))?;
let results = index
.search(&query_vector, top_k)
.map_err(|e| anyhow::anyhow!("HNSW search failed: {e}"))?;
let mut matches = Vec::new();
for (vector_id, distance) in results {
let metadata_json: Option<String> = conn
.query_row(
"SELECT metadata FROM hnsw_vectors WHERE id = ?1",
[vector_id],
|row| row.get(0),
)
.optional()
.unwrap_or(None);
let metadata_json = match metadata_json {
Some(m) => m,
None => continue,
};
let metadata: serde_json::Value =
serde_json::from_str(&metadata_json).unwrap_or(serde_json::json!({}));
if let Some(entity_id) = metadata["entity_id"].as_i64() {
matches.push((entity_id, distance));
}
}
Ok(Some(matches))
}
pub fn resolve_semantic_to_symbol(
db_path: &std::path::Path,
query: &str,
) -> anyhow::Result<Option<crate::graph::magellan_integration::SymbolInfo>> {
let matches = resolve_semantic_query(db_path, query, 5)?;
let matches = match matches {
Some(m) => m,
None => return Ok(None),
};
for (entity_id, _distance) in matches {
if let Some(info) = resolve_symbol_by_entity_id(db_path, entity_id)? {
return Ok(Some(info));
}
}
Ok(None)
}
pub fn resolve_semantic_to_symbols(
db_path: &std::path::Path,
query: &str,
top_k: usize,
) -> anyhow::Result<Option<Vec<crate::graph::magellan_integration::SymbolInfo>>> {
let matches = resolve_semantic_query(db_path, query, top_k)?;
let matches = match matches {
Some(m) => m,
None => return Ok(None),
};
let mut symbols = Vec::new();
for (entity_id, _distance) in matches {
if let Some(info) = resolve_symbol_by_entity_id(db_path, entity_id)? {
symbols.push(info);
}
}
if symbols.is_empty() {
Ok(None)
} else {
Ok(Some(symbols))
}
}
pub fn resolve_symbol_by_entity_id(
db_path: &std::path::Path,
entity_id: i64,
) -> anyhow::Result<Option<crate::graph::magellan_integration::SymbolInfo>> {
use crate::graph::magellan_integration::SymbolInfo;
let conn =
Connection::open(db_path).map_err(|e| anyhow::anyhow!("Failed to open DB: {e}"))?;
let row: Option<(String, String, String)> = conn
.query_row(
"SELECT name, file_path, data FROM graph_entities WHERE id = ?1 AND kind = 'Symbol'",
[entity_id],
|row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
))
},
)
.optional()
.map_err(|e| anyhow::anyhow!("Failed to query symbol by entity ID: {e}"))?;
let (name, file_path, data_json) = match row {
Some(r) => r,
None => return Ok(None),
};
let data: serde_json::Value = serde_json::from_str(&data_json)
.map_err(|e| anyhow::anyhow!("Failed to parse symbol data JSON: {e}"))?;
let byte_start = data
.get("byte_start")
.and_then(|v| v.as_u64())
.unwrap_or(0) as usize;
let byte_end = data
.get("byte_end")
.and_then(|v| v.as_u64())
.unwrap_or(0) as usize;
let kind = data
.get("kind")
.and_then(|v| v.as_str())
.unwrap_or("Unknown")
.to_string();
let start_line = data.get("start_line").and_then(|v| v.as_u64()).map(|l| l as usize);
let end_line = data.get("end_line").and_then(|v| v.as_u64()).map(|l| l as usize);
Ok(Some(SymbolInfo {
entity_id,
name,
file_path,
kind,
byte_start,
byte_end,
start_line,
end_line,
}))
}