parent
0d1ae51bf0
commit
e2a8913778
@ -1,110 +0,0 @@
|
||||
import logging
|
||||
from typing import Optional, List, Dict
|
||||
from aiomysql import DictCursor
|
||||
|
||||
class KbDao:
|
||||
def __init__(self, mysql_pool):
|
||||
self.mysql_pool = mysql_pool
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
async def create_kb(self, kb: Dict) -> int:
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor(DictCursor) as cur:
|
||||
await cur.execute(
|
||||
"INSERT INTO t_ai_kb(kb_name, short_name) VALUES(%s, %s)",
|
||||
(kb['name'], kb.get('description', '')))
|
||||
await conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
async def get_kb(self, kb_id: int) -> Optional[Dict]:
|
||||
"""获取知识库详情"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor(DictCursor) as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM t_ai_kb WHERE id = %s", (kb_id,))
|
||||
return await cur.fetchone()
|
||||
|
||||
async def update_kb(self, kb_id: int, kb: Dict) -> bool:
|
||||
"""更新知识库信息"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"UPDATE t_ai_kb SET name = %s, description = %s WHERE id = %s",
|
||||
(kb['name'], kb['description'], kb_id))
|
||||
await conn.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
async def delete_kb(self, kb_id: int) -> bool:
|
||||
"""删除知识库"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"DELETE FROM t_ai_kb WHERE id = %s", (kb_id,))
|
||||
await conn.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
async def create_kb_file(self, file: Dict) -> int:
|
||||
"""创建知识库文件记录"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor(DictCursor) as cur:
|
||||
await cur.execute(
|
||||
"""INSERT INTO t_ai_kb_files
|
||||
(kb_id, file_name, file_path, file_size, file_type, state)
|
||||
VALUES(%s, %s, %s, %s, %s, %s)""",
|
||||
(file['kb_id'], file['file_name'], file['file_path'],
|
||||
file['file_size'], file['file_type'], file['state']))
|
||||
await conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
async def get_kb_file(self, file_id: int) -> Optional[Dict]:
|
||||
"""获取文件详情"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor(DictCursor) as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM t_ai_kb_files WHERE id = %s", (file_id,))
|
||||
return await cur.fetchone()
|
||||
|
||||
async def update_kb_file(self, file_id: int, file: Dict) -> bool:
|
||||
"""更新文件信息"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"""UPDATE t_ai_kb_files SET
|
||||
kb_id = %s, file_name = %s, file_path = %s,
|
||||
file_size = %s, file_type = %s, state = %s
|
||||
WHERE id = %s""",
|
||||
(file['kb_id'], file['file_name'], file['file_path'],
|
||||
file['file_size'], file['file_type'], file['state'], file_id))
|
||||
await conn.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
async def delete_kb_file(self, file_id: int) -> bool:
|
||||
"""删除文件记录"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"DELETE FROM t_ai_kb_files WHERE id = %s", (file_id,))
|
||||
await conn.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
async def handle_upload(self, kb_id: int, file) -> Dict:
|
||||
"""处理文件上传"""
|
||||
# 文件保存逻辑
|
||||
# 数据库记录创建
|
||||
# 返回文件信息
|
||||
pass
|
||||
|
||||
async def get_unprocessed_files(self) -> List[Dict]:
|
||||
"""获取未处理文件列表"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor(DictCursor) as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM t_ai_kb_files WHERE state = 0")
|
||||
return await cur.fetchall()
|
||||
|
||||
async def list_kbs(self) -> List[Dict]:
|
||||
"""获取所有知识库列表"""
|
||||
async with self.mysql_pool.acquire() as conn:
|
||||
async with conn.cursor(DictCursor) as cur:
|
||||
await cur.execute("SELECT * FROM t_ai_kb")
|
||||
return await cur.fetchall()
|
||||
Binary file not shown.
@ -1,55 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
class KbBase(BaseModel):
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
|
||||
class KbCreate(KbBase):
|
||||
pass
|
||||
|
||||
class KbUpdate(KbBase):
|
||||
pass
|
||||
|
||||
class Kb(KbBase):
|
||||
id: int
|
||||
create_time: datetime
|
||||
update_time: datetime
|
||||
|
||||
class Config:
|
||||
orm_mode = True
|
||||
|
||||
class KbFileBase(BaseModel):
|
||||
kb_id: int
|
||||
file_name: str
|
||||
file_path: str
|
||||
file_size: int
|
||||
file_type: str
|
||||
state: int = 0
|
||||
|
||||
class KbFileCreate(KbFileBase):
|
||||
pass
|
||||
|
||||
class KbFileUpdate(KbFileBase):
|
||||
pass
|
||||
|
||||
class KbFile(KbFileBase):
|
||||
id: int
|
||||
create_time: datetime
|
||||
update_time: datetime
|
||||
|
||||
class Config:
|
||||
orm_mode = True
|
||||
|
||||
class KbModel(BaseModel):
|
||||
kb_name: str
|
||||
short_name: str
|
||||
is_delete: Optional[int] = 0
|
||||
|
||||
class KbFileModel(BaseModel):
|
||||
file_name: str
|
||||
ext_name: str
|
||||
kb_id: int
|
||||
is_delete: Optional[int] = 0
|
||||
state: Optional[int] = 0
|
||||
Binary file not shown.
Loading…
Reference in new issue