fix: manage_scheduled_task + search_web 注入 + 删除 LM Studio

- manage_scheduled_task: INSERT 加 user_id='', 检查 psql 退出码
- search_web: 用 python json.dumps 构造请求体(防止注入)
- 删除 lmstudio.rs / LmStudioConfig,仅保留 DeepSeek
- config.json 移除 lmstudio 段
This commit is contained in:
2026-06-02 13:51:46 +08:00
parent 7ab5052642
commit a549bd3b2f
8 changed files with 29 additions and 226 deletions
-5
View File
@@ -9,11 +9,6 @@
"api_key": "${DEEPSEEK_API_KEY}",
"model": "${DEEPSEEK_MODEL}",
"base_url": "${DEEPSEEK_BASE_URL}"
},
"lmstudio": {
"base_url": "${LM_STUDIO_BASE_URL}",
"model": "${LM_STUDIO_MODEL}",
"api_key": "${LM_STUDIO_API_KEY}"
}
}
}
@@ -15,23 +15,30 @@ case "$ACTION" in
add)
[ -z "$CMD" ] && echo '{"error":"add 需要 command 参数"}' && exit 0
psql "$DB_URL" -v name="$NAME" -v desc="$DESC" -v cmd="$CMD" -v interval="$INTERVAL" << 'SQL'
INSERT INTO scheduled_tasks (id, name, description, command, interval_seconds, enabled, next_run_at)
VALUES (gen_random_uuid(), :'name', :'desc', :'cmd', :interval, true, NOW() + interval ':interval seconds')
INSERT INTO scheduled_tasks (id, name, description, command, interval_seconds, enabled, user_id, next_run_at)
VALUES (gen_random_uuid(), :'name', :'desc', :'cmd', :interval, true, '', NOW() + interval ':interval seconds')
ON CONFLICT (name) DO UPDATE SET command=EXCLUDED.command, interval_seconds=EXCLUDED.interval_seconds, enabled=true, description=EXCLUDED.description
SQL
echo '{"ok":true,"action":"add","name":"'"$NAME"'"}'
if [ $? -eq 0 ]; then
echo '{"ok":true,"action":"add","name":"'"$NAME"'"}'
else
echo '{"error":"添加失败"}'
fi
;;
delete)
psql "$DB_URL" -v name="$NAME" -c "DELETE FROM scheduled_tasks WHERE name = :'name'" 2>/dev/null
echo '{"ok":true,"action":"delete","name":"'"$NAME"'"}'
psql "$DB_URL" -v name="$NAME" -c "DELETE FROM scheduled_tasks WHERE name = :'name'" 2>/dev/null && \
echo '{"ok":true,"action":"delete","name":"'"$NAME"'"}' || \
echo '{"error":"删除失败"}'
;;
enable)
psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = true WHERE name = :'name'" 2>/dev/null
echo '{"ok":true,"action":"enable","name":"'"$NAME"'"}'
psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = true WHERE name = :'name'" 2>/dev/null && \
echo '{"ok":true,"action":"enable","name":"'"$NAME"'"}' || \
echo '{"error":"启用失败"}'
;;
disable)
psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = false WHERE name = :'name'" 2>/dev/null
echo '{"ok":true,"action":"disable","name":"'"$NAME"'"}'
psql "$DB_URL" -v name="$NAME" -c "UPDATE scheduled_tasks SET enabled = false WHERE name = :'name'" 2>/dev/null && \
echo '{"ok":true,"action":"disable","name":"'"$NAME"'"}' || \
echo '{"error":"禁用失败"}'
;;
*)
echo '{"error":"未知操作: '"$ACTION"'(支持: add, delete, enable, disable"}'
+10 -3
View File
@@ -9,9 +9,16 @@ COUNT=$(echo "$input" | python3 -c "import sys,json; print(json.load(sys.stdin).
API_KEY="${TAVILY_API_KEY}"
[ -z "$API_KEY" ] && echo '{"error":"请设置 TAVILY_API_KEY"}' && exit 0
RESP=$(curl -s -X POST "https://api.tavily.com/search" \
-H "Content-Type: application/json" \
-d "{\"api_key\":\"$API_KEY\",\"query\":\"$QUERY\",\"max_results\":$COUNT,\"include_answer\":true}")
export SEARCH_QUERY="$QUERY" SEARCH_COUNT="$COUNT"
RESP=$(python3 -c "
import urllib.request, json, os
key = os.environ['TAVILY_API_KEY']
query = os.environ['SEARCH_QUERY']
count = int(os.environ.get('SEARCH_COUNT', 5))
body = json.dumps({'api_key': key, 'query': query, 'max_results': count, 'include_answer': True}).encode()
req = urllib.request.Request('https://api.tavily.com/search', data=body, headers={'Content-Type': 'application/json'})
print(urllib.request.urlopen(req).read().decode())
" 2>/dev/null)
export RESP QUERY COUNT
python3 -c "
-35
View File
@@ -15,17 +15,11 @@ pub struct AppConfig {
#[derive(Debug, Clone, Deserialize)]
pub struct LlmConfig {
/// LLM 提供商:deepseek | lmstudio
#[serde(default = "default_llm_provider")]
pub provider: String,
/// DeepSeek 配置
#[serde(default)]
pub deepseek: DeepSeekConfig,
/// LM Studio 配置
#[serde(default)]
pub lmstudio: LmStudioConfig,
}
#[derive(Debug, Clone, Deserialize)]
@@ -38,16 +32,6 @@ pub struct DeepSeekConfig {
pub base_url: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct LmStudioConfig {
#[serde(default = "default_lmstudio_base_url")]
pub base_url: String,
#[serde(default = "default_lmstudio_model")]
pub model: String,
#[serde(default)]
pub api_key: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct StorageConfig {
/// 状态文件目录
@@ -76,14 +60,6 @@ fn default_deepseek_base_url() -> String {
env_or_default("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")
}
fn default_lmstudio_base_url() -> String {
env_or_default("LM_STUDIO_BASE_URL", "http://localhost:1234/v1")
}
fn default_lmstudio_model() -> String {
env_or_default("LM_STUDIO_MODEL", "local-model")
}
fn default_state_dir() -> String {
".data/weixin-ilink".to_string()
}
@@ -97,7 +73,6 @@ impl Default for LlmConfig {
Self {
provider: default_llm_provider(),
deepseek: DeepSeekConfig::default(),
lmstudio: LmStudioConfig::default(),
}
}
}
@@ -112,16 +87,6 @@ impl Default for DeepSeekConfig {
}
}
impl Default for LmStudioConfig {
fn default() -> Self {
Self {
base_url: default_lmstudio_base_url(),
model: default_lmstudio_model(),
api_key: String::new(),
}
}
}
impl Default for StorageConfig {
fn default() -> Self {
Self {
-157
View File
@@ -1,157 +0,0 @@
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
use async_trait::async_trait;
use reqwest::Client as HttpClient;
use tokio::sync::mpsc;
/// LM Studio 提供商(OpenAI 兼容接口)
pub struct LmStudioProvider {
http: HttpClient,
base_url: String,
api_key: String,
}
impl LmStudioProvider {
pub fn new() -> Result<Self, String> {
let base_url = std::env::var("LM_STUDIO_BASE_URL")
.unwrap_or_else(|_| "http://localhost:1234/v1".to_string());
// 标准化 base_url: 确保有协议和 /v1 路径
let base_url = normalize_base_url(&base_url);
let api_key = std::env::var("LM_STUDIO_API_KEY").unwrap_or_default();
Ok(Self {
http: HttpClient::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
base_url,
api_key,
})
}
}
#[async_trait]
impl LlmProvider for LmStudioProvider {
fn name(&self) -> &str {
"lmstudio"
}
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String> {
let url = format!("{}/chat/completions", self.base_url);
let body = serde_json::json!({
"model": config.model,
"messages": messages,
"temperature": config.temperature,
"max_tokens": config.max_tokens,
"stream": true,
});
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
let http = self.http.clone();
let api_key = self.api_key.clone();
tokio::spawn(async move {
if let Err(e) = stream_openai(http, &url, &api_key, body, tx.clone()).await {
let _ = tx.send(StreamChunk::Error(e)).await;
}
});
Ok(rx)
}
}
/// 与 DeepSeek 公用相同的 SSE 流式处理,但去除 thinking 相关字段
async fn stream_openai(
http: HttpClient,
url: &str,
api_key: &str,
body: serde_json::Value,
tx: StreamSender,
) -> Result<(), String> {
let mut req = http.post(url).header("Content-Type", "application/json");
if !api_key.is_empty() {
req = req.header("Authorization", format!("Bearer {}", api_key));
}
let response = req
.json(&body)
.send()
.await
.map_err(|e| format!("请求失败: {}", e))?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
return Err(format!("HTTP {}: {}", status, text));
}
let mut full_text = String::new();
let mut buf = String::new();
let mut stream = response.bytes_stream();
use futures_util::StreamExt;
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
let text = String::from_utf8_lossy(&chunk);
buf.push_str(&text);
while let Some(line_end) = buf.find('\n') {
let line = buf[..line_end].to_string();
buf = buf[line_end + 1..].to_string();
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
// LM Studio 使用标准 OpenAI SSE 格式
if let Some(parsed) = parse_chat_chunk(trimmed) {
match parsed {
ParsedChunk::Text(t) => {
full_text.push_str(&t);
let _ = tx.send(StreamChunk::Text(t)).await;
}
ParsedChunk::Reasoning(r) => {
let _ = tx.send(StreamChunk::Reasoning(r)).await;
}
_ => {}
}
}
}
}
let _ = tx
.send(StreamChunk::Done {
text: full_text,
reasoning: None,
tool_calls: None,
usage: None,
})
.await;
Ok(())
}
/// 确保 base_url 有协议前缀和 /v1 路径
fn normalize_base_url(raw: &str) -> String {
let raw = raw.trim();
let with_protocol = if raw.starts_with("http://") || raw.starts_with("https://") {
raw.to_string()
} else {
format!("http://{}", raw)
};
// 去掉末尾 /
let trimmed = with_protocol.trim_end_matches('/');
// 确保路径以 /v1 结尾
if trimmed.ends_with("/v1") {
trimmed.to_string()
} else {
format!("{}/v1", trimmed)
}
}
-1
View File
@@ -1,6 +1,5 @@
pub mod conversation;
pub mod deepseek;
pub mod lmstudio;
pub mod provider;
pub mod types;
+2 -10
View File
@@ -27,16 +27,8 @@ pub trait LlmProvider: Send + Sync {
pub type BoxedProvider = Arc<dyn LlmProvider>;
/// 从配置创建恰当的提供商
pub fn create_provider(config: &ConversationConfig) -> Result<BoxedProvider, String> {
let provider = std::env::var("LLM_PROVIDER")
.unwrap_or_else(|_| "deepseek".to_string())
.to_lowercase();
match provider.as_str() {
"deepseek" => Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?)),
"lmstudio" => Ok(Arc::new(super::lmstudio::LmStudioProvider::new()?)),
_ => Err(format!("不支持的 LLM 提供商: {}", provider)),
}
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?))
}
// ─── 内部:SSE 解析工具 ───
+1 -6
View File
@@ -230,12 +230,7 @@ async fn cmd_listen(
let conversation = if enable_llm {
let sys_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT")
.unwrap_or_else(|_| DEFAULT_SYSTEM_PROMPT.to_string());
let provider = std::env::var("LLM_PROVIDER").unwrap_or_else(|_| "deepseek".to_string());
let model = if provider == "lmstudio" {
config.llm.lmstudio.model.clone()
} else {
config.llm.deepseek.model.clone()
};
let model = config.llm.deepseek.model.clone();
let mut cfg = ConversationConfig {
system_prompt: sys_prompt,