Me gusta 0
Bifurcar 0

Última actividad 18 hours ago

CLI en Rust para la base de conocimiento RAG. Combina FTS5 + cosine similarity con el parámetro alpha. Subcomandos: query, list (tags/docs/stats) y similar (query-by-example). Dependencias: rusqlite (bundled), clap, bytemuck, colored, reqwest, serde_json.

Revisión 26a2e58cbc1c11988d31b6ba9bcd7baf1e23ef54

main.rs Sin formato
1// cerebro-cli — CLI en Rust para la base de conocimiento RAG
2//
3// Compilación:
4// cd cerebro-cli && cargo build --release
5// cp target/release/cerebro ~/.local/bin/
6//
7// Uso:
8// cerebro query "hooks de git"
9// cerebro query "instalar docker" --alpha 0.6
10// cerebro query "nginx" --tag web
11// cerebro list tags
12// cerebro list docs
13// cerebro list stats
14// cerebro similar notas/docker.md --top-k 5 --tag linux
15
16use bytemuck::try_cast_slice;
17use clap::{Parser, Subcommand};
18use colored::*;
19use rusqlite::{Connection, Result as SqlResult};
20use std::path::{Path, PathBuf};
21
22const DB_PATH: &str = "~/.cerebro/rag_conocimiento.db";
23const EMBEDDING_DIM: usize = 1024;
24const EMBEDDING_BYTES: usize = EMBEDDING_DIM * 4;
25
26/// Elimina stopwords del español para la consulta FTS5
27fn limpiar_fts(texto: &str) -> String {
28 let stopwords = ["de", "la", "que", "el", "en", "y", "a", "los", "del", "se",
29 "las", "por", "un", "para", "con", "no", "una", "su", "al", "lo",
30 "como", "mas", "más", "pero", "sus", "le", "ya", "o", "este", "si",
31 "sí", "porque", "esta", "entre", "cuando", "muy", "sin", "sobre",
32 "tambien", "también", "me", "hasta", "hay", "donde", "quien", "desde",
33 "todo", "nos", "durante", "todos", "uno", "les", "ni", "contra",
34 "otros", "ese", "eso", "ante", "ellos", "e", "esto", "mi", "antes",
35 "algunos", "unos", "yo", "otro", "otras", "otra", "tanto", "esa",
36 "estos", "mucho", "quienes", "nada", "muchos", "cual", "poco",
37 "ella", "estar", "estas", "algo", "nosotros", "mis", "tu", "tus",
38 "ellas", "os", "esos", "esas", "estoy", "estan", "están", "estaba",
39 "estaban", "fue", "fueron", "es", "son", "ser", "era", "eran",
40 "ha", "han", "he", "hemos", "habia", "había", "habian", "habían",
41 "hubo", "hi", "hice", "hicimos", "puede", "pueden", "puedo",
42 "podemos", "dela", "tras", "bajo", "ademas", "además", "solo",
43 "sólo", "tan", "bien", "hacer", "tener", "tengo", "tiene",
44 "tenemos", "ir", "voy", "va", "van", "vamos", "hace", "hacen",
45 "soy", "eres", "somos", "aun", "aún", "acerca", "asi", "así",
46 "buen", "buena", "buenos", "buenas", "cada", "casi", "cualquier",
47 "dado", "dar", "decir", "dijo", "don", "dos", "ejemplo", "ello",
48 "embargo", "estamos", "estuvo", "estuvieron", "fuera", "fui",
49 "gran", "grande", "grandes", "haber", "hacia", "hayan", "hizo",
50 "hubiera", "hubiesen", "luego", "mientras", "mismo", "misma",
51 "mismos", "necesita", "necesitan", "parte", "poca", "pocos",
52 "pues", "resulta", "resultan", "sea", "sean", "según", "segun",
53 "sera", "será", "sería", "seria", "siempre", "sino", "tener",
54 "tenia", "tenía", "tenido", "tiene", "tienen", "toda", "todo",
55 "todas", "todos", "tuvo", "unas", "ustedes", "varios", "varias",
56 "veces", "ver", "vez", "te", "él", "vosotros",
57 ];
58 let lower = texto.to_lowercase();
59 let filtrado: Vec<&str> = lower
60 .split_whitespace()
61 .filter(|tok| !tok.is_empty() && !stopwords.contains(tok))
62 .collect();
63 if filtrado.is_empty() { texto.to_string() }
64 else { filtrado.join(" ") }
65}
66
67#[derive(Parser)]
68#[command(name = "cerebro", version, about = "RAG knowledge base CLI")]
69struct Cli {
70 #[command(subcommand)]
71 command: Commands,
72}
73
74#[derive(Subcommand)]
75enum Commands {
76 Query {
77 texto: Option<String>,
78 #[arg(short, long, default_value_t = 0.4)]
79 alpha: f32,
80 #[arg(short = k, long, default_value_t = 5)]
81 top_k: u32,
82 #[arg(long)]
83 tag: Option<String>,
84 },
85 List { tipo: String },
86 Similar {
87 archivo: PathBuf,
88 #[arg(short = k, long, default_value_t = 5)]
89 top_k: u32,
90 #[arg(long)]
91 tag: Option<String>,
92 },
93}
94
95fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
96 let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
97 let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
98 let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
99 dot / (norm_a * norm_b + 1e-10)
100}
101
102fn sigmoid_fts(score: f32) -> f32 {
103 1.0 / (1.0 + (-score / 10.0).exp())
104}
105
106fn obtener_embedding(texto: &str) -> Option<Vec<f32>> {
107 let client = reqwest::blocking::Client::builder()
108 .timeout(std::time::Duration::from_secs(15)).build().ok()?;
109 let body = serde_json::json!({"model": "bge-m3", "prompt": texto});
110 let resp = client.post("http://localhost:11434/api/embeddings")
111 .json(&body).send().ok()?;
112 let json: serde_json::Value = resp.json().ok()?;
113 let arr = json.get("embedding")?.as_array()?;
114 let vec: Vec<f32> = arr.iter()
115 .filter_map(|v| v.as_f64().map(|f| f as f32)).collect();
116 if vec.len() == EMBEDDING_DIM { Some(vec) } else { None }
117}
118
119fn ejecutar_query(texto: &str, alpha: f32, top_k: u32, tag: Option<&str>) {
120 let db_path = expand_path(DB_PATH);
121 if !Path::new(&db_path).exists() { eprintln!("BD no encontrada"); return; }
122 let conn = open_db(&db_path).unwrap();
123
124 let mut where_clauses = vec!["chunks_fts MATCH ?1".to_string()];
125 let query_fts = limpiar_fts(texto);
126 let mut param_values: Vec<rusqlite::types::Value> =
127 vec![rusqlite::types::Value::Text(query_fts)];
128 if let Some(t) = tag {
129 where_clauses.push("c.tags LIKE ?2".to_string());
130 param_values.push(rusqlite::types::Value::Text(format!("%{}%", t)));
131 }
132 let fts_candidates = top_k.saturating_mul(40).clamp(50, 500);
133 param_values.push(rusqlite::types::Value::Integer(fts_candidates as i64));
134
135 let sql = format!(
136 "SELECT c.id, c.content, e.vector, c.doc_id, bm25(chunks_fts) as fts_score, \
137 c.title, c.doc_path FROM chunks_fts JOIN chunks c ON chunks_fts.rowid = c.id \
138 JOIN embeddings e ON e.chunk_id = c.id WHERE {} ORDER BY fts_score DESC LIMIT {}",
139 where_clauses.join(" AND "),
140 where_clauses.len() + 1
141 );
142
143 // ... (query execution logic with cosine similarity + dedup)
144 let query_emb = obtener_embedding(texto);
145 println!("Consulta: {} | Alpha: {}", texto, alpha);
146 println!("Implementación completa en https://gist.atareao.es/");
147}
148
149fn main() {
150 let cli = Cli::parse();
151 match &cli.command {
152 Commands::Query { texto, alpha, top_k, tag } => {
153 ejecutar_query(texto.as_deref().unwrap_or(""), *alpha, *top_k, tag.as_deref());
154 }
155 Commands::List { tipo } => {
156 println!("Listando: {} (implementación completa en el gist)", tipo);
157 }
158 Commands::Similar { archivo, top_k, tag } => {
159 println!("Similar: {} (implementación completa en el gist)", archivo.display());
160 }
161 }
162}
163
164fn expand_path(path: &str) -> String {
165 if path.starts_with(~) {
166 let home = std::env::var("HOME").unwrap_or_default();
167 path.replacen(~, &home, 1)
168 } else { path.to_string() }
169}
170
171fn open_db(path: &str) -> SqlResult<Connection> {
172 let expanded = expand_path(path);
173 let conn = Connection::open(&expanded)?;
174 conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA query_only=1;")?;
175 Ok(conn)
176}

Desarrollado por Opengist ⋅ Load: 88ms⋅

Español