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)
.with_context(|| format!("failed to compute embedding for '{operation_id}'"))?;
delete.execute(rusqlite::params![operation_id])?;
insert
.execute(rusqlite::params![operation_id, vector_to_le_bytes(&vector)])
.with_context(|| format!("failed to insert embedding for '{operation_id}'"))?;
}
drop(select);
drop(delete);
drop(insert);
check_completeness(&conn, path)?;
conn.execute("VACUUM", [])
.with_context(|| format!("failed to VACUUM '{}'", path.display()))?;
Ok(count)
}
fn check_completeness(conn: &rusqlite::Connection, path: &Path) -> anyhow::Result<()> {
let endpoints_count: usize =
conn.query_row("SELECT COUNT(*) FROM endpoints", [], |row| row.get(0))?;
let semantic_count: usize =
conn.query_row("SELECT COUNT(*) FROM semantic_endpoints", [], |row| {
row.get(0)
})?;
if endpoints_count != semantic_count {
let mut stmt = conn.prepare(
"SELECT operation_id FROM endpoints
WHERE operation_id NOT IN (SELECT operation_id FROM semantic_endpoints)",
)?;
let missing: Vec<String> = stmt
.query_map([], |row| row.get(0))?
.collect::<Result<_, _>>()?;
anyhow::bail!(
"'{}' is incomplete: {endpoints_count} endpoint(s) but only {semantic_count} \
semantic_endpoints row(s); missing operation_id(s): {}",
path.display(),
missing.join(", ")
);
}
Ok(())
}
fn targets() -> Vec<PathBuf> {
let mut args = std::env::args().skip(1);
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 main() -> anyhow::Result<()> {
for path in targets() {
ensure_raw_db(&path)?;
let count = populate_one(&path)?;
recompress_and_remove_raw(&path)?;
println!(
"populated embeddings for {count} operation(s) in '{}'",
path.display()
);
}
Ok(())
}