mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-02 03:55:55 +00:00
fix: ruff check errors
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user