277 lines
9.6 KiB
Rust
277 lines
9.6 KiB
Rust
|
mod config;
|
||
|
|
mod database;
|
||
|
|
mod spider;
|
||
|
|
mod analyzer;
|
||
|
|
mod ui;
|
||
|
|
|
||
|
|
use config::{Config, ConfigManager};
|
||
|
|
use database::DatabaseManager;
|
||
|
|
use spider::SpiderManager;
|
||
|
|
use analyzer::LLMAnalyzer;
|
||
|
|
use ui::{AppState, draw_indicator, draw_waveform, get_score_label};
|
||
|
|
|
||
|
|
use eframe::egui;
|
||
|
|
use std::sync::Arc;
|
||
|
|
use std::thread;
|
||
|
|
use std::time::Duration;
|
||
|
|
|
||
|
|
struct GubaApp {
|
||
|
|
config_manager: ConfigManager,
|
||
|
|
db: DatabaseManager,
|
||
|
|
spider: SpiderManager,
|
||
|
|
analyzer: LLMAnalyzer,
|
||
|
|
state: AppState,
|
||
|
|
stock_data: Arc<parking_lot::Mutex<Vec<(String, f64)>>>,
|
||
|
|
config_open: bool,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl GubaApp {
|
||
|
|
fn new(cc: &eframe::CreationContext<'_>, config_manager: ConfigManager) -> Self {
|
||
|
|
let config = config_manager.get();
|
||
|
|
|
||
|
|
let db = DatabaseManager::new(&config.database.path)
|
||
|
|
.expect("Failed to initialize database");
|
||
|
|
|
||
|
|
let spider = SpiderManager::new(config.spider.clone());
|
||
|
|
|
||
|
|
let mut analyzer = LLMAnalyzer::new(config.llm_api.clone());
|
||
|
|
|
||
|
|
let state = AppState::new();
|
||
|
|
let stock_data = Arc::new(parking_lot::Mutex::new(Vec::new()));
|
||
|
|
|
||
|
|
let db_clone = Arc::new(database::DatabaseManager::new(&config.database.path)
|
||
|
|
.expect("Failed to create database clone"));
|
||
|
|
let spider_clone = Arc::new(spider::SpiderManager::new(config.spider.clone()));
|
||
|
|
let mut analyzer_clone = analyzer::LLMAnalyzer::new(config.llm_api.clone());
|
||
|
|
|
||
|
|
let state_clone = Arc::new(ui::AppState::new());
|
||
|
|
let stock_data_clone = stock_data.clone();
|
||
|
|
|
||
|
|
thread::spawn(move || {
|
||
|
|
run_background_task(db_clone, spider_clone, &mut analyzer_clone, state_clone, stock_data_clone);
|
||
|
|
});
|
||
|
|
|
||
|
|
Self {
|
||
|
|
config_manager,
|
||
|
|
db,
|
||
|
|
spider,
|
||
|
|
analyzer,
|
||
|
|
state,
|
||
|
|
stock_data,
|
||
|
|
config_open: false,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn run_background_task(
|
||
|
|
db: Arc<DatabaseManager>,
|
||
|
|
spider: Arc<SpiderManager>,
|
||
|
|
analyzer: &mut LLMAnalyzer,
|
||
|
|
state: Arc<ui::AppState>,
|
||
|
|
stock_data: Arc<parking_lot::Mutex<Vec<(String, f64)>>>,
|
||
|
|
) {
|
||
|
|
let mut no_content_count = 0i32;
|
||
|
|
let mut fetch_interval = 15u64;
|
||
|
|
|
||
|
|
loop {
|
||
|
|
if !state.running.load(Ordering::SeqCst) {
|
||
|
|
thread::sleep(Duration::from_secs(1));
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
|
||
|
|
{
|
||
|
|
let mut status = state.status_text.lock();
|
||
|
|
*status = "正在爬取评论...".to_string();
|
||
|
|
}
|
||
|
|
|
||
|
|
let comments = match spider.fetch() {
|
||
|
|
Ok(c) => c,
|
||
|
|
Err(e) => {
|
||
|
|
let mut status = state.status_text.lock();
|
||
|
|
*status = format!("爬取失败: {}", e);
|
||
|
|
thread::sleep(Duration::from_secs(5));
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
if comments.is_empty() {
|
||
|
|
no_content_count += 1;
|
||
|
|
let interval = fetch_interval * (1 + no_content_count.min(4) as u64);
|
||
|
|
let mut status = state.status_text.lock();
|
||
|
|
*status = format!("无新内容，{}秒后重试", interval);
|
||
|
|
thread::sleep(Duration::from_secs(interval));
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
|
||
|
|
no_content_count = 0;
|
||
|
|
state.fetch_count.fetch_add(1, Ordering::SeqCst);
|
||
|
|
|
||
|
|
{
|
||
|
|
let mut status = state.status_text.lock();
|
||
|
|
*status = format!("获取到 {} 条评论", comments.len());
|
||
|
|
}
|
||
|
|
|
||
|
|
match db.add_comments_batch(&comments) {
|
||
|
|
Ok(new_ids) if !new_ids.is_empty() => {
|
||
|
|
let thresholds = {
|
||
|
|
let config = ConfigManager::new("config.json");
|
||
|
|
config.get().ui.thresholds
|
||
|
|
};
|
||
|
|
|
||
|
|
for id in &new_ids {
|
||
|
|
if let Ok(unanalyzed) = db.get_unanalyzed_comments(1) {
|
||
|
|
if let Some(comment) = unanalyzed.first() {
|
||
|
|
match analyzer.analyze(&comment.content) {
|
||
|
|
Ok((score, label)) => {
|
||
|
|
let _ = db.mark_analyzed(comment.id, score, &label);
|
||
|
|
state.analysis_count.fetch_add(1, Ordering::SeqCst);
|
||
|
|
|
||
|
|
let scores = db.get_all_scores(100).unwrap_or_default();
|
||
|
|
if !scores.is_empty() {
|
||
|
|
let avg: i32 = scores.iter().sum::<i32>() / scores.len() as i32;
|
||
|
|
state.current_score.store(avg, Ordering::SeqCst);
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut status = state.status_text.lock();
|
||
|
|
*status = format!("分析完成: {}分 - {}", score, label);
|
||
|
|
}
|
||
|
|
Err(e) => {
|
||
|
|
let _ = db.mark_analyzed(comment.id, 50, "分析异常");
|
||
|
|
let mut status = state.status_text.lock();
|
||
|
|
*status = format!("分析失败: {}", e);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
thread::sleep(Duration::from_secs(1));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
_ => {}
|
||
|
|
}
|
||
|
|
|
||
|
|
if let Ok(data) = spider.fetch_sse_stock_data() {
|
||
|
|
if data.value > 0.0 {
|
||
|
|
let mut stocks = stock_data.lock();
|
||
|
|
stocks.push((data.time.clone(), data.value));
|
||
|
|
if stocks.len() > 100 {
|
||
|
|
stocks.remove(0);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let interval = fetch_interval * (1 + no_content_count.min(4) as u64);
|
||
|
|
thread::sleep(Duration::from_secs(interval));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl eframe::App for GubaApp {
|
||
|
|
fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
|
||
|
|
egui::CentralPanel::default().show(ctx, |ui| {
|
||
|
|
ui.heading("股吧人气指示器");
|
||
|
|
ui.horizontal(|ui| {
|
||
|
|
if ui.button(if self.state.running.load(Ordering::SeqCst) { "停止" } else { "开始" }).clicked() {
|
||
|
|
let running = self.state.running.load(Ordering::SeqCst);
|
||
|
|
self.state.running.store(!running, Ordering::SeqCst);
|
||
|
|
}
|
||
|
|
|
||
|
|
if ui.button("刷新").clicked() {
|
||
|
|
self.state.no_content_count.store(0, Ordering::SeqCst);
|
||
|
|
}
|
||
|
|
|
||
|
|
if ui.button("配置").clicked() {
|
||
|
|
self.config_open = true;
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
ui.separator();
|
||
|
|
|
||
|
|
let config = self.config_manager.get();
|
||
|
|
let score = self.state.current_score.load(Ordering::SeqCst);
|
||
|
|
let cold = config.ui.thresholds.cold;
|
||
|
|
let warm = config.ui.thresholds.warm;
|
||
|
|
|
||
|
|
ui.horizontal(|ui| {
|
||
|
|
ui.label("当前情绪:");
|
||
|
|
draw_indicator(ui, score, cold, warm);
|
||
|
|
});
|
||
|
|
|
||
|
|
ui.separator();
|
||
|
|
|
||
|
|
ui.label("上证指数走势:");
|
||
|
|
let stock_data = self.stock_data.lock();
|
||
|
|
let data: Vec<(String, f64)> = stock_data.clone();
|
||
|
|
drop(stock_data);
|
||
|
|
|
||
|
|
egui::ScrollArea::vertical().max_height(200.0).show(ui, |ui| {
|
||
|
|
draw_waveform(ui, &data, ui.available_width(), 150.0);
|
||
|
|
});
|
||
|
|
|
||
|
|
ui.separator();
|
||
|
|
|
||
|
|
ui.horizontal(|ui| {
|
||
|
|
ui.label("状态: ");
|
||
|
|
let status = self.state.status_text.lock();
|
||
|
|
ui.label(status.clone());
|
||
|
|
});
|
||
|
|
|
||
|
|
ui.horizontal(|ui| {
|
||
|
|
ui.label(format!("爬取次数: {}", self.state.fetch_count.load(Ordering::SeqCst)));
|
||
|
|
ui.label(format!("分析次数: {}", self.state.analysis_count.load(Ordering::SeqCst)));
|
||
|
|
});
|
||
|
|
});
|
||
|
|
|
||
|
|
if self.config_open {
|
||
|
|
egui::Window::new("配置").open(&mut self.config_open).show(ctx, |ui| {
|
||
|
|
let config = self.config_manager.get();
|
||
|
|
|
||
|
|
ui.group(|ui| {
|
||
|
|
ui.label("API配置:");
|
||
|
|
ui.text_edit_singleline(&mut self.config_manager.get().llm_api.api_key.clone());
|
||
|
|
});
|
||
|
|
|
||
|
|
ui.group(|ui| {
|
||
|
|
ui.label("爬虫配置:");
|
||
|
|
ui.text_edit_singleline(&mut self.config_manager.get().spider.target_url.clone());
|
||
|
|
});
|
||
|
|
|
||
|
|
ui.group(|ui| {
|
||
|
|
ui.label("阈值设置:");
|
||
|
|
ui.horizontal(|ui| {
|
||
|
|
ui.label("冷阈值:");
|
||
|
|
ui.add(egui::DragValue::new(&mut self.config_manager.get().ui.thresholds.cold).range(0..=100));
|
||
|
|
ui.label("热阈值:");
|
||
|
|
ui.add(egui::DragValue::new(&mut self.config_manager.get().ui.thresholds.warm).range(0..=100));
|
||
|
|
});
|
||
|
|
});
|
||
|
|
|
||
|
|
if ui.button("保存").clicked() {
|
||
|
|
let _ = self.config_manager.save();
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn main() -> eframe::Result<()> {
|
||
|
|
let config_manager = ConfigManager::new("config.json");
|
||
|
|
|
||
|
|
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
|
||
|
|
.init();
|
||
|
|
|
||
|
|
log::info!("股吧人气指示器启动");
|
||
|
|
|
||
|
|
let options = eframe::NativeOptions {
|
||
|
|
viewport: egui::ViewportBuilder::default()
|
||
|
|
.with_inner_size([400.0, 600.0])
|
||
|
|
.with_min_inner_size([300.0, 400.0])
|
||
|
|
.with_resizable(true),
|
||
|
|
..Default::default()
|
||
|
|
};
|
||
|
|
|
||
|
|
eframe::run_native(
|
||
|
|
"股吧人气指示器",
|
||
|
|
options,
|
||
|
|
Box::new(|cc| Ok(Box::new(GubaApp::new(cc, config_manager)))),
|
||
|
|
)
|
||
|
|
}
|