app/database.py

import functools
import io
import sqlite3

import numpy as np


class ItemDatabase:
    def __init__(self, db_path):
        sqlite3.register_converter("embedding", self._text_to_numpy_array)
        self._db_path = db_path

    @staticmethod
    def _text_to_numpy_array(text):
        out = io.BytesIO(text)
        out.seek(0)
        return np.load(out)

    def _connect(self):
        return sqlite3.connect(
            self._db_path, detect_types=sqlite3.PARSE_DECLTYPES)

    def search_items(self, query, n_items=10):
        with self._connect() as conn:
            c = conn.cursor()
            c.execute(f"select item_id from items where title like '%{query}%'")
            rows = c.fetchall()[:n_items]
            return [row[0] for row in rows]

    @functools.lru_cache(maxsize=2**14)
    def get_item(self, item_id):
        with self._connect() as conn:
            c = conn.cursor()
            c.row_factory = sqlite3.Row
            c.execute(f"select * from items where item_id like '{item_id}'")
            return c.fetchone()