fix: modify rag database

This commit is contained in:
WangCham
2025-07-05 17:53:11 +08:00
committed by Junyan Qin
parent 34fe8b324d
commit 552fee9bac
2 changed files with 5 additions and 4 deletions
+2 -2
View File
@@ -63,7 +63,7 @@ class RAG_Manager:
) )
async def create_knowledge_base(self, kb_name: str, kb_description: str ,): async def create_knowledge_base(self, kb_name: str, kb_description: str, embedding_model: str = "", top_k: int = 5):
""" """
Creates a new knowledge base with the given name and description. Creates a new knowledge base with the given name and description.
If a knowledge base with the same name already exists, it returns that one. If a knowledge base with the same name already exists, it returns that one.
@@ -82,7 +82,7 @@ class RAG_Manager:
def _add_kb_sync(): def _add_kb_sync():
session = SessionLocal() session = SessionLocal()
try: try:
new_kb = KnowledgeBase(name=kb_name, description=kb_description) new_kb = KnowledgeBase(name=kb_name, description=kb_description, embedding_model=embedding_model, top_k=top_k)
session.add(new_kb) session.add(new_kb)
session.commit() session.commit()
session.refresh(new_kb) session.refresh(new_kb)
+2 -1
View File
@@ -11,7 +11,8 @@ class KnowledgeBase(Base):
name = Column(String, index=True) name = Column(String, index=True)
description = Column(Text) description = Column(Text)
created_at = Column(DateTime, default=datetime.utcnow) created_at = Column(DateTime, default=datetime.utcnow)
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): class File(Base):