"""
清理测试数据脚本
"""

from dotenv import load_dotenv
from pathlib import Path

_env_path = Path(__file__).parent.parent / ".env"
load_dotenv(str(_env_path))

from sqlalchemy import text
from app.db.session import engine

conn = engine.connect()

# 清空知识库
conn.execute(text("DELETE FROM knowledge_embeddings"))
conn.execute(text("DELETE FROM knowledge_items"))
conn.commit()

# 查行数
ki = conn.execute(text("SELECT COUNT(*) FROM knowledge_items")).scalar()
ke = conn.execute(text("SELECT COUNT(*) FROM knowledge_embeddings")).scalar()

print(f"knowledge_items rows: {ki}")
print(f"knowledge_embeddings rows: {ke}")

conn.close()