updated
This commit is contained in:
@@ -1,50 +0,0 @@
|
||||
use std::fs::{self, File};
|
||||
use std::io::{self, BufRead, BufReader, BufWriter, Write};
|
||||
use std::path::Path;
|
||||
|
||||
pub fn merge_english_jsonl() -> io::Result<()> {
|
||||
let input_dir = "data/LLM_English";
|
||||
let output_dir = "output";
|
||||
let output_file = "english_master.jsonl";
|
||||
|
||||
fs::create_dir_all(output_dir)?;
|
||||
|
||||
let output_path = Path::new(output_dir).join(output_file);
|
||||
let mut writer = BufWriter::new(File::create(&output_path)?);
|
||||
|
||||
let files = [
|
||||
"english_explain.jsonl",
|
||||
"english_reasoning.jsonl",
|
||||
"english_summary.jsonl",
|
||||
"instruction_following.jsonl",
|
||||
];
|
||||
|
||||
let mut total_lines: usize = 0;
|
||||
|
||||
for filename in &files {
|
||||
let input_path = Path::new(input_dir).join(filename);
|
||||
|
||||
if !input_path.exists() {
|
||||
println!("⚠️ Warning: {} not found, skipping...", filename);
|
||||
continue;
|
||||
}
|
||||
|
||||
println!("📄 Merging: {}", filename);
|
||||
|
||||
let file = File::open(&input_path)?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if !trimmed.is_empty() {
|
||||
writeln!(writer, "{}", trimmed)?;
|
||||
total_lines += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
writer.flush()?;
|
||||
println!("✅ Successfully merged {} lines into → {}", total_lines, output_path.display());
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,77 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::{Device, Tensor, safetensors, DType, backend::Backend};
|
||||
use candle_nn::VarBuilder;
|
||||
use candle_transformers::models::bert::{BertModel, Config};
|
||||
use tokenizers::Tokenizer;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub struct Embedder {
|
||||
model: BertModel,
|
||||
tokenizer: Tokenizer,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl Embedder {
|
||||
pub fn new() -> Result<Self> {
|
||||
let device = Device::Cpu;
|
||||
let cache_dir = std::path::Path::new("models/all-MiniLM-L6-v2");
|
||||
std::fs::create_dir_all(cache_dir)?;
|
||||
|
||||
let model_file = cache_dir.join("model.safetensors");
|
||||
let config_file = cache_dir.join("config.json");
|
||||
let tokenizer_file = cache_dir.join("tokenizer.json");
|
||||
|
||||
if !model_file.exists() {
|
||||
println!("📥 Downloading all-MiniLM-L6-v2 model (80 MB)...");
|
||||
download_file(
|
||||
"https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/model.safetensors",
|
||||
&model_file,
|
||||
)?;
|
||||
download_file(
|
||||
"https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/config.json",
|
||||
&config_file,
|
||||
)?;
|
||||
download_file(
|
||||
"https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/tokenizer.json",
|
||||
&tokenizer_file,
|
||||
)?;
|
||||
}
|
||||
|
||||
let config: Config = serde_json::from_str(&std::fs::read_to_string(config_file)?)?;
|
||||
let tokenizer = Tokenizer::from_file(tokenizer_file).map_err(anyhow::Error::msg)?;
|
||||
|
||||
// Load safetensors file into a map of tensors
|
||||
let tensors = unsafe { safetensors::load(&model_file, &device) }?;
|
||||
// Create a backend from the tensors
|
||||
let backend = Arc::new(Backend::from_tensors(tensors, &device)?);
|
||||
let vb = VarBuilder::from_backend(backend);
|
||||
let model = BertModel::load(vb, &config)?;
|
||||
|
||||
Ok(Self { model, tokenizer, device })
|
||||
}
|
||||
|
||||
pub fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
|
||||
let mut embeddings = Vec::with_capacity(texts.len());
|
||||
for text in texts {
|
||||
let tokens = self
|
||||
.tokenizer
|
||||
.encode(text, true)
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
let token_ids = Tensor::new(tokens.get_ids(), &self.device)?.unsqueeze(0)?;
|
||||
let attention_mask = Tensor::new(tokens.get_attention_mask(), &self.device)?.unsqueeze(0)?;
|
||||
let output = self.model.forward(&token_ids, &attention_mask, None)?;
|
||||
let emb = output.mean(1)?;
|
||||
let vec = emb.to_vec2()?[0].clone();
|
||||
embeddings.push(vec);
|
||||
}
|
||||
Ok(embeddings)
|
||||
}
|
||||
}
|
||||
|
||||
fn download_file(url: &str, dest: &std::path::Path) -> Result<()> {
|
||||
let response = reqwest::blocking::get(url)?;
|
||||
let mut file = std::fs::File::create(dest)?;
|
||||
let content = response.bytes()?;
|
||||
std::io::copy(&mut content.as_ref(), &mut file)?;
|
||||
Ok(())
|
||||
}
|
||||
+243
-196
@@ -1,10 +1,13 @@
|
||||
use eframe::egui::{self, Color32, RichText};
|
||||
use serde_json::Value;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::mpsc;
|
||||
use std::thread;
|
||||
use std::time::Instant;
|
||||
|
||||
pub static WEB_BROWSER_MODE: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
fn pretty_json(text: &str) -> String {
|
||||
match serde_json::from_str::<Value>(text) {
|
||||
Ok(v) => serde_json::to_string_pretty(&v).unwrap_or_else(|_| text.to_string()),
|
||||
@@ -294,6 +297,8 @@ struct AirustGui {
|
||||
sidebar_sources: Vec<(String, u64)>,
|
||||
stats_pending: bool,
|
||||
|
||||
should_close: bool,
|
||||
|
||||
tx: mpsc::Sender<Msg>,
|
||||
rx: mpsc::Receiver<Msg>,
|
||||
}
|
||||
@@ -371,6 +376,7 @@ impl AirustGui {
|
||||
sidebar_total: "—".to_string(),
|
||||
sidebar_sources: Vec::new(),
|
||||
stats_pending: false,
|
||||
should_close: false,
|
||||
|
||||
tx,
|
||||
rx,
|
||||
@@ -603,6 +609,10 @@ impl eframe::App for AirustGui {
|
||||
AppMode::Desktop => self.ui_desktop(ui),
|
||||
}
|
||||
|
||||
if self.should_close {
|
||||
ui.ctx().send_viewport_cmd(egui::ViewportCommand::Close);
|
||||
}
|
||||
|
||||
ui.ctx().request_repaint_after(std::time::Duration::from_millis(16));
|
||||
}
|
||||
}
|
||||
@@ -630,7 +640,8 @@ impl AirustGui {
|
||||
egui::Button::new(RichText::new("🌐 Web Browser").size(18.0)),
|
||||
).clicked() {
|
||||
let _ = open::that(format!("https://{}", "127.0.0.1:3000"));
|
||||
self.mode = AppMode::Desktop;
|
||||
WEB_BROWSER_MODE.store(true, Ordering::SeqCst);
|
||||
self.should_close = true;
|
||||
}
|
||||
ui.add_space(spacing);
|
||||
if ui.add_sized(
|
||||
@@ -649,59 +660,72 @@ impl AirustGui {
|
||||
|
||||
impl AirustGui {
|
||||
fn ui_desktop(&mut self, ui: &mut egui::Ui) {
|
||||
let (_, _, _, _, accent, text) = self.theme.colors();
|
||||
let (_, surface, surface2, border, accent, text) = self.theme.colors();
|
||||
let dim = text.linear_multiply(0.6);
|
||||
let sidebar_w = 340.0;
|
||||
|
||||
// ── Sidebar ──
|
||||
egui::Panel::left("sidebar")
|
||||
.resizable(false)
|
||||
.default_size(250.0)
|
||||
.show_inside(ui, |ui| {
|
||||
// Sidebar header
|
||||
ui.heading(RichText::new("AiRust").color(accent).size(18.0));
|
||||
ui.label(RichText::new("Cyber Operator v0.2.0").size(11.0).color(text.linear_multiply(0.6)));
|
||||
ui.separator();
|
||||
// ── Sidebar (fixed 340px, manual layout) ──
|
||||
let sider = egui::Rect::from_min_size(ui.max_rect().left_top(), egui::vec2(sidebar_w, ui.max_rect().height()));
|
||||
let mut su = ui.new_child(egui::UiBuilder::new().id_salt("sb").max_rect(sider).layout(egui::Layout::top_down(egui::Align::Min)));
|
||||
su.painter().rect_filled(sider, 0.0, surface);
|
||||
|
||||
// System Status
|
||||
ui.label(RichText::new("SYSTEM STATUS").size(11.0).color(text.linear_multiply(0.6)).strong());
|
||||
ui.horizontal(|ui| {
|
||||
ui.group(|ui| {
|
||||
ui.set_min_width(ui.available_width());
|
||||
egui::ScrollArea::vertical().id_salt("sb_scroll").show(&mut su, |ui| {
|
||||
ui.set_min_width(sidebar_w - 8.0);
|
||||
ui.add_space(16.0);
|
||||
|
||||
// Header
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("AI").size(18.0).color(accent).strong()); ui.label(RichText::new("RUST").size(18.0).color(text)); });
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("Cyber Operator v0.2.0").size(11.0).color(dim)); });
|
||||
ui.add_space(12.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.separator(); });
|
||||
|
||||
// System Status
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("SYSTEM STATUS").size(11.0).color(dim).strong()); });
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(16.0);
|
||||
let cw = ((ui.available_width() - 8.0) / 2.0).max(1.0);
|
||||
let mut first = true;
|
||||
for (val, label) in [(&self.sidebar_total, "Entries"), (&self.sidebar_sources.len().to_string(), "Sources")] {
|
||||
if !first { ui.add_space(8.0); } else { first = false; }
|
||||
egui::Frame::NONE.fill(surface2).stroke((1.0, border)).corner_radius(4).show(ui, |ui| {
|
||||
ui.set_min_size(egui::vec2(cw, 52.0));
|
||||
ui.vertical_centered(|ui| {
|
||||
ui.label(RichText::new(&self.sidebar_total).size(20.0).color(accent).strong());
|
||||
ui.label(RichText::new("Entries").size(10.0).color(text.linear_multiply(0.6)));
|
||||
ui.label(RichText::new(val).size(20.0).color(accent).strong());
|
||||
ui.label(RichText::new(label).size(10.0).color(dim));
|
||||
});
|
||||
});
|
||||
ui.group(|ui| {
|
||||
ui.set_min_width(ui.available_width());
|
||||
ui.vertical_centered(|ui| {
|
||||
ui.label(RichText::new(&self.sidebar_sources.len().to_string()).size(20.0).color(accent).strong());
|
||||
ui.label(RichText::new("Sources").size(10.0).color(text.linear_multiply(0.6)));
|
||||
}
|
||||
});
|
||||
|
||||
// Sources
|
||||
ui.add_space(10.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("SOURCES").size(11.0).color(dim).strong()); });
|
||||
ui.add_space(4.0);
|
||||
egui::ScrollArea::vertical().id_salt("sb_src").max_height(120.0).show(ui, |ui| {
|
||||
for (src, cnt) in &self.sidebar_sources {
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(16.0);
|
||||
ui.label(RichText::new(src).size(12.0));
|
||||
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
|
||||
let _ = ui.allocate_space(egui::vec2(8.0, 0.0));
|
||||
ui.label(RichText::new(cnt.to_string()).size(12.0).color(accent).strong());
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
if self.sidebar_sources.is_empty() {
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("No sources").size(11.0).color(dim)); });
|
||||
}
|
||||
});
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.separator(); });
|
||||
|
||||
// Knowledge Sources
|
||||
ui.add_space(8.0);
|
||||
ui.label(RichText::new("SOURCES").size(11.0).color(text.linear_multiply(0.6)).strong());
|
||||
egui::ScrollArea::vertical().max_height(150.0).show(ui, |ui| {
|
||||
for (src, cnt) in &self.sidebar_sources {
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(src).size(12.0));
|
||||
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
|
||||
ui.label(RichText::new(cnt.to_string()).size(12.0).color(accent).strong());
|
||||
});
|
||||
});
|
||||
}
|
||||
if self.sidebar_sources.is_empty() {
|
||||
ui.label(RichText::new("No sources").size(11.0).color(text.linear_multiply(0.5)));
|
||||
}
|
||||
});
|
||||
|
||||
ui.separator();
|
||||
|
||||
// Theme
|
||||
ui.label(RichText::new("THEME").size(11.0).color(text.linear_multiply(0.6)).strong());
|
||||
egui::ComboBox::from_id_salt("sidebar_theme")
|
||||
// Theme
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("THEME").size(11.0).color(dim).strong()); });
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(16.0);
|
||||
egui::ComboBox::from_id_salt("sb_theme")
|
||||
.selected_text(self.theme.name())
|
||||
.width(ui.available_width())
|
||||
.show_ui(ui, |ui| {
|
||||
@@ -712,98 +736,90 @@ impl AirustGui {
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.separator(); });
|
||||
|
||||
ui.separator();
|
||||
|
||||
// Target
|
||||
ui.label(RichText::new("TARGET").size(11.0).color(text.linear_multiply(0.6)).strong());
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_sized(
|
||||
[ui.available_width() - 50.0, 24.0],
|
||||
egui::TextEdit::singleline(&mut self.tgt_val).hint_text("IP / hostname..."),
|
||||
);
|
||||
if ui.button("Set").clicked() {
|
||||
let val = self.tgt_val.clone();
|
||||
self.spawn("tgt_status", move |url| {
|
||||
match api_post(url, "/api/target/set", serde_json::json!({"target": val})) {
|
||||
Ok(r) => ("tgt_status".into(), GuiMsg::TargetResult(r)),
|
||||
Err(e) => ("tgt_status".into(), GuiMsg::TargetResult(format!("Error: {}", e))),
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
if !self.tgt_status.is_empty() {
|
||||
ui.label(RichText::new(&self.tgt_status).size(10.0).color(text.linear_multiply(0.6)));
|
||||
}
|
||||
|
||||
ui.separator();
|
||||
|
||||
// Chat Sessions
|
||||
ui.label(RichText::new("CHAT SESSIONS").size(11.0).color(text.linear_multiply(0.6)).strong());
|
||||
if ui.button("+ New Chat").clicked() {
|
||||
self.do_new_session();
|
||||
}
|
||||
egui::ScrollArea::vertical().max_height(180.0).show(ui, |ui| {
|
||||
self.render_session_list(ui);
|
||||
});
|
||||
if ui.button("Delete All Chats").clicked() {
|
||||
let id = format!("sess_{}", chrono_now_ms());
|
||||
self.sessions = vec![ChatSession {
|
||||
id: id.clone(), name: "Chat 1".to_string(), messages: Vec::new(),
|
||||
}];
|
||||
self.active_session_id = id;
|
||||
}
|
||||
|
||||
ui.separator();
|
||||
|
||||
// Quick Actions
|
||||
ui.label(RichText::new("QUICK ACTIONS").size(11.0).color(text.linear_multiply(0.6)).strong());
|
||||
if ui.button("Clear Chat").clicked() {
|
||||
if let Some(s) = self.active_session_mut() {
|
||||
s.messages.clear();
|
||||
}
|
||||
}
|
||||
if ui.button("Purge Knowledge").clicked() {
|
||||
self.spawn("kn_stats", move |url| {
|
||||
match api_post(url, "/api/knowledge/purge", serde_json::json!({})) {
|
||||
Ok(r) => ("kn_stats".into(), GuiMsg::Set(r)),
|
||||
Err(e) => ("kn_stats".into(), GuiMsg::Set(format!("Error: {}", e))),
|
||||
// Target
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("TARGET").size(11.0).color(dim).strong()); });
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(16.0);
|
||||
let iw = (ui.available_width() - 50.0 - 16.0).max(0.0);
|
||||
ui.add_sized([iw, 24.0], egui::TextEdit::singleline(&mut self.tgt_val).hint_text("IP / hostname..."));
|
||||
if ui.button("Set").clicked() {
|
||||
let val = self.tgt_val.clone();
|
||||
self.spawn("tgt_status", move |url| {
|
||||
match api_post(url, "/api/target/set", serde_json::json!({"target": val})) {
|
||||
Ok(r) => ("tgt_status".into(), GuiMsg::TargetResult(r)),
|
||||
Err(e) => ("tgt_status".into(), GuiMsg::TargetResult(format!("Error: {}", e))),
|
||||
}
|
||||
});
|
||||
}
|
||||
if ui.button("Kill Server").clicked() {
|
||||
let url = self.server_url.clone();
|
||||
thread::spawn(move || { let _ = api_post(&url, "/api/shutdown", serde_json::json!({})); });
|
||||
}
|
||||
});
|
||||
if !self.tgt_status.is_empty() {
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new(&self.tgt_status).size(10.0).color(dim)); });
|
||||
}
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.separator(); });
|
||||
|
||||
// ── Main area ──
|
||||
egui::Panel::top("tab_bar").show_inside(ui, |ui| {
|
||||
// Chat Sessions
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("CHAT SESSIONS").size(11.0).color(dim).strong()); });
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); if ui.button("+ New Chat").clicked() { self.do_new_session(); } });
|
||||
egui::ScrollArea::vertical().id_salt("sb_sess").max_height(150.0).show(ui, |ui| { self.render_session_list(ui); });
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); if ui.button("Delete All Chats").clicked() { let id = format!("sess_{}", chrono_now_ms()); self.sessions = vec![ChatSession { id: id.clone(), name: "Chat 1".to_string(), messages: Vec::new() }]; self.active_session_id = id; } });
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.separator(); });
|
||||
|
||||
// Quick Actions
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); ui.label(RichText::new("QUICK ACTIONS").size(11.0).color(dim).strong()); });
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); if ui.button("Clear Chat").clicked() { if let Some(s) = self.active_session_mut() { s.messages.clear(); } } });
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); if ui.button("Purge Knowledge").clicked() { self.spawn("kn_stats", move |url| { match api_post(url, "/api/knowledge/purge", serde_json::json!({})) { Ok(r) => ("kn_stats".into(), GuiMsg::Set(r)), Err(e) => ("kn_stats".into(), GuiMsg::Set(format!("Error: {}", e))), } }); } });
|
||||
ui.horizontal(|ui| { ui.add_space(16.0); if ui.button("Kill Server").clicked() { let url = self.server_url.clone(); thread::spawn(move || { let _ = api_post(&url, "/api/shutdown", serde_json::json!({})); }); } });
|
||||
ui.add_space(12.0);
|
||||
});
|
||||
|
||||
// ── Main area (to the right of sidebar) ──
|
||||
let main_x = sidebar_w;
|
||||
let main_w = (ui.max_rect().width() - main_x).max(0.0);
|
||||
let main_rect = egui::Rect::from_min_size(egui::pos2(main_x, ui.max_rect().top()), egui::vec2(main_w, ui.max_rect().height()));
|
||||
let mut mu = ui.new_child(egui::UiBuilder::new().id_salt("main").max_rect(main_rect).layout(egui::Layout::top_down(egui::Align::Min)));
|
||||
|
||||
// Tab bar
|
||||
let tab_h = 36.0;
|
||||
egui::Frame::NONE.fill(surface).stroke((1.0, border)).show(&mut mu, |ui| {
|
||||
ui.set_min_size(egui::vec2(main_w, tab_h));
|
||||
ui.horizontal(|ui| {
|
||||
let active = &mut self.active_tab;
|
||||
for tab in Tab::all() {
|
||||
let selected = *tab == *active;
|
||||
if ui.selectable_label(selected, RichText::new(tab.label()).size(13.0)).clicked() {
|
||||
*active = *tab;
|
||||
let sel = *tab == self.active_tab;
|
||||
let c = if sel { accent } else { dim };
|
||||
if ui.selectable_label(sel, RichText::new(tab.label()).size(13.0).color(c)).clicked() {
|
||||
self.active_tab = *tab;
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
egui::CentralPanel::default().show_inside(ui, |ui| {
|
||||
match self.active_tab {
|
||||
Tab::Chat => self.ui_chat(ui),
|
||||
Tab::Execute => self.ui_exec(ui),
|
||||
Tab::Scripts => self.ui_scripts(ui),
|
||||
Tab::SelfLearn => self.ui_self_learn(ui),
|
||||
Tab::ScriptStudio => self.ui_script_studio(ui),
|
||||
Tab::PdfImport => self.ui_pdf_import(ui),
|
||||
Tab::Ocr => self.ui_ocr(ui),
|
||||
Tab::Books => self.ui_books(ui),
|
||||
Tab::Datasets => self.ui_datasets(ui),
|
||||
Tab::Knowledge => self.ui_knowledge(ui),
|
||||
}
|
||||
});
|
||||
// Content area (below tab bar)
|
||||
let content_rect = egui::Rect::from_min_size(
|
||||
egui::pos2(main_x, ui.max_rect().top() + tab_h + 1.0),
|
||||
egui::vec2(main_w, (ui.max_rect().height() - tab_h - 1.0).max(0.0)),
|
||||
);
|
||||
let mut cu = ui.new_child(egui::UiBuilder::new().id_salt("content").max_rect(content_rect).layout(egui::Layout::top_down(egui::Align::Min)));
|
||||
match self.active_tab {
|
||||
Tab::Chat => self.ui_chat(&mut cu),
|
||||
Tab::Execute => self.ui_exec(&mut cu),
|
||||
Tab::Scripts => self.ui_scripts(&mut cu),
|
||||
Tab::SelfLearn => self.ui_self_learn(&mut cu),
|
||||
Tab::ScriptStudio => self.ui_script_studio(&mut cu),
|
||||
Tab::PdfImport => self.ui_pdf_import(&mut cu),
|
||||
Tab::Ocr => self.ui_ocr(&mut cu),
|
||||
Tab::Books => self.ui_books(&mut cu),
|
||||
Tab::Datasets => self.ui_datasets(&mut cu),
|
||||
Tab::Knowledge => self.ui_knowledge(&mut cu),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -812,92 +828,123 @@ impl AirustGui {
|
||||
impl AirustGui {
|
||||
// ── Chat Tab ──
|
||||
fn ui_chat(&mut self, ui: &mut egui::Ui) {
|
||||
let avail = ui.available_height();
|
||||
let input_h = 40.0;
|
||||
let (_, _, _, _, _, text) = self.theme.colors();
|
||||
let av = ui.available_size();
|
||||
let input_h = 60.0;
|
||||
let msg_h = (av.y - input_h).max(100.0);
|
||||
let (_, surface, _, border, accent, text) = self.theme.colors();
|
||||
let base = ui.max_rect().left_top();
|
||||
let full_w = av.x;
|
||||
|
||||
// Messages
|
||||
egui::ScrollArea::vertical()
|
||||
.max_height((avail - input_h - 10.0).max(100.0))
|
||||
// Messages — styled like web UI, fills space above input
|
||||
let msg_rect = egui::Rect::from_min_size(base, egui::vec2(full_w, msg_h));
|
||||
let mut mu = ui.new_child(egui::UiBuilder::new().id_salt("chat_msgs_out").max_rect(msg_rect).layout(egui::Layout::top_down(egui::Align::Min)));
|
||||
egui::ScrollArea::vertical().id_salt("chat_msgs")
|
||||
.auto_shrink([false, false])
|
||||
.stick_to_bottom(true)
|
||||
.show(ui, |ui| {
|
||||
.show(&mut mu, |ui| {
|
||||
ui.add_space(8.0);
|
||||
if let Some(session) = self.active_session() {
|
||||
for msg in &session.messages {
|
||||
let role_color = match msg.role.as_str() {
|
||||
"user" => text.linear_multiply(0.8),
|
||||
"assistant" => self.theme.colors().4,
|
||||
"assistant" => accent,
|
||||
_ => text.linear_multiply(0.5),
|
||||
};
|
||||
ui.label(RichText::new(format!("[{}]", msg.role)).size(11.0).color(role_color).strong());
|
||||
if msg.content.is_empty() && msg.role == "assistant" && self.chat_streaming {
|
||||
ui.label(RichText::new("...").size(13.0).color(text.linear_multiply(0.5)));
|
||||
let role_label = match msg.role.as_str() {
|
||||
"user" => "user",
|
||||
"assistant" => "assistant",
|
||||
_ => msg.role.as_str(),
|
||||
};
|
||||
ui.add(egui::Label::new(RichText::new(role_label).size(11.0).color(role_color).strong().monospace()));
|
||||
let empty_streaming = msg.role == "assistant" && msg.content.is_empty() && self.chat_streaming;
|
||||
if empty_streaming {
|
||||
egui::Frame::NONE.fill(surface).stroke((1.0, border)).corner_radius(6).show(ui, |ui| {
|
||||
ui.label(RichText::new("...").size(13.0).color(text.linear_multiply(0.5)));
|
||||
});
|
||||
} else {
|
||||
ui.label(RichText::new(&msg.content).size(13.0));
|
||||
egui::Frame::NONE.fill(surface).stroke((1.0, border)).corner_radius(6).show(ui, |ui| {
|
||||
ui.label(RichText::new(&msg.content).size(13.0).color(text));
|
||||
});
|
||||
if msg.role == "assistant" {
|
||||
let resp = ui.response();
|
||||
let mut rect = resp.rect;
|
||||
rect.max.x = rect.min.x + 3.0;
|
||||
rect.min.x -= 1.0;
|
||||
ui.painter().rect_filled(rect, egui::CornerRadius::same(2), accent);
|
||||
}
|
||||
}
|
||||
ui.add_space(6.0);
|
||||
ui.add_space(10.0);
|
||||
}
|
||||
if session.messages.is_empty() {
|
||||
ui.label(RichText::new("Start a conversation...").size(12.0).color(text.linear_multiply(0.5)));
|
||||
ui.add(egui::Label::new(RichText::new("Start a conversation...").size(12.0).color(text.linear_multiply(0.5))));
|
||||
}
|
||||
}
|
||||
ui.add_space(8.0);
|
||||
});
|
||||
|
||||
// Input area — anchored at bottom
|
||||
let inp_rect = egui::Rect::from_min_size(egui::pos2(base.x, base.y + msg_h), egui::vec2(full_w, input_h));
|
||||
let mut iu = ui.new_child(egui::UiBuilder::new().id_salt("chat_inp_out").max_rect(inp_rect).layout(egui::Layout::top_down(egui::Align::Min)));
|
||||
egui::Frame::NONE.fill(surface).stroke((1.0, border)).show(&mut iu, |ui| {
|
||||
ui.add_space(12.0);
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(16.0);
|
||||
let inp_w = (ui.available_width() - 70.0 - 16.0 - 10.0).max(0.0);
|
||||
let resp = ui.add_sized(
|
||||
[inp_w, 36.0],
|
||||
egui::TextEdit::singleline(&mut self.chat_input)
|
||||
.hint_text("Type your message...")
|
||||
.desired_width(f32::INFINITY),
|
||||
);
|
||||
let send = ui.add(egui::Button::new(RichText::new("Send").size(12.0)).min_size(egui::vec2(70.0, 36.0)));
|
||||
let disabled = self.chat_streaming;
|
||||
if !disabled && (send.clicked() || (resp.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)))) {
|
||||
let msg = self.chat_input.trim().to_string();
|
||||
if !msg.is_empty() {
|
||||
if let Some(s) = self.active_session_mut() {
|
||||
s.messages.push(ChatMsg { role: "user".into(), content: msg.clone() });
|
||||
s.messages.push(ChatMsg { role: "assistant".into(), content: String::new() });
|
||||
}
|
||||
self.chat_streaming = true;
|
||||
let tx = self.tx.clone();
|
||||
let url = self.server_url.clone();
|
||||
thread::spawn(move || {
|
||||
let url_f = format!("{}/api/chat", url);
|
||||
match http_client()
|
||||
.post(&url_f)
|
||||
.json(&serde_json::json!({"message": msg}))
|
||||
.send()
|
||||
{
|
||||
Ok(resp) => {
|
||||
let reader = BufReader::new(resp);
|
||||
let mut reply = String::new();
|
||||
for line in reader.lines() {
|
||||
match line {
|
||||
Ok(line) => {
|
||||
if line.starts_with("data: ") {
|
||||
let data = &line[6..];
|
||||
if data == "[DONE]" { break; }
|
||||
reply.push_str(data);
|
||||
let _ = tx.send(("_chat_stream".into(), GuiMsg::ChatStreamChunk(reply.clone())));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(("_chat_resp".into(), GuiMsg::Set(format!("Error reading stream: {}", e))));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = tx.send(("_chat_stream".into(), GuiMsg::ChatStreamDone));
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(("_chat_resp".into(), GuiMsg::Set(format!("Error: {}", e))));
|
||||
}
|
||||
}
|
||||
});
|
||||
self.chat_input.clear();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Input
|
||||
ui.separator();
|
||||
let btn_w = 70.0;
|
||||
ui.horizontal(|ui| {
|
||||
let resp = ui.add_sized(
|
||||
[ui.available_width() - btn_w - 10.0, 30.0],
|
||||
egui::TextEdit::singleline(&mut self.chat_input).hint_text("Type your message..."),
|
||||
);
|
||||
let send = ui.button("Send");
|
||||
let disabled = self.chat_streaming;
|
||||
if !disabled && (send.clicked() || (resp.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)))) {
|
||||
let msg = self.chat_input.trim().to_string();
|
||||
if !msg.is_empty() {
|
||||
if let Some(s) = self.active_session_mut() {
|
||||
s.messages.push(ChatMsg { role: "user".into(), content: msg.clone() });
|
||||
s.messages.push(ChatMsg { role: "assistant".into(), content: String::new() });
|
||||
}
|
||||
self.chat_streaming = true;
|
||||
let tx = self.tx.clone();
|
||||
let url = self.server_url.clone();
|
||||
thread::spawn(move || {
|
||||
let url_f = format!("{}/api/chat", url);
|
||||
match http_client()
|
||||
.post(&url_f)
|
||||
.json(&serde_json::json!({"message": msg}))
|
||||
.send()
|
||||
{
|
||||
Ok(resp) => {
|
||||
let reader = BufReader::new(resp);
|
||||
let mut reply = String::new();
|
||||
for line in reader.lines() {
|
||||
match line {
|
||||
Ok(line) => {
|
||||
if line.starts_with("data: ") {
|
||||
let data = &line[6..];
|
||||
if data == "[DONE]" { break; }
|
||||
reply.push_str(data);
|
||||
let _ = tx.send(("_chat_stream".into(), GuiMsg::ChatStreamChunk(reply.clone())));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(("_chat_resp".into(), GuiMsg::Set(format!("Error reading stream: {}", e))));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = tx.send(("_chat_stream".into(), GuiMsg::ChatStreamDone));
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(("_chat_resp".into(), GuiMsg::Set(format!("Error: {}", e))));
|
||||
}
|
||||
}
|
||||
});
|
||||
self.chat_input.clear();
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -1,40 +0,0 @@
|
||||
use std::fs::File;
|
||||
use std::io::{ BufRead, BufReader};
|
||||
|
||||
pub fn load_knowledge() -> String {
|
||||
let path = "output/english_master.jsonl";
|
||||
|
||||
let file = match File::open(path) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
eprintln!("❌ Could not open {}: {}", path, e);
|
||||
return "Knowledge base not found.".to_string();
|
||||
}
|
||||
};
|
||||
|
||||
let reader = BufReader::new(file);
|
||||
let mut knowledge = String::from("=== ENGLISH MASTERY KNOWLEDGE BASE ===\n\n");
|
||||
let mut count = 0;
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = match line {
|
||||
Ok(l) => l,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Just add the raw line as-is (since it's JSONL)
|
||||
knowledge.push_str(trimmed);
|
||||
knowledge.push_str("\n\n");
|
||||
count += 1;
|
||||
}
|
||||
|
||||
knowledge.push_str(&format!("\n=== TOTAL: {} EXAMPLES LOADED ===\n", count));
|
||||
println!("📚 Successfully loaded {} examples from english_master.jsonl", count);
|
||||
|
||||
knowledge
|
||||
}
|
||||
+20
-23
@@ -9,12 +9,11 @@ use std::collections::hash_map::DefaultHasher;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::fs;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::io::{self, Write};
|
||||
use std::path::Path;
|
||||
use std::sync::LazyLock;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
pub static OLLAMA_SEM: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(3));
|
||||
pub static OLLAMA_SEM: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(5));
|
||||
|
||||
const INDEX_CACHE_PATH: &str = "output/index_cache.json";
|
||||
const ARCHIVE_THRESHOLD: usize = 5000;
|
||||
@@ -216,9 +215,8 @@ impl KnowledgeIndex {
|
||||
self.entries.len() + self.archived.as_ref().map(|a| a.entries.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.is_empty()
|
||||
pub fn embeddings_ready(&self) -> bool {
|
||||
!self.embeddings.is_empty()
|
||||
}
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
@@ -297,25 +295,24 @@ impl KnowledgeIndex {
|
||||
}
|
||||
}
|
||||
|
||||
println!("\nBuilding knowledge embeddings...");
|
||||
self.embeddings = Vec::with_capacity(self.entries.len());
|
||||
println!("\nBuilding knowledge embeddings ({} entries)...", self.entries.len());
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
|
||||
for (i, entry) in self.entries.iter().enumerate() {
|
||||
print!("\r [{}/{}]", i + 1, self.entries.len());
|
||||
io::stdout().flush()?;
|
||||
let texts: Vec<String> = self
|
||||
.entries
|
||||
.iter()
|
||||
.map(|e| format!("{} {}", e.instruction, e.output))
|
||||
.collect();
|
||||
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let text = format!("{} {}", entry.instruction, entry.output);
|
||||
let req = GenerateEmbeddingsRequest::new(
|
||||
self.embed_model.clone(),
|
||||
EmbeddingsInput::Single(text),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
let resp = ollama.generate_embeddings(req).await?;
|
||||
let emb = resp.embeddings.into_iter().next().unwrap_or_default();
|
||||
self.embeddings.push(emb);
|
||||
}
|
||||
println!("\nDone. ({} vectors)", self.embeddings.len());
|
||||
let req = GenerateEmbeddingsRequest::new(
|
||||
self.embed_model.clone(),
|
||||
EmbeddingsInput::Multiple(texts),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
let resp = ollama.generate_embeddings(req).await?;
|
||||
|
||||
self.embeddings = resp.embeddings;
|
||||
println!("Done. ({} vectors)", self.embeddings.len());
|
||||
|
||||
self.maybe_archive();
|
||||
self.reindex();
|
||||
@@ -746,7 +743,7 @@ pub fn format_knowledge(entries: &[&KnowledgeEntry]) -> String {
|
||||
|
||||
let mut kb = String::new();
|
||||
for (i, e) in entries.iter().enumerate() {
|
||||
if i >= 5 {
|
||||
if i >= 8 {
|
||||
break;
|
||||
}
|
||||
let confidence = source_importance(&e.source);
|
||||
|
||||
+279
-51
@@ -21,8 +21,10 @@ use crate::web_server::{start_server, AppState};
|
||||
use anyhow::Result;
|
||||
use ollama_rs::Ollama;
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
fn main() -> Result<()> {
|
||||
@@ -32,51 +34,97 @@ fn main() -> Result<()> {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
return rt.block_on(run_pdf_import(&args[2..]));
|
||||
}
|
||||
if args.len() >= 3 && args[1].eq_ignore_ascii_case("learn-topics") {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
return rt.block_on(run_learn_topics(&args[2]));
|
||||
}
|
||||
if args.len() >= 2 && args[1].eq_ignore_ascii_case("cli") {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
return rt.block_on(run_cli());
|
||||
}
|
||||
if args.len() >= 2 && args[1].eq_ignore_ascii_case("export-training") {
|
||||
return export_training_data();
|
||||
}
|
||||
|
||||
if args.iter().any(|a| a == "--browser") {
|
||||
if args.iter().any(|a| a == "--desktop") {
|
||||
run_gui_mode_main_thread()
|
||||
} else {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
rt.block_on(run_browser_mode())
|
||||
} else {
|
||||
run_gui_mode_main_thread()
|
||||
}
|
||||
}
|
||||
|
||||
fn discover_jsonl_files(root: &str) -> Vec<(String, String)> {
|
||||
let mut files = Vec::new();
|
||||
let root_path = Path::new(root);
|
||||
if !root_path.exists() {
|
||||
return files;
|
||||
}
|
||||
collect_jsonl(root_path, root_path, &mut files);
|
||||
files
|
||||
}
|
||||
|
||||
fn collect_jsonl(base: &Path, dir: &Path, files: &mut Vec<(String, String)>) {
|
||||
if let Ok(entries) = fs::read_dir(dir) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
collect_jsonl(base, &path, files);
|
||||
} else if path.extension().map_or(false, |e| e == "jsonl") {
|
||||
let full = path.to_string_lossy().to_string();
|
||||
let label = if let Some(parent) = path.parent() {
|
||||
if parent == base {
|
||||
path.file_stem()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
} else {
|
||||
parent
|
||||
.file_name()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
}
|
||||
} else {
|
||||
path.file_stem()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
};
|
||||
files.push((full, label));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_state() -> Result<AppState> {
|
||||
println!("AiRust — Cyber Operator v0.2.0");
|
||||
|
||||
let mut idx = KnowledgeIndex::new("nomic-embed-text");
|
||||
let mut idx = KnowledgeIndex::new("dolphin-llama3:8b");
|
||||
|
||||
if let Ok(samples) = load_english("output/english_master.jsonl") {
|
||||
let clean = crate::pipeline::preprocess::remove_empty(samples);
|
||||
idx.extend_from_samples(clean, "english-master");
|
||||
let mut total = 0usize;
|
||||
let mut discovered = 0usize;
|
||||
let jsonl_files = discover_jsonl_files("output");
|
||||
for (path, label) in &jsonl_files {
|
||||
if let Ok(samples) = load_english(path) {
|
||||
let clean = crate::pipeline::preprocess::remove_empty(samples);
|
||||
let n = clean.len();
|
||||
idx.extend_from_samples(clean, label);
|
||||
total += n;
|
||||
discovered += 1;
|
||||
} else {
|
||||
println!(" ⚠ Could not load: {}", path);
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(samples) = load_english("output/random_new_knowledge.jsonl") {
|
||||
let clean = crate::pipeline::preprocess::remove_empty(samples);
|
||||
idx.extend_from_samples(clean, "dynamic");
|
||||
}
|
||||
|
||||
if let Ok(samples) = load_english("output/cyber_tools.jsonl") {
|
||||
let clean = crate::pipeline::preprocess::remove_empty(samples);
|
||||
idx.extend_from_samples(clean, "cyber-tools");
|
||||
}
|
||||
|
||||
println!("Loaded {} knowledge entries.", idx.len());
|
||||
println!("Loaded {} entries from {} file(s).", total, discovered);
|
||||
|
||||
let ollama = Arc::new(Ollama::default());
|
||||
|
||||
match ollama.list_local_models().await {
|
||||
Ok(models) => {
|
||||
let names: Vec<_> = models.iter().map(|m| m.name.clone()).collect();
|
||||
if !names.contains(&"sadiq-bd/llama3.2-3b-uncensored".to_string()) {
|
||||
println!("Warning: pull sadiq-bd/llama3.2-3b-uncensored");
|
||||
}
|
||||
if !names.contains(&"nomic-embed-text".to_string()) {
|
||||
println!("Warning: pull nomic-embed-text");
|
||||
if !names.contains(&"dolphin-llama3:8b".to_string()) {
|
||||
println!("Warning: pull dolphin-llama3:8b");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -87,27 +135,30 @@ async fn build_state() -> Result<AppState> {
|
||||
history: Arc::new(Mutex::new(Vec::new())),
|
||||
target: Arc::new(Mutex::new(None)),
|
||||
shutdown: Arc::new(tokio::sync::Notify::new()),
|
||||
model_cache: Arc::new(Mutex::new(web_server::ModelCache {
|
||||
names: Vec::new(),
|
||||
fetched_at: Instant::now(),
|
||||
})),
|
||||
embed_cache: Arc::new(Mutex::new(EmbeddingCache::new())),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
idx.build_embeddings(&ollama).await?;
|
||||
let knowledge = Arc::new(Mutex::new(idx));
|
||||
|
||||
// Build embeddings in background — server starts immediately
|
||||
// Holds the lock during build (only slow on first run; cached after)
|
||||
let bg_knowledge = knowledge.clone();
|
||||
let bg_ollama = ollama.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut k = bg_knowledge.lock().await;
|
||||
if let Err(e) = k.build_embeddings(&bg_ollama).await {
|
||||
eprintln!("Background embedding build: {e}");
|
||||
}
|
||||
});
|
||||
|
||||
Ok(AppState {
|
||||
ollama,
|
||||
knowledge: Arc::new(Mutex::new(idx)),
|
||||
knowledge,
|
||||
history: Arc::new(Mutex::new(Vec::new())),
|
||||
target: Arc::new(Mutex::new(None)),
|
||||
shutdown: Arc::new(tokio::sync::Notify::new()),
|
||||
model_cache: Arc::new(Mutex::new(web_server::ModelCache {
|
||||
names: Vec::new(),
|
||||
fetched_at: Instant::now(),
|
||||
})),
|
||||
embed_cache: Arc::new(Mutex::new(EmbeddingCache::new())),
|
||||
})
|
||||
}
|
||||
@@ -139,14 +190,23 @@ fn run_gui_mode_main_thread() -> Result<()> {
|
||||
let http_url = "http://127.0.0.1:3001".to_string();
|
||||
crate::gui::run_gui(&http_url);
|
||||
|
||||
server_handle.abort();
|
||||
if crate::gui::WEB_BROWSER_MODE.load(Ordering::SeqCst) {
|
||||
// Web Browser button was clicked — keep server running, wait for shutdown from web UI
|
||||
eprintln!("\n[*] Server running at https://127.0.0.1:3000");
|
||||
eprintln!("[*] Shut down from Web UI: Quick Actions > Kill Server");
|
||||
let _ = rt.block_on(server_handle);
|
||||
} else {
|
||||
// X button was clicked — kill server and exit
|
||||
server_handle.abort();
|
||||
}
|
||||
let _ = rt.shutdown_timeout(std::time::Duration::from_secs(3));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_cli() -> Result<()> {
|
||||
println!("AiRust CLI mode");
|
||||
|
||||
let mut idx = KnowledgeIndex::new("nomic-embed-text");
|
||||
let mut idx = KnowledgeIndex::new("dolphin-llama3:8b");
|
||||
|
||||
if let Ok(samples) = load_english("output/english_master.jsonl") {
|
||||
let clean = crate::pipeline::preprocess::remove_empty(samples);
|
||||
@@ -158,6 +218,10 @@ async fn run_cli() -> Result<()> {
|
||||
idx.extend_from_samples(clean, "cyber-tools");
|
||||
}
|
||||
|
||||
if let Ok(samples) = load_english("output/pdf_knowledge.jsonl") {
|
||||
idx.extend_from_samples(samples, "pdf-import");
|
||||
}
|
||||
|
||||
println!("Loaded {} entries.", idx.len());
|
||||
|
||||
let ollama = Ollama::default();
|
||||
@@ -187,7 +251,7 @@ async fn run_cli() -> Result<()> {
|
||||
}
|
||||
|
||||
let emb_req = ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
"dolphin-llama3:8b".to_string(),
|
||||
ollama_rs::generation::embeddings::request::EmbeddingsInput::Single(input.to_string()),
|
||||
)
|
||||
.keep_alive(ollama_rs::generation::parameters::KeepAlive::Indefinitely);
|
||||
@@ -218,10 +282,10 @@ async fn run_cli() -> Result<()> {
|
||||
msgs.extend(capped);
|
||||
msgs.push(ChatMessage::user(input.to_string()));
|
||||
|
||||
let req = ChatMessageRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), msgs)
|
||||
let req = ChatMessageRequest::new("dolphin-llama3:8b".to_string(), msgs)
|
||||
.options(
|
||||
ModelOptions::default()
|
||||
.num_ctx(4096)
|
||||
.num_ctx(8192)
|
||||
.num_thread(8)
|
||||
.num_predict(1024)
|
||||
.temperature(0.7)
|
||||
@@ -246,6 +310,143 @@ async fn run_cli() -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_learn_topics(topic_file: &str) -> Result<()> {
|
||||
use crate::data_ls::schema::EnglishSample;
|
||||
use crate::data_ls::loader::load_english;
|
||||
use crate::training::websearch::{append_sample_to_jsonl, web_search_with_context};
|
||||
use std::collections::HashSet;
|
||||
use std::io::Write;
|
||||
|
||||
let content = std::fs::read_to_string(topic_file)?;
|
||||
let topics: Vec<&str> = content.lines().filter(|l| !l.trim().is_empty()).collect();
|
||||
|
||||
if topics.is_empty() {
|
||||
eprintln!("No topics found in {}", topic_file);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let out_path = std::path::Path::new("output/random_new_knowledge.jsonl");
|
||||
let abs_path = std::fs::canonicalize(&out_path).unwrap_or_else(|_| out_path.to_path_buf());
|
||||
|
||||
// Load already-known topics so we can skip them (resume after crash)
|
||||
let known: HashSet<String> = load_english(out_path.to_str().unwrap())
|
||||
.map(|s| s.into_iter().map(|e| e.instruction).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
let pending: Vec<&&str> = topics.iter().filter(|t| !known.contains(t.trim())).collect();
|
||||
|
||||
if pending.len() != topics.len() {
|
||||
println!("Skipping {} already-learned topics.", topics.len() - pending.len());
|
||||
}
|
||||
|
||||
if pending.is_empty() {
|
||||
println!("All topics already learned. Nothing to do.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("Learning {} topics from {}", pending.len(), topic_file);
|
||||
println!("Using dolphin-llama3:8b");
|
||||
println!("Output: {}", abs_path.display());
|
||||
println!();
|
||||
|
||||
// Load all existing knowledge for context
|
||||
let mut all_knowledge: Vec<EnglishSample> = Vec::new();
|
||||
let kb_files = discover_jsonl_files("output");
|
||||
for (path, _) in &kb_files {
|
||||
if let Ok(samples) = load_english(path) {
|
||||
all_knowledge.extend(samples);
|
||||
}
|
||||
}
|
||||
println!("Loaded {} knowledge entries for context.", all_knowledge.len());
|
||||
|
||||
let ollama = Ollama::default();
|
||||
let model = "dolphin-llama3:8b";
|
||||
|
||||
std::fs::create_dir_all("output")?;
|
||||
|
||||
let mut learned = 0;
|
||||
let mut failed = 0;
|
||||
|
||||
// Build a word index for simple keyword matching
|
||||
fn topic_words(topic: &str) -> Vec<String> {
|
||||
topic
|
||||
.to_lowercase()
|
||||
.split(|c: char| !c.is_alphanumeric())
|
||||
.filter(|w| w.len() > 3)
|
||||
.map(|w| w.to_string())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn match_score(topic_words: &[String], sample: &EnglishSample) -> usize {
|
||||
let haystack = format!("{} {} {}", sample.instruction, sample.output, sample.reasoning).to_lowercase();
|
||||
topic_words.iter().filter(|w| haystack.contains(w.as_str())).count()
|
||||
}
|
||||
|
||||
for (i, topic) in pending.iter().enumerate() {
|
||||
let topic = topic.trim();
|
||||
print!("[{}/{}] {}... ", i + 1, pending.len(), topic);
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
// Find up to 5 most relevant knowledge entries
|
||||
let tw = topic_words(topic);
|
||||
let mut scored: Vec<(usize, &EnglishSample)> = all_knowledge
|
||||
.iter()
|
||||
.map(|s| (match_score(&tw, s), s))
|
||||
.collect();
|
||||
scored.sort_by(|a, b| b.0.cmp(&a.0));
|
||||
scored.truncate(5);
|
||||
|
||||
let context: String = scored
|
||||
.iter()
|
||||
.filter(|(score, _)| *score > 0)
|
||||
.enumerate()
|
||||
.map(|(idx, (_, s))| {
|
||||
format!(
|
||||
"[{}] Q: {}\nA: {}\n\n",
|
||||
idx + 1,
|
||||
s.instruction,
|
||||
s.output
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
match if context.is_empty() {
|
||||
web_search_with_context(topic, "", &ollama, model).await
|
||||
} else {
|
||||
web_search_with_context(topic, &context, &ollama, model).await
|
||||
} {
|
||||
Ok(answer) => {
|
||||
let sample = EnglishSample {
|
||||
r#type: "instruction_following".into(),
|
||||
instruction: topic.to_string(),
|
||||
input: String::new(),
|
||||
reasoning: format!("AI-generated knowledge about {}", topic),
|
||||
output: answer.trim().to_string(),
|
||||
};
|
||||
if let Err(e) = append_sample_to_jsonl(&sample, out_path.to_str().unwrap()) {
|
||||
eprintln!(" write failed: {e}");
|
||||
failed += 1;
|
||||
} else {
|
||||
// Verify the write by counting lines
|
||||
let count = load_english(out_path.to_str().unwrap())
|
||||
.map(|s| s.len())
|
||||
.unwrap_or(0);
|
||||
println!(" OK ({} entries total)", count);
|
||||
learned += 1;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!(" FAILED: {e}");
|
||||
failed += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!("\nDone. Learned: {}, Failed: {}", learned, failed);
|
||||
println!("All entries in {}", abs_path.display());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_pdf_import(args: &[String]) -> Result<()> {
|
||||
use crate::pdf_import::import_pdf_path;
|
||||
use anyhow::anyhow;
|
||||
@@ -312,17 +513,7 @@ async fn run_pdf_import(args: &[String]) -> Result<()> {
|
||||
return Err(anyhow!("Exists: {}", out_path.display()));
|
||||
}
|
||||
|
||||
let models = ollama.list_local_models().await.unwrap_or_default();
|
||||
let names: Vec<_> = models.iter().map(|m| m.name.clone()).collect();
|
||||
let model = if names.contains(&"qwen2.5-coder:1.5b".to_string()) {
|
||||
"qwen2.5-coder:1.5b"
|
||||
} else if names.contains(&"sadiq-bd/llama3.2-3b-uncensored".to_string()) {
|
||||
"sadiq-bd/llama3.2-3b-uncensored"
|
||||
} else {
|
||||
return Err(anyhow!(
|
||||
"No models available. Pull sadiq-bd/llama3.2-3b-uncensored or qwen2.5-coder:1.5b"
|
||||
));
|
||||
};
|
||||
let model = "dolphin-llama3:8b";
|
||||
|
||||
println!(
|
||||
"Distilling {} -> {} [{}]",
|
||||
@@ -355,3 +546,40 @@ async fn run_pdf_import(args: &[String]) -> Result<()> {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn export_training_data() -> Result<()> {
|
||||
use std::io::Write;
|
||||
let files = [
|
||||
"output/english_master.jsonl",
|
||||
"output/random_new_knowledge.jsonl",
|
||||
"output/cyber_tools.jsonl",
|
||||
"output/pdf_knowledge.jsonl",
|
||||
];
|
||||
|
||||
let mut all = Vec::new();
|
||||
for path in &files {
|
||||
if let Ok(samples) = load_english(path) {
|
||||
all.extend(samples);
|
||||
}
|
||||
}
|
||||
|
||||
if all.is_empty() {
|
||||
eprintln!("No knowledge data found in output/*.jsonl");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let out_path = "output/training_data.jsonl";
|
||||
let mut file = std::fs::File::create(out_path)?;
|
||||
for s in &all {
|
||||
let alpaca = serde_json::json!({
|
||||
"instruction": s.instruction,
|
||||
"input": s.input,
|
||||
"output": s.output,
|
||||
});
|
||||
writeln!(file, "{}", alpaca)?;
|
||||
}
|
||||
|
||||
println!("Exported {} samples to {}", all.len(), out_path);
|
||||
println!("Ready for fine-tuning tools (unsloth, llama.cpp, Axolotl, etc.)");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+1
-1
@@ -35,7 +35,7 @@ fn find_pdf_files(input_dir: &Path) -> Result<Vec<PathBuf>> {
|
||||
const MAX_CHUNK_CHARS: usize = 10000;
|
||||
|
||||
/// Split large PDF text into manageable chunks while preserving whole paragraphs.
|
||||
fn split_text_into_chunks(text: &str, max_chars: usize) -> Vec<String> {
|
||||
pub fn split_text_into_chunks(text: &str, max_chars: usize) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let mut current = String::new();
|
||||
|
||||
|
||||
@@ -1,50 +0,0 @@
|
||||
use std::fs::{self, File};
|
||||
use std::io::{self, BufRead, BufReader, BufWriter, Write};
|
||||
use std::path::Path;
|
||||
|
||||
pub fn merge_english_jsonl() -> io::Result<()> {
|
||||
let input_dir = "data/LLM_English";
|
||||
let output_dir = "output";
|
||||
let output_file = "english_master.jsonl";
|
||||
|
||||
fs::create_dir_all(output_dir)?;
|
||||
|
||||
let output_path = Path::new(output_dir).join(output_file);
|
||||
let mut writer = BufWriter::new(File::create(&output_path)?);
|
||||
|
||||
let files = [
|
||||
"english_explain.jsonl",
|
||||
"english_reasoning.jsonl",
|
||||
"english_summary.jsonl",
|
||||
"instruction_following.jsonl",
|
||||
];
|
||||
|
||||
let mut total_lines: usize = 0;
|
||||
|
||||
for filename in &files {
|
||||
let input_path = Path::new(input_dir).join(filename);
|
||||
|
||||
if !input_path.exists() {
|
||||
println!("⚠️ Warning: {} not found, skipping...", filename);
|
||||
continue;
|
||||
}
|
||||
|
||||
println!("📄 Merging: {}", filename);
|
||||
|
||||
let file = File::open(&input_path)?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if !trimmed.is_empty() {
|
||||
writeln!(writer, "{}", trimmed)?;
|
||||
total_lines += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
writer.flush()?;
|
||||
println!("✅ Successfully merged {} lines into → {}", total_lines, output_path.display());
|
||||
Ok(())
|
||||
}
|
||||
+137
-285
@@ -1,5 +1,4 @@
|
||||
// src/training/websearch.rs
|
||||
#![allow(dead_code)]
|
||||
|
||||
use crate::data_ls::schema::EnglishSample;
|
||||
use crate::knowledge_index::OLLAMA_SEM;
|
||||
@@ -12,13 +11,27 @@ use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbedd
|
||||
use ollama_rs::generation::parameters::{FormatType, KeepAlive};
|
||||
use ollama_rs::models::ModelOptions;
|
||||
use ollama_rs::Ollama;
|
||||
use reqwest::{header, Client};
|
||||
use scraper::{Html, Selector};
|
||||
use serde_json::Value;
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
use tokio::time::{timeout, Duration};
|
||||
use urlencoding::encode;
|
||||
use tokio::time::Duration;
|
||||
|
||||
/// Convert a JSON Value to EnglishSample, handling `output` being either a string or an object.
|
||||
fn value_to_english_sample(v: &Value) -> Result<EnglishSample> {
|
||||
let typ = v.get("type").and_then(|t| t.as_str()).unwrap_or("instruction_following").to_string();
|
||||
let instruction = v.get("instruction").and_then(|t| t.as_str()).unwrap_or("").to_string();
|
||||
let input = v.get("input").and_then(|t| t.as_str()).unwrap_or("").to_string();
|
||||
let reasoning = v.get("reasoning").and_then(|t| t.as_str()).unwrap_or("").to_string();
|
||||
let output = match v.get("output") {
|
||||
Some(Value::String(s)) => s.clone(),
|
||||
Some(Value::Object(m)) => {
|
||||
m.get("content").and_then(|c| c.as_str()).unwrap_or("").to_string()
|
||||
}
|
||||
_ => String::new(),
|
||||
};
|
||||
Ok(EnglishSample { r#type: typ, instruction, input, reasoning, output })
|
||||
}
|
||||
|
||||
fn extract_json_from_response(text: &str) -> Option<String> {
|
||||
let trimmed = text.trim();
|
||||
@@ -117,200 +130,68 @@ fn validate_sample(sample: &EnglishSample) -> Result<()> {
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SearchEngine {
|
||||
Dolphin,
|
||||
DuckDuckGo,
|
||||
Deepseek,
|
||||
Qwen,
|
||||
Grok,
|
||||
AiSearch,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SearchEngine {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"{}",
|
||||
match self {
|
||||
SearchEngine::Dolphin => "Dolphin Chat",
|
||||
SearchEngine::DuckDuckGo => "DuckDuckGo",
|
||||
SearchEngine::Deepseek => "Deepseek",
|
||||
SearchEngine::Qwen => "Qwen",
|
||||
SearchEngine::Grok => "Grok",
|
||||
}
|
||||
)
|
||||
write!(f, "AI Search")
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchEngine {
|
||||
pub fn search_url(&self, query: &str) -> String {
|
||||
match self {
|
||||
SearchEngine::Dolphin => format!("https://chat.dphn.ai/?q={}", encode(query)),
|
||||
SearchEngine::DuckDuckGo => {
|
||||
format!("https://html.duckduckgo.com/html/?q={}", encode(query))
|
||||
}
|
||||
SearchEngine::Deepseek => format!("https://deepseek.ai/search?q={}", encode(query)),
|
||||
SearchEngine::Qwen => format!("https://qwen.ai/search?q={}", encode(query)),
|
||||
SearchEngine::Grok => format!("https://grok.com/search?q={}", encode(query)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn create_browser_client() -> Result<Client, reqwest::Error> {
|
||||
let mut headers = header::HeaderMap::new();
|
||||
headers.insert(
|
||||
header::ACCEPT,
|
||||
header::HeaderValue::from_static(
|
||||
"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
|
||||
),
|
||||
);
|
||||
headers.insert(
|
||||
header::ACCEPT_LANGUAGE,
|
||||
header::HeaderValue::from_static("en-US,en;q=0.9"),
|
||||
);
|
||||
headers.insert(
|
||||
header::CONNECTION,
|
||||
header::HeaderValue::from_static("keep-alive"),
|
||||
);
|
||||
headers.insert(
|
||||
header::USER_AGENT,
|
||||
header::HeaderValue::from_static("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"),
|
||||
);
|
||||
Client::builder().default_headers(headers).build()
|
||||
}
|
||||
|
||||
pub async fn web_search(query: &str, engine: SearchEngine) -> Result<String> {
|
||||
let url = engine.search_url(query);
|
||||
let client = create_browser_client()?;
|
||||
let body = client
|
||||
.get(&url)
|
||||
.header(header::REFERER, "https://chat.dphn.ai/")
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?
|
||||
.text()
|
||||
.await?;
|
||||
let doc = Html::parse_document(&body);
|
||||
|
||||
let snippets: Vec<String> = match engine {
|
||||
SearchEngine::Dolphin => {
|
||||
let selector = Selector::parse("p, div, span").unwrap();
|
||||
let mut snippets: Vec<String> = doc
|
||||
.select(&selector)
|
||||
.filter_map(|el| {
|
||||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||||
if text.len() >= 20 && text.len() <= 300 {
|
||||
Some(text)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(5)
|
||||
.collect();
|
||||
|
||||
if snippets.is_empty() {
|
||||
let body_text = doc.root_element().text().collect::<Vec<_>>().join(" ");
|
||||
snippets = body_text
|
||||
.split_terminator(|c| c == '.' || c == '!' || c == '?')
|
||||
.map(str::trim)
|
||||
.filter(|text| text.len() >= 20 && text.len() <= 300)
|
||||
.take(5)
|
||||
.map(String::from)
|
||||
.collect();
|
||||
}
|
||||
|
||||
snippets
|
||||
}
|
||||
SearchEngine::DuckDuckGo => {
|
||||
let selector = Selector::parse(".result__snippet").unwrap();
|
||||
doc.select(&selector)
|
||||
.take(5)
|
||||
.map(|el| el.inner_html())
|
||||
.collect()
|
||||
}
|
||||
SearchEngine::Deepseek => {
|
||||
let selector = Selector::parse(".snippet, .search-result, .result, p").unwrap();
|
||||
doc.select(&selector)
|
||||
.filter_map(|el| {
|
||||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||||
if text.len() >= 20 && text.len() <= 300 {
|
||||
Some(text)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(5)
|
||||
.collect()
|
||||
}
|
||||
SearchEngine::Qwen => {
|
||||
let selector = Selector::parse(".search-result, .result, p").unwrap();
|
||||
doc.select(&selector)
|
||||
.filter_map(|el| {
|
||||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||||
if text.len() >= 20 && text.len() <= 300 {
|
||||
Some(text)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(5)
|
||||
.collect()
|
||||
}
|
||||
SearchEngine::Grok => {
|
||||
let selector = Selector::parse(".search-result, .result, p").unwrap();
|
||||
doc.select(&selector)
|
||||
.filter_map(|el| {
|
||||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||||
if text.len() >= 20 && text.len() <= 300 {
|
||||
Some(text)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(5)
|
||||
.collect()
|
||||
}
|
||||
};
|
||||
|
||||
if snippets.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"{} returned no usable results from {}",
|
||||
engine,
|
||||
url
|
||||
));
|
||||
}
|
||||
|
||||
Ok(format!("{}\n\nSearch page: {}", snippets.join(" "), url))
|
||||
}
|
||||
|
||||
/// Use LLM to convert search result into a structured EnglishSample
|
||||
pub async fn derive_sample_from_search(
|
||||
pub async fn web_search(
|
||||
query: &str,
|
||||
snippets: &str,
|
||||
engine: SearchEngine,
|
||||
ollama: Option<&Ollama>,
|
||||
model_name: Option<&str>,
|
||||
) -> Result<String> {
|
||||
let _ = engine;
|
||||
let ollama = ollama.ok_or_else(|| anyhow!("AI search requires Ollama client"))?;
|
||||
let model = model_name.unwrap_or("dolphin-llama3:8b");
|
||||
web_search_inner(query, None, &ollama, model).await
|
||||
}
|
||||
|
||||
pub async fn web_search_with_context(
|
||||
query: &str,
|
||||
context: &str,
|
||||
ollama: &Ollama,
|
||||
model_name: &str,
|
||||
) -> Result<EnglishSample> {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let prompt = format!(
|
||||
"Convert the following search result into a structured learning sample in JSON format.\n\
|
||||
Use this exact schema:\n\
|
||||
{{\"type\": \"instruction_following\", \"instruction\": \"{}\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}\n\n\
|
||||
Query: {}\nSearch result: {}\n\nOnly output valid JSON, no extra text. Fill in the reasoning and output based on the search result. The instruction should be the query itself or a rephrased version. The type is always \"instruction_following\".",
|
||||
query, query, snippets
|
||||
);
|
||||
let req = ChatMessageRequest::new(model_name.to_string(), vec![ChatMessage::user(prompt)])
|
||||
.options(
|
||||
ModelOptions::default()
|
||||
.num_ctx(8192)
|
||||
.num_thread(8)
|
||||
.num_predict(1024)
|
||||
.temperature(0.85),
|
||||
) -> Result<String> {
|
||||
web_search_inner(query, Some(context), ollama, model_name).await
|
||||
}
|
||||
|
||||
async fn web_search_inner(
|
||||
query: &str,
|
||||
context: Option<&str>,
|
||||
ollama: &Ollama,
|
||||
model: &str,
|
||||
) -> Result<String> {
|
||||
let prompt = if let Some(ctx) = context {
|
||||
format!(
|
||||
"Here is relevant reference material:\n{}\n\nBased on the reference material above, answer the following query accurately. Use specific details from the reference material. If the reference material doesn't fully answer the query, supplement with your own knowledge.\n\nQuery: {}\n\nProvide a detailed, accurate answer.",
|
||||
ctx, query
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"Answer the following query with a detailed, accurate response. Include specific facts, examples, and explanations.\n\nQuery: {}\n\nProvide a comprehensive answer based on your knowledge. Be thorough and specific.",
|
||||
query
|
||||
)
|
||||
};
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let opts = ModelOptions::default()
|
||||
.num_ctx(8192)
|
||||
.num_thread(8)
|
||||
.num_predict(2048)
|
||||
.temperature(0.7);
|
||||
let msg = ChatMessageRequest::new(model.to_string(), vec![ChatMessage::user(prompt)])
|
||||
.options(opts)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
let resp = timeout(Duration::from_secs(120), ollama.send_chat_messages(req))
|
||||
.await
|
||||
.map_err(|_| anyhow!("Sample generation timed out after 120s"))??;
|
||||
let json_str = resp.message.content;
|
||||
let sample: EnglishSample = serde_json::from_str(&json_str)?;
|
||||
Ok(sample)
|
||||
let result = match tokio::time::timeout(Duration::from_secs(600), ollama.send_chat_messages(msg)).await {
|
||||
Ok(Ok(resp)) => resp.message.content,
|
||||
Ok(Err(e)) => return Err(anyhow!("AI search error: {}", e)),
|
||||
Err(_) => return Err(anyhow!("AI search timed out")),
|
||||
};
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn derive_sample_from_chunk(
|
||||
@@ -320,66 +201,29 @@ async fn derive_sample_from_chunk(
|
||||
total: usize,
|
||||
ollama: &Ollama,
|
||||
model_name: &str,
|
||||
) -> Result<EnglishSample> {
|
||||
) -> Result<Vec<EnglishSample>> {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let prompt = format!(
|
||||
"Extract ONE learning sample from the following text in {}.\n\
|
||||
Use this exact schema: {{\"type\": \"instruction_following\", \"instruction\": \"...\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}.\n\
|
||||
"Extract ALL learning samples from the following text in {}.\n\
|
||||
Return a JSON array of objects. \
|
||||
Each object uses this schema: {{\"type\": \"instruction_following\", \"instruction\": \"...\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}.\n\
|
||||
The instruction must be a specific technical question or task based *only* on the text.\n\
|
||||
The output must be a direct answer or summary (no empty strings unless truly absent).\n\
|
||||
Do NOT include any meta-dialogue, do NOT describe the format, and do NOT output anything except the JSON object.\n\nText:\n{}",
|
||||
The output must be a direct answer or summary.\n\
|
||||
Extract every distinct piece of knowledge — do not leave anything out.\n\
|
||||
Do NOT include any meta-dialogue and do NOT output anything except the JSON array.\n\nText:\n{}",
|
||||
pdf_name, chunk
|
||||
);
|
||||
|
||||
let opts = ModelOptions::default()
|
||||
.num_ctx(8192)
|
||||
.num_thread(8)
|
||||
.num_predict(2048)
|
||||
.num_predict(4096)
|
||||
.temperature(0.85);
|
||||
let request = GenerationRequest::new(model_name.to_string(), prompt.clone())
|
||||
.format(FormatType::Json)
|
||||
.options(opts.clone());
|
||||
|
||||
let resp = match ollama.generate(request).await {
|
||||
Ok(r) => r,
|
||||
Err(e) if model_name == "deepseek-r1:7b" => {
|
||||
eprintln!(
|
||||
"⚠️ deepseek-r1:7b failed, falling back to qwen2.5-coder:1.5b. Error: {}",
|
||||
e
|
||||
);
|
||||
let fallback_req =
|
||||
GenerationRequest::new("qwen2.5-coder:1.5b".to_string(), prompt.clone())
|
||||
.format(FormatType::Json)
|
||||
.options(opts.clone());
|
||||
match ollama.generate(fallback_req).await {
|
||||
Ok(r) => r,
|
||||
Err(_) => {
|
||||
eprintln!(
|
||||
"⚠️ qwen2.5-coder:1.5b also failed; falling back to sadiq-bd/llama3.2-3b-uncensored."
|
||||
);
|
||||
let fallback_req = GenerationRequest::new(
|
||||
"sadiq-bd/llama3.2-3b-uncensored".to_string(),
|
||||
prompt,
|
||||
)
|
||||
.format(FormatType::Json)
|
||||
.options(opts);
|
||||
ollama.generate(fallback_req).await?
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) if model_name == "qwen2.5-coder:1.5b" => {
|
||||
eprintln!(
|
||||
"⚠️ qwen2.5-coder:1.5b failed, falling back to sadiq-bd/llama3.2-3b-uncensored. Error: {}",
|
||||
e
|
||||
);
|
||||
let fallback_req =
|
||||
GenerationRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), prompt)
|
||||
.format(FormatType::Json)
|
||||
.options(opts);
|
||||
ollama.generate(fallback_req).await?
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
let resp = ollama.generate(request).await?;
|
||||
|
||||
let raw_response = resp.response;
|
||||
let mut json_candidate = extract_json_from_response(&raw_response)
|
||||
@@ -388,61 +232,69 @@ async fn derive_sample_from_chunk(
|
||||
json_candidate = repair_json_braces(&json_candidate);
|
||||
}
|
||||
|
||||
let sample: EnglishSample = match serde_json::from_str::<EnglishSample>(&json_candidate) {
|
||||
Ok(mut s) => {
|
||||
s.output = clean_output(&s.output);
|
||||
s
|
||||
}
|
||||
Err(err) => {
|
||||
let repaired_candidate = repair_json_braces(&json_candidate);
|
||||
if repaired_candidate != json_candidate {
|
||||
serde_json::from_str(&repaired_candidate).map_err(|second_err| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to parse JSON from PDF chunk {} of {}: {}\nRetry parse error: {}\nCandidate: {}\nRepaired: {}\nRaw response: {}",
|
||||
index + 1,
|
||||
total,
|
||||
err,
|
||||
second_err,
|
||||
json_candidate,
|
||||
repaired_candidate,
|
||||
raw_response
|
||||
)
|
||||
})?
|
||||
} else if let Some(cleaned) = extract_json_from_response(&raw_response) {
|
||||
serde_json::from_str(&cleaned).map_err(|second_err| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to parse JSON from PDF chunk {} of {}: {}\nRetry parse error: {}\nCleaned response: {}\nRaw response: {}",
|
||||
index + 1,
|
||||
total,
|
||||
err,
|
||||
second_err,
|
||||
cleaned,
|
||||
raw_response
|
||||
)
|
||||
})?
|
||||
} else {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to parse JSON from PDF chunk {} of {}.\nRaw response:\n{}",
|
||||
index + 1,
|
||||
total,
|
||||
raw_response
|
||||
));
|
||||
}
|
||||
let samples = parse_samples_from_json(&json_candidate, index, total, &raw_response)?;
|
||||
|
||||
for (i, sample) in samples.iter().enumerate() {
|
||||
validate_sample(sample).map_err(|validation_err| {
|
||||
anyhow::anyhow!(
|
||||
"Validation failed for PDF chunk {} of {} sample {}: {}\nRaw response: {}",
|
||||
index + 1,
|
||||
total,
|
||||
i + 1,
|
||||
validation_err,
|
||||
raw_response
|
||||
)
|
||||
})?;
|
||||
}
|
||||
|
||||
Ok(samples)
|
||||
}
|
||||
|
||||
fn parse_samples_from_json(
|
||||
json_candidate: &str,
|
||||
index: usize,
|
||||
total: usize,
|
||||
raw_response: &str,
|
||||
) -> Result<Vec<EnglishSample>> {
|
||||
let try_parse = |s: &str| -> Option<Vec<EnglishSample>> {
|
||||
let v: Value = serde_json::from_str(s).ok()?;
|
||||
let items: Vec<&Value> = match &v {
|
||||
Value::Array(arr) => arr.iter().collect(),
|
||||
_ => vec![&v],
|
||||
};
|
||||
let mut out = Vec::with_capacity(items.len());
|
||||
for item in items {
|
||||
let mut sample = value_to_english_sample(item).ok()?;
|
||||
sample.output = clean_output(&sample.output);
|
||||
out.push(sample);
|
||||
}
|
||||
Some(out)
|
||||
};
|
||||
|
||||
validate_sample(&sample).map_err(|validation_err| {
|
||||
anyhow::anyhow!(
|
||||
"Validation failed for PDF chunk {} of {}: {}\nCandidate: {}\nRaw response: {}",
|
||||
index + 1,
|
||||
total,
|
||||
validation_err,
|
||||
json_candidate,
|
||||
raw_response
|
||||
)
|
||||
})?;
|
||||
if let Some(samples) = try_parse(json_candidate) {
|
||||
return Ok(samples);
|
||||
}
|
||||
|
||||
Ok(sample)
|
||||
let repaired = repair_json_braces(json_candidate);
|
||||
if repaired != json_candidate {
|
||||
if let Some(samples) = try_parse(&repaired) {
|
||||
return Ok(samples);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(cleaned) = extract_json_from_response(raw_response) {
|
||||
if let Some(samples) = try_parse(&cleaned) {
|
||||
return Ok(samples);
|
||||
}
|
||||
}
|
||||
|
||||
Err(anyhow::anyhow!(
|
||||
"Failed to parse JSON from PDF chunk {} of {}.\nCandidate: {}\nRaw response: {}",
|
||||
index + 1,
|
||||
total,
|
||||
json_candidate,
|
||||
raw_response
|
||||
))
|
||||
}
|
||||
|
||||
/// Convert a large PDF text into a sequence of structured EnglishSample objects.
|
||||
@@ -475,7 +327,7 @@ pub async fn derive_samples_from_pdf_text(
|
||||
|
||||
let batch_results = join_all(futures).await;
|
||||
for result in batch_results {
|
||||
all_samples.push(result?);
|
||||
all_samples.extend(result?);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+271
-11
@@ -209,7 +209,7 @@ body {
|
||||
/* Themed chat backgrounds */
|
||||
[data-theme="cyberpunk"] .chat-messages::before {
|
||||
content: '';
|
||||
position: absolute;
|
||||
position: fixed;
|
||||
inset: 0;
|
||||
pointer-events: none;
|
||||
background-image:
|
||||
@@ -222,7 +222,7 @@ body {
|
||||
|
||||
[data-theme="venom"] .chat-messages::before {
|
||||
content: '';
|
||||
position: absolute;
|
||||
position: fixed;
|
||||
inset: 0;
|
||||
pointer-events: none;
|
||||
background-image:
|
||||
@@ -240,7 +240,7 @@ body {
|
||||
}
|
||||
[data-theme="rick"] .chat-messages::before {
|
||||
content: '';
|
||||
position: absolute;
|
||||
position: fixed;
|
||||
inset: 0;
|
||||
pointer-events: none;
|
||||
background-image:
|
||||
@@ -277,14 +277,74 @@ body {
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 6px;
|
||||
padding: 12px 16px;
|
||||
white-space: pre-wrap;
|
||||
line-height: 1.5;
|
||||
line-height: 1.6;
|
||||
}
|
||||
|
||||
.message.assistant .content {
|
||||
border-left: 3px solid var(--accent);
|
||||
}
|
||||
|
||||
.message .content h2 { font-size: 15px; margin: 28px 0 14px; color: var(--accent); }
|
||||
.message .content h3 { font-size: 13px; margin: 22px 0 10px; color: var(--accent); }
|
||||
.message .content p { margin: 0 0 14px; }
|
||||
.message .content p:last-child { margin-bottom: 0; }
|
||||
.message .content ul, .message .content ol { margin: 6px 0 14px; padding-left: 20px; }
|
||||
.message .content li { margin-bottom: 6px; }
|
||||
.message .content pre {
|
||||
background: var(--bg);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 4px;
|
||||
padding: 10px 12px;
|
||||
margin: 8px 0;
|
||||
overflow-x: auto;
|
||||
font-size: 12px;
|
||||
line-height: 1.5;
|
||||
position: relative;
|
||||
}
|
||||
.message .content pre code {
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
tab-size: 4;
|
||||
}
|
||||
.message .content pre .copy-code-btn {
|
||||
position: absolute;
|
||||
top: 4px;
|
||||
right: 4px;
|
||||
font-size: 10px;
|
||||
padding: 2px 8px;
|
||||
background: var(--surface);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 3px;
|
||||
color: var(--text-dim);
|
||||
cursor: pointer;
|
||||
opacity: 0.5;
|
||||
z-index: 2;
|
||||
}
|
||||
.message .content pre .copy-code-btn:hover {
|
||||
opacity: 1;
|
||||
color: var(--accent);
|
||||
}
|
||||
/* syntax highlighting */
|
||||
.message .content .syn-comment { color: #6a9955; }
|
||||
.message .content .syn-string { color: #ce9178; }
|
||||
.message .content .syn-keyword { color: #569cd6; font-weight: bold; }
|
||||
.message .content .syn-number { color: #b5cea8; }
|
||||
.message .content .syn-builtin { color: #dcdcaa; }
|
||||
.message .content .syn-type { color: #4ec9b0; }
|
||||
.message .content code {
|
||||
background: var(--bg);
|
||||
padding: 1px 4px;
|
||||
border-radius: 3px;
|
||||
font-size: 12px;
|
||||
}
|
||||
.message .content pre code {
|
||||
background: none;
|
||||
padding: 0;
|
||||
border-radius: 0;
|
||||
}
|
||||
.message .content hr { border: none; border-top: 1px solid var(--border); margin: 12px 0; }
|
||||
.message .content strong { color: var(--accent); }
|
||||
|
||||
.message .meta {
|
||||
font-size: 10px;
|
||||
color: var(--text-dim);
|
||||
@@ -598,6 +658,12 @@ select:focus { border-color: var(--accent); }
|
||||
<textarea class="chat-input" id="chat-input" rows="1" placeholder="Type your message..." onkeydown="handleChatKey(event)"></textarea>
|
||||
<button class="btn primary" onclick="sendChat()">Send</button>
|
||||
</div>
|
||||
<div style="display:flex;gap:8px;align-items:center;padding:4px 0;font-size:11px;color:var(--text-dim);">
|
||||
<label style="display:flex;align-items:center;gap:4px;cursor:pointer;">
|
||||
<input type="checkbox" id="rick-voice" checked> 🎙 Rick Voice
|
||||
</label>
|
||||
<button class="btn" id="stop-speech-btn" style="font-size:10px;padding:2px 8px;display:none;" onclick="stopSpeech()">Stop</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -817,10 +883,11 @@ function addMessage(role, content, meta = '') {
|
||||
const container = document.getElementById('chat-messages');
|
||||
const msg = document.createElement('div');
|
||||
msg.className = 'message ' + role;
|
||||
const safeContent = role === 'assistant' ? sanitizeHtml(content) : escapeHtml(content);
|
||||
msg.innerHTML = `
|
||||
<div class="role">${role}</div>
|
||||
<div class="content">${escapeHtml(content)}</div>
|
||||
${meta ? `<div class="meta">${meta}</div>` : ''}
|
||||
<div class="content">${safeContent}</div>
|
||||
${meta ? `<div class="meta">${escapeHtml(meta)}</div>` : ''}
|
||||
`;
|
||||
container.appendChild(msg);
|
||||
container.scrollTop = container.scrollHeight;
|
||||
@@ -852,7 +919,8 @@ function saveChat() {
|
||||
const msgs = [];
|
||||
document.querySelectorAll('#chat-messages .message').forEach(el => {
|
||||
const role = el.classList.contains('user') ? 'user' : el.classList.contains('assistant') ? 'assistant' : 'system';
|
||||
const content = el.querySelector('.content')?.textContent || '';
|
||||
const contentEl = el.querySelector('.content');
|
||||
const content = contentEl ? (role === 'assistant' ? contentEl.innerHTML : contentEl.textContent) : '';
|
||||
const meta = el.querySelector('.meta')?.textContent || '';
|
||||
msgs.push({ role, content, meta });
|
||||
});
|
||||
@@ -964,7 +1032,9 @@ function restoreChat() {
|
||||
for (const m of session.messages) {
|
||||
const el = document.createElement('div');
|
||||
el.className = 'message ' + m.role;
|
||||
el.innerHTML = `<div class="role">${m.role}</div><div class="content">${escapeHtml(m.content)}</div>${m.meta ? `<div class="meta">${escapeHtml(m.meta)}</div>` : ''}`;
|
||||
const safeContent = m.role === 'assistant' ? sanitizeHtml(m.content) : escapeHtml(m.content);
|
||||
el.innerHTML = `<div class="role">${m.role}</div><div class="content">${safeContent}</div>${m.meta ? `<div class="meta">${escapeHtml(m.meta)}</div>` : ''}`;
|
||||
if (m.role === 'assistant') addCodeCopyButtons(el.querySelector('.content'));
|
||||
container.appendChild(el);
|
||||
}
|
||||
container.scrollTop = container.scrollHeight;
|
||||
@@ -974,6 +1044,135 @@ function escapeHtml(text) {
|
||||
return text.replace(/&/g, '&').replace(/</g, '<').replace(/>/g, '>');
|
||||
}
|
||||
|
||||
function sanitizeHtml(raw) {
|
||||
// Step 1: convert markdown patterns to HTML before DOM parsing
|
||||
let h = raw;
|
||||
const codeBlocks = [];
|
||||
|
||||
// Insert newline before code blocks that run into text: "text\`\`\`python" → "text\n\`\`\`python"
|
||||
h = h.replace(/([^\n])```/g, '$1\n```');
|
||||
|
||||
// Extract code blocks — handle both proper (```python\ncode```) and malformed
|
||||
// (```pythoncode``` or ```pythoncode``, no newline after language)
|
||||
h = h.replace(/```(\w*)\n?([\s\S]*?)``(?:`)?/g, (_, lang, code) => {
|
||||
const idx = codeBlocks.length;
|
||||
code = fixCodeNewlines(code);
|
||||
// Store raw code — DOM handles escaping later
|
||||
codeBlocks.push({ lang: lang || '', code: code.trim() });
|
||||
return `\n%%CODEBLOCK${idx}%%\n`;
|
||||
});
|
||||
|
||||
// Insert newline before concatenated list items: "text* **Item**" or "text* Item"
|
||||
h = h.replace(/([^\n])(\* )(?=\*?\*?[A-Z])/g, '$1\n$2');
|
||||
h = h.replace(/([^\n])(- )(?=\S)/g, '$1\n$2');
|
||||
// Insert newline before concatenated numbered items: "text1. Item"
|
||||
h = h.replace(/([^\n])(\d+\. )(?=\*?\*?[A-Z])/g, '$1\n$2');
|
||||
|
||||
// Convert **bold** (must come before *italic*)
|
||||
h = h.replace(/\*\*(.+?)\*\*/g, '<strong>$1</strong>');
|
||||
// Convert *italic* (but not inside words like *list markers or *emphasis*)
|
||||
h = h.replace(/(?<!\w)\*(.+?)\*(?!\w)/g, '<em>$1</em>');
|
||||
|
||||
// Convert `inline code` (skip placeholders)
|
||||
h = h.replace(/`([^`]+)`/g, (m, c) => c.startsWith('%') && c.endsWith('%') ? m : `<code>${c}</code>`);
|
||||
|
||||
// Convert headings: ###, ##, #
|
||||
h = h.replace(/^### (.+)$/gm, '<h3>$1</h3>');
|
||||
h = h.replace(/^## (.+)$/gm, '<h2>$1</h2>');
|
||||
h = h.replace(/^# (.+)$/gm, '<h2>$1</h2>');
|
||||
|
||||
// Horizontal rules
|
||||
h = h.replace(/^-{3,}$/gm, '<hr>');
|
||||
|
||||
// Convert list items: lines starting with * or -
|
||||
h = h.replace(/^[*] (.+)$/gm, '<li>$1</li>');
|
||||
h = h.replace(/^[-] (.+)$/gm, '<li>$1</li>');
|
||||
|
||||
// Wrap consecutive <li> in <ul>
|
||||
h = h.replace(/((?:<li>.*<\/li>\n?)+)/g, '<ul>$1</ul>');
|
||||
|
||||
// Wrap remaining lines in <p> (skip already-wrapped elements and code block placeholders)
|
||||
h = h.replace(/^(?!\s*$|<[hulpo]|<li|<pre|%%CODE)(.+)$/gm, '<p>$1</p>');
|
||||
|
||||
// Restore code blocks (give them breathing room)
|
||||
h = h.replace(/%%CODEBLOCK(\d+)%%/g, (_, idx) => {
|
||||
const block = codeBlocks[parseInt(idx)];
|
||||
return `\n<pre><code class="language-${block.lang}">${block.code}</code></pre>\n`;
|
||||
});
|
||||
|
||||
// Step 2: fix missing spaces in text
|
||||
const doc = new DOMParser().parseFromString('<div>' + h + '</div>', 'text/html');
|
||||
const root = doc.body.firstElementChild;
|
||||
if (!root) return escapeHtml(raw);
|
||||
|
||||
// Remove scripts and event handlers
|
||||
root.querySelectorAll('script').forEach(s => s.remove());
|
||||
root.querySelectorAll('*').forEach(el => {
|
||||
for (const attr of el.getAttributeNames()) {
|
||||
if (attr.startsWith('on') || (attr === 'href' && el.getAttribute('href').toLowerCase().startsWith('javascript:')))
|
||||
el.removeAttribute(attr);
|
||||
}
|
||||
});
|
||||
|
||||
// Fix missing spaces in text nodes (skip code/pre)
|
||||
const walker = doc.createTreeWalker(root, 4, null, false);
|
||||
const textNodes = [];
|
||||
while (walker.nextNode()) textNodes.push(walker.currentNode);
|
||||
for (const node of textNodes) {
|
||||
let parent = node.parentNode;
|
||||
while (parent) {
|
||||
if (parent.tagName === 'CODE' || parent.tagName === 'PRE') break;
|
||||
parent = parent.parentNode;
|
||||
}
|
||||
if (parent) continue;
|
||||
let t = node.textContent;
|
||||
t = t.replace(/([a-z])([A-Z][a-z])/g, '$1 $2');
|
||||
t = t.replace(/([.:!?])([A-Za-z])/g, '$1 $2');
|
||||
node.textContent = t;
|
||||
}
|
||||
// Syntax highlight code blocks on the DOM (safe — textContent is auto-escaped by DOMParser)
|
||||
root.querySelectorAll('pre > code').forEach(el => {
|
||||
const raw = el.textContent;
|
||||
const lang = el.className.replace('language-', '');
|
||||
el.innerHTML = highlightCode(escapeHtml(raw), lang);
|
||||
});
|
||||
return root.innerHTML;
|
||||
}
|
||||
|
||||
// === Speech / TTS (Rick Voice) ===
|
||||
let currentUtterance = null;
|
||||
|
||||
function stopSpeech() {
|
||||
if (currentUtterance) {
|
||||
window.speechSynthesis.cancel();
|
||||
currentUtterance = null;
|
||||
}
|
||||
document.getElementById('stop-speech-btn').style.display = 'none';
|
||||
}
|
||||
|
||||
function speakText(text) {
|
||||
stopSpeech();
|
||||
const plain = text.replace(/<[^>]*>/g, '').replace(/```[\s\S]*?```/g, '').trim();
|
||||
if (!plain) return;
|
||||
const utterance = new SpeechSynthesisUtterance(plain);
|
||||
utterance.rate = 0.75;
|
||||
utterance.pitch = 0.3;
|
||||
utterance.volume = 1.0;
|
||||
const voices = window.speechSynthesis.getVoices();
|
||||
const en = voices.find(v => v.lang.startsWith('en'));
|
||||
if (en) utterance.voice = en;
|
||||
currentUtterance = utterance;
|
||||
utterance.onend = () => { currentUtterance = null; document.getElementById('stop-speech-btn').style.display = 'none'; };
|
||||
document.getElementById('stop-speech-btn').style.display = 'inline-block';
|
||||
window.speechSynthesis.speak(utterance);
|
||||
}
|
||||
|
||||
function speakMsg(btn) {
|
||||
const msgEl = btn.closest('.message');
|
||||
const content = msgEl.querySelector('.content');
|
||||
speakText(content.textContent || content.innerText);
|
||||
}
|
||||
|
||||
async function sendChat() {
|
||||
const input = document.getElementById('chat-input');
|
||||
const msg = input.value.trim();
|
||||
@@ -985,10 +1184,11 @@ async function sendChat() {
|
||||
const container = document.getElementById('chat-messages');
|
||||
const msgEl = document.createElement('div');
|
||||
msgEl.className = 'message assistant';
|
||||
msgEl.innerHTML = '<div class="role">assistant</div><div class="content"></div><div class="meta"></div>';
|
||||
msgEl.innerHTML = '<div class="role">assistant</div><div class="content"></div><div class="meta"></div><button class="speak-btn" onclick="speakMsg(this)" style="display:none;font-size:10px;padding:2px 8px;margin-top:4px;background:var(--surface2);border:1px solid var(--border);color:var(--accent);border-radius:4px;cursor:pointer;">🔊 Speak</button>';
|
||||
container.appendChild(msgEl);
|
||||
const contentEl = msgEl.querySelector('.content');
|
||||
const metaEl = msgEl.querySelector('.meta');
|
||||
const speakBtn = msgEl.querySelector('.speak-btn');
|
||||
|
||||
try {
|
||||
const res = await fetch(API + '/api/chat', {
|
||||
@@ -1031,10 +1231,15 @@ async function sendChat() {
|
||||
try {
|
||||
const m = JSON.parse(data);
|
||||
metaEl.textContent = m.duration_ms + 'ms | ' + m.knowledge_hits + ' knowledge hits';
|
||||
speakBtn.style.display = 'inline-block';
|
||||
if (document.getElementById('rick-voice').checked) {
|
||||
speakMsg(speakBtn);
|
||||
}
|
||||
} catch (_) {}
|
||||
} else {
|
||||
reply += data;
|
||||
contentEl.textContent = reply;
|
||||
contentEl.innerHTML = sanitizeHtml(reply);
|
||||
addCodeCopyButtons(contentEl);
|
||||
container.scrollTop = container.scrollHeight;
|
||||
}
|
||||
}
|
||||
@@ -1470,6 +1675,61 @@ function copyScript() {
|
||||
});
|
||||
}
|
||||
|
||||
function copyCodeBlock(btn) {
|
||||
const code = btn.parentElement.querySelector('code');
|
||||
if (!code) return;
|
||||
navigator.clipboard.writeText(code.textContent).then(() => {
|
||||
const orig = btn.textContent;
|
||||
btn.textContent = 'Copied!';
|
||||
setTimeout(() => btn.textContent = orig, 2000);
|
||||
});
|
||||
}
|
||||
|
||||
function addCodeCopyButtons(container) {
|
||||
container.querySelectorAll('pre').forEach(pre => {
|
||||
if (pre.querySelector('.copy-code-btn')) return;
|
||||
const btn = document.createElement('button');
|
||||
btn.className = 'copy-code-btn';
|
||||
btn.textContent = 'Copy';
|
||||
btn.onclick = () => copyCodeBlock(btn);
|
||||
pre.appendChild(btn);
|
||||
});
|
||||
}
|
||||
|
||||
function fixCodeNewlines(code) {
|
||||
let c = code;
|
||||
// After # comment followed immediately by $variable (most common gap)
|
||||
c = c.replace(/(#[^\n]*)(\$\w)/g, '$1\n$2');
|
||||
// After } followed by non-whitespace (e.g. }for, }while)
|
||||
c = c.replace(/}(\s*)(?![\s\n])(\S)/g, '}\n$1$2');
|
||||
// After ) followed by a capital-letter word or $var (likely next statement)
|
||||
c = c.replace(/\)(\s*)(\$\w|[A-Z]\w+)/g, ')\n$1$2');
|
||||
// Keyword concatenated with previous word (no word boundary — model joins them)
|
||||
c = c.replace(/(\w)(import|from|def|class|return|function|const|let|var|pub|fn|use|mod|while|for|if|with|try|except)\b/g, '$1\n$2');
|
||||
// Semicolon not already followed by newline
|
||||
c = c.replace(/;\s*(?!\n)/g, ';\n');
|
||||
// Collapse multiple blank lines
|
||||
c = c.replace(/\n{3,}/g, '\n\n');
|
||||
return c;
|
||||
}
|
||||
|
||||
function highlightCode(code, lang) {
|
||||
// Generic syntax highlighting — works for most languages
|
||||
// Comments (single-line, must be done first so they aren't processed further)
|
||||
code = code.replace(/(#[^\n]*)/g, '<span class="syn-comment">$1</span>');
|
||||
code = code.replace(/(\/\/[^\n]*)/g, '<span class="syn-comment">$1</span>');
|
||||
// Strings (double-quoted)
|
||||
code = code.replace(/"([^"\\]*(?:\\.[^"\\]*)*)"/g, '<span class="syn-string">"$1"</span>');
|
||||
// Strings (single-quoted)
|
||||
code = code.replace(/'([^'\\]*(?:\\.[^'\\]*)*)'/g, '<span class="syn-string">\'$1\'</span>');
|
||||
// Numbers
|
||||
code = code.replace(/\b(\d+\.?\d*)\b/g, '<span class="syn-number">$1</span>');
|
||||
// Keywords
|
||||
const keywords = /\b(for|foreach|while|if|else|elif|switch|case|break|continue|return|function|class|import|from|export|def|var|let|const|using|namespace|public|private|static|void|int|string|bool|float|double|null|true|false|try|catch|finally|throw|new|async|await|yield|in|do|done|then|end|print|echo|write)\b/g;
|
||||
code = code.replace(keywords, '<span class="syn-keyword">$1</span>');
|
||||
return code;
|
||||
}
|
||||
|
||||
async function runGeneratedScript() {
|
||||
if (!generatedCode) return;
|
||||
const output = document.getElementById('studio-output');
|
||||
|
||||
+113
-271
@@ -16,7 +16,6 @@ use serde::{Deserialize, Serialize};
|
||||
use std::convert::Infallible;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::time::{timeout, Duration};
|
||||
use tokio_stream::StreamExt;
|
||||
@@ -28,15 +27,9 @@ use crate::knowledge_index::{
|
||||
format_knowledge, EmbeddingCache, KnowledgeIndex, SearchProfile, OLLAMA_SEM,
|
||||
};
|
||||
use crate::training::websearch::{
|
||||
append_sample_to_jsonl, compute_sample_embedding, derive_sample_from_search, web_search,
|
||||
SearchEngine,
|
||||
append_sample_to_jsonl, compute_sample_embedding, web_search, SearchEngine,
|
||||
};
|
||||
|
||||
pub struct ModelCache {
|
||||
pub names: Vec<String>,
|
||||
pub fetched_at: Instant,
|
||||
}
|
||||
|
||||
const SYSTEM_PROMPT: &str = include_str!("../prompts/system.md");
|
||||
const SCRIPTGEN_PROMPT: &str = "CURRENT TOOL KNOWLEDGE:\n\
|
||||
{knowledge}\n\n\
|
||||
@@ -70,10 +63,40 @@ pub struct AppState {
|
||||
pub history: Arc<Mutex<Vec<ChatMessage>>>,
|
||||
pub target: Arc<Mutex<Option<String>>>,
|
||||
pub shutdown: Arc<tokio::sync::Notify>,
|
||||
pub model_cache: Arc<Mutex<ModelCache>>,
|
||||
pub embed_cache: Arc<Mutex<EmbeddingCache>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
/// Get embedding for text using the model, with LRU cache.
|
||||
/// Returns empty vec on failure.
|
||||
pub async fn embed_text(&self, text: &str) -> Vec<f32> {
|
||||
{
|
||||
let cache = self.embed_cache.lock().await;
|
||||
if let Some(e) = cache.get(text) {
|
||||
return e.clone();
|
||||
}
|
||||
}
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let req = GenerateEmbeddingsRequest::new(
|
||||
"dolphin-llama3:8b".to_string(),
|
||||
EmbeddingsInput::Single(text.to_string()),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
let emb = self
|
||||
.ollama
|
||||
.generate_embeddings(req)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|r| r.embeddings.into_iter().next())
|
||||
.unwrap_or_default();
|
||||
if !emb.is_empty() {
|
||||
let mut cache = self.embed_cache.lock().await;
|
||||
cache.insert(text, emb.clone());
|
||||
}
|
||||
emb
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct ChatRequest {
|
||||
pub message: String,
|
||||
@@ -107,6 +130,7 @@ pub struct ExecResponse {
|
||||
pub struct KnowledgeStats {
|
||||
pub total: usize,
|
||||
pub sources: Vec<(String, usize)>,
|
||||
pub embeddings_ready: bool,
|
||||
}
|
||||
|
||||
pub async fn start_server(addr: &str, state: AppState, open_browser: bool) -> anyhow::Result<()> {
|
||||
@@ -218,37 +242,8 @@ async fn handle_chat_sync(
|
||||
let start = std::time::Instant::now();
|
||||
let query = req.message.clone();
|
||||
|
||||
let q_emb = {
|
||||
let cache = state.embed_cache.lock().await;
|
||||
cache.get(&query).cloned()
|
||||
};
|
||||
let q_emb = match q_emb {
|
||||
Some(e) => Some(e),
|
||||
None => {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let emb_req = GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
EmbeddingsInput::Single(query.clone()),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
match timeout(
|
||||
Duration::from_secs(10),
|
||||
state.ollama.generate_embeddings(emb_req),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(r)) => {
|
||||
let e = r.embeddings.into_iter().next();
|
||||
if let Some(ref emb) = e {
|
||||
let mut cache = state.embed_cache.lock().await;
|
||||
cache.insert(&query, emb.clone());
|
||||
}
|
||||
e
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
};
|
||||
let q_emb = state.embed_text(&query).await;
|
||||
let q_emb = if q_emb.is_empty() { None } else { Some(q_emb) };
|
||||
|
||||
let profile = SearchProfile::from_query(&query);
|
||||
let (knowledge_text, hit_count) = match q_emb {
|
||||
@@ -281,11 +276,11 @@ async fn handle_chat_sync(
|
||||
let query_copy = query.clone();
|
||||
msgs.push(ChatMessage::user(query));
|
||||
|
||||
let chat_req = ChatMessageRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), msgs)
|
||||
let chat_req = ChatMessageRequest::new("dolphin-llama3:8b".to_string(), msgs)
|
||||
.options(
|
||||
ModelOptions::default()
|
||||
.num_predict(-1)
|
||||
.num_ctx(32768)
|
||||
.num_predict(1024)
|
||||
.num_ctx(8192)
|
||||
.repeat_penalty(1.1),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
@@ -347,40 +342,8 @@ async fn handle_chat(
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
|
||||
|
||||
// Embedding search — lock is NOT held during the network call
|
||||
let q_emb = {
|
||||
let cache = state.embed_cache.lock().await;
|
||||
cache.get(&query).cloned()
|
||||
};
|
||||
let q_emb = match q_emb {
|
||||
Some(e) => Some(e),
|
||||
None => {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let emb_req =
|
||||
ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
ollama_rs::generation::embeddings::request::EmbeddingsInput::Single(
|
||||
query.clone(),
|
||||
),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
match timeout(
|
||||
Duration::from_secs(10),
|
||||
state.ollama.generate_embeddings(emb_req),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(r)) => {
|
||||
let e = r.embeddings.into_iter().next();
|
||||
if let Some(ref emb) = e {
|
||||
let mut cache = state.embed_cache.lock().await;
|
||||
cache.insert(&query, emb.clone());
|
||||
}
|
||||
e
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
};
|
||||
let q_emb = state.embed_text(&query).await;
|
||||
let q_emb = if q_emb.is_empty() { None } else { Some(q_emb) };
|
||||
|
||||
let profile = SearchProfile::from_query(&query);
|
||||
let (knowledge_text, hit_count) = match q_emb {
|
||||
@@ -412,16 +375,16 @@ async fn handle_chat(
|
||||
let query_for_hist = query.clone();
|
||||
msgs.push(ChatMessage::user(query));
|
||||
|
||||
let chat_req = ChatMessageRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), msgs)
|
||||
let chat_req = ChatMessageRequest::new("dolphin-llama3:8b".to_string(), msgs)
|
||||
.options(
|
||||
ModelOptions::default()
|
||||
.num_ctx(32768)
|
||||
.num_ctx(8192)
|
||||
.num_thread(8)
|
||||
.num_predict(-1)
|
||||
.temperature(0.85)
|
||||
.num_predict(1024)
|
||||
.temperature(0.7)
|
||||
.repeat_penalty(1.05)
|
||||
.top_k(40)
|
||||
.top_p(0.95),
|
||||
.top_p(0.9),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
|
||||
@@ -496,21 +459,6 @@ async fn handle_script(Json(req): Json<ScriptRequest>) -> Result<Json<ExecRespon
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_cached_models(state: &AppState) -> Vec<String> {
|
||||
let mut cache = state.model_cache.lock().await;
|
||||
if cache.fetched_at.elapsed() < Duration::from_secs(60) && !cache.names.is_empty() {
|
||||
return cache.names.clone();
|
||||
}
|
||||
match state.ollama.list_local_models().await {
|
||||
Ok(models) => {
|
||||
cache.names = models.iter().map(|m| m.name.clone()).collect();
|
||||
cache.fetched_at = Instant::now();
|
||||
cache.names.clone()
|
||||
}
|
||||
Err(_) => cache.names.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_detect_lang(Json(req): Json<DetectLangRequest>) -> Json<serde_json::Value> {
|
||||
let lang = detect_language(&req.code);
|
||||
Json(serde_json::json!({
|
||||
@@ -548,33 +496,7 @@ async fn handle_add_knowledge(
|
||||
let output = req.get("output").and_then(|v| v.as_str()).unwrap_or("");
|
||||
|
||||
let know_text = format!("{} {}", instruction, output);
|
||||
let emb = {
|
||||
let cache = state.embed_cache.lock().await;
|
||||
cache.get(&know_text).cloned()
|
||||
};
|
||||
let emb = match emb {
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let emb_req =
|
||||
ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
ollama_rs::generation::embeddings::request::EmbeddingsInput::Single(
|
||||
know_text.clone(),
|
||||
),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
let e = match state.ollama.generate_embeddings(emb_req).await {
|
||||
Ok(r) => r.embeddings.into_iter().next().unwrap_or_default(),
|
||||
Err(_) => vec![],
|
||||
};
|
||||
if !e.is_empty() {
|
||||
let mut cache = state.embed_cache.lock().await;
|
||||
cache.insert(&know_text, e.clone());
|
||||
}
|
||||
e
|
||||
}
|
||||
};
|
||||
let emb = state.embed_text(&know_text).await;
|
||||
|
||||
let mut knowledge = state.knowledge.lock().await;
|
||||
let dup = if emb.is_empty() {
|
||||
@@ -600,35 +522,7 @@ async fn handle_search(
|
||||
) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
let query = req.get("query").and_then(|v| v.as_str()).unwrap_or("");
|
||||
|
||||
let q_emb = {
|
||||
let cache = state.embed_cache.lock().await;
|
||||
cache.get(query).cloned()
|
||||
};
|
||||
let q_emb = match q_emb {
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let emb_req =
|
||||
ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
ollama_rs::generation::embeddings::request::EmbeddingsInput::Single(
|
||||
query.to_string(),
|
||||
),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
match state.ollama.generate_embeddings(emb_req).await {
|
||||
Ok(r) => {
|
||||
let e = r.embeddings.into_iter().next().unwrap_or_default();
|
||||
if !e.is_empty() {
|
||||
let mut cache = state.embed_cache.lock().await;
|
||||
cache.insert(query, e.clone());
|
||||
}
|
||||
e
|
||||
}
|
||||
Err(_) => vec![],
|
||||
}
|
||||
}
|
||||
};
|
||||
let q_emb = state.embed_text(query).await;
|
||||
|
||||
let profile = SearchProfile::from_query(query);
|
||||
let results: Vec<_> = if q_emb.is_empty() {
|
||||
@@ -659,6 +553,7 @@ async fn handle_knowledge_stats(State(state): State<AppState>) -> Json<Knowledge
|
||||
Json(KnowledgeStats {
|
||||
total: knowledge.len(),
|
||||
sources,
|
||||
embeddings_ready: knowledge.embeddings_ready(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -671,7 +566,9 @@ async fn handle_learn_search(
|
||||
_ => return Err(StatusCode::BAD_REQUEST),
|
||||
};
|
||||
|
||||
let snippets = match web_search(&query, SearchEngine::DuckDuckGo).await {
|
||||
let model_name = "dolphin-llama3:8b";
|
||||
|
||||
let answer = match web_search(&query, SearchEngine::AiSearch, Some(&state.ollama), Some(model_name)).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
return Ok(Json(
|
||||
@@ -680,23 +577,12 @@ async fn handle_learn_search(
|
||||
}
|
||||
};
|
||||
|
||||
let model_name = {
|
||||
let names = get_cached_models(&state).await;
|
||||
if names.contains(&"qwen2.5-coder:1.5b".to_string()) {
|
||||
"qwen2.5-coder:1.5b"
|
||||
} else {
|
||||
"sadiq-bd/llama3.2-3b-uncensored"
|
||||
}
|
||||
};
|
||||
|
||||
let sample = match derive_sample_from_search(&query, &snippets, &state.ollama, model_name).await
|
||||
{
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
return Ok(Json(
|
||||
serde_json::json!({"success": false, "error": format!("Derive failed: {}", e)}),
|
||||
))
|
||||
}
|
||||
let sample = EnglishSample {
|
||||
r#type: "instruction_following".into(),
|
||||
instruction: query.clone(),
|
||||
input: String::new(),
|
||||
reasoning: format!("AI answer based on known knowledge about {}", query),
|
||||
output: answer.trim().to_string(),
|
||||
};
|
||||
|
||||
if let Err(e) = append_sample_to_jsonl(&sample, "output/random_new_knowledge.jsonl") {
|
||||
@@ -715,7 +601,7 @@ async fn handle_learn_search(
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let e = compute_sample_embedding(&sample, "nomic-embed-text", &state.ollama)
|
||||
let e = compute_sample_embedding(&sample, "dolphin-llama3:8b", &state.ollama)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
if !e.is_empty() {
|
||||
@@ -756,35 +642,27 @@ async fn handle_learn_auto(
|
||||
"cloud security best practices 2026",
|
||||
];
|
||||
|
||||
let names = get_cached_models(&state).await;
|
||||
let model_name = if names.contains(&"qwen2.5-coder:1.5b".to_string()) {
|
||||
"qwen2.5-coder:1.5b"
|
||||
} else {
|
||||
"sadiq-bd/llama3.2-3b-uncensored"
|
||||
};
|
||||
let model_name = "dolphin-llama3:8b";
|
||||
|
||||
let tasks: Vec<_> = topics
|
||||
.into_iter()
|
||||
.map(|topic| {
|
||||
let state = state.clone();
|
||||
async move {
|
||||
let snippets = match web_search(topic, SearchEngine::DuckDuckGo).await {
|
||||
let answer = match web_search(topic, SearchEngine::AiSearch, Some(&state.ollama), Some(model_name)).await {
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
return serde_json::json!({"topic": topic, "error": "Web search failed"})
|
||||
}
|
||||
};
|
||||
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let sample =
|
||||
match derive_sample_from_search(topic, &snippets, &state.ollama, model_name)
|
||||
.await
|
||||
{
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
return serde_json::json!({"topic": topic, "error": e.to_string()})
|
||||
}
|
||||
};
|
||||
let sample = EnglishSample {
|
||||
r#type: "instruction_following".into(),
|
||||
instruction: topic.to_string(),
|
||||
input: String::new(),
|
||||
reasoning: format!("AI answer about {}", topic),
|
||||
output: answer.trim().to_string(),
|
||||
};
|
||||
|
||||
if let Err(e) = append_sample_to_jsonl(&sample, "output/random_new_knowledge.jsonl")
|
||||
{
|
||||
@@ -803,7 +681,7 @@ async fn handle_learn_auto(
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let e =
|
||||
compute_sample_embedding(&sample, "nomic-embed-text", &state.ollama)
|
||||
compute_sample_embedding(&sample, "dolphin-llama3:8b", &state.ollama)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
if !e.is_empty() {
|
||||
@@ -874,40 +752,7 @@ async fn handle_scriptgen(
|
||||
return Json(serde_json::json!({"success": false, "error": "No description provided."}));
|
||||
}
|
||||
|
||||
let q_emb = {
|
||||
let cache = state.embed_cache.lock().await;
|
||||
cache.get(&description).cloned()
|
||||
};
|
||||
let q_emb = match q_emb {
|
||||
Some(e) => e,
|
||||
None => {
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let emb_req =
|
||||
ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
ollama_rs::generation::embeddings::request::EmbeddingsInput::Single(
|
||||
description.clone(),
|
||||
),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
match timeout(
|
||||
Duration::from_secs(10),
|
||||
state.ollama.generate_embeddings(emb_req),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(r)) => {
|
||||
let e = r.embeddings.into_iter().next().unwrap_or_default();
|
||||
if !e.is_empty() {
|
||||
let mut cache = state.embed_cache.lock().await;
|
||||
cache.insert(&description, e.clone());
|
||||
}
|
||||
e
|
||||
}
|
||||
_ => vec![],
|
||||
}
|
||||
}
|
||||
};
|
||||
let q_emb = state.embed_text(&description).await;
|
||||
|
||||
let tool_knowledge = if q_emb.is_empty() {
|
||||
"Cybersecurity tools knowledge loaded.".to_string()
|
||||
@@ -934,19 +779,19 @@ async fn handle_scriptgen(
|
||||
ChatMessage::user(description.clone()),
|
||||
];
|
||||
|
||||
let chat_req = ChatMessageRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), msgs)
|
||||
let chat_req = ChatMessageRequest::new("dolphin-llama3:8b".to_string(), msgs)
|
||||
.options(
|
||||
ModelOptions::default()
|
||||
.num_ctx(32768)
|
||||
.num_ctx(4096)
|
||||
.num_thread(8)
|
||||
.num_predict(-1)
|
||||
.temperature(0.85)
|
||||
.num_predict(2048)
|
||||
.temperature(0.7)
|
||||
.repeat_penalty(1.05),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
match timeout(
|
||||
Duration::from_secs(600),
|
||||
Duration::from_secs(300),
|
||||
state.ollama.send_chat_messages(chat_req),
|
||||
)
|
||||
.await
|
||||
@@ -1018,17 +863,8 @@ async fn handle_pdf_import(
|
||||
let out_path = Path::new("output").join(&out_name);
|
||||
std::fs::create_dir_all("output").map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
|
||||
let names = get_cached_models(&state).await;
|
||||
let model = if names.contains(&"qwen2.5-coder:1.5b".to_string()) {
|
||||
"qwen2.5-coder:1.5b"
|
||||
} else if names.contains(&"sadiq-bd/llama3.2-3b-uncensored".to_string()) {
|
||||
"sadiq-bd/llama3.2-3b-uncensored"
|
||||
} else {
|
||||
"sadiq-bd/llama3.2-3b-uncensored"
|
||||
};
|
||||
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let done = match import_pdf_path(&temp_path, &out_path, &state.ollama, model, false).await {
|
||||
let done = match import_pdf_path(&temp_path, &out_path, &state.ollama, "dolphin-llama3:8b", false).await {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&temp_path);
|
||||
@@ -1041,12 +877,27 @@ async fn handle_pdf_import(
|
||||
|
||||
let _ = std::fs::remove_file(&temp_path);
|
||||
|
||||
let book_name = filename
|
||||
.rsplit_once('.')
|
||||
.map(|(name, _)| name.to_string())
|
||||
.unwrap_or_else(|| filename.clone());
|
||||
|
||||
let mut knowledge = state.knowledge.lock().await;
|
||||
let source_name = format!("pdf-import:{}", filename);
|
||||
if let Ok(samples) = crate::data_ls::loader::load_english(out_path.to_str().unwrap()) {
|
||||
let clean = crate::pipeline::preprocess::remove_empty(samples);
|
||||
let tagged: Vec<_> = clean
|
||||
.into_iter()
|
||||
.map(|mut s| {
|
||||
s.instruction = format!("[{}] {}", book_name, s.instruction);
|
||||
s
|
||||
})
|
||||
.collect();
|
||||
for s in &tagged {
|
||||
let _ = append_sample_to_jsonl(s, "output/pdf_knowledge.jsonl");
|
||||
}
|
||||
let existing = knowledge.len();
|
||||
knowledge.extend_from_samples(clean, &source_name);
|
||||
knowledge.extend_from_samples(tagged, &source_name);
|
||||
if knowledge.len() > existing {
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
eprintln!("save_cache failed: {e}");
|
||||
@@ -1186,22 +1037,7 @@ async fn handle_epub_import(
|
||||
Err(e) => return Ok(Json(serde_json::json!({"success": false, "error": e}))),
|
||||
};
|
||||
|
||||
let chunk_size = 5000;
|
||||
let chunks: Vec<String> = {
|
||||
let mut chunks = Vec::new();
|
||||
let mut current = String::with_capacity(chunk_size);
|
||||
for c in text.chars() {
|
||||
current.push(c);
|
||||
if current.len() >= chunk_size {
|
||||
chunks.push(current);
|
||||
current = String::with_capacity(chunk_size);
|
||||
}
|
||||
}
|
||||
if !current.is_empty() {
|
||||
chunks.push(current);
|
||||
}
|
||||
chunks
|
||||
};
|
||||
let chunks = crate::pdf_import::split_text_into_chunks(&text, 10000);
|
||||
|
||||
let out_name = filename.replace(".epub", ".jsonl");
|
||||
let out_path = std::path::Path::new("output").join(&out_name);
|
||||
@@ -1209,19 +1045,12 @@ async fn handle_epub_import(
|
||||
let mut out_file =
|
||||
std::fs::File::create(&out_path).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
|
||||
let names = get_cached_models(&state).await;
|
||||
let model = if names.contains(&"qwen2.5-coder:1.5b".to_string()) {
|
||||
"qwen2.5-coder:1.5b"
|
||||
} else {
|
||||
"sadiq-bd/llama3.2-3b-uncensored"
|
||||
};
|
||||
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let samples = match crate::training::websearch::derive_samples_from_pdf_text(
|
||||
std::path::Path::new(&filename),
|
||||
chunks,
|
||||
&state.ollama,
|
||||
model,
|
||||
"dolphin-llama3:8b",
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -1234,11 +1063,24 @@ async fn handle_epub_import(
|
||||
};
|
||||
drop(_permit);
|
||||
|
||||
let book_name = filename
|
||||
.rsplit_once('.')
|
||||
.map(|(name, _)| name.to_string())
|
||||
.unwrap_or_else(|| filename.clone());
|
||||
|
||||
let tagged: Vec<_> = samples
|
||||
.into_iter()
|
||||
.map(|mut s| {
|
||||
s.instruction = format!("[{}] {}", book_name, s.instruction);
|
||||
s
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut count = 0;
|
||||
for s in &samples {
|
||||
for s in &tagged {
|
||||
let line = serde_json::to_string(s).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
writeln!(out_file, "{}", line).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
if let Err(e) = append_sample_to_jsonl(s, "output/random_new_knowledge.jsonl") {
|
||||
if let Err(e) = append_sample_to_jsonl(s, "output/pdf_knowledge.jsonl") {
|
||||
eprintln!("append_sample failed: {e}");
|
||||
}
|
||||
count += 1;
|
||||
@@ -1246,7 +1088,7 @@ async fn handle_epub_import(
|
||||
|
||||
let mut knowledge = state.knowledge.lock().await;
|
||||
knowledge.extend_from_samples(
|
||||
crate::pipeline::preprocess::remove_empty(samples),
|
||||
crate::pipeline::preprocess::remove_empty(tagged),
|
||||
&format!("epub:{}", filename),
|
||||
);
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
@@ -1326,14 +1168,14 @@ async fn handle_dataset_generate(
|
||||
);
|
||||
|
||||
let req_msg = ChatMessageRequest::new(
|
||||
"sadiq-bd/llama3.2-3b-uncensored".to_string(),
|
||||
"dolphin-llama3:8b".to_string(),
|
||||
vec![ChatMessage::user(prompt)],
|
||||
)
|
||||
.options(
|
||||
ModelOptions::default()
|
||||
.num_ctx(32768)
|
||||
.num_ctx(4096)
|
||||
.num_thread(8)
|
||||
.num_predict(-1)
|
||||
.num_predict(1024)
|
||||
.temperature(0.7),
|
||||
)
|
||||
.keep_alive(KeepAlive::Indefinitely);
|
||||
@@ -1409,7 +1251,7 @@ async fn handle_dataset_generate(
|
||||
let _permit = OLLAMA_SEM.acquire().await.unwrap();
|
||||
let emb_req =
|
||||
ollama_rs::generation::embeddings::request::GenerateEmbeddingsRequest::new(
|
||||
"nomic-embed-text".to_string(),
|
||||
"dolphin-llama3:8b".to_string(),
|
||||
ollama_rs::generation::embeddings::request::EmbeddingsInput::Single(
|
||||
text.clone(),
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user