fix: ruff check errors

This commit is contained in:
Junyan Qin
2026-01-23 13:30:44 +08:00
parent c90f2d6a12
commit e60cb6ad0e
4 changed files with 52 additions and 84 deletions

View File

@@ -93,7 +93,7 @@ class RuntimeBot:
pipeline_uuid=self.bot_entity.use_pipeline_uuid,
)
else:
await self.logger.info(f'Pipeline skipped for person message due to webhook response')
await self.logger.info('Pipeline skipped for person message due to webhook response')
async def on_group_message(
event: platform_events.GroupMessage,
@@ -136,7 +136,7 @@ class RuntimeBot:
pipeline_uuid=self.bot_entity.use_pipeline_uuid,
)
else:
await self.logger.info(f'Pipeline skipped for group message due to webhook response')
await self.logger.info('Pipeline skipped for group message due to webhook response')
self.adapter.register_listener(platform_events.FriendMessage, on_friend_message)
self.adapter.register_listener(platform_events.GroupMessage, on_group_message)

View File

@@ -56,7 +56,7 @@ class WebhookPusher:
# Check if any webhook responded with skip_pipeline=true
for result in results:
if isinstance(result, dict) and result.get('skip_pipeline') is True:
self.logger.info(f'Webhook responded with skip_pipeline=true, skipping pipeline for person message')
self.logger.info('Webhook responded with skip_pipeline=true, skipping pipeline for person message')
return True
return False
@@ -103,7 +103,7 @@ class WebhookPusher:
# Check if any webhook responded with skip_pipeline=true
for result in results:
if isinstance(result, dict) and result.get('skip_pipeline') is True:
self.logger.info(f'Webhook responded with skip_pipeline=true, skipping pipeline for group message')
self.logger.info('Webhook responded with skip_pipeline=true, skipping pipeline for group message')
return True
return False

View File

@@ -1,19 +1,18 @@
from __future__ import annotations
import asyncio
from typing import Any, Dict
from sqlalchemy import create_engine, text, Column, String, Text
from sqlalchemy.orm import declarative_base, sessionmaker, Session
from sqlalchemy.orm import declarative_base
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from pgvector.sqlalchemy import Vector
from langbot.pkg.vector.vdb import VectorDatabase
from langbot.pkg.core import app
import uuid
Base = declarative_base()
class PgVectorEntry(Base):
"""SQLAlchemy model for pgvector entries"""
__tablename__ = 'langbot_vectors'
id = Column(String, primary_key=True)
@@ -31,11 +30,11 @@ class PgVectorDatabase(VectorDatabase):
self,
ap: app.Application,
connection_string: str = None,
host: str = "localhost",
host: str = 'localhost',
port: int = 5432,
database: str = "langbot",
user: str = "postgres",
password: str = "postgres"
database: str = 'langbot',
user: str = 'postgres',
password: str = 'postgres',
):
"""Initialize pgvector database
@@ -54,14 +53,10 @@ class PgVectorDatabase(VectorDatabase):
if connection_string:
self.connection_string = connection_string
else:
self.connection_string = (
f"postgresql+psycopg://{user}:{password}@{host}:{port}/{database}"
)
self.connection_string = f'postgresql+psycopg://{user}:{password}@{host}:{port}/{database}'
self.async_connection_string = self.connection_string.replace(
"postgresql://", "postgresql+asyncpg://"
).replace(
"postgresql+psycopg://", "postgresql+asyncpg://"
self.async_connection_string = self.connection_string.replace('postgresql://', 'postgresql+asyncpg://').replace(
'postgresql+psycopg://', 'postgresql+asyncpg://'
)
self.engine = None
@@ -75,35 +70,25 @@ class PgVectorDatabase(VectorDatabase):
"""Initialize database connection and create tables"""
try:
# Create async engine for async operations
self.async_engine = create_async_engine(
self.async_connection_string,
echo=False,
pool_pre_ping=True
)
self.AsyncSessionLocal = async_sessionmaker(
self.async_engine,
class_=AsyncSession,
expire_on_commit=False
)
self.async_engine = create_async_engine(self.async_connection_string, echo=False, pool_pre_ping=True)
self.AsyncSessionLocal = async_sessionmaker(self.async_engine, class_=AsyncSession, expire_on_commit=False)
# Create sync engine for table creation
sync_connection_string = self.connection_string.replace(
"postgresql+asyncpg://", "postgresql+psycopg://"
)
sync_connection_string = self.connection_string.replace('postgresql+asyncpg://', 'postgresql+psycopg://')
self.engine = create_engine(sync_connection_string, echo=False)
# Create pgvector extension and tables
with self.engine.connect() as conn:
# Enable pgvector extension
conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
conn.commit()
# Create tables
Base.metadata.create_all(self.engine)
self.ap.logger.info(f"Connected to PostgreSQL with pgvector")
self.ap.logger.info('Connected to PostgreSQL with pgvector')
except Exception as e:
self.ap.logger.error(f"Failed to connect to PostgreSQL: {e}")
self.ap.logger.error(f'Failed to connect to PostgreSQL: {e}')
raise
async def get_or_create_collection(self, collection: str):
@@ -144,24 +129,20 @@ class PgVectorDatabase(VectorDatabase):
id=vector_id,
collection=collection,
embedding=embeddings_list[i],
text=metadata.get("text", ""),
file_id=metadata.get("file_id", ""),
chunk_uuid=metadata.get("uuid", "")
text=metadata.get('text', ''),
file_id=metadata.get('file_id', ''),
chunk_uuid=metadata.get('uuid', ''),
)
session.add(entry)
await session.commit()
self.ap.logger.info(
f"Added {len(ids)} embeddings to pgvector collection '{collection}'"
)
self.ap.logger.info(f"Added {len(ids)} embeddings to pgvector collection '{collection}'")
except Exception as e:
await session.rollback()
self.ap.logger.error(f"Error adding embeddings to pgvector: {e}")
self.ap.logger.error(f'Error adding embeddings to pgvector: {e}')
raise
async def search(
self, collection: str, query_embedding: list[float], k: int = 5
) -> Dict[str, Any]:
async def search(self, collection: str, query_embedding: list[float], k: int = 5) -> Dict[str, Any]:
"""Search for similar vectors using cosine distance
Args:
@@ -177,7 +158,7 @@ class PgVectorDatabase(VectorDatabase):
async with self.AsyncSessionLocal() as session:
try:
# Use cosine distance for similarity search
from sqlalchemy import select, func
from sqlalchemy import select
# Query for similar vectors
stmt = (
@@ -186,7 +167,7 @@ class PgVectorDatabase(VectorDatabase):
PgVectorEntry.text,
PgVectorEntry.file_id,
PgVectorEntry.chunk_uuid,
PgVectorEntry.embedding.cosine_distance(query_embedding).label('distance')
PgVectorEntry.embedding.cosine_distance(query_embedding).label('distance'),
)
.filter(PgVectorEntry.collection == collection)
.order_by(PgVectorEntry.embedding.cosine_distance(query_embedding))
@@ -204,25 +185,17 @@ class PgVectorDatabase(VectorDatabase):
for row in rows:
ids.append(row.id)
distances.append(float(row.distance))
metadatas.append({
"text": row.text or "",
"file_id": row.file_id or "",
"uuid": row.chunk_uuid or ""
})
metadatas.append(
{'text': row.text or '', 'file_id': row.file_id or '', 'uuid': row.chunk_uuid or ''}
)
result_dict = {
"ids": [ids],
"distances": [distances],
"metadatas": [metadatas]
}
result_dict = {'ids': [ids], 'distances': [distances], 'metadatas': [metadatas]}
self.ap.logger.info(
f"pgvector search in '{collection}' returned {len(ids)} results"
)
self.ap.logger.info(f"pgvector search in '{collection}' returned {len(ids)} results")
return result_dict
except Exception as e:
self.ap.logger.error(f"Error searching pgvector: {e}")
self.ap.logger.error(f'Error searching pgvector: {e}')
raise
async def delete_by_file_id(self, collection: str, file_id: str) -> None:
@@ -239,8 +212,7 @@ class PgVectorDatabase(VectorDatabase):
from sqlalchemy import delete
stmt = delete(PgVectorEntry).where(
PgVectorEntry.collection == collection,
PgVectorEntry.file_id == file_id
PgVectorEntry.collection == collection, PgVectorEntry.file_id == file_id
)
await session.execute(stmt)
await session.commit()
@@ -250,7 +222,7 @@ class PgVectorDatabase(VectorDatabase):
)
except Exception as e:
await session.rollback()
self.ap.logger.error(f"Error deleting from pgvector: {e}")
self.ap.logger.error(f'Error deleting from pgvector: {e}')
raise
async def delete_collection(self, collection: str):
@@ -266,16 +238,14 @@ class PgVectorDatabase(VectorDatabase):
try:
from sqlalchemy import delete
stmt = delete(PgVectorEntry).where(
PgVectorEntry.collection == collection
)
stmt = delete(PgVectorEntry).where(PgVectorEntry.collection == collection)
await session.execute(stmt)
await session.commit()
self.ap.logger.info(f"Deleted pgvector collection '{collection}'")
except Exception as e:
await session.rollback()
self.ap.logger.error(f"Error deleting pgvector collection: {e}")
self.ap.logger.error(f'Error deleting pgvector collection: {e}')
raise
async def close(self):

View File

@@ -3,10 +3,8 @@ from __future__ import annotations
import asyncio
from typing import Any, Dict, List
import sqlalchemy
from langbot.pkg.core import app
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.vector.vdb import VectorDatabase
try:
@@ -87,14 +85,16 @@ class SeekDBVectorDatabase(VectorDatabase):
self._collections: Dict[str, Any] = {}
self._collection_configs: Dict[str, HNSWConfiguration] = {}
self._escape_table = str.maketrans({
'\x00': '',
'\\': '\\\\',
'"': '\\"',
'\n': '\\n',
'\r': '\\r',
'\t': '\\t',
})
self._escape_table = str.maketrans(
{
'\x00': '',
'\\': '\\\\',
'"': '\\"',
'\n': '\\n',
'\r': '\\r',
'\t': '\\t',
}
)
async def _get_or_create_collection_internal(self, collection: str, vector_size: int = None) -> Any:
"""Internal method to get or create a collection with proper configuration."""
@@ -133,8 +133,10 @@ class SeekDBVectorDatabase(VectorDatabase):
def _clean_metadata(self, meta: Dict[str, Any]) -> Dict[str, Any]:
"""SeekDB metadata doesn't support \\ and ", insert will error 3104"""
return {
k: v.translate(self._escape_table) if isinstance(v, str)
else v if v is None or isinstance(v, (int, float, bool))
k: v.translate(self._escape_table)
if isinstance(v, str)
else v
if v is None or isinstance(v, (int, float, bool))
else str(v)
for k, v in meta.items()
if v is not None
@@ -145,11 +147,7 @@ class SeekDBVectorDatabase(VectorDatabase):
return await self._get_or_create_collection_internal(collection)
async def add_embeddings(
self,
collection: str,
ids: List[str],
embeddings_list: List[List[float]],
metadatas: List[Dict[str, Any]]
self, collection: str, ids: List[str], embeddings_list: List[List[float]], metadatas: List[Dict[str, Any]]
) -> None:
"""Add vector embeddings to the specified collection.