sql_trainer/
lib.rs

1use std::result::Result;
2
3use std::path::Path;
4
5use rusqlite::Connection;
6use rusqlite::types::ValueRef;
7
8use tabled::{
9    builder::Builder,
10    settings::{Alignment, Style, object::Columns},
11};
12
13use ollama_rs::generation::chat::ChatMessage;
14use ollama_rs::generation::chat::request::ChatMessageRequest;
15use ollama_rs::{Ollama, error::OllamaError};
16
17use regex::Regex;
18
19pub fn file_exists(filename: &str) -> bool {
20    let path = Path::new(filename);
21
22    path.exists()
23}
24
25pub fn execute_user_select(conn: &Connection, query: &str) -> Result<(), rusqlite::Error> {
26    let mut stmt = conn.prepare(query)?;
27    let col_count = stmt.column_count();
28
29    let col_names: Vec<String> = stmt
30        .column_names()
31        .iter()
32        .map(|s| String::from(*s))
33        .collect::<Vec<_>>();
34
35    let mut rows = stmt.query([])?;
36
37    let mut builder = Builder::default();
38    builder.push_record(col_names.iter().cloned());
39
40    while let Some(row) = rows.next()? {
41        let mut cells = Vec::new();
42        for i in 0..col_count {
43            let value_ref = row.get_ref(i)?;
44
45            let display_value = match value_ref {
46                ValueRef::Null => "NULL".to_string(),
47                ValueRef::Integer(i) => i.to_string(),
48                ValueRef::Real(f) => f.to_string(),
49                ValueRef::Text(t) => String::from_utf8_lossy(t).to_string(),
50                ValueRef::Blob(b) => format!("[BLOB, len={}]", b.len()),
51            };
52            cells.push(display_value);
53        }
54        builder.push_record(cells);
55    }
56
57    let mut table = builder.build();
58    table.with(Style::modern());
59    table.modify(Columns::first(), Alignment::right());
60    println!("{table}");
61
62    Ok(())
63}
64
65pub fn create_sqlite_file(
66    filename: &str,
67    sql_statements: &[&str],
68) -> Result<(), Box<dyn std::error::Error>> {
69    let conn = Connection::open(filename)?;
70
71    for i in sql_statements {
72        conn.execute(i, [])?;
73    }
74
75    Ok(())
76}
77
78async fn send_message_to_ollama(message: &str) -> Result<String, ollama_rs::error::OllamaError> {
79    let mut ollama = Ollama::default();
80    let model = "qwen3:latest".to_string();
81    let mut history = vec![];
82
83    let res = ollama
84        .send_chat_messages_with_history(
85            &mut history,
86            ChatMessageRequest::new(model, vec![ChatMessage::user(String::from(message))]),
87        )
88        .await;
89
90    Ok(res.unwrap().message.content)
91}
92
93fn format_output(output: &str) -> String {
94    let re = Regex::new(r"(?s)<think>.*?</think>").unwrap();
95    re.replace_all(output, "").to_string()
96}
97
98pub async fn generate_db(
99    topic: &str,
100    conn: &Connection,
101) -> Result<(), ollama_rs::error::OllamaError> {
102    let prompt = format!(
103        "Create sql instructions for a sqlite database about {topic}. Please ONLY output sql Instructions. Nothing else. The output gets directly converted into a .sqlite file. Be creative with the topic and create sample data, not just tables. Also be sure to create more lots of diverse entries in the tables so they aren't that empty, but dont repeat yourself. Make sure the syntax is correct and safe and do not use strings in actual data, as this causes parsing problems."
104    );
105    let mut message = send_message_to_ollama(&prompt).await?;
106    message = format_output(&message).to_string();
107
108    println!("{:?}", message);
109
110    conn.execute_batch(&message)
111        .expect("Failed to execute insert queries");
112
113    Ok(())
114}
115
116pub async fn ask_sql_question(message: &str) -> Result<String, OllamaError> {
117    let prompt = String::from(
118        "Please answer the following question as an sql-expert in sqlite. Be concise, short and only mention relevant things. Message: ",
119    ) + message;
120    let output = send_message_to_ollama(&prompt).await?;
121    Ok(format_output(&output))
122}
sql_trainer/lib.rs

sql_trainer/
lib.rs
