subx_cli/services/ai/
prompts.rs

1use crate::Result;
2use crate::error::SubXError;
3use crate::services::ai::{AnalysisRequest, ConfidenceScore, MatchResult, VerificationRequest};
4use serde_json;
5
6/// Prompt builder trait for AI providers.
7pub trait PromptBuilder {
8    /// Build analysis prompt.
9    fn build_analysis_prompt(&self, request: &AnalysisRequest) -> String {
10        build_analysis_prompt_base(request)
11    }
12
13    /// Build verification prompt.
14    fn build_verification_prompt(&self, request: &VerificationRequest) -> String {
15        build_verification_prompt_base(request)
16    }
17
18    /// System message for analysis prompt.
19    fn get_analysis_system_message() -> &'static str {
20        "You are a professional subtitle matching assistant that can analyze the correspondence between video and subtitle files."
21    }
22
23    /// System message for verification prompt.
24    fn get_verification_system_message() -> &'static str {
25        "Please evaluate the confidence level of subtitle matching and provide a score between 0-1."
26    }
27}
28
29/// Response parsing trait for AI providers.
30pub trait ResponseParser {
31    /// Parse match result.
32    fn parse_match_result(&self, response: &str) -> Result<MatchResult> {
33        parse_match_result_base(response)
34    }
35
36    /// Parse confidence score.
37    fn parse_confidence_score(&self, response: &str) -> Result<ConfidenceScore> {
38        parse_confidence_score_base(response)
39    }
40}
41
42/// Build analysis prompt for AI providers.
43pub fn build_analysis_prompt_base(request: &AnalysisRequest) -> String {
44    let mut prompt = String::new();
45    prompt.push_str(
46        "Please analyze the matching relationship between the following video and subtitle files. Each file has a unique ID that you must use in your response.\n\n",
47    );
48    prompt.push_str("Video files:\n");
49    for video in &request.video_files {
50        prompt.push_str(&format!("- {}\n", video));
51    }
52    prompt.push_str("\nSubtitle files:\n");
53    for subtitle in &request.subtitle_files {
54        prompt.push_str(&format!("- {}\n", subtitle));
55    }
56    if !request.content_samples.is_empty() {
57        prompt.push_str("\nSubtitle content preview:\n");
58        for sample in &request.content_samples {
59            prompt.push_str(&format!("File: {}\n", sample.filename));
60            prompt.push_str(&format!("Content: {}\n\n", sample.content_preview));
61        }
62    }
63    prompt.push_str(
64        "Please provide matching suggestions based on filename patterns, content similarity, and other factors.\n\
65Response format must be JSON using the file IDs:\n\
66{\n\
67  \"matches\": [\n\
68    {\n\
69      \"video_file_id\": \"file_abc123456789abcd\",\n\
70      \"subtitle_file_id\": \"file_def456789abcdef0\",\n\
71      \"confidence\": 0.95,\n\
72      \"match_factors\": [\"filename_similarity\", \"content_correlation\"]\n\
73    }\n\
74  ],\n\
75  \"confidence\": 0.9,\n\
76  \"reasoning\": \"Explanation for the matching decisions\"\n\
77}",
78    );
79    prompt
80}
81
82/// Parse matching results from AI response.
83pub fn parse_match_result_base(response: &str) -> Result<MatchResult> {
84    let json_start = response.find('{').unwrap_or(0);
85    let json_end = response.rfind('}').map(|i| i + 1).unwrap_or(response.len());
86    let json_str = &response[json_start..json_end];
87    serde_json::from_str(json_str)
88        .map_err(|e| SubXError::AiService(format!("AI response parsing failed: {}", e)))
89}
90
91/// Build verification prompt for AI providers.
92pub fn build_verification_prompt_base(request: &VerificationRequest) -> String {
93    let mut prompt = String::new();
94    prompt.push_str(
95        "Please evaluate the confidence level based on the following matching information:\n",
96    );
97    prompt.push_str(&format!("Video file: {}\n", request.video_file));
98    prompt.push_str(&format!("Subtitle file: {}\n", request.subtitle_file));
99    prompt.push_str("Matching factors:\n");
100    for factor in &request.match_factors {
101        prompt.push_str(&format!("- {}\n", factor));
102    }
103    prompt.push_str(
104        "\nPlease respond in JSON format as follows:\n{\"score\": 0.9,\"factors\": [\"...\"]}",
105    );
106    prompt
107}
108
109/// Parse confidence score from AI response.
110pub fn parse_confidence_score_base(response: &str) -> Result<ConfidenceScore> {
111    let json_start = response.find('{').unwrap_or(0);
112    let json_end = response.rfind('}').map(|i| i + 1).unwrap_or(response.len());
113    let json_str = &response[json_start..json_end];
114    serde_json::from_str(json_str)
115        .map_err(|e| SubXError::AiService(format!("AI confidence parsing failed: {}", e)))
116}
117
118#[cfg(test)]
119mod tests {
120
121    use crate::services::ai::prompts::{PromptBuilder, ResponseParser};
122    use crate::services::ai::{AnalysisRequest, OpenAIClient};
123
124    #[test]
125    fn test_ai_prompt_with_file_ids_english() {
126        let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
127        let request = AnalysisRequest {
128            video_files: vec!["ID:file_abc123456789abcd | Name:movie.mkv | Path:movie.mkv".into()],
129            subtitle_files: vec![
130                "ID:file_def456789abcdef0 | Name:movie.srt | Path:movie.srt".into(),
131            ],
132            content_samples: vec![],
133        };
134
135        let prompt = client.build_analysis_prompt(&request);
136
137        assert!(prompt.contains("ID:file_abc123456789abcd"));
138        assert!(prompt.contains("video_file_id"));
139        assert!(prompt.contains("subtitle_file_id"));
140        assert!(prompt.contains("Please analyze the matching"));
141        assert!(prompt.contains("unique ID"));
142        assert!(prompt.contains("Response format must be JSON"));
143        assert!(!prompt.contains("請分析"));
144        assert!(!prompt.contains("影片檔案"));
145        assert!(!prompt.contains("字幕檔案"));
146    }
147
148    #[test]
149    fn test_parse_match_result_with_ids() {
150        let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
151        let json_resp = r#"{
152            "matches": [{
153                "video_file_id": "file_abc123456789abcd",
154                "subtitle_file_id": "file_def456789abcdef0",
155                "confidence": 0.95,
156                "match_factors": ["filename_similarity"]
157            }],
158            "confidence": 0.9,
159            "reasoning": "Strong match based on filename patterns"
160        }"#;
161
162        let result = client.parse_match_result(json_resp).unwrap();
163        assert_eq!(result.matches.len(), 1);
164        assert_eq!(result.matches[0].video_file_id, "file_abc123456789abcd");
165        assert_eq!(result.matches[0].subtitle_file_id, "file_def456789abcdef0");
166        assert_eq!(result.matches[0].confidence, 0.95);
167        assert_eq!(result.matches[0].match_factors[0], "filename_similarity");
168    }
169
170    #[test]
171    fn test_ai_prompt_structure_consistency() {
172        let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
173        let request = AnalysisRequest {
174            video_files: vec![
175                "ID:file_video1 | Name:video1.mkv | Path:season1/video1.mkv".into(),
176                "ID:file_video2 | Name:video2.mkv | Path:season1/video2.mkv".into(),
177            ],
178            subtitle_files: vec![
179                "ID:file_sub1 | Name:sub1.srt | Path:season1/sub1.srt".into(),
180                "ID:file_sub2 | Name:sub2.srt | Path:season1/sub2.srt".into(),
181            ],
182            content_samples: vec![],
183        };
184
185        let prompt = client.build_analysis_prompt(&request);
186
187        assert!(prompt.contains("ID:file_video1"));
188        assert!(prompt.contains("ID:file_video2"));
189        assert!(prompt.contains("ID:file_sub1"));
190        assert!(prompt.contains("ID:file_sub2"));
191        assert!(prompt.contains("Video files:"));
192        assert!(prompt.contains("Subtitle files:"));
193        assert!(prompt.contains("Response format must be JSON"));
194    }
195
196    #[test]
197    fn test_parse_confidence_score() {
198        let client = OpenAIClient::new("test_key".into(), "gpt-4.1".into(), 0.1, 1000, 0, 0);
199        let json_resp = r#"{
200            "score": 0.88,
201            "factors": ["filename_similarity", "content_correlation"]
202        }"#;
203
204        let result = client.parse_confidence_score(json_resp).unwrap();
205        assert_eq!(result.score, 0.88);
206        assert_eq!(
207            result.factors,
208            vec![
209                "filename_similarity".to_string(),
210                "content_correlation".to_string()
211            ]
212        );
213    }
214}