Mi piace 0
Forka 0

Ultima attività 17 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.

atareao's Avatar atareao ha revisionato questo gist 17 hours ago. Vai alla revisione

1 file changed, 176 insertions

main.rs(file creato)

@@ -0,0 +1,176 @@
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 +
16 + use bytemuck::try_cast_slice;
17 + use clap::{Parser, Subcommand};
18 + use colored::*;
19 + use rusqlite::{Connection, Result as SqlResult};
20 + use std::path::{Path, PathBuf};
21 +
22 + const DB_PATH: &str = "~/.cerebro/rag_conocimiento.db";
23 + const EMBEDDING_DIM: usize = 1024;
24 + const EMBEDDING_BYTES: usize = EMBEDDING_DIM * 4;
25 +
26 + /// Elimina stopwords del español para la consulta FTS5
27 + fn 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")]
69 + struct Cli {
70 + #[command(subcommand)]
71 + command: Commands,
72 + }
73 +
74 + #[derive(Subcommand)]
75 + enum 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 +
95 + fn 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 +
102 + fn sigmoid_fts(score: f32) -> f32 {
103 + 1.0 / (1.0 + (-score / 10.0).exp())
104 + }
105 +
106 + fn 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 +
119 + fn 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 +
149 + fn 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 +
164 + fn 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 +
171 + fn 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 + }
Più nuovi Più vecchi

Creato da Opengist ⋅ Load: 33ms⋅

Italiano