Skip to content

Commit

Permalink
it
Browse files Browse the repository at this point in the history
  • Loading branch information
louis030195 committed Feb 25, 2025
1 parent 5305ec9 commit 9c6e403
Show file tree
Hide file tree
Showing 4 changed files with 162 additions and 66 deletions.
113 changes: 64 additions & 49 deletions screenpipe-server/src/db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -664,7 +664,7 @@ impl DatabaseManager {

// Split query parts between frame metadata and OCR content
if !query.is_empty() {
ocr_fts_parts.push(query.to_owned()); // Just use the query directly
ocr_fts_parts.push(query); // Just use the query directly
}
if let Some(app) = app_name {
if !app.is_empty() {
Expand Down Expand Up @@ -695,9 +695,9 @@ impl DatabaseManager {

let where_clause = match (frame_query.is_empty(), ocr_query.is_empty()) {
(true, true) => "WHERE 1=1".to_owned(),
(false, true) => "WHERE frames_fts MATCH ?".to_owned(),
(true, false) => "WHERE ocr_text_fts MATCH ?".to_owned(),
(false, false) => "WHERE frames_fts MATCH ? AND ocr_text_fts MATCH ?".to_owned(),
(false, true) => "WHERE frames_fts MATCH ?1".to_owned(),
(true, false) => "WHERE ocr_text_fts MATCH ?1".to_owned(),
(false, false) => "WHERE frames_fts MATCH ?1 AND ocr_text_fts MATCH ?2".to_owned(),
};

let sql = format!(
Expand All @@ -721,29 +721,46 @@ impl DatabaseManager {
LEFT JOIN vision_tags ON frames.id = vision_tags.vision_id
LEFT JOIN tags ON vision_tags.tag_id = tags.id
{}
AND (?2 IS NULL OR frames.timestamp >= ?2)
AND (?3 IS NULL OR frames.timestamp <= ?3)
AND (?4 IS NULL OR COALESCE(ocr_text.text_length, LENGTH(ocr_text.text)) >= ?4)
AND (?5 IS NULL OR COALESCE(ocr_text.text_length, LENGTH(ocr_text.text)) <= ?5)
AND (?6 IS NULL OR frames.name LIKE '%' || ?6 || '%' COLLATE NOCASE)
AND (?3 IS NULL OR frames.timestamp >= ?3)
AND (?4 IS NULL OR frames.timestamp <= ?4)
AND (?5 IS NULL OR COALESCE(ocr_text.text_length, LENGTH(ocr_text.text)) >= ?5)
AND (?6 IS NULL OR COALESCE(ocr_text.text_length, LENGTH(ocr_text.text)) <= ?6)
AND (?7 IS NULL OR frames.name LIKE '%' || ?7 || '%' COLLATE NOCASE)
GROUP BY ocr_text.frame_id
ORDER BY frames.timestamp DESC
LIMIT ?7 OFFSET ?8
LIMIT ?8 OFFSET ?9
"#,
base_sql, where_clause
);

let raw_results: Vec<OCRResultRaw> = sqlx::query_as(&sql)
.bind(if frame_query.is_empty() && ocr_query.is_empty() {
"*".to_owned()
} else {
// Use ocr_query if frame_query is empty
if frame_query.is_empty() {
ocr_query
} else {
frame_query
}
})
println!("sql: {}", sql);
println!("frame_query: {}", frame_query);
println!("ocr_query: {}", ocr_query);
println!("start_time: {:?}", start_time);
println!("end_time: {:?}", end_time);
println!("min_length: {:?}", min_length);
println!("max_length: {:?}", max_length);
println!("frame_name: {:?}", frame_name);
println!("where_clause: {:?}", where_clause);
let mut query_builder = sqlx::query_as(&sql);

// Properly bind parameters based on which queries are present
match (frame_query.is_empty(), ocr_query.is_empty()) {
(true, true) => {
query_builder = query_builder.bind("*".to_owned());
}
(false, true) => {
query_builder = query_builder.bind(&frame_query);
}
(true, false) => {
query_builder = query_builder.bind(&ocr_query);
}
(false, false) => {
query_builder = query_builder.bind(&frame_query).bind(&ocr_query);
}
}

let raw_results: Vec<OCRResultRaw> = query_builder
.bind(start_time)
.bind(end_time)
.bind(min_length.map(|l| l as i64))
Expand All @@ -754,6 +771,8 @@ impl DatabaseManager {
.fetch_all(&self.pool)
.await?;

println!("raw_results: {:?}", raw_results);

Ok(raw_results
.into_iter()
.map(|raw| OCRResult {
Expand Down Expand Up @@ -1163,18 +1182,6 @@ impl DatabaseManager {
.fetch_one(&self.pool)
.await?
}
ContentType::All => {
sqlx::query_scalar(&sql)
.bind(ocr_query)
.bind(start_time)
.bind(end_time)
.bind(min_length.map(|l| l as i64))
.bind(max_length.map(|l| l as i64))
.bind(frame_name)
.bind(json_array)
.fetch_one(&self.pool)
.await?
}
_ => {
sqlx::query_scalar(&sql)
.bind(query)
Expand Down Expand Up @@ -1453,13 +1460,13 @@ impl DatabaseManager {
f.timestamp,
f.offset_index,
ot.text,
ot.app_name,
ot.window_name,
f.app_name,
f.window_name,
vc.device_name as screen_device,
vc.file_path as video_path,
strftime('%Y-%m-%d %H:%M', f.timestamp) as minute_group,
ROW_NUMBER() OVER (
PARTITION BY strftime('%Y-%m-%d %H:%M', f.timestamp), ot.app_name, vc.device_name
PARTITION BY strftime('%Y-%m-%d %H:%M', f.timestamp), f.app_name, vc.device_name
ORDER BY f.timestamp DESC
) as rn
FROM frames f
Expand Down Expand Up @@ -1578,10 +1585,10 @@ impl DatabaseManager {
fts_parts.push(query.to_owned());
}
if let Some(app) = app_name {
fts_parts.push(format!("app:\"{}\"", app));
fts_parts.push(format!("app:{}", app));
}
if let Some(window) = window_name {
fts_parts.push(format!("window:\"{}\"", window));
fts_parts.push(format!("window:{}", window));
}
let combined_query = fts_parts.join(" ");

Expand All @@ -1603,13 +1610,13 @@ impl DatabaseManager {
ui_monitoring.id,
ui_monitoring.text_output,
ui_monitoring.timestamp,
ui_monitoring.app,
ui_monitoring.window,
ui_monitoring.app as app_name,
ui_monitoring.window as window_name,
ui_monitoring.initial_traversal_at,
video_chunks.file_path,
frames.offset_index,
frames.browser_url,
frames.name as frame_name
frames.name as frame_name,
frames.browser_url
FROM {}
LEFT JOIN frames ON
frames.timestamp BETWEEN
Expand All @@ -1619,15 +1626,23 @@ impl DatabaseManager {
{}
AND (?2 IS NULL OR ui_monitoring.timestamp >= ?2)
AND (?3 IS NULL OR ui_monitoring.timestamp <= ?3)
GROUP BY ui_monitoring.id
ORDER BY ui_monitoring.timestamp DESC
LIMIT ?4 OFFSET ?5
"#,
base_sql, where_clause
);

println!("sql: {}", sql);
println!("combined_query: {}", combined_query);
println!("start_time: {:?}", start_time);
println!("end_time: {:?}", end_time);
println!("limit: {}", limit);
println!("offset: {}", offset);

sqlx::query_as(&sql)
.bind(if combined_query.is_empty() {
"".to_owned()
"*".to_owned()
} else {
combined_query
})
Expand Down Expand Up @@ -2046,9 +2061,9 @@ impl DatabaseManager {
video_chunks.file_path,
frames.offset_index,
frames.name as frame_name,
ocr_text.app_name,
frames.app_name,
ocr_text.ocr_engine,
ocr_text.window_name,
frames.window_name,
GROUP_CONCAT(tags.name, ',') as tags,
frames.browser_url
FROM embedding_matches
Expand Down Expand Up @@ -2226,11 +2241,11 @@ impl DatabaseManager {
conditions.push("f.timestamp <= ?");
}

// Add app names condition if provided
// Add app names condition if provided - updated to use f.app_name
if let Some(apps) = &app_names {
if !apps.is_empty() {
let placeholders = vec!["?"; apps.len()].join(",");
let app_condition = format!("o.app_name IN ({})", placeholders);
let app_condition = format!("f.app_name IN ({})", placeholders);
owned_conditions.push(app_condition);
conditions.push(owned_conditions.last().unwrap().as_str());
}
Expand Down Expand Up @@ -2261,8 +2276,8 @@ SELECT
f.id,
f.timestamp,
f.browser_url as url,
o.app_name,
o.window_name,
f.app_name,
f.window_name,
o.text as ocr_text,
o.text_json
FROM frames f
Expand Down
12 changes: 10 additions & 2 deletions screenpipe-server/src/db_types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,14 @@ pub enum SearchResult {
UI(UiContent),
}

#[derive(FromRow, Debug)]
pub struct Frame {
pub id: i64,
pub timestamp: DateTime<Utc>,
pub browser_url: String,
pub app_name: String,
pub window_name: String,
}
#[derive(FromRow, Debug)]
pub struct OCRResultRaw {
pub frame_id: i64,
Expand Down Expand Up @@ -128,9 +136,9 @@ pub struct UiContent {
#[sqlx(rename = "text_output")]
pub text: String,
pub timestamp: DateTime<Utc>,
#[sqlx(rename = "app")]
#[sqlx(rename = "app_name")]
pub app_name: String,
#[sqlx(rename = "window")]
#[sqlx(rename = "window_name")]
pub window_name: String,
pub initial_traversal_at: Option<DateTime<Utc>>,
pub file_path: String,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,8 @@ CREATE VIRTUAL TABLE IF NOT EXISTS frames_fts USING fts5(
);

-- Create triggers for frames
CREATE TRIGGER IF NOT EXISTS frames_ai AFTER INSERT ON frames
WHEN (NEW.browser_url IS NOT NULL AND NEW.browser_url != '')
OR (NEW.app_name IS NOT NULL AND NEW.app_name != '')
OR (NEW.window_name IS NOT NULL AND NEW.window_name != '')
OR (NEW.focused IS NOT NULL)
BEGIN
INSERT OR IGNORE INTO frames_fts(rowid, browser_url, app_name, window_name, focused)
CREATE TRIGGER IF NOT EXISTS frames_ai AFTER INSERT ON frames BEGIN
INSERT INTO frames_fts(rowid, browser_url, app_name, window_name, focused)
VALUES (
NEW.id,
COALESCE(NEW.browser_url, ''),
Expand All @@ -40,12 +35,14 @@ WHEN (NEW.browser_url IS NOT NULL AND NEW.browser_url != '')
OR (NEW.window_name IS NOT NULL AND NEW.window_name != '')
OR (NEW.focused IS NOT NULL)
BEGIN
UPDATE frames_fts
SET browser_url = COALESCE(NEW.browser_url, ''),
app_name = COALESCE(NEW.app_name, ''),
window_name = COALESCE(NEW.window_name, ''),
focused = COALESCE(NEW.focused, 0)
WHERE rowid = OLD.id;
INSERT OR REPLACE INTO frames_fts(rowid, browser_url, app_name, window_name, focused)
VALUES (
NEW.id,
COALESCE(NEW.browser_url, ''),
COALESCE(NEW.app_name, ''),
COALESCE(NEW.window_name, ''),
COALESCE(NEW.focused, 0)
);
END;

CREATE TRIGGER IF NOT EXISTS frames_ad AFTER DELETE ON frames
Expand Down
Loading

0 comments on commit 9c6e403

Please sign in to comment.