刘国栋
454 lines
18 KiB
Python
454 lines
18 KiB
Python
"""
|
||
轻量级数据持久化模块 - MVP版本
|
||
支持 SQLite 和 JSON 两种存储方式
|
||
"""
|
||
import sqlite3
|
||
import json
|
||
import os
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional, Any
|
||
import pandas as pd
|
||
|
||
|
||
class DataStorage:
|
||
"""统一的数据存储接口，支持SQLite和JSON两种后端"""
|
||
|
||
def __init__(self, storage_type: str = "sqlite", db_path: str = "geo_data.db"):
|
||
"""
|
||
Args:
|
||
storage_type: "sqlite" 或 "json"
|
||
db_path: SQLite数据库路径，或JSON文件目录
|
||
"""
|
||
self.storage_type = storage_type
|
||
self.db_path = db_path
|
||
|
||
if storage_type == "sqlite":
|
||
self._init_sqlite()
|
||
else:
|
||
self._init_json()
|
||
|
||
def _init_sqlite(self):
|
||
"""初始化SQLite数据库"""
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
|
||
# 关键词表
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS keywords (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
keyword TEXT NOT NULL,
|
||
brand TEXT,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
|
||
# 内容表（生成的文章）
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS articles (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
keyword TEXT,
|
||
platform TEXT,
|
||
content TEXT,
|
||
filename TEXT,
|
||
brand TEXT,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
|
||
# 优化记录表
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS optimizations (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
original_content TEXT,
|
||
optimized_content TEXT,
|
||
changes TEXT,
|
||
platform TEXT,
|
||
brand TEXT,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
|
||
# 验证结果表
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS verify_results (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
query TEXT,
|
||
brand TEXT,
|
||
verify_model TEXT,
|
||
mention_count INTEGER,
|
||
mention_position TEXT,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
def _init_json(self):
|
||
"""初始化JSON存储目录"""
|
||
Path(self.db_path).mkdir(parents=True, exist_ok=True)
|
||
|
||
# ==================== 关键词相关 ====================
|
||
|
||
def save_keywords(self, keywords: List[str], brand: str):
|
||
"""保存关键词列表"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
for keyword in keywords:
|
||
cursor.execute(
|
||
"INSERT INTO keywords (keyword, brand) VALUES (?, ?)",
|
||
(keyword, brand)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
else:
|
||
# JSON方式：追加到文件
|
||
json_file = Path(self.db_path) / "keywords.json"
|
||
data = []
|
||
if json_file.exists():
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
for keyword in keywords:
|
||
data.append({
|
||
"keyword": keyword,
|
||
"brand": brand,
|
||
"created_at": datetime.now().isoformat()
|
||
})
|
||
|
||
with open(json_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
def get_keywords(self, brand: Optional[str] = None) -> List[str]:
|
||
"""获取关键词列表"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
if brand:
|
||
cursor.execute("SELECT keyword FROM keywords WHERE brand = ?", (brand,))
|
||
else:
|
||
cursor.execute("SELECT keyword FROM keywords")
|
||
keywords = [row[0] for row in cursor.fetchall()]
|
||
conn.close()
|
||
return keywords
|
||
else:
|
||
json_file = Path(self.db_path) / "keywords.json"
|
||
if not json_file.exists():
|
||
return []
|
||
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
if brand:
|
||
return [item["keyword"] for item in data if item.get("brand") == brand]
|
||
return [item["keyword"] for item in data]
|
||
|
||
# ==================== 文章内容相关 ====================
|
||
|
||
def save_article(self, keyword: str, platform: str, content: str,
|
||
filename: str, brand: str):
|
||
"""保存生成的文章"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
cursor.execute("""
|
||
INSERT INTO articles (keyword, platform, content, filename, brand)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
""", (keyword, platform, content, filename, brand))
|
||
conn.commit()
|
||
conn.close()
|
||
else:
|
||
json_file = Path(self.db_path) / "articles.json"
|
||
data = []
|
||
if json_file.exists():
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
data.append({
|
||
"keyword": keyword,
|
||
"platform": platform,
|
||
"content": content,
|
||
"filename": filename,
|
||
"brand": brand,
|
||
"created_at": datetime.now().isoformat()
|
||
})
|
||
|
||
with open(json_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
def get_articles(self, brand: Optional[str] = None,
|
||
platform: Optional[str] = None) -> List[Dict]:
|
||
"""获取文章列表"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
if brand and platform:
|
||
df = pd.read_sql_query(
|
||
"SELECT * FROM articles WHERE brand = ? AND platform = ?",
|
||
conn, params=(brand, platform)
|
||
)
|
||
elif brand:
|
||
df = pd.read_sql_query(
|
||
"SELECT * FROM articles WHERE brand = ?",
|
||
conn, params=(brand,)
|
||
)
|
||
else:
|
||
df = pd.read_sql_query("SELECT * FROM articles", conn)
|
||
conn.close()
|
||
return df.to_dict('records')
|
||
else:
|
||
json_file = Path(self.db_path) / "articles.json"
|
||
if not json_file.exists():
|
||
return []
|
||
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
if brand and platform:
|
||
return [item for item in data
|
||
if item.get("brand") == brand and item.get("platform") == platform]
|
||
elif brand:
|
||
return [item for item in data if item.get("brand") == brand]
|
||
return data
|
||
|
||
# ==================== 优化记录相关 ====================
|
||
|
||
def save_optimization(self, original_content: str, optimized_content: str,
|
||
changes: str, platform: str, brand: str):
|
||
"""保存优化记录"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
cursor.execute("""
|
||
INSERT INTO optimizations
|
||
(original_content, optimized_content, changes, platform, brand)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
""", (original_content, optimized_content, changes, platform, brand))
|
||
conn.commit()
|
||
conn.close()
|
||
else:
|
||
json_file = Path(self.db_path) / "optimizations.json"
|
||
data = []
|
||
if json_file.exists():
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
data.append({
|
||
"original_content": original_content,
|
||
"optimized_content": optimized_content,
|
||
"changes": changes,
|
||
"platform": platform,
|
||
"brand": brand,
|
||
"created_at": datetime.now().isoformat()
|
||
})
|
||
|
||
with open(json_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
def get_optimizations(self, brand: Optional[str] = None) -> List[Dict]:
|
||
"""获取优化记录"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
if brand:
|
||
df = pd.read_sql_query(
|
||
"SELECT * FROM optimizations WHERE brand = ? ORDER BY created_at DESC",
|
||
conn, params=(brand,)
|
||
)
|
||
else:
|
||
df = pd.read_sql_query(
|
||
"SELECT * FROM optimizations ORDER BY created_at DESC",
|
||
conn
|
||
)
|
||
conn.close()
|
||
return df.to_dict('records')
|
||
else:
|
||
json_file = Path(self.db_path) / "optimizations.json"
|
||
if not json_file.exists():
|
||
return []
|
||
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
if brand:
|
||
return [item for item in data if item.get("brand") == brand]
|
||
return data
|
||
|
||
# ==================== 验证结果相关 ====================
|
||
|
||
def save_verify_results(self, results: List[Dict]):
|
||
"""批量保存验证结果"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
for result in results:
|
||
cursor.execute("""
|
||
INSERT INTO verify_results
|
||
(query, brand, verify_model, mention_count, mention_position)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
""", (
|
||
result.get("问题"),
|
||
result.get("品牌"),
|
||
result.get("验证模型"),
|
||
result.get("提及次数"),
|
||
result.get("位置")
|
||
))
|
||
conn.commit()
|
||
conn.close()
|
||
else:
|
||
json_file = Path(self.db_path) / "verify_results.json"
|
||
data = []
|
||
if json_file.exists():
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
for result in results:
|
||
data.append({
|
||
"query": result.get("问题"),
|
||
"brand": result.get("品牌"),
|
||
"verify_model": result.get("验证模型"),
|
||
"mention_count": result.get("提及次数"),
|
||
"mention_position": result.get("位置"),
|
||
"created_at": datetime.now().isoformat()
|
||
})
|
||
|
||
with open(json_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
def get_verify_results(self, brand: Optional[str] = None, include_timestamp: bool = False) -> pd.DataFrame:
|
||
"""获取验证结果（返回DataFrame）
|
||
|
||
Args:
|
||
brand: 品牌名称，如果为None则返回所有品牌
|
||
include_timestamp: 是否包含时间戳字段
|
||
"""
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
if include_timestamp:
|
||
if brand:
|
||
df = pd.read_sql_query(
|
||
"""SELECT query as "问题", brand as "品牌", verify_model as "验证模型",
|
||
mention_count as "提及次数", mention_position as "位置",
|
||
created_at as "验证时间"
|
||
FROM verify_results WHERE brand = ? ORDER BY created_at DESC""",
|
||
conn, params=(brand,)
|
||
)
|
||
else:
|
||
df = pd.read_sql_query(
|
||
"""SELECT query as "问题", brand as "品牌", verify_model as "验证模型",
|
||
mention_count as "提及次数", mention_position as "位置",
|
||
created_at as "验证时间"
|
||
FROM verify_results ORDER BY created_at DESC""",
|
||
conn
|
||
)
|
||
else:
|
||
if brand:
|
||
df = pd.read_sql_query(
|
||
"""SELECT query as "问题", brand as "品牌", verify_model as "验证模型",
|
||
mention_count as "提及次数", mention_position as "位置"
|
||
FROM verify_results WHERE brand = ?""",
|
||
conn, params=(brand,)
|
||
)
|
||
else:
|
||
df = pd.read_sql_query(
|
||
"""SELECT query as "问题", brand as "品牌", verify_model as "验证模型",
|
||
mention_count as "提及次数", mention_position as "位置"
|
||
FROM verify_results""",
|
||
conn
|
||
)
|
||
conn.close()
|
||
if include_timestamp and not df.empty and "验证时间" in df.columns:
|
||
df["验证时间"] = pd.to_datetime(df["验证时间"])
|
||
return df
|
||
else:
|
||
json_file = Path(self.db_path) / "verify_results.json"
|
||
if not json_file.exists():
|
||
return pd.DataFrame()
|
||
|
||
with open(json_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
if brand:
|
||
data = [item for item in data if item.get("brand") == brand]
|
||
|
||
# 转换为DataFrame格式
|
||
records = []
|
||
for item in data:
|
||
record = {
|
||
"问题": item.get("query"),
|
||
"品牌": item.get("brand"),
|
||
"验证模型": item.get("verify_model"),
|
||
"提及次数": item.get("mention_count"),
|
||
"位置": item.get("mention_position")
|
||
}
|
||
if include_timestamp and "created_at" in item:
|
||
record["验证时间"] = pd.to_datetime(item.get("created_at"))
|
||
records.append(record)
|
||
|
||
df = pd.DataFrame(records)
|
||
if include_timestamp and not df.empty and "验证时间" in df.columns:
|
||
df = df.sort_values("验证时间", ascending=False)
|
||
return df
|
||
|
||
# ==================== 统计功能 ====================
|
||
|
||
def get_stats(self, brand: Optional[str] = None) -> Dict[str, Any]:
|
||
"""获取统计数据"""
|
||
stats = {}
|
||
|
||
if self.storage_type == "sqlite":
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
|
||
# 关键词数量
|
||
if brand:
|
||
cursor.execute("SELECT COUNT(*) FROM keywords WHERE brand = ?", (brand,))
|
||
else:
|
||
cursor.execute("SELECT COUNT(*) FROM keywords")
|
||
stats["keywords_count"] = cursor.fetchone()[0]
|
||
|
||
# 文章数量
|
||
if brand:
|
||
cursor.execute("SELECT COUNT(*) FROM articles WHERE brand = ?", (brand,))
|
||
else:
|
||
cursor.execute("SELECT COUNT(*) FROM articles")
|
||
stats["articles_count"] = cursor.fetchone()[0]
|
||
|
||
# 优化记录数量
|
||
if brand:
|
||
cursor.execute("SELECT COUNT(*) FROM optimizations WHERE brand = ?", (brand,))
|
||
else:
|
||
cursor.execute("SELECT COUNT(*) FROM optimizations")
|
||
stats["optimizations_count"] = cursor.fetchone()[0]
|
||
|
||
# 验证结果数量
|
||
if brand:
|
||
cursor.execute("SELECT COUNT(*) FROM verify_results WHERE brand = ?", (brand,))
|
||
else:
|
||
cursor.execute("SELECT COUNT(*) FROM verify_results")
|
||
stats["verify_results_count"] = cursor.fetchone()[0]
|
||
|
||
conn.close()
|
||
else:
|
||
# JSON方式统计
|
||
keywords_file = Path(self.db_path) / "keywords.json"
|
||
articles_file = Path(self.db_path) / "articles.json"
|
||
optimizations_file = Path(self.db_path) / "optimizations.json"
|
||
verify_file = Path(self.db_path) / "verify_results.json"
|
||
|
||
def count_json(file_path, brand_filter=None):
|
||
if not file_path.exists():
|
||
return 0
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
if brand_filter:
|
||
return len([item for item in data if item.get("brand") == brand_filter])
|
||
return len(data)
|
||
|
||
stats["keywords_count"] = count_json(keywords_file, brand)
|
||
stats["articles_count"] = count_json(articles_file, brand)
|
||
stats["optimizations_count"] = count_json(optimizations_file, brand)
|
||
stats["verify_results_count"] = count_json(verify_file, brand)
|
||
|
||
return stats
|