use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};
use anyhow::Context;
use github_mcp::data::store::{VERSION_STORE_FILES, open_store_read_write};
use github_mcp::services::embedding_service::embed;
const ZSTD_LEVEL: i32 = 19;
const EMBEDDING_DIM: usize = 384;
fn zst_sibling(path: &Path) -> PathBuf {
let mut name = path.as_os_str().to_owned();
name.push(".zst");
PathBuf::from(name)
}
fn ensure_raw_db(path: &Path) -> anyhow::Result<()> {
if path.exists() {
return Ok(());
}
let zst_path = zst_sibling(path);
let compressed = std::fs::read(&zst_path).with_context(|| {
format!(
"neither '{}' nor '{}' exists",
path.display(),
zst_path.display()
)
})?;
let decompressed = zstd::stream::decode_all(compressed.as_slice())
.with_context(|| format!("failed to decompress '{}'", zst_path.display()))?;
std::fs::write(path, decompressed)
.with_context(|| format!("failed to write decompressed '{}'", path.display()))?;
Ok(())
}
fn recompress_and_remove_raw(path: &Path) -> anyhow::Result<()> {
let raw =
std::fs::read(path).with_context(|| format!("failed to read '{}'", path.display()))?;
let compressed = zstd::stream::encode_all(raw.as_slice(), ZSTD_LEVEL)
.with_context(|| format!("failed to zstd-compress '{}'", path.display()))?;
let zst_path = zst_sibling(path);
let mut file = File::create(&zst_path)
.with_context(|| format!("failed to create '{}'", zst_path.display()))?;
file.write_all(&compressed)
.with_context(|| format!("failed to write '{}'", zst_path.display()))?;
std::fs::remove_file(path).with_context(|| {
format!(
"failed to remove raw '{}' after recompressing",
path.display()
)
})?;
Ok(())
}
struct EndpointRow {
operation_id: String,
path: String,
method: String,
summary: Option<String>,
description: Option<String>,
}
fn vector_to_le_bytes(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|value| value.to_le_bytes())
.collect()
}
fn populate_one(path: &Path) -> anyhow::Result<usize> {
let conn = open_store_read_write(path)?;
conn.execute("DROP TABLE IF EXISTS semantic_endpoints", [])
.with_context(|| {
format!(
"failed to drop 'semantic_endpoints' in '{}'",
path.display()
)
})?;
conn.execute(
&format!(
"CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
operation_id TEXT PRIMARY KEY,
embedding FLOAT[{EMBEDDING_DIM}]
)"
),
[],
)
.with_context(|| {
format!(
"failed to recreate 'semantic_endpoints' in '{}'",
path.display()
)
})?;
let mut select =
conn.prepare("SELECT operation_id, path, method, summary, description FROM endpoints")?;
let rows: Vec<EndpointRow> = select
.query_map([], |row| {
Ok(EndpointRow {
operation_id: row.get(0)?,
path: row.get(1)?,
method: row.get(2)?,
summary: row.get(3)?,
description: row.get(4)?,
})
})?
.collect::<Result<_, _>>()?;
let mut delete = conn.prepare("DELETE FROM semantic_endpoints WHERE operation_id = ?1")?;
let mut insert =
conn.prepare("INSERT INTO semantic_endpoints (operation_id, embedding) VALUES (?1, ?2)")?;
let count = rows.len();
for row in rows {
let operation_id = row.operation_id.clone();
let text = [
Some(row.method),
Some(row.path),
row.summary,
row.description,
]
.into_iter()
.flatten()
.collect::<Vec<_>>()
.join(" ");
let vector = embed(&text).map_err(|err| {
anyhow::anyhow!(
"failed to compute embedding for '{operation_id}' in '{}': {err}",
path.display()
)
})?;
delete.execute(rusqlite::params![row.operation_id])?;
insert.execute(rusqlite::params![
row.operation_id,
vector_to_le_bytes(&vector)
])?;
}
conn.execute("VACUUM", [])
.with_context(|| format!("failed to VACUUM '{}'", path.display()))?;
Ok(count)
}
fn missing_operation_ids(path: &Path) -> anyhow::Result<Vec<String>> {
let conn = open_store_read_write(path)?;
let mut select = conn.prepare(
"SELECT e.operation_id FROM endpoints e \
LEFT JOIN semantic_endpoints s ON s.operation_id = e.operation_id \
WHERE s.operation_id IS NULL",
)?;
let missing: Vec<String> = select
.query_map([], |row| row.get(0))?
.collect::<Result<_, _>>()?;
Ok(missing)
}
fn targets_from(mut args: impl Iterator<Item = String>) -> Vec<PathBuf> {
match args.next().as_deref() {
Some("--all") | None => VERSION_STORE_FILES
.iter()
.map(|(_, file)| PathBuf::from(file))
.collect(),
Some(path) => vec![PathBuf::from(path)],
}
}
fn targets() -> Vec<PathBuf> {
targets_from(std::env::args().skip(1))
}
fn main() -> anyhow::Result<()> {
let mut had_mismatch = false;
for path in targets() {
ensure_raw_db(&path)?;
let count = populate_one(&path)?;
println!(
"populated embeddings for {count} operation(s) in '{}'",
path.display()
);
let missing = missing_operation_ids(&path)?;
if !missing.is_empty() {
had_mismatch = true;
eprintln!(
"ERROR: '{}' has {} endpoint(s) missing from semantic_endpoints (row-count \
parity check failed): {}",
path.display(),
missing.len(),
missing.join(", ")
);
}
recompress_and_remove_raw(&path)?;
}
if had_mismatch {
anyhow::bail!(
"embedding population incomplete: one or more stores have semantic_endpoints \
row counts below their endpoints row counts (see errors above)"
);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn target_selection_covers_defaults_all_and_explicit_paths_without_mutating_stores() {
let expected = VERSION_STORE_FILES
.iter()
.map(|(_, file)| PathBuf::from(file))
.collect::<Vec<_>>();
assert_eq!(targets_from(std::iter::empty()), expected);
assert_eq!(targets_from(["--all".to_string()].into_iter()), expected);
assert_eq!(
targets_from(["copy.db".to_string()].into_iter()),
vec![PathBuf::from("copy.db")]
);
}
}