Merge branch 'rc/new-plugin' into refactor/new-plugin-system

pkg/entity/persistence/bot.py

@@ -4,7 +4,7 @@ from .base import Base
 
 
 class Bot(Base):
-    """机器人"""
+    """Bot"""
 
     __tablename__ = 'bots'
 

pkg/entity/persistence/metadata.py

@@ -12,7 +12,7 @@ initial_metadata = [
 
 
 class Metadata(Base):
-    """数据库元数据"""
+    """Database metadata"""
 
     __tablename__ = 'metadata'
 

pkg/entity/persistence/model.py

@@ -4,7 +4,7 @@ from .base import Base
 
 
 class LLMModel(Base):
-    """LLM 模型"""
+    """LLM model"""
 
     __tablename__ = 'llm_models'
 
@@ -23,3 +23,24 @@ class LLMModel(Base):
         server_default=sqlalchemy.func.now(),
         onupdate=sqlalchemy.func.now(),
     )
+
+
+class EmbeddingModel(Base):
+    """Embedding 模型"""
+
+    __tablename__ = 'embedding_models'
+
+    uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True, unique=True)
+    name = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
+    description = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
+    requester = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
+    requester_config = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default={})
+    api_keys = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
+    extra_args = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default={})
+    created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now())
+    updated_at = sqlalchemy.Column(
+        sqlalchemy.DateTime,
+        nullable=False,
+        server_default=sqlalchemy.func.now(),
+        onupdate=sqlalchemy.func.now(),
+    )

pkg/entity/persistence/pipeline.py

@@ -4,7 +4,7 @@ from .base import Base
 
 
 class LegacyPipeline(Base):
-    """旧版流水线"""
+    """Legacy pipeline"""
 
     __tablename__ = 'legacy_pipelines'
 
@@ -20,13 +20,12 @@ class LegacyPipeline(Base):
     )
     for_version = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
     is_default = sqlalchemy.Column(sqlalchemy.Boolean, nullable=False, default=False)
-
     stages = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
     config = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
 
 
 class PipelineRunRecord(Base):
-    """流水线运行记录"""
+    """Pipeline run record"""
 
     __tablename__ = 'pipeline_run_records'
 
@@ -43,3 +42,4 @@ class PipelineRunRecord(Base):
     started_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
     finished_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
     result = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
+    knowledge_base_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)

pkg/entity/persistence/plugin.py

@@ -4,7 +4,7 @@ from .base import Base
 
 
 class PluginSetting(Base):
-    """插件配置"""
+    """Plugin setting"""
 
     __tablename__ = 'plugin_settings'
 

pkg/entity/persistence/rag.py

@@ -0,0 +1,50 @@
+import sqlalchemy
+from .base import Base
+
+# Base = declarative_base()
+# DATABASE_URL = os.getenv('DATABASE_URL', 'sqlite:///./rag_knowledge.db')
+# print("Using database URL:", DATABASE_URL)
+
+
+# engine = create_engine(DATABASE_URL, connect_args={'check_same_thread': False})
+
+# SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+# def create_db_and_tables():
+#     """Creates all database tables defined in the Base."""
+#     Base.metadata.create_all(bind=engine)
+#     print('Database tables created or already exist.')
+
+
+class KnowledgeBase(Base):
+    __tablename__ = 'knowledge_bases'
+    uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True, unique=True)
+    name = sqlalchemy.Column(sqlalchemy.String, index=True)
+    description = sqlalchemy.Column(sqlalchemy.Text)
+    created_at = sqlalchemy.Column(sqlalchemy.DateTime, default=sqlalchemy.func.now())
+    embedding_model_uuid = sqlalchemy.Column(sqlalchemy.String, default='')
+    top_k = sqlalchemy.Column(sqlalchemy.Integer, default=5)
+
+
+class File(Base):
+    __tablename__ = 'knowledge_base_files'
+    uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True, unique=True)
+    kb_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
+    file_name = sqlalchemy.Column(sqlalchemy.String)
+    extension = sqlalchemy.Column(sqlalchemy.String)
+    created_at = sqlalchemy.Column(sqlalchemy.DateTime, default=sqlalchemy.func.now())
+    status = sqlalchemy.Column(sqlalchemy.String, default='pending')  # pending, processing, completed, failed
+
+
+class Chunk(Base):
+    __tablename__ = 'knowledge_base_chunks'
+    uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True, unique=True)
+    file_id = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
+    text = sqlalchemy.Column(sqlalchemy.Text)
+
+
+# class Vector(Base):
+#     __tablename__ = 'knowledge_base_vectors'
+#     uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True, unique=True)
+#     chunk_id = sqlalchemy.Column(sqlalchemy.String, nullable=True)
+#     embedding = sqlalchemy.Column(sqlalchemy.LargeBinary)

pkg/entity/persistence/vector.py

@@ -0,0 +1,13 @@
+from sqlalchemy import Column, Integer, ForeignKey, LargeBinary
+from sqlalchemy.orm import declarative_base, relationship
+
+Base = declarative_base()
+
+
+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
+
+    chunk = relationship('Chunk', back_populates='vector')

pkg/entity/rag/retriever.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import pydantic
+
+from typing import Any
+
+
+class RetrieveResultEntry(pydantic.BaseModel):
+    id: str
+
+    metadata: dict[str, Any]
+
+    distance: float