xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-06-09 07:16:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: ruff check --fix

Browse Source
This commit is contained in:
Junyan Qin
2025-07-05 21:56:54 +08:00
parent 39c062f73e
commit 8d28ace252
23 changed files with 647 additions and 737 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
pkg/rag/knowledge/services/database.py
View File

@@ -1,19 +1,20 @@
from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, ForeignKey, LargeBinary
from sqlalchemy.orm import declarative_base, sessionmaker, relationship
from datetime import datetime
import numpy as np # 用于处理从LargeBinary转换回来的embedding
Base = declarative_base()
class KnowledgeBase(Base):
__tablename__ = 'kb'
id = Column(Integer, primary_key=True, index=True)
name = Column(String, index=True)
description = Column(Text)
created_at = Column(DateTime, default=datetime.utcnow)
embedding_model = Column(String, default="") # 默认嵌入模型
embedding_model = Column(String, default='') # 默认嵌入模型
top_k = Column(Integer, default=5) # 默认返回的top_k数量
files = relationship("File", back_populates="knowledge_base")
files = relationship('File', back_populates='knowledge_base')
class File(Base):
__tablename__ = 'file'
@@ -24,8 +25,9 @@ class File(Base):
created_at = Column(DateTime, default=datetime.utcnow)
file_type = Column(String)
status = Column(Integer, default=0) # 0: 未处理, 1: 处理中, 2: 已处理, 3: 错误
knowledge_base = relationship("KnowledgeBase", back_populates="files")
chunks = relationship("Chunk", back_populates="file")
knowledge_base = relationship('KnowledgeBase', back_populates='files')
chunks = relationship('Chunk', back_populates='file')
class Chunk(Base):
__tablename__ = 'chunks'
@@ -33,26 +35,30 @@ class Chunk(Base):
file_id = Column(Integer, ForeignKey('file.id'))
text = Column(Text)
file = relationship("File", back_populates="chunks")
vector = relationship("Vector", uselist=False, back_populates="chunk") # One-to-one
file = relationship('File', back_populates='chunks')
vector = relationship('Vector', uselist=False, back_populates='chunk') # One-to-one
class Vector(Base):
__tablename__ = 'vectors'
id = Column(Integer, primary_key=True, index=True)
chunk_id = Column(Integer, ForeignKey('chunks.id'), unique=True)
embedding = Column(LargeBinary) # Store embeddings as binary
embedding = Column(LargeBinary) # Store embeddings as binary
chunk = relationship('Chunk', back_populates='vector')
chunk = relationship("Chunk", back_populates="vector")
# 数据库连接
DATABASE_URL = "sqlite:///./knowledge_base.db" # 生产环境请更换为 PostgreSQL/MySQL
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False} if "sqlite" in DATABASE_URL else {})
DATABASE_URL = 'sqlite:///./knowledge_base.db' # 生产环境请更换为 PostgreSQL/MySQL
engine = create_engine(DATABASE_URL, connect_args={'check_same_thread': False} if 'sqlite' in DATABASE_URL else {})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# 创建所有表 (可以在应用启动时执行一次)
def create_db_and_tables():
Base.metadata.create_all(bind=engine)
print("Database tables created/checked.")
print('Database tables created/checked.')
# 定义嵌入维度（请根据你实际使用的模型调整）
EMBEDDING_DIM = 1024
EMBEDDING_DIM = 1024