tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-07 11:45:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(db): extract fts and vec function

Browse Source
This commit is contained in:
arkohut 2024-10-12 23:59:06 +08:00
parent f793438d41
commit 974f056dc2
1 changed files with 69 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

130
memos/models.py
View File

@ -313,6 +313,73 @@ def update_entity_last_scan_at_for_metadata(mapper, connection, target):
session.commit() session.commit()
async def update_or_insert_entities_vec(connection, target_id, embedding):
try:
# First, try to update the existing row
result = connection.execute(
text(
"UPDATE entities_vec SET embedding = :embedding WHERE rowid = :id"
),
{
"id": target_id,
"embedding": serialize_float32(embedding),
},
)
# If no row was updated (i.e., the row doesn't exist), then insert a new row
if result.rowcount == 0:
connection.execute(
text(
"INSERT INTO entities_vec (rowid, embedding) VALUES (:id, :embedding)"
),
{
"id": target_id,
"embedding": serialize_float32(embedding),
},
)
except Exception as e:
print(f"Error updating entities_vec: {e}")
def update_or_insert_entities_fts(connection, target_id, filepath, tags, metadata):
try:
# First, try to update the existing row
result = connection.execute(
text(
"""
UPDATE entities_fts
SET filepath = :filepath, tags = :tags, metadata = :metadata
WHERE id = :id
"""
),
{
"id": target_id,
"filepath": filepath,
"tags": tags,
"metadata": metadata,
},
)
# If no row was updated (i.e., the row doesn't exist), then insert a new row
if result.rowcount == 0:
connection.execute(
text(
"""
INSERT INTO entities_fts(id, filepath, tags, metadata)
VALUES(:id, :filepath, :tags, :metadata)
"""
),
{
"id": target_id,
"filepath": filepath,
"tags": tags,
"metadata": metadata,
},
)
except Exception as e:
print(f"Error updating entities_fts: {e}")
async def update_fts_and_vec(mapper, connection, target): async def update_fts_and_vec(mapper, connection, target):
session = Session(bind=connection) session = Session(bind=connection)
@ -344,42 +411,7 @@ async def update_fts_and_vec(mapper, connection, target):
) )
# Update FTS table # Update FTS table
try: update_or_insert_entities_fts(connection, target.id, target.filepath, tags, metadata)
# First, try to update the existing row
result = connection.execute(
text(
"""
UPDATE entities_fts
SET filepath = :filepath, tags = :tags, metadata = :metadata
WHERE id = :id
"""
),
{
"id": target.id,
"filepath": target.filepath,
"tags": tags,
"metadata": metadata,
},
)
# If no row was updated (i.e., the row doesn't exist), then insert a new row
if result.rowcount == 0:
connection.execute(
text(
"""
INSERT INTO entities_fts(id, filepath, tags, metadata)
VALUES(:id, :filepath, :tags, :metadata)
"""
),
{
"id": target.id,
"filepath": target.filepath,
"tags": tags,
"metadata": metadata,
},
)
except Exception as e:
print(f"Error updating entities_fts: {e}")
# Prepare vector data # Prepare vector data
metadata_text = "\n".join( metadata_text = "\n".join(
@ -399,31 +431,7 @@ async def update_fts_and_vec(mapper, connection, target):
# Update vector table # Update vector table
if embedding: if embedding:
try: await update_or_insert_entities_vec(connection, target.id, embedding)
# First, try to update the existing row
result = connection.execute(
text(
"UPDATE entities_vec SET embedding = :embedding WHERE rowid = :id"
),
{
"id": target.id,
"embedding": serialize_float32(embedding),
},
)
# If no row was updated (i.e., the row doesn't exist), then insert a new row
if result.rowcount == 0:
connection.execute(
text(
"INSERT INTO entities_vec (rowid, embedding) VALUES (:id, :embedding)"
),
{
"id": target.id,
"embedding": serialize_float32(embedding),
},
)
except Exception as e:
print(f"Error updating entities_vec: {e}")
session.commit() session.commit()