File size: 1,768 Bytes
b9a59c9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import sqlite3
import os

DB_PATH = os.path.join(os.path.dirname(__file__), "papers.db")

def get_db_connection():
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
        CREATE VIRTUAL TABLE IF NOT EXISTS papers USING fts5(
            title,
            abstract,
            author,
            year UNINDEXED,
            url UNINDEXED,
            venue UNINDEXED
        )
    ''')
    conn.commit()
    conn.close()

def search_papers(keywords: list, sort_by: str = "relevance"):
    conn = get_db_connection()
    cursor = conn.cursor()
    
    clean_keywords = []
    for kw in keywords:
        # Sanitize keyword
        kw = kw.replace('"', '').replace("'", "")
        # Handle -, +, &, etc. by keeping them inside quotes for FTS5
        clean_keywords.append(f'"{kw}"')
        
    match_query = " OR ".join(clean_keywords)
    fts_query = f"title : ({match_query})"
    
    if sort_by == "year":
        query = '''
            SELECT title, abstract, author, year, url, venue, rank 
            FROM papers 
            WHERE papers MATCH ? 
            ORDER BY year DESC
            LIMIT 50
        '''
    else:
        query = '''
            SELECT title, abstract, author, year, url, venue, rank 
            FROM papers 
            WHERE papers MATCH ? 
            ORDER BY rank
            LIMIT 50
        '''
        
    try:
        cursor.execute(query, (fts_query,))
        results = cursor.fetchall()
        return [dict(row) for row in results]
    except sqlite3.Error as e:
        print(f"Database error: {e}")
        return []
    finally:
        conn.close()