xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-22 20:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: ruff check errors

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/platform/botmgr.py
+2 -2
View File
@@ -93,7 +93,7 @@ class RuntimeBot:
pipeline_uuid=self.bot_entity.use_pipeline_uuid, pipeline_uuid=self.bot_entity.use_pipeline_uuid,
) )
else: 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( async def on_group_message(
event: platform_events.GroupMessage, event: platform_events.GroupMessage,
@@ -136,7 +136,7 @@ class RuntimeBot:
pipeline_uuid=self.bot_entity.use_pipeline_uuid, pipeline_uuid=self.bot_entity.use_pipeline_uuid,
) )
else: 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.FriendMessage, on_friend_message)
self.adapter.register_listener(platform_events.GroupMessage, on_group_message) self.adapter.register_listener(platform_events.GroupMessage, on_group_message)
src/langbot/pkg/platform/webhook_pusher.py
+2 -2
View File
@@ -56,7 +56,7 @@ class WebhookPusher:
# Check if any webhook responded with skip_pipeline=true # Check if any webhook responded with skip_pipeline=true
for result in results: for result in results:
if isinstance(result, dict) and result.get('skip_pipeline') is True: 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 True
return False return False
@@ -103,7 +103,7 @@ class WebhookPusher:
# Check if any webhook responded with skip_pipeline=true # Check if any webhook responded with skip_pipeline=true
for result in results: for result in results:
if isinstance(result, dict) and result.get('skip_pipeline') is True: 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 True
return False return False
src/langbot/pkg/vector/vdbs/pgvector_db.py
+33 -63
View File
@@ -1,19 +1,18 @@
from __future__ import annotations from __future__ import annotations
import asyncio
from typing import Any, Dict from typing import Any, Dict
from sqlalchemy import create_engine, text, Column, String, Text 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 sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from pgvector.sqlalchemy import Vector from pgvector.sqlalchemy import Vector
from langbot.pkg.vector.vdb import VectorDatabase from langbot.pkg.vector.vdb import VectorDatabase
from langbot.pkg.core import app from langbot.pkg.core import app
import uuid
Base = declarative_base() Base = declarative_base()
class PgVectorEntry(Base): class PgVectorEntry(Base):
"""SQLAlchemy model for pgvector entries""" """SQLAlchemy model for pgvector entries"""
__tablename__ = 'langbot_vectors' __tablename__ = 'langbot_vectors'
id = Column(String, primary_key=True) id = Column(String, primary_key=True)
@@ -31,11 +30,11 @@ class PgVectorDatabase(VectorDatabase):
self, self,
ap: app.Application, ap: app.Application,
connection_string: str = None, connection_string: str = None,
host: str = "localhost", host: str = 'localhost',
port: int = 5432, port: int = 5432,
database: str = "langbot", database: str = 'langbot',
user: str = "postgres", user: str = 'postgres',
password: str = "postgres" password: str = 'postgres',
): ):
"""Initialize pgvector database """Initialize pgvector database
@@ -54,14 +53,10 @@ class PgVectorDatabase(VectorDatabase):
if connection_string: if connection_string:
self.connection_string = connection_string self.connection_string = connection_string
else: else:
self.connection_string = ( self.connection_string = f'postgresql+psycopg://{user}:{password}@{host}:{port}/{database}'
f"postgresql+psycopg://{user}:{password}@{host}:{port}/{database}"
)
self.async_connection_string = self.connection_string.replace( self.async_connection_string = self.connection_string.replace('postgresql://', 'postgresql+asyncpg://').replace(
"postgresql://", "postgresql+asyncpg://" 'postgresql+psycopg://', 'postgresql+asyncpg://'
).replace(
"postgresql+psycopg://", "postgresql+asyncpg://"
) )
self.engine = None self.engine = None
@@ -75,35 +70,25 @@ class PgVectorDatabase(VectorDatabase):
"""Initialize database connection and create tables""" """Initialize database connection and create tables"""
try: try:
# Create async engine for async operations # Create async engine for async operations
self.async_engine = create_async_engine( self.async_engine = create_async_engine(self.async_connection_string, echo=False, pool_pre_ping=True)
self.async_connection_string, self.AsyncSessionLocal = async_sessionmaker(self.async_engine, class_=AsyncSession, expire_on_commit=False)
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 # Create sync engine for table creation
sync_connection_string = self.connection_string.replace( sync_connection_string = self.connection_string.replace('postgresql+asyncpg://', 'postgresql+psycopg://')
"postgresql+asyncpg://", "postgresql+psycopg://"
)
self.engine = create_engine(sync_connection_string, echo=False) self.engine = create_engine(sync_connection_string, echo=False)
# Create pgvector extension and tables # Create pgvector extension and tables
with self.engine.connect() as conn: with self.engine.connect() as conn:
# Enable pgvector extension # Enable pgvector extension
conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
conn.commit() conn.commit()
# Create tables # Create tables
Base.metadata.create_all(self.engine) 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: 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 raise
async def get_or_create_collection(self, collection: str): async def get_or_create_collection(self, collection: str):
@@ -144,24 +129,20 @@ class PgVectorDatabase(VectorDatabase):
id=vector_id, id=vector_id,
collection=collection, collection=collection,
embedding=embeddings_list[i], embedding=embeddings_list[i],
text=metadata.get("text", ""), text=metadata.get('text', ''),
file_id=metadata.get("file_id", ""), file_id=metadata.get('file_id', ''),
chunk_uuid=metadata.get("uuid", "") chunk_uuid=metadata.get('uuid', ''),
) )
session.add(entry) session.add(entry)
await session.commit() await session.commit()
self.ap.logger.info( self.ap.logger.info(f"Added {len(ids)} embeddings to pgvector collection '{collection}'")
f"Added {len(ids)} embeddings to pgvector collection '{collection}'"
)
except Exception as e: except Exception as e:
await session.rollback() 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 raise
async def search( async def search(self, collection: str, query_embedding: list[float], k: int = 5) -> Dict[str, Any]:
self, collection: str, query_embedding: list[float], k: int = 5
) -> Dict[str, Any]:
"""Search for similar vectors using cosine distance """Search for similar vectors using cosine distance
Args: Args:
@@ -177,7 +158,7 @@ class PgVectorDatabase(VectorDatabase):
async with self.AsyncSessionLocal() as session: async with self.AsyncSessionLocal() as session:
try: try:
# Use cosine distance for similarity search # Use cosine distance for similarity search
from sqlalchemy import select, func from sqlalchemy import select
# Query for similar vectors # Query for similar vectors
stmt = ( stmt = (
@@ -186,7 +167,7 @@ class PgVectorDatabase(VectorDatabase):
PgVectorEntry.text, PgVectorEntry.text,
PgVectorEntry.file_id, PgVectorEntry.file_id,
PgVectorEntry.chunk_uuid, PgVectorEntry.chunk_uuid,
PgVectorEntry.embedding.cosine_distance(query_embedding).label('distance') PgVectorEntry.embedding.cosine_distance(query_embedding).label('distance'),
) )
.filter(PgVectorEntry.collection == collection) .filter(PgVectorEntry.collection == collection)
.order_by(PgVectorEntry.embedding.cosine_distance(query_embedding)) .order_by(PgVectorEntry.embedding.cosine_distance(query_embedding))
@@ -204,25 +185,17 @@ class PgVectorDatabase(VectorDatabase):
for row in rows: for row in rows:
ids.append(row.id) ids.append(row.id)
distances.append(float(row.distance)) distances.append(float(row.distance))
metadatas.append({ metadatas.append(
"text": row.text or "", {'text': row.text or '', 'file_id': row.file_id or '', 'uuid': row.chunk_uuid or ''}
"file_id": row.file_id or "", )
"uuid": row.chunk_uuid or ""
})
result_dict = { result_dict = {'ids': [ids], 'distances': [distances], 'metadatas': [metadatas]}
"ids": [ids],
"distances": [distances],
"metadatas": [metadatas]
}
self.ap.logger.info( self.ap.logger.info(f"pgvector search in '{collection}' returned {len(ids)} results")
f"pgvector search in '{collection}' returned {len(ids)} results"
)
return result_dict return result_dict
except Exception as e: except Exception as e:
self.ap.logger.error(f"Error searching pgvector: {e}") self.ap.logger.error(f'Error searching pgvector: {e}')
raise raise
async def delete_by_file_id(self, collection: str, file_id: str) -> None: async def delete_by_file_id(self, collection: str, file_id: str) -> None:
@@ -239,8 +212,7 @@ class PgVectorDatabase(VectorDatabase):
from sqlalchemy import delete from sqlalchemy import delete
stmt = delete(PgVectorEntry).where( stmt = delete(PgVectorEntry).where(
PgVectorEntry.collection == collection, PgVectorEntry.collection == collection, PgVectorEntry.file_id == file_id
PgVectorEntry.file_id == file_id
) )
await session.execute(stmt) await session.execute(stmt)
await session.commit() await session.commit()
@@ -250,7 +222,7 @@ class PgVectorDatabase(VectorDatabase):
) )
except Exception as e: except Exception as e:
await session.rollback() await session.rollback()
self.ap.logger.error(f"Error deleting from pgvector: {e}") self.ap.logger.error(f'Error deleting from pgvector: {e}')
raise raise
async def delete_collection(self, collection: str): async def delete_collection(self, collection: str):
@@ -266,16 +238,14 @@ class PgVectorDatabase(VectorDatabase):
try: try:
from sqlalchemy import delete from sqlalchemy import delete
stmt = delete(PgVectorEntry).where( stmt = delete(PgVectorEntry).where(PgVectorEntry.collection == collection)
PgVectorEntry.collection == collection
)
await session.execute(stmt) await session.execute(stmt)
await session.commit() await session.commit()
self.ap.logger.info(f"Deleted pgvector collection '{collection}'") self.ap.logger.info(f"Deleted pgvector collection '{collection}'")
except Exception as e: except Exception as e:
await session.rollback() await session.rollback()
self.ap.logger.error(f"Error deleting pgvector collection: {e}") self.ap.logger.error(f'Error deleting pgvector collection: {e}')
raise raise
async def close(self): async def close(self):
src/langbot/pkg/vector/vdbs/seekdb.py
+15 -17
View File
@@ -3,10 +3,8 @@ from __future__ import annotations
import asyncio import asyncio
from typing import Any, Dict, List from typing import Any, Dict, List
import sqlalchemy
from langbot.pkg.core import app from langbot.pkg.core import app
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.vector.vdb import VectorDatabase from langbot.pkg.vector.vdb import VectorDatabase
try: try:
@@ -87,14 +85,16 @@ class SeekDBVectorDatabase(VectorDatabase):
self._collections: Dict[str, Any] = {} self._collections: Dict[str, Any] = {}
self._collection_configs: Dict[str, HNSWConfiguration] = {} self._collection_configs: Dict[str, HNSWConfiguration] = {}
self._escape_table = str.maketrans({ self._escape_table = str.maketrans(
'\x00': '', {
'\\': '\\\\', '\x00': '',
'"': '\\"', '\\': '\\\\',
'\n': '\\n', '"': '\\"',
'\r': '\\r', '\n': '\\n',
'\t': '\\t', '\r': '\\r',
}) '\t': '\\t',
}
)
async def _get_or_create_collection_internal(self, collection: str, vector_size: int = None) -> Any: 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.""" """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]: def _clean_metadata(self, meta: Dict[str, Any]) -> Dict[str, Any]:
"""SeekDB metadata doesn't support \\ and ", insert will error 3104""" """SeekDB metadata doesn't support \\ and ", insert will error 3104"""
return { return {
k: v.translate(self._escape_table) if isinstance(v, str) k: v.translate(self._escape_table)
else v if v is None or isinstance(v, (int, float, bool)) if isinstance(v, str)
else v
if v is None or isinstance(v, (int, float, bool))
else str(v) else str(v)
for k, v in meta.items() for k, v in meta.items()
if v is not None if v is not None
@@ -145,11 +147,7 @@ class SeekDBVectorDatabase(VectorDatabase):
return await self._get_or_create_collection_internal(collection) return await self._get_or_create_collection_internal(collection)
async def add_embeddings( async def add_embeddings(
self, self, collection: str, ids: List[str], embeddings_list: List[List[float]], metadatas: List[Dict[str, Any]]
collection: str,
ids: List[str],
embeddings_list: List[List[float]],
metadatas: List[Dict[str, Any]]
) -> None: ) -> None:
"""Add vector embeddings to the specified collection. """Add vector embeddings to the specified collection.