MaxKB/apps/knowledge/models/knowledge.py
CaptainB 3205ff2a65
Some checks failed
sync2gitee / repo-sync (push) Has been cancelled
feat: add knowledge management
2025-04-25 14:58:22 +08:00

76 lines
3.0 KiB
Python

import uuid_utils.compat as uuid
from django.db import models
from django.db.models.signals import pre_delete
from django.dispatch import receiver
from mptt.fields import TreeForeignKey
from mptt.models import MPTTModel
from common.db.sql_execute import select_one
from common.mixins.app_model_mixin import AppModelMixin
from models_provider.models import Model
from users.models import User
class KnowledgeType(models.IntegerChoices):
base = 0, '通用类型'
web = 1, 'web站点类型'
lark = 2, '飞书类型'
yuque = 3, '语雀类型'
def default_model():
# todo : 这里需要从数据库中获取默认的模型
return uuid.UUID('42f63a3d-427e-11ef-b3ec-a8a1595801ab')
class KnowledgeModule(MPTTModel, AppModelMixin):
id = models.CharField(primary_key=True, max_length=64, editable=False, verbose_name="主键id")
name = models.CharField(max_length=64, verbose_name="文件夹名称")
user = models.ForeignKey(User, on_delete=models.CASCADE, verbose_name="用户id")
workspace_id = models.CharField(max_length=64, verbose_name="工作空间id", default="default", db_index=True)
parent = TreeForeignKey('self', on_delete=models.CASCADE, null=True, blank=True, related_name='children')
class Meta:
db_table = "knowledge_module"
class MPTTMeta:
order_insertion_by = ['name']
class Knowledge(AppModelMixin):
id = models.UUIDField(primary_key=True, max_length=128, default=uuid.uuid7, editable=False, verbose_name="主键id")
name = models.CharField(max_length=150, verbose_name="知识库名称")
desc = models.CharField(max_length=256, verbose_name="描述")
user = models.ForeignKey(User, on_delete=models.DO_NOTHING, verbose_name="所属用户")
type = models.IntegerField(verbose_name='类型', choices=KnowledgeType.choices, default=KnowledgeType.base)
embedding_mode = models.ForeignKey(Model, on_delete=models.DO_NOTHING, verbose_name="向量模型",
default=default_model)
meta = models.JSONField(verbose_name="元数据", default=dict)
class Meta:
db_table = "knowledge"
class File(AppModelMixin):
id = models.UUIDField(primary_key=True, max_length=128, default=uuid.uuid7, editable=False, verbose_name="主键id")
file_name = models.CharField(max_length=256, verbose_name="文件名称", default="")
loid = models.IntegerField(verbose_name="loid")
meta = models.JSONField(verbose_name="文件关联数据", default=dict)
class Meta:
db_table = "file"
def save(self, bytea=None, force_insert=False, force_update=False, using=None, update_fields=None):
result = select_one("SELECT lo_from_bytea(%s, %s::bytea) as loid", [0, bytea])
self.loid = result['loid']
super().save()
def get_bytes(self):
result = select_one(f'SELECT lo_get({self.loid}) as "data"', [])
return result['data']
@receiver(pre_delete, sender=File)
def on_delete_file(sender, instance, **kwargs):
select_one(f'SELECT lo_unlink({instance.loid})', [])