mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-01-25 14:54:26 +00:00
91d6edc4eb
1.create knowledge space 2.list knowledge space 3.create knowledge document 4.list knowledge document 5.save document chunks 6.sync embedding document
83 lines
3.5 KiB
Python
83 lines
3.5 KiB
Python
from datetime import datetime
|
|
|
|
from sqlalchemy import Column, Integer, String, DateTime, create_engine
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
|
|
from pilot.configs.config import Config
|
|
|
|
from pilot.server.knowledge.request.knowledge_request import KnowledgeSpaceRequest
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
CFG = Config()
|
|
Base = declarative_base()
|
|
class KnowledgeSpaceEntity(Base):
|
|
__tablename__ = 'knowledge_space'
|
|
id = Column(Integer, primary_key=True)
|
|
name = Column(String(100))
|
|
vector_type = Column(String(100))
|
|
desc = Column(String(100))
|
|
owner = Column(String(100))
|
|
gmt_created = Column(DateTime)
|
|
gmt_modified = Column(DateTime)
|
|
|
|
def __repr__(self):
|
|
return f"KnowledgeSpaceEntity(id={self.id}, name='{self.name}', vector_type='{self.vector_type}', desc='{self.desc}', owner='{self.owner}', gmt_created='{self.gmt_created}', gmt_modified='{self.gmt_modified}')"
|
|
|
|
|
|
class KnowledgeSpaceDao:
|
|
def __init__(self):
|
|
database = "knowledge_management"
|
|
self.db_engine = create_engine(f'mysql+pymysql://{CFG.LOCAL_DB_USER}:{CFG.LOCAL_DB_PASSWORD}@{CFG.LOCAL_DB_HOST}:{CFG.LOCAL_DB_PORT}/{database}', echo=True)
|
|
self.Session = sessionmaker(bind=self.db_engine)
|
|
|
|
def create_knowledge_space(self, space:KnowledgeSpaceRequest):
|
|
session = self.Session()
|
|
knowledge_space = KnowledgeSpaceEntity(
|
|
name=space.name,
|
|
vector_type=space.vector_type,
|
|
desc=space.desc,
|
|
owner=space.owner,
|
|
gmt_created=datetime.now(),
|
|
gmt_modified=datetime.now()
|
|
)
|
|
session.add(knowledge_space)
|
|
session.commit()
|
|
|
|
session.close()
|
|
|
|
def get_knowledge_space(self, query:KnowledgeSpaceEntity):
|
|
session = self.Session()
|
|
knowledge_spaces = session.query(KnowledgeSpaceEntity)
|
|
if query.id is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.id == query.id)
|
|
if query.name is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.name == query.name)
|
|
if query.vector_type is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.vector_type == query.vector_type)
|
|
if query.desc is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.desc == query.desc)
|
|
if query.owner is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.owner == query.owner)
|
|
if query.gmt_created is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.gmt_created == query.gmt_created)
|
|
if query.gmt_modified is not None:
|
|
knowledge_spaces = knowledge_spaces.filter(KnowledgeSpaceEntity.gmt_modified == query.gmt_modified)
|
|
|
|
knowledge_spaces = knowledge_spaces.order_by(KnowledgeSpaceEntity.gmt_created.desc())
|
|
result = knowledge_spaces.all()
|
|
return result
|
|
|
|
def update_knowledge_space(self, space_id:int, space:KnowledgeSpaceEntity):
|
|
cursor = self.conn.cursor()
|
|
query = "UPDATE knowledge_space SET name = %s, vector_type = %s, desc = %s, owner = %s WHERE id = %s"
|
|
cursor.execute(query, (space.name, space.vector_type, space.desc, space.owner, space_id))
|
|
self.conn.commit()
|
|
cursor.close()
|
|
|
|
def delete_knowledge_space(self, space_id:int):
|
|
cursor = self.conn.cursor()
|
|
query = "DELETE FROM knowledge_space WHERE id = %s"
|
|
cursor.execute(query, (space_id,))
|
|
self.conn.commit()
|
|
cursor.close()
|