use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::{CollectorError, CollectorResult, ExceptionRecord};
#[async_trait]
pub trait LlmChannel: Send + Sync {
async fn send(&self, system_prompt: &str, user_prompt: &str) -> CollectorResult<String>;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassificationResult {
pub signature: String,
pub action: ClassificationAction,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ClassificationAction {
Duplicate {
issue_number: u32,
},
NewIssue {
title: String,
body: String,
},
}
pub struct LlmClassifier {
channel: Arc<dyn LlmChannel>,
timeout_secs: u64,
}
impl std::fmt::Debug for LlmClassifier {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LlmClassifier")
.field("timeout_secs", &self.timeout_secs)
.finish_non_exhaustive()
}
}
#[derive(Debug, Serialize)]
struct ExceptionInput {
signature: String,
component: String,
message: String,
stacktrace: String,
}
#[derive(Debug, Serialize)]
struct ExistingIssue {
number: u32,
title: String,
body: String,
}
impl LlmClassifier {
#[must_use]
pub fn new(channel: Arc<dyn LlmChannel>) -> Self {
Self {
channel,
timeout_secs: 30,
}
}
pub async fn classify_batch(
&self,
records: &[ExceptionRecord],
existing_issues: &[ExistingIssueData],
) -> CollectorResult<Vec<ClassificationResult>> {
let system_prompt = build_system_prompt();
let user_prompt = build_user_prompt(records, existing_issues);
let response = self.channel.send(&system_prompt, &user_prompt).await;
let text = match response {
Ok(text) => text,
Err(e) => {
let results = fallback_classify(records);
return if results.is_empty() {
Err(e)
} else {
Ok(results)
};
}
};
let Ok(results) = parse_llm_response(&text) else {
let retry = self.channel.send(&system_prompt, &user_prompt).await?;
return parse_llm_response(&retry).map_err(|_| CollectorError::Reporter {
reason: "LLM response not valid JSON after retry".to_string(),
});
};
Ok(results)
}
}
#[cfg(feature = "http-llm")]
#[must_use]
pub fn default_classifier() -> LlmClassifier {
use crate::channel::HttpLlmChannel;
LlmClassifier::new(Arc::new(HttpLlmChannel::from_env()))
}
#[derive(Debug, Clone)]
pub struct ExistingIssueData {
pub number: u32,
pub title: String,
pub body: String,
}
const SYSTEM_PROMPT: &str =
"\
You are a TokenFleet exception classification assistant. Your job is to determine whether each new \
exception matches an existing open issue (semantic duplicate) or is truly new. For \
duplicates: reference the existing issue number. For new issues: generate a concise title \
and a structured body with: exception summary, component, module path, location, and stack \
trace. Respond ONLY with a JSON array of objects, each with keys: \"signature\", \"action\" \
(\"duplicate\" or \"new_issue\"), \"issue_number\" (for duplicates), \"title\" and \"body\" \
(for new issues).";
fn build_system_prompt() -> String {
SYSTEM_PROMPT.to_string()
}
fn build_user_prompt(records: &[ExceptionRecord], existing_issues: &[ExistingIssueData]) -> String {
let exceptions: Vec<ExceptionInput> = records
.iter()
.map(|r| ExceptionInput {
signature: r.dedup_signature.clone(),
component: r.component.clone(),
message: r.message.clone(),
stacktrace: r.stacktrace.clone(),
})
.collect();
let issues: Vec<ExistingIssue> = existing_issues
.iter()
.map(|i| ExistingIssue {
number: i.number,
title: i.title.clone(),
body: i.body.clone(),
})
.collect();
let exceptions_json = serde_json::to_string(&exceptions).unwrap_or_default();
let issues_json = serde_json::to_string(&issues).unwrap_or_default();
format!(
"## Existing Open Issues (JSON)\n{issues_json}\n\n## New Exceptions \
(JSON)\n{exceptions_json}"
)
}
fn parse_llm_response(text: &str) -> CollectorResult<Vec<ClassificationResult>> {
let json_text = text
.trim()
.strip_prefix("```json")
.and_then(|s| s.strip_suffix("```"))
.map_or(text.trim(), str::trim);
let results: Vec<ClassificationResult> =
serde_json::from_str(json_text).map_err(CollectorError::SerdeJson)?;
Ok(results)
}
fn fallback_classify(records: &[ExceptionRecord]) -> Vec<ClassificationResult> {
records
.iter()
.map(|r| ClassificationResult {
signature: r.dedup_signature.clone(),
action: ClassificationAction::NewIssue {
title: format!(
"[{}] {}: {}",
r.component,
r.kind,
&r.message[..r.message.len().min(80)]
),
body: build_fallback_body(r),
},
})
.collect()
}
fn build_fallback_body(record: &ExceptionRecord) -> String {
format!(
"## 异常摘要\n{message}\n\n## 堆栈\n```\n{stacktrace}\n```\n\n## 影响范围\n- 组件: \
{component}\n- 模块: {module}\n- 位置: {location}\n- 时间: {timestamp}\n\n---\n> 由 \
TokenFleet Exception Collector 自动生成(降级模式)",
message = record.message,
stacktrace = record.stacktrace,
component = record.component,
module = record.module_path.as_deref().unwrap_or("unknown"),
location = record.location.as_deref().unwrap_or("unknown"),
timestamp = record.timestamp.format("%Y-%m-%d %H:%M:%S UTC"),
)
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::unwrap_in_result,
clippy::expect_used,
clippy::panic,
clippy::pedantic,
clippy::indexing_slicing,
reason = "test module relaxes production lint strictness"
)]
mod tests {
use super::*;
use crate::{ExceptionKind, ExceptionRecord};
struct MockChannel {
response: String,
should_fail: bool,
}
#[async_trait]
impl LlmChannel for MockChannel {
async fn send(&self, _system: &str, _user: &str) -> CollectorResult<String> {
if self.should_fail {
Err(CollectorError::Reporter {
reason: "mock channel failure".to_string(),
})
} else {
Ok(self.response.clone())
}
}
}
fn make_record(signature: &str, message: &str) -> ExceptionRecord {
let mut record =
ExceptionRecord::new("test-comp", ExceptionKind::ErrorLog, message, "stack trace");
record.dedup_signature = signature.to_string();
record
}
fn duplicate_response(sig: &str, issue: u32) -> String {
format!(r#"[{{"signature":"{sig}","action":{{"issue_number":{issue}}}}}]"#)
}
fn new_issue_response(sig: &str) -> String {
format!(r#"[{{"signature":"{sig}","action":{{"title":"Test Issue","body":"Test body"}}}}]"#)
}
#[test]
fn test_should_detect_duplicate_via_llm() {
let response = duplicate_response("abc123", 42);
let channel = Arc::new(MockChannel {
response,
should_fail: false,
});
let classifier = LlmClassifier::new(channel);
let records = vec![make_record("abc123", "test error")];
let existing = vec![ExistingIssueData {
number: 42,
title: "Existing issue".to_string(),
body: "Old body".to_string(),
}];
let rt = tokio::runtime::Runtime::new().unwrap();
let results = rt
.block_on(classifier.classify_batch(&records, &existing))
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].signature, "abc123");
match &results[0].action {
ClassificationAction::Duplicate { issue_number } => {
assert_eq!(*issue_number, 42);
}
_ => panic!("expected Duplicate"),
}
}
#[test]
fn test_should_generate_new_issue_via_llm() {
let response = new_issue_response("xyz789");
let channel = Arc::new(MockChannel {
response,
should_fail: false,
});
let classifier = LlmClassifier::new(channel);
let records = vec![make_record("xyz789", "new error")];
let rt = tokio::runtime::Runtime::new().unwrap();
let results = rt
.block_on(classifier.classify_batch(&records, &[]))
.unwrap();
assert_eq!(results.len(), 1);
match &results[0].action {
ClassificationAction::NewIssue { title, body } => {
assert_eq!(title, "Test Issue");
assert_eq!(body, "Test body");
}
_ => panic!("expected NewIssue"),
}
}
#[test]
fn test_should_fallback_on_llm_failure() {
let channel = Arc::new(MockChannel {
response: String::new(),
should_fail: true,
});
let classifier = LlmClassifier::new(channel);
let records = vec![make_record("sig1", "test error")];
let rt = tokio::runtime::Runtime::new().unwrap();
let results = rt
.block_on(classifier.classify_batch(&records, &[]))
.unwrap();
assert_eq!(results.len(), 1);
match &results[0].action {
ClassificationAction::NewIssue { title, .. } => {
assert!(title.contains("test error"));
}
_ => panic!("expected fallback NewIssue"),
}
}
#[test]
fn test_should_fallback_on_non_json_response() {
let channel = Arc::new(MockChannel {
response: "not valid json at all".to_string(),
should_fail: false,
});
let classifier = LlmClassifier::new(channel);
let records = vec![make_record("sig1", "test")];
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(classifier.classify_batch(&records, &[]));
assert!(result.is_err());
}
#[test]
fn test_parse_llm_response_should_handle_json_with_code_fence() {
let text = r###"```json
[{"signature":"abc","action":{"title":"T","body":"B"}}]
```"###;
let results = parse_llm_response(text).unwrap();
assert_eq!(results.len(), 1);
}
#[test]
fn test_parse_llm_response_should_handle_plain_json() {
let text = r#"[{"signature":"abc","action":{"title":"T","body":"B"}}]"#;
let results = parse_llm_response(text).unwrap();
assert_eq!(results.len(), 1);
}
#[test]
fn test_parse_llm_response_should_error_on_invalid() {
let text = "not json";
assert!(parse_llm_response(text).is_err());
}
#[test]
fn test_build_fallback_body_should_include_component() {
let mut record = make_record("sig", "test message");
record.module_path = Some("mod".to_string());
record.location = Some("file.rs:1".to_string());
let body = build_fallback_body(&record);
assert!(body.contains("test message"));
assert!(body.contains("test-comp"));
assert!(body.contains("降级模式"));
}
#[test]
fn test_system_prompt_should_be_non_empty() {
let prompt = build_system_prompt();
assert!(!prompt.is_empty());
assert!(prompt.contains("TokenFleet"));
}
#[test]
fn test_user_prompt_should_include_records() {
let records = vec![make_record("sig-a", "error one")];
let issues = vec![ExistingIssueData {
number: 1,
title: "Old".to_string(),
body: "Old body".to_string(),
}];
let prompt = build_user_prompt(&records, &issues);
assert!(prompt.contains("sig-a"));
assert!(prompt.contains("error one"));
assert!(prompt.contains("Old"));
assert!(prompt.contains("Existing Open Issues"));
}
#[test]
fn test_classification_action_serde_duplicate() {
let action = ClassificationAction::Duplicate { issue_number: 5 };
let json = serde_json::to_string(&action).unwrap();
assert!(json.contains("issue_number"));
assert!(json.contains("5"));
}
#[test]
fn test_classification_action_serde_new_issue() {
let action = ClassificationAction::NewIssue {
title: "T".to_string(),
body: "B".to_string(),
};
let json = serde_json::to_string(&action).unwrap();
assert!(json.contains("title"));
assert!(json.contains("B"));
}
}