You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
138 lines
5.0 KiB
Python
138 lines
5.0 KiB
Python
from __future__ import annotations
|
|
import hashlib
|
|
import asyncpg
|
|
import numpy as np
|
|
from pgvector.sqlalchemy import Vector
|
|
from pgvector.asyncpg import register_vector
|
|
import sqlalchemy
|
|
from sqlalchemy.orm import mapped_column, relationship, Mapped, deferred, defer
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
import config
|
|
from api.model.base import BaseHelper, BaseModel
|
|
from service.database import DatabaseService
|
|
|
|
class TitleIndexModel(BaseModel):
|
|
__tablename__ = "embedding_search_title_index"
|
|
|
|
id: Mapped[int] = mapped_column(sqlalchemy.Integer, primary_key=True, autoincrement=True)
|
|
sha1: Mapped[str] = mapped_column(sqlalchemy.String(255), index=True)
|
|
title: Mapped[str] = mapped_column(sqlalchemy.String(255), index=True)
|
|
page_id: Mapped[int] = mapped_column(sqlalchemy.Integer, index=True)
|
|
collection_id: Mapped[int] = mapped_column(sqlalchemy.Integer, index=True)
|
|
indexed_rev_id: Mapped[int] = mapped_column(sqlalchemy.Integer, nullable=True)
|
|
latest_rev_id: Mapped[int] = mapped_column(sqlalchemy.Integer)
|
|
embedding: Mapped[np.ndarray] = deferred(mapped_column(Vector(config.EMBEDDING_VECTOR_SIZE), nullable=True))
|
|
|
|
embedding_index = sqlalchemy.Index("embedding_search_title_index_embedding_idx", embedding,
|
|
postgresql_using='ivfflat',
|
|
postgresql_ops={'embedding': 'vector_cosine_ops'})
|
|
|
|
class TitleIndexHelper(BaseHelper):
|
|
__tablename__ = "embedding_search_title_index"
|
|
|
|
columns = [
|
|
"id",
|
|
"sha1",
|
|
"title",
|
|
"page_id",
|
|
"collection_id",
|
|
"indexed_rev_id",
|
|
"latest_rev_id"
|
|
"embedding",
|
|
]
|
|
|
|
def __init__(self, dbs: DatabaseService):
|
|
super().__init__(dbs)
|
|
|
|
async def __aenter__(self):
|
|
if not self.initialized:
|
|
self.dbpool = self.dbs.pool.acquire()
|
|
self.dbi = await self.dbpool.__aenter__()
|
|
|
|
await register_vector(self.dbi)
|
|
|
|
return await super().__aenter__()
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
await self.dbpool.__aexit__(exc_type, exc, tb)
|
|
return await super().__aexit__(exc_type, exc, tb)
|
|
|
|
def get_columns(self, exclude=[]):
|
|
if len(exclude) == 0:
|
|
return ", ".join(self.columns)
|
|
|
|
return ", ".join([col for col in self.columns if col not in exclude])
|
|
|
|
"""
|
|
Add a title to the index
|
|
"""
|
|
async def add(self, obj: TitleIndexModel):
|
|
obj.sha1 = hashlib.sha1(obj.title.encode("utf-8")).hexdigest()
|
|
|
|
if await self.find_by_sha1(obj.sha1) is None:
|
|
self.session.add(obj)
|
|
await self.session.commit()
|
|
await self.session.refresh(obj)
|
|
return obj
|
|
|
|
return False
|
|
|
|
"""
|
|
Remove a title from the index
|
|
"""
|
|
async def remove(self, title: str):
|
|
title_sha1 = hashlib.sha1(title.encode("utf-8")).hexdigest()
|
|
stmt = sqlalchemy.delete(TitleIndexModel).where(TitleIndexModel.sha1 == title_sha1)
|
|
await self.session.execute(stmt)
|
|
|
|
async def update(self, obj: TitleIndexModel):
|
|
obj.sha1 = hashlib.sha1(obj.title.encode("utf-8")).hexdigest()
|
|
await self.session.merge(obj)
|
|
await self.session.commit()
|
|
await self.session.refresh(obj)
|
|
return obj
|
|
|
|
async def refresh(self, obj: TitleIndexModel):
|
|
await self.session.refresh(obj)
|
|
return obj
|
|
|
|
"""
|
|
Search for titles by consine similary
|
|
"""
|
|
async def search_title_embedding(self, embedding: np.ndarray, limit: int = 10):
|
|
ret = self.dbi.fetch("""SELECT %s, embedding <-> $1 AS distance
|
|
FROM embedding_search_title_index
|
|
ORDER BY distance DESC
|
|
LIMIT %d""" % (self.get_columns(exclude=['embedding']), limit),
|
|
embedding)
|
|
|
|
return ret
|
|
|
|
"""
|
|
Find a title in the index
|
|
"""
|
|
async def find_by_title(self, title: str):
|
|
sha1 = hashlib.sha1(title.encode("utf-8")).hexdigest()
|
|
return await self.find_by_sha1(sha1)
|
|
|
|
async def find_by_sha1(self, sha1: str):
|
|
stmt = sqlalchemy.select(TitleIndexModel).where(TitleIndexModel.sha1 == sha1)
|
|
return await self.session.scalar(stmt)
|
|
|
|
async def find_by_page_id(self, page_id: int):
|
|
stmt = sqlalchemy.select(TitleIndexModel).where(TitleIndexModel.page_id == page_id)
|
|
return await self.session.scalar(stmt)
|
|
|
|
async def find_list_by_collection_id(self, collection_id: int) -> sqlalchemy.ScalarResult[TitleIndexModel]:
|
|
stmt = sqlalchemy.select(TitleIndexModel).where(TitleIndexModel.collection_id == collection_id)
|
|
return await self.session.scalars(stmt)
|
|
|
|
async def get_need_update_index_list(self, collection_id: int) -> list[TitleIndexModel]:
|
|
page_list = await self.find_list_by_collection_id(collection_id)
|
|
result: list[TitleIndexModel] = []
|
|
for page_info in page_list:
|
|
if page_info.indexed_rev_id != page_info.latest_rev_id:
|
|
result.append(page_info)
|
|
|
|
return result |