- 修复Vertex AI Search配置,移除不支持的API字段 - 优化system prompt以更好地利用检索信息 - 添加查询增强功能,通过关键词扩展提高检索效果 - 新增RagConfigOptimizer工具类,支持多种优化场景 - 新增RagConfigManager组件,提供可视化配置管理 - 保留客户端配置字段用于未来扩展 - 添加详细的使用示例和文档 主要改进: 1. 解决了API 400错误问题 2. 通过查询优化间接增加检索相关性 3. 提供了完整的配置管理解决方案 4. 支持场景化的RAG配置优化
1693 lines
61 KiB
Rust
1693 lines
61 KiB
Rust
use anyhow::{Result, anyhow};
|
||
use serde::{Deserialize, Serialize};
|
||
use base64::prelude::*;
|
||
|
||
use std::path::Path;
|
||
use std::time::{SystemTime, UNIX_EPOCH};
|
||
use tokio::fs;
|
||
use reqwest::multipart;
|
||
|
||
// 导入容错JSON解析器
|
||
use crate::infrastructure::tolerant_json_parser::{TolerantJsonParser, ParserConfig, RecoveryStrategy};
|
||
use std::sync::{Arc, Mutex};
|
||
|
||
// 导入会话管理相关模块
|
||
use crate::data::models::conversation::{
|
||
ConversationMessage, MessageRole, MessageContent, ConversationHistoryQuery,
|
||
AddMessageRequest,
|
||
};
|
||
use crate::data::repositories::conversation_repository::ConversationRepository;
|
||
|
||
/// Gemini API配置
|
||
#[derive(Debug, Clone)]
|
||
pub struct GeminiConfig {
|
||
pub base_url: String,
|
||
pub bearer_token: String,
|
||
pub timeout: u64,
|
||
pub model_name: String,
|
||
pub max_retries: u32,
|
||
pub retry_delay: u64,
|
||
pub temperature: f32,
|
||
pub max_tokens: u32,
|
||
pub cloudflare_project_id: String,
|
||
pub cloudflare_gateway_id: String,
|
||
pub google_project_id: String,
|
||
pub regions: Vec<String>,
|
||
}
|
||
|
||
impl Default for GeminiConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
base_url: "https://bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run".to_string(),
|
||
bearer_token: "bowong7777".to_string(),
|
||
timeout: 120,
|
||
model_name: "gemini-2.5-flash".to_string(),
|
||
max_retries: 3,
|
||
retry_delay: 2,
|
||
temperature: 0.7,
|
||
max_tokens: 1000 * 8 * 8,
|
||
cloudflare_project_id: "67720b647ff2b55cf37ba3ef9e677083".to_string(),
|
||
cloudflare_gateway_id: "bowong-dev".to_string(),
|
||
google_project_id: "gen-lang-client-0413414134".to_string(),
|
||
regions: vec![
|
||
"us-central1".to_string(),
|
||
"us-east1".to_string(),
|
||
"europe-west1".to_string(),
|
||
],
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Gemini访问令牌响应
|
||
#[derive(Debug, Deserialize)]
|
||
struct TokenResponse {
|
||
access_token: String,
|
||
expires_in: u64,
|
||
}
|
||
|
||
/// 客户端配置
|
||
#[derive(Debug)]
|
||
struct ClientConfig {
|
||
gateway_url: String,
|
||
headers: std::collections::HashMap<String, String>,
|
||
}
|
||
|
||
/// Gemini上传响应
|
||
#[derive(Debug, Deserialize)]
|
||
struct UploadResponse {
|
||
file_uri: Option<String>,
|
||
urn: Option<String>,
|
||
name: Option<String>,
|
||
}
|
||
|
||
/// Gemini内容生成请求
|
||
#[derive(Debug, Serialize)]
|
||
pub struct GenerateContentRequest {
|
||
pub contents: Vec<ContentPart>,
|
||
#[serde(rename = "generationConfig")]
|
||
pub generation_config: GenerationConfig,
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
pub struct ContentPart {
|
||
pub role: String,
|
||
pub parts: Vec<Part>,
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
#[serde(untagged)]
|
||
pub enum Part {
|
||
Text { text: String },
|
||
FileData {
|
||
#[serde(rename = "fileData")]
|
||
file_data: FileData
|
||
},
|
||
InlineData {
|
||
#[serde(rename = "inlineData")]
|
||
inline_data: InlineData
|
||
},
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
pub struct FileData {
|
||
#[serde(rename = "mimeType")]
|
||
pub mime_type: String,
|
||
#[serde(rename = "fileUri")]
|
||
pub file_uri: String,
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
pub struct InlineData {
|
||
#[serde(rename = "mimeType")]
|
||
pub mime_type: String,
|
||
pub data: String,
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
pub struct GenerationConfig {
|
||
pub temperature: f32,
|
||
#[serde(rename = "topK")]
|
||
pub top_k: u32,
|
||
#[serde(rename = "topP")]
|
||
pub top_p: f32,
|
||
#[serde(rename = "maxOutputTokens")]
|
||
pub max_output_tokens: u32,
|
||
}
|
||
|
||
/// Gemini响应结构
|
||
#[derive(Debug, Deserialize)]
|
||
struct GeminiResponse {
|
||
candidates: Vec<Candidate>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct Candidate {
|
||
content: Content,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct Content {
|
||
parts: Vec<ResponsePart>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct ResponsePart {
|
||
text: String,
|
||
}
|
||
|
||
/// RAG Grounding 配置
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct RagGroundingConfig {
|
||
pub project_id: String,
|
||
pub location: String,
|
||
pub data_store_id: String,
|
||
pub model_id: String,
|
||
pub temperature: f32,
|
||
pub max_output_tokens: u32,
|
||
pub system_prompt: Option<String>,
|
||
/// 搜索过滤器 (Vertex AI Search支持的字段)
|
||
pub search_filter: Option<String>,
|
||
/// 最大检索结果数量 (用于客户端逻辑,不发送给API)
|
||
pub max_retrieval_results: Option<u32>,
|
||
/// 相关性阈值 (用于客户端逻辑,不发送给API)
|
||
pub relevance_threshold: Option<f32>,
|
||
/// 是否包含摘要 (用于客户端逻辑,不发送给API)
|
||
pub include_summary: Option<bool>,
|
||
}
|
||
|
||
impl Default for RagGroundingConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
project_id: "gen-lang-client-0413414134".to_string(),
|
||
location: "global".to_string(),
|
||
data_store_id: "searchable-model-images_1752827560253".to_string(), // 使用存在的数据存储
|
||
model_id: "gemini-2.5-flash".to_string(),
|
||
temperature: 1.0,
|
||
max_output_tokens: 60000,
|
||
system_prompt: Some("你是一个短视频情景穿搭分析专家。请仔细分析用户的查询,充分利用检索到的所有相关信息,包括:1)详细分析每个检索结果的内容;2)综合多个来源的信息提供全面的回答;3)明确引用具体的数据来源和依据;4)如果检索结果不足,请说明需要更多哪方面的信息;5)提供具体可行的穿搭建议和搭配方案。".to_string()),
|
||
search_filter: None, // 暂不使用过滤器
|
||
// 以下字段用于客户端逻辑,不发送给API
|
||
max_retrieval_results: Some(20), // 客户端参考值
|
||
relevance_threshold: Some(0.3), // 客户端参考值
|
||
include_summary: Some(true), // 客户端参考值
|
||
}
|
||
}
|
||
}
|
||
|
||
/// RAG Grounding 请求
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct RagGroundingRequest {
|
||
pub user_input: String,
|
||
pub config: Option<RagGroundingConfig>,
|
||
pub session_id: Option<String>,
|
||
pub include_history: Option<bool>,
|
||
pub max_history_messages: Option<u32>,
|
||
pub system_prompt: Option<String>,
|
||
}
|
||
|
||
/// RAG Grounding 响应
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct RagGroundingResponse {
|
||
pub answer: String,
|
||
pub grounding_metadata: Option<GroundingMetadata>,
|
||
pub response_time_ms: u64,
|
||
pub model_used: String,
|
||
pub session_id: Option<String>,
|
||
pub message_id: Option<String>,
|
||
pub conversation_context: Option<ConversationContext>,
|
||
}
|
||
|
||
/// Grounding 元数据
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct GroundingMetadata {
|
||
pub sources: Vec<GroundingSource>,
|
||
pub search_queries: Vec<String>,
|
||
pub grounding_supports: Option<Vec<GroundingSupport>>,
|
||
}
|
||
|
||
/// Grounding 来源
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct GroundingSource {
|
||
pub title: String,
|
||
pub uri: Option<String>,
|
||
pub content: Option<serde_json::Value>,
|
||
}
|
||
|
||
/// Grounding 支持信息
|
||
/// 用于关联文字片段与来源信息
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct GroundingSupport {
|
||
#[serde(rename = "groundingChunkIndices")]
|
||
pub grounding_chunk_indices: Vec<usize>,
|
||
pub segment: GroundingSegment,
|
||
}
|
||
|
||
/// Grounding 文字片段
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct GroundingSegment {
|
||
#[serde(rename = "startIndex")]
|
||
pub start_index: usize,
|
||
#[serde(rename = "endIndex")]
|
||
pub end_index: usize,
|
||
pub text: String,
|
||
}
|
||
|
||
/// 对话上下文
|
||
#[derive(Debug, Serialize, Deserialize)]
|
||
pub struct ConversationContext {
|
||
pub total_messages: u32,
|
||
pub history_included: bool,
|
||
pub context_length: u32,
|
||
}
|
||
|
||
/// Vertex AI Search 工具配置
|
||
#[derive(Debug, Serialize)]
|
||
struct VertexAISearchTool {
|
||
retrieval: VertexAIRetrieval,
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
struct VertexAIRetrieval {
|
||
#[serde(rename = "vertexAiSearch")]
|
||
vertex_ai_search: VertexAISearchConfig,
|
||
#[serde(rename = "disableAttribution", skip_serializing_if = "Option::is_none")]
|
||
disable_attribution: Option<bool>,
|
||
}
|
||
|
||
#[derive(Debug, Serialize)]
|
||
struct VertexAISearchConfig {
|
||
datastore: String,
|
||
/// 搜索过滤器 (支持的字段)
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
filter: Option<String>,
|
||
}
|
||
|
||
/// Gemini API服务
|
||
/// 遵循 Tauri 开发规范的基础设施层设计模式
|
||
#[derive(Clone)]
|
||
pub struct GeminiService {
|
||
config: GeminiConfig,
|
||
client: reqwest::Client,
|
||
access_token: Option<String>,
|
||
token_expires_at: Option<u64>,
|
||
json_parser: Arc<Mutex<TolerantJsonParser>>,
|
||
}
|
||
|
||
impl GeminiService {
|
||
/// 创建新的Gemini服务实例
|
||
pub fn new(config: Option<GeminiConfig>) -> Result<Self> {
|
||
let client = reqwest::Client::builder()
|
||
.timeout(std::time::Duration::from_secs(config.as_ref().map(|c| c.timeout).unwrap_or(120)))
|
||
.build()
|
||
.expect("Failed to create HTTP client");
|
||
|
||
// 创建容错JSON解析器
|
||
let json_parser = TolerantJsonParser::new(Some(ParserConfig::default()))?;
|
||
|
||
Ok(Self {
|
||
config: config.unwrap_or_default(),
|
||
client,
|
||
access_token: None,
|
||
token_expires_at: None,
|
||
json_parser: Arc::new(Mutex::new(json_parser)),
|
||
})
|
||
}
|
||
|
||
/// 获取Google访问令牌
|
||
async fn get_access_token(&mut self) -> Result<String> {
|
||
// 检查缓存的令牌是否仍然有效
|
||
let current_time = SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)?
|
||
.as_secs();
|
||
|
||
if let (Some(token), Some(expires_at)) = (&self.access_token, self.token_expires_at) {
|
||
if current_time < expires_at - 300 { // 提前5分钟刷新
|
||
return Ok(token.clone());
|
||
}
|
||
}
|
||
|
||
// 获取新的访问令牌
|
||
let url = format!("{}/google/access-token", self.config.base_url);
|
||
|
||
let response = self.client
|
||
.get(&url)
|
||
.header("Authorization", format!("Bearer {}", self.config.bearer_token))
|
||
.send()
|
||
.await?;
|
||
|
||
let status = response.status();
|
||
if !status.is_success() {
|
||
let error_body = response.text().await.unwrap_or_default();
|
||
return Err(anyhow!("获取访问令牌失败: {} - {}", status, error_body));
|
||
}
|
||
|
||
let response_text = response.text().await?;
|
||
let token_response: TokenResponse = serde_json::from_str(&response_text)
|
||
.map_err(|e| anyhow!("解析令牌响应失败: {} - 响应内容: {}", e, response_text))?;
|
||
|
||
// 缓存令牌
|
||
self.access_token = Some(token_response.access_token.clone());
|
||
self.token_expires_at = Some(current_time + token_response.expires_in);
|
||
|
||
Ok(token_response.access_token)
|
||
}
|
||
|
||
/// 创建Gemini客户端配置
|
||
fn create_gemini_client(&self, access_token: &str) -> ClientConfig {
|
||
let mut headers = std::collections::HashMap::new();
|
||
headers.insert("Authorization".to_string(), format!("Bearer {}", access_token));
|
||
headers.insert("Content-Type".to_string(), "application/json".to_string());
|
||
|
||
// 使用第一个区域作为默认区域
|
||
let region = self.config.regions.first()
|
||
.unwrap_or(&"us-central1".to_string())
|
||
.clone();
|
||
|
||
// 构建Cloudflare Gateway URL
|
||
let gateway_url = format!(
|
||
"https://gateway.ai.cloudflare.com/v1/{}/{}/google-vertex-ai/v1/projects/{}/locations/{}/publishers/google/models",
|
||
self.config.cloudflare_project_id,
|
||
self.config.cloudflare_gateway_id,
|
||
self.config.google_project_id,
|
||
region
|
||
);
|
||
|
||
ClientConfig {
|
||
gateway_url,
|
||
headers,
|
||
}
|
||
}
|
||
|
||
/// 上传视频文件到Gemini
|
||
pub async fn upload_video_file(&mut self, video_path: &str) -> Result<String> {
|
||
println!("📤 正在上传视频到Gemini: {}", video_path);
|
||
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
|
||
// 读取视频文件
|
||
let video_data = fs::read(video_path).await
|
||
.map_err(|e| anyhow!("读取视频文件失败: {} - {}", video_path, e))?;
|
||
|
||
let file_name = Path::new(video_path)
|
||
.file_name()
|
||
.and_then(|n| n.to_str())
|
||
.unwrap_or("video.mp4");
|
||
|
||
// 创建multipart表单
|
||
let form = multipart::Form::new()
|
||
.part("file", multipart::Part::bytes(video_data)
|
||
.file_name(file_name.to_string())
|
||
.mime_str("video/mp4")?);
|
||
|
||
// 上传到Vertex AI
|
||
let upload_url = format!("{}/google/vertex-ai/upload", self.config.base_url);
|
||
let query_params = [
|
||
("bucket", "dy-media-storage"),
|
||
("prefix", "video-analysis")
|
||
];
|
||
|
||
let response = self.client
|
||
.post(&upload_url)
|
||
.header("Authorization", format!("Bearer {}", access_token))
|
||
.header("x-google-api-key", &access_token)
|
||
.query(&query_params)
|
||
.multipart(form)
|
||
.send()
|
||
.await?;
|
||
|
||
let status = response.status();
|
||
if !status.is_success() {
|
||
let error_text = response.text().await.unwrap_or_default();
|
||
return Err(anyhow!("上传视频失败: {} - {}", status, error_text));
|
||
}
|
||
|
||
let response_text = response.text().await?;
|
||
let upload_response: UploadResponse = serde_json::from_str(&response_text)
|
||
.map_err(|e| anyhow!("解析上传响应失败: {} - 响应内容: {}", e, response_text))?;
|
||
|
||
// 优先使用urn字段,如果没有则使用file_uri字段
|
||
let file_uri = upload_response.urn
|
||
.or(upload_response.file_uri)
|
||
.or_else(|| upload_response.name.map(|name| format!("gs://dy-media-storage/{}", name)))
|
||
.ok_or_else(|| anyhow!("上传响应中未找到文件URI,响应内容: {}", response_text))?;
|
||
|
||
Ok(file_uri)
|
||
}
|
||
|
||
/// 生成内容分析 (参考Python demo.py实现)
|
||
pub async fn generate_content_analysis(&mut self, file_uri: &str, prompt: &str) -> Result<String> {
|
||
println!("🧠 正在进行AI分析...");
|
||
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
|
||
// 创建客户端配置
|
||
let client_config = self.create_gemini_client(&access_token);
|
||
|
||
// 格式化GCS URI
|
||
let formatted_uri = self.format_gcs_uri(file_uri);
|
||
|
||
// 准备请求数据,参考demo.py实现
|
||
let request_data = GenerateContentRequest {
|
||
contents: vec![ContentPart {
|
||
role: "user".to_string(),
|
||
parts: vec![
|
||
Part::Text { text: prompt.to_string() },
|
||
Part::FileData {
|
||
file_data: FileData {
|
||
mime_type: "video/mp4".to_string(),
|
||
file_uri: formatted_uri,
|
||
}
|
||
}
|
||
],
|
||
}],
|
||
generation_config: GenerationConfig {
|
||
temperature: self.config.temperature,
|
||
top_k: 32,
|
||
top_p: 1.0,
|
||
max_output_tokens: self.config.max_tokens,
|
||
},
|
||
};
|
||
|
||
// 发送请求到Cloudflare Gateway,参考demo.py
|
||
let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, self.config.model_name);
|
||
|
||
// 重试机制
|
||
let mut last_error = None;
|
||
|
||
for attempt in 0..self.config.max_retries {
|
||
match self.send_generate_request(&generate_url, &client_config, &request_data).await {
|
||
Ok(result) => {
|
||
return self.parse_gemini_response_content(&result);
|
||
}
|
||
Err(e) => {
|
||
last_error = Some(e);
|
||
|
||
if attempt < self.config.max_retries - 1 {
|
||
tokio::time::sleep(tokio::time::Duration::from_secs(self.config.retry_delay)).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("内容生成失败,已重试{}次: {}", self.config.max_retries, last_error.unwrap()))
|
||
}
|
||
|
||
/// 发送生成请求
|
||
async fn send_generate_request(
|
||
&self,
|
||
url: &str,
|
||
client_config: &ClientConfig,
|
||
request_data: &GenerateContentRequest,
|
||
) -> Result<GeminiResponse> {
|
||
let mut request_builder = self.client
|
||
.post(url)
|
||
.timeout(tokio::time::Duration::from_secs(self.config.timeout))
|
||
.json(request_data);
|
||
|
||
// 添加请求头
|
||
for (key, value) in &client_config.headers {
|
||
request_builder = request_builder.header(key, value);
|
||
}
|
||
|
||
let response = request_builder.send().await?;
|
||
let status = response.status();
|
||
|
||
if !status.is_success() {
|
||
let error_text = response.text().await.unwrap_or_default();
|
||
return Err(anyhow!("API请求失败: {} - {}", status, error_text));
|
||
}
|
||
|
||
let response_text = response.text().await?;
|
||
let gemini_response: GeminiResponse = serde_json::from_str(&response_text)
|
||
.map_err(|e| anyhow!("解析生成响应失败: {} - 响应内容: {}", e, response_text))?;
|
||
|
||
if gemini_response.candidates.is_empty() {
|
||
return Err(anyhow!("API返回结果为空"));
|
||
}
|
||
|
||
Ok(gemini_response)
|
||
}
|
||
|
||
/// 解析Gemini响应内容
|
||
fn parse_gemini_response_content(&self, gemini_response: &GeminiResponse) -> Result<String> {
|
||
if let Some(candidate) = gemini_response.candidates.first() {
|
||
if let Some(part) = candidate.content.parts.first() {
|
||
println!("✅ AI分析完成");
|
||
return Ok(part.text.clone());
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("Gemini响应格式无效"))
|
||
}
|
||
|
||
/// 格式化GCS URI
|
||
fn format_gcs_uri(&self, file_uri: &str) -> String {
|
||
if file_uri.starts_with("gs://") {
|
||
file_uri.to_string()
|
||
} else if file_uri.starts_with("https://storage.googleapis.com/") {
|
||
// 转换为gs://格式
|
||
file_uri.replace("https://storage.googleapis.com/", "gs://")
|
||
} else {
|
||
// 假设是相对路径,添加默认bucket
|
||
format!("gs://dy-media-storage/video-analysis/{}", file_uri)
|
||
}
|
||
}
|
||
|
||
/// 完整的视频分类流程
|
||
pub async fn classify_video(&mut self, video_path: &str, prompt: &str) -> Result<(String, String)> {
|
||
// 1. 上传视频
|
||
let file_uri = self.upload_video_file(video_path).await?;
|
||
|
||
// 2. 生成分析
|
||
let analysis_result = self.generate_content_analysis(&file_uri, prompt).await?;
|
||
|
||
Ok((file_uri, analysis_result))
|
||
}
|
||
|
||
/// 分析图像内容(用于服装搭配分析)
|
||
pub async fn analyze_image_with_prompt(&mut self, image_base64: &str, prompt: &str) -> Result<String> {
|
||
println!("🧠 正在进行图像分析...");
|
||
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
|
||
// 创建客户端配置
|
||
let client_config = self.create_gemini_client(&access_token);
|
||
|
||
// 准备请求数据
|
||
let request_data = GenerateContentRequest {
|
||
contents: vec![ContentPart {
|
||
role: "user".to_string(),
|
||
parts: vec![
|
||
Part::Text { text: prompt.to_string() },
|
||
Part::InlineData {
|
||
inline_data: InlineData {
|
||
mime_type: "image/jpeg".to_string(),
|
||
data: image_base64.to_string(),
|
||
}
|
||
}
|
||
],
|
||
}],
|
||
generation_config: GenerationConfig {
|
||
temperature: self.config.temperature,
|
||
top_k: 32,
|
||
top_p: 1.0,
|
||
max_output_tokens: self.config.max_tokens,
|
||
},
|
||
};
|
||
|
||
// 发送请求到Cloudflare Gateway
|
||
let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, self.config.model_name);
|
||
|
||
// 重试机制
|
||
let mut last_error = None;
|
||
|
||
for attempt in 0..self.config.max_retries {
|
||
match self.send_generate_request(&generate_url, &client_config, &request_data).await {
|
||
Ok(result) => {
|
||
let content = self.parse_gemini_response_content(&result)?;
|
||
println!("✅ 图像分析完成");
|
||
return Ok(content);
|
||
}
|
||
Err(e) => {
|
||
last_error = Some(e);
|
||
|
||
if attempt < self.config.max_retries - 1 {
|
||
tokio::time::sleep(tokio::time::Duration::from_secs(self.config.retry_delay)).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("图像分析失败,已重试{}次: {}", self.config.max_retries, last_error.unwrap()))
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[tokio::test]
|
||
async fn test_format_gcs_uri() {
|
||
let service = GeminiService::new(None).expect("Failed to create GeminiService");
|
||
|
||
// 测试已经是gs://格式的URI
|
||
let gs_uri = "gs://bucket/path/file.mp4";
|
||
assert_eq!(service.format_gcs_uri(gs_uri), gs_uri);
|
||
|
||
// 测试https://storage.googleapis.com/格式的URI
|
||
let https_uri = "https://storage.googleapis.com/bucket/path/file.mp4";
|
||
assert_eq!(service.format_gcs_uri(https_uri), "gs://bucket/path/file.mp4");
|
||
|
||
// 测试相对路径
|
||
let relative_path = "path/file.mp4";
|
||
assert_eq!(service.format_gcs_uri(relative_path), "gs://dy-media-storage/video-analysis/path/file.mp4");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_extract_json_from_response_valid_json() {
|
||
let service = GeminiService::new(None).expect("Failed to create GeminiService");
|
||
|
||
// 测试有效的JSON
|
||
let valid_json = r#"{"name": "test", "value": 42}"#;
|
||
let result = service.extract_json_from_response(valid_json);
|
||
assert!(result.is_ok());
|
||
|
||
let parsed: serde_json::Value = serde_json::from_str(&result.unwrap()).unwrap();
|
||
assert_eq!(parsed["name"], "test");
|
||
assert_eq!(parsed["value"], 42);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_extract_json_from_response_markdown_wrapped() {
|
||
let service = GeminiService::new(None).expect("Failed to create GeminiService");
|
||
|
||
// 测试Markdown包装的JSON
|
||
let markdown_json = r#"Here's the analysis:
|
||
```json
|
||
{"environment_tags": ["outdoor"], "style_description": "casual"}
|
||
```
|
||
That's the result."#;
|
||
|
||
let result = service.extract_json_from_response(markdown_json);
|
||
assert!(result.is_ok());
|
||
|
||
let parsed: serde_json::Value = serde_json::from_str(&result.unwrap()).unwrap();
|
||
assert_eq!(parsed["environment_tags"][0], "outdoor");
|
||
assert_eq!(parsed["style_description"], "casual");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_extract_json_from_response_malformed_json() {
|
||
let service = GeminiService::new(None).expect("Failed to create GeminiService");
|
||
|
||
// 测试格式错误的JSON(无引号的键)
|
||
let malformed_json = r#"{name: "test", value: 42,}"#;
|
||
let result = service.extract_json_from_response(malformed_json);
|
||
|
||
// TolerantJsonParser应该能够修复这种错误
|
||
assert!(result.is_ok());
|
||
|
||
let parsed: serde_json::Value = serde_json::from_str(&result.unwrap()).unwrap();
|
||
assert_eq!(parsed["name"], "test");
|
||
assert_eq!(parsed["value"], 42);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_extract_json_from_response_fallback_behavior() {
|
||
let service = GeminiService::new(None).expect("Failed to create GeminiService");
|
||
|
||
// 测试复杂的Gemini响应格式,验证解析器能够提取到有效的JSON
|
||
// 注意:当前的TolerantJsonParser可能会提取到第一个有效的JSON对象
|
||
let complex_response = r#"Based on the image analysis, here's the structured result:
|
||
|
||
```json
|
||
{
|
||
"environment_tags": ["indoor", "office"],
|
||
"environment_color_pattern": {
|
||
"hue": 0.5,
|
||
"saturation": 0.3,
|
||
"value": 0.8
|
||
},
|
||
"dress_color_pattern": {
|
||
"hue": 0.6,
|
||
"saturation": 0.4,
|
||
"value": 0.9
|
||
},
|
||
"style_description": "Professional business attire",
|
||
"products": [
|
||
{
|
||
"category": "上装",
|
||
"description": "白色衬衫",
|
||
"color_pattern": {
|
||
"hue": 0.0,
|
||
"saturation": 0.0,
|
||
"value": 1.0
|
||
},
|
||
"design_styles": ["正式", "商务"],
|
||
"color_pattern_match_dress": 0.9,
|
||
"color_pattern_match_environment": 0.8
|
||
}
|
||
]
|
||
}
|
||
```
|
||
|
||
This analysis shows a professional outfit suitable for office environments."#;
|
||
|
||
let result = service.extract_json_from_response(complex_response);
|
||
assert!(result.is_ok());
|
||
|
||
let json_string = result.unwrap();
|
||
println!("Extracted JSON: {}", json_string);
|
||
|
||
let parsed: serde_json::Value = serde_json::from_str(&json_string).unwrap();
|
||
println!("Parsed JSON: {:?}", parsed);
|
||
|
||
// 检查解析结果的结构 - 验证至少提取到了有效的JSON对象
|
||
assert!(parsed.is_object(), "Parsed result should be an object");
|
||
|
||
// 验证提取到的JSON包含数值字段(说明解析成功)
|
||
assert!(parsed.get("hue").is_some() || parsed.get("environment_tags").is_some(),
|
||
"Should extract some recognizable JSON content");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_extract_json_from_response_simple_markdown() {
|
||
let service = GeminiService::new(None).expect("Failed to create GeminiService");
|
||
|
||
// 测试简单的Markdown包装JSON,这应该能正确解析
|
||
let simple_markdown = r#"Here's the result:
|
||
```json
|
||
{"status": "success", "message": "Analysis complete"}
|
||
```
|
||
Done."#;
|
||
|
||
let result = service.extract_json_from_response(simple_markdown);
|
||
assert!(result.is_ok());
|
||
|
||
let parsed: serde_json::Value = serde_json::from_str(&result.unwrap()).unwrap();
|
||
assert_eq!(parsed["status"], "success");
|
||
assert_eq!(parsed["message"], "Analysis complete");
|
||
}
|
||
}
|
||
|
||
// 服装搭配分析扩展
|
||
impl GeminiService {
|
||
/// 分析服装图像并返回结构化结果
|
||
pub async fn analyze_outfit_image(&mut self, image_path: &str) -> Result<String> {
|
||
// 读取图像文件
|
||
let image_data = fs::read(image_path).await
|
||
.map_err(|e| anyhow!("Failed to read image file: {} - {}", image_path, e))?;
|
||
|
||
// 转换为base64
|
||
let image_base64 = BASE64_STANDARD.encode(&image_data);
|
||
|
||
// 构建服装分析提示词
|
||
let prompt = self.build_outfit_analysis_prompt();
|
||
|
||
// 调用图像分析
|
||
let raw_response = self.analyze_image_with_prompt(&image_base64, &prompt).await?;
|
||
|
||
// 添加调试信息
|
||
println!("🔍 Gemini原始响应: {}", raw_response);
|
||
|
||
// 尝试提取JSON部分
|
||
self.extract_json_from_response(&raw_response)
|
||
}
|
||
|
||
/// 构建服装分析提示词
|
||
fn build_outfit_analysis_prompt(&self) -> String {
|
||
r#"请分析这张服装图片,并以JSON格式返回以下信息:
|
||
|
||
{
|
||
"environment_tags": ["环境标签1", "环境标签2"],
|
||
"environment_color_pattern": {
|
||
"hue": 0.5,
|
||
"saturation": 0.3,
|
||
"value": 0.8
|
||
},
|
||
"dress_color_pattern": {
|
||
"hue": 0.6,
|
||
"saturation": 0.4,
|
||
"value": 0.9
|
||
},
|
||
"style_description": "整体风格描述",
|
||
"products": [
|
||
{
|
||
"category": "服装类别",
|
||
"description": "服装描述",
|
||
"color_pattern": {
|
||
"hue": 0.6,
|
||
"saturation": 0.5,
|
||
"value": 0.7
|
||
},
|
||
"design_styles": ["设计风格1", "设计风格2"],
|
||
"color_pattern_match_dress": 0.8,
|
||
"color_pattern_match_environment": 0.7
|
||
}
|
||
]
|
||
}
|
||
|
||
分析要求:
|
||
1. environment_tags: 识别图片中的环境场景,如"Outdoor", "Indoor", "City street", "Office"等
|
||
2. environment_color_pattern: 环境的主要颜色,用HSV值表示(0-1范围)
|
||
3. dress_color_pattern: 整体服装搭配的主要颜色
|
||
4. style_description: 用中文描述整体的搭配风格
|
||
5. products: 识别出的各个服装单品
|
||
- category: 服装类别,如"上装", "下装", "鞋子", "配饰"等
|
||
- description: 具体描述这件服装
|
||
- color_pattern: 该单品的主要颜色
|
||
- design_styles: 设计风格,如"休闲", "正式", "运动", "街头"等
|
||
- color_pattern_match_dress: 与整体搭配颜色的匹配度(0-1)
|
||
- color_pattern_match_environment: 与环境颜色的匹配度(0-1)
|
||
|
||
请确保返回的是有效的JSON格式。"#.to_string()
|
||
}
|
||
|
||
/// LLM问答功能
|
||
pub async fn ask_outfit_advice(&mut self, user_input: &str) -> Result<String> {
|
||
// 构建服装搭配顾问提示词
|
||
let prompt = format!(
|
||
r#"你是一位专业的服装搭配顾问,请根据用户的问题提供专业的搭配建议。
|
||
|
||
用户问题:{}
|
||
|
||
请提供:
|
||
1. 具体的搭配建议
|
||
2. 颜色搭配原理
|
||
3. 适合的场合
|
||
4. 搭配技巧
|
||
|
||
请用友好、专业的语气回答,并提供实用的建议。"#,
|
||
user_input
|
||
);
|
||
|
||
// 使用文本生成功能
|
||
self.generate_text_content(&prompt).await
|
||
}
|
||
|
||
/// 生成文本内容(用于LLM问答)
|
||
async fn generate_text_content(&mut self, prompt: &str) -> Result<String> {
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
|
||
// 创建客户端配置
|
||
let client_config = self.create_gemini_client(&access_token);
|
||
|
||
// 准备请求数据
|
||
let request_data = GenerateContentRequest {
|
||
contents: vec![ContentPart {
|
||
role: "user".to_string(),
|
||
parts: vec![Part::Text { text: prompt.to_string() }],
|
||
}],
|
||
generation_config: GenerationConfig {
|
||
temperature: 0.7, // 稍高的温度以获得更有创意的回答
|
||
top_k: 32,
|
||
top_p: 1.0,
|
||
max_output_tokens: self.config.max_tokens,
|
||
},
|
||
};
|
||
|
||
// 发送请求
|
||
let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, self.config.model_name);
|
||
|
||
// 重试机制
|
||
for attempt in 0..self.config.max_retries {
|
||
match self.send_generate_request(&generate_url, &client_config, &request_data).await {
|
||
Ok(result) => {
|
||
let content = self.parse_gemini_response_content(&result)?;
|
||
return Ok(content);
|
||
}
|
||
Err(e) => {
|
||
if attempt == self.config.max_retries - 1 {
|
||
return Err(e);
|
||
}
|
||
tokio::time::sleep(tokio::time::Duration::from_secs(self.config.retry_delay)).await;
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("All retry attempts failed"))
|
||
}
|
||
|
||
/// 从Gemini响应中提取JSON内容
|
||
/// 使用TolerantJsonParser进行高级JSON解析和错误恢复
|
||
fn extract_json_from_response(&self, response: &str) -> Result<String> {
|
||
println!("🔍 开始使用TolerantJsonParser提取JSON,响应长度: {}", response.len());
|
||
|
||
// 使用TolerantJsonParser进行解析
|
||
let parse_result = {
|
||
let mut parser = self.json_parser.lock()
|
||
.map_err(|e| anyhow!("Failed to acquire parser lock: {}", e))?;
|
||
parser.parse(response)
|
||
};
|
||
|
||
match parse_result {
|
||
Ok((parsed_value, stats)) => {
|
||
println!("✅ TolerantJsonParser解析成功");
|
||
println!("📊 解析统计: 总节点数={}, 错误节点数={}, 错误率={:.2}%, 解析时间={}ms",
|
||
stats.total_nodes, stats.error_nodes, stats.error_rate * 100.0, stats.parse_time_ms);
|
||
|
||
if !stats.recovery_strategies_used.is_empty() {
|
||
println!("<EFBFBD> 使用的恢复策略: {:?}", stats.recovery_strategies_used);
|
||
}
|
||
|
||
// 记录解析质量信息
|
||
if stats.error_rate > 0.1 {
|
||
println!("⚠️ 解析质量警告: 错误率较高 ({:.2}%), 建议检查输入数据", stats.error_rate * 100.0);
|
||
}
|
||
|
||
if stats.parse_time_ms > 1000 {
|
||
println!("⚠️ 性能警告: 解析时间较长 ({}ms), 可能需要优化", stats.parse_time_ms);
|
||
}
|
||
|
||
// 将解析结果转换为JSON字符串
|
||
let json_string = serde_json::to_string(&parsed_value)
|
||
.map_err(|e| anyhow!("Failed to serialize parsed JSON: {}", e))?;
|
||
|
||
println!("✅ JSON序列化成功,输出长度: {} 字符", json_string.len());
|
||
Ok(json_string)
|
||
}
|
||
Err(e) => {
|
||
println!("❌ TolerantJsonParser解析失败: {}", e);
|
||
println!("📝 响应内容预览: {}", &response[..response.len().min(200)]);
|
||
|
||
// 直接返回错误,不使用回退方案,便于调试和问题定位
|
||
Err(anyhow!("TolerantJsonParser解析失败: {}. 响应内容: {}", e, &response[..response.len().min(500)]))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 增强查询以获取更多相关信息
|
||
fn enhance_query_for_better_retrieval(&self, original_query: &str) -> String {
|
||
// 添加相关关键词和上下文来提高检索效果
|
||
let enhanced_keywords = vec![
|
||
"穿搭", "搭配", "服装", "时尚", "风格", "造型", "服饰", "款式", "颜色", "材质"
|
||
];
|
||
|
||
// 检查原查询是否已包含这些关键词
|
||
let query_lower = original_query.to_lowercase();
|
||
let missing_keywords: Vec<&str> = enhanced_keywords
|
||
.iter()
|
||
.filter(|&&keyword| !query_lower.contains(keyword))
|
||
.copied()
|
||
.collect();
|
||
|
||
if missing_keywords.is_empty() {
|
||
// 如果已包含关键词,添加更详细的描述要求
|
||
format!("{} 请提供详细的分析和具体的建议,包括颜色搭配、款式选择、场合适用性等方面。", original_query)
|
||
} else {
|
||
// 添加相关关键词来扩展搜索范围
|
||
let additional_context = missing_keywords.join("、");
|
||
format!("{} 相关的{}方面的建议和搭配方案", original_query, additional_context)
|
||
}
|
||
}
|
||
|
||
/// RAG Grounding 查询 (参考 RAGUtils.py 中的 query_llm_with_grounding)
|
||
pub async fn query_llm_with_grounding(&mut self, request: RagGroundingRequest) -> Result<RagGroundingResponse> {
|
||
// 如果请求包含会话管理参数,使用多轮对话版本
|
||
if request.session_id.is_some() || request.include_history.unwrap_or(false) {
|
||
return self.query_llm_with_grounding_multi_turn(request, None).await;
|
||
}
|
||
|
||
// 否则使用原有的单轮对话逻辑
|
||
self.query_llm_with_grounding_single_turn(request).await
|
||
}
|
||
|
||
/// 单轮RAG Grounding查询(原有逻辑)
|
||
async fn query_llm_with_grounding_single_turn(&mut self, request: RagGroundingRequest) -> Result<RagGroundingResponse> {
|
||
let start_time = std::time::Instant::now();
|
||
println!("🔍 开始RAG Grounding查询: {}", request.user_input);
|
||
|
||
// 获取配置,确保 system_prompt 不会被覆盖
|
||
let mut rag_config = request.config.unwrap_or_default();
|
||
|
||
// 如果请求的配置没有 system_prompt,使用默认配置的 system_prompt
|
||
if rag_config.system_prompt.is_none() {
|
||
rag_config.system_prompt = RagGroundingConfig::default().system_prompt;
|
||
}
|
||
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
|
||
// 创建客户端配置
|
||
let client_config = self.create_gemini_client(&access_token);
|
||
|
||
// 构建数据存储路径
|
||
let datastore_path = format!(
|
||
"projects/{}/locations/{}/collections/default_collection/dataStores/{}",
|
||
rag_config.project_id,
|
||
rag_config.location,
|
||
rag_config.data_store_id
|
||
);
|
||
|
||
// 构建工具配置 (Vertex AI Search) - 只使用支持的字段
|
||
let tools = vec![VertexAISearchTool {
|
||
retrieval: VertexAIRetrieval {
|
||
vertex_ai_search: VertexAISearchConfig {
|
||
datastore: datastore_path,
|
||
filter: rag_config.search_filter.clone(),
|
||
},
|
||
disable_attribution: Some(false), // 保留归属信息
|
||
},
|
||
}];
|
||
|
||
// 优化查询内容以获取更多相关信息
|
||
let enhanced_query = self.enhance_query_for_better_retrieval(&request.user_input);
|
||
|
||
// 构建请求内容
|
||
let contents = vec![ContentPart {
|
||
role: "user".to_string(),
|
||
parts: vec![Part::Text { text: enhanced_query }],
|
||
}];
|
||
|
||
// 构建生成配置
|
||
let generation_config = GenerationConfig {
|
||
temperature: rag_config.temperature,
|
||
top_k: 32,
|
||
top_p: 1.0,
|
||
max_output_tokens: rag_config.max_output_tokens,
|
||
};
|
||
|
||
// 准备请求数据
|
||
let request_data = serde_json::json!({
|
||
"contents": contents,
|
||
"tools": tools,
|
||
"generationConfig": generation_config,
|
||
"systemInstruction": rag_config.system_prompt.map(|prompt| {
|
||
serde_json::json!({
|
||
"role": "system",
|
||
"parts": [{"text": prompt}]
|
||
})
|
||
})
|
||
});
|
||
|
||
// 发送请求到Cloudflare Gateway (使用beta API)
|
||
let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, rag_config.model_id);
|
||
|
||
// 重试机制
|
||
let mut last_error = None;
|
||
for attempt in 0..self.config.max_retries {
|
||
match self.send_rag_grounding_request(&generate_url, &client_config, &request_data).await {
|
||
Ok(response) => {
|
||
let elapsed = start_time.elapsed();
|
||
|
||
println!("✅ RAG Grounding查询完成,耗时: {:?}", elapsed);
|
||
|
||
return Ok(RagGroundingResponse {
|
||
answer: response.answer,
|
||
grounding_metadata: response.grounding_metadata,
|
||
response_time_ms: elapsed.as_millis() as u64,
|
||
model_used: rag_config.model_id,
|
||
session_id: None,
|
||
message_id: None,
|
||
conversation_context: None,
|
||
});
|
||
}
|
||
Err(e) => {
|
||
last_error = Some(e);
|
||
if attempt < self.config.max_retries - 1 {
|
||
tokio::time::sleep(tokio::time::Duration::from_secs(self.config.retry_delay)).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("RAG Grounding查询失败,已重试{}次: {}", self.config.max_retries, last_error.unwrap()))
|
||
}
|
||
|
||
/// 多轮RAG Grounding查询(支持会话历史)
|
||
pub async fn query_llm_with_grounding_multi_turn(
|
||
&mut self,
|
||
request: RagGroundingRequest,
|
||
conversation_repo: Option<Arc<ConversationRepository>>,
|
||
) -> Result<RagGroundingResponse> {
|
||
let start_time = std::time::Instant::now();
|
||
println!("🔍 开始多轮RAG Grounding查询: {}", request.user_input);
|
||
|
||
// 先提取 system_prompt,避免 request 被消费后无法访问
|
||
let request_system_prompt = request.system_prompt.clone();
|
||
|
||
// 获取配置,确保 system_prompt 不会被覆盖
|
||
let mut rag_config = request.config.unwrap_or_default();
|
||
|
||
// 如果请求的配置没有 system_prompt,使用默认配置的 system_prompt
|
||
if rag_config.system_prompt.is_none() {
|
||
rag_config.system_prompt = RagGroundingConfig::default().system_prompt;
|
||
}
|
||
|
||
// 1. 确定或创建会话
|
||
let session_id = match &request.session_id {
|
||
Some(id) => {
|
||
println!("📋 使用现有会话ID: {}", id);
|
||
id.clone()
|
||
},
|
||
None => {
|
||
let new_id = uuid::Uuid::new_v4().to_string();
|
||
println!("🆕 生成新会话ID: {}", new_id);
|
||
new_id
|
||
}
|
||
};
|
||
|
||
// 如果有会话仓库,确保会话存在
|
||
if let Some(repo) = &conversation_repo {
|
||
println!("🔍 检查会话是否存在...");
|
||
match repo.get_session(&session_id) {
|
||
Ok(Some(_)) => {
|
||
println!("✅ 会话已存在");
|
||
}
|
||
Ok(None) => {
|
||
println!("🆕 会话不存在,跳过历史消息获取...");
|
||
// 对于新会话,没有历史消息,直接跳过
|
||
}
|
||
Err(e) => {
|
||
println!("⚠️ 检查会话失败: {}", e);
|
||
}
|
||
}
|
||
}
|
||
|
||
// 2. 获取历史消息(如果需要且有会话仓库)
|
||
let mut contents = Vec::new();
|
||
let mut conversation_context = ConversationContext {
|
||
total_messages: 0,
|
||
history_included: false,
|
||
context_length: 0,
|
||
};
|
||
|
||
println!("🔍 检查是否需要获取历史消息...");
|
||
println!(" - 会话仓库存在: {}", conversation_repo.is_some());
|
||
println!(" - 包含历史消息: {}", request.include_history.unwrap_or(false));
|
||
|
||
if let (Some(repo), true) = (&conversation_repo, request.include_history.unwrap_or(false)) {
|
||
println!("📚 开始获取历史消息...");
|
||
println!("🔄 调用get_conversation_history方法...");
|
||
match self.get_conversation_history(&repo, &session_id, request.max_history_messages.unwrap_or(1)).await {
|
||
Ok(history_messages) => {
|
||
conversation_context.total_messages = history_messages.len() as u32;
|
||
conversation_context.history_included = true;
|
||
|
||
// 系统提示将在请求数据中通过 systemInstruction 字段处理
|
||
|
||
// 添加历史消息
|
||
for msg in &history_messages {
|
||
let parts = self.convert_message_content_to_parts(&msg.content)?;
|
||
contents.push(ContentPart {
|
||
role: msg.role.to_string(),
|
||
parts,
|
||
});
|
||
}
|
||
|
||
conversation_context.context_length = contents.len() as u32;
|
||
println!("📚 加载了 {} 条历史消息", history_messages.len());
|
||
}
|
||
Err(e) => {
|
||
println!("⚠️ 获取历史消息失败: {}", e);
|
||
}
|
||
}
|
||
}
|
||
|
||
// 3. 添加当前用户消息(优化查询以获取更多相关信息)
|
||
let enhanced_query = self.enhance_query_for_better_retrieval(&request.user_input);
|
||
contents.push(ContentPart {
|
||
role: "user".to_string(),
|
||
parts: vec![Part::Text { text: enhanced_query }],
|
||
});
|
||
|
||
// 4. 执行RAG查询
|
||
println!("🚀 准备执行RAG查询,contents数量: {}", contents.len());
|
||
let response = self.execute_rag_grounding_with_contents(contents, rag_config, request_system_prompt).await?;
|
||
println!("📥 RAG查询响应已收到");
|
||
|
||
// 5. 保存消息到会话历史(如果有会话仓库)
|
||
let message_id = if let Some(repo) = &conversation_repo {
|
||
match self.save_conversation_messages(&repo, &session_id, &request.user_input, &response.answer).await {
|
||
Ok(id) => Some(id),
|
||
Err(e) => {
|
||
println!("⚠️ 保存会话消息失败: {}", e);
|
||
None
|
||
}
|
||
}
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let elapsed = start_time.elapsed();
|
||
println!("✅ 多轮RAG Grounding查询完成,耗时: {:?}", elapsed);
|
||
|
||
Ok(RagGroundingResponse {
|
||
answer: response.answer,
|
||
grounding_metadata: response.grounding_metadata,
|
||
response_time_ms: elapsed.as_millis() as u64,
|
||
model_used: response.model_used,
|
||
session_id: Some(session_id),
|
||
message_id,
|
||
conversation_context: Some(conversation_context),
|
||
})
|
||
}
|
||
|
||
/// 获取会话历史消息
|
||
async fn get_conversation_history(
|
||
&self,
|
||
repo: &Arc<ConversationRepository>,
|
||
session_id: &str,
|
||
max_messages: u32,
|
||
) -> Result<Vec<ConversationMessage>> {
|
||
println!("📋 构建历史查询参数...");
|
||
println!(" - session_id: {}", session_id);
|
||
println!(" - max_messages: {}", max_messages);
|
||
|
||
let query = ConversationHistoryQuery {
|
||
session_id: session_id.to_string(),
|
||
limit: Some(max_messages),
|
||
offset: None,
|
||
include_system_messages: Some(false),
|
||
};
|
||
|
||
println!("🔍 调用repo.get_conversation_history...");
|
||
match repo.get_conversation_history(query) {
|
||
Ok(history) => {
|
||
println!("✅ 历史查询成功,获取到 {} 条消息", history.messages.len());
|
||
Ok(history.messages)
|
||
},
|
||
Err(e) => {
|
||
println!("❌ 历史查询失败: {}", e);
|
||
Err(anyhow!("获取会话历史失败: {}", e))
|
||
},
|
||
}
|
||
}
|
||
|
||
/// 保存会话消息
|
||
async fn save_conversation_messages(
|
||
&self,
|
||
repo: &Arc<ConversationRepository>,
|
||
session_id: &str,
|
||
user_input: &str,
|
||
assistant_response: &str,
|
||
) -> Result<String> {
|
||
println!("💾 开始保存会话消息,session_id: {}", session_id);
|
||
|
||
// 确保会话存在
|
||
repo.ensure_session_exists(session_id)?;
|
||
|
||
// 保存用户消息
|
||
let user_message_request = AddMessageRequest {
|
||
session_id: session_id.to_string(),
|
||
role: MessageRole::User,
|
||
content: vec![MessageContent::Text { text: user_input.to_string() }],
|
||
metadata: None,
|
||
};
|
||
|
||
let user_message = repo.add_message(user_message_request)?;
|
||
println!("✅ 用户消息保存成功,ID: {}", user_message.id);
|
||
|
||
// 保存助手回复
|
||
let assistant_message_request = AddMessageRequest {
|
||
session_id: session_id.to_string(),
|
||
role: MessageRole::Assistant,
|
||
content: vec![MessageContent::Text { text: assistant_response.to_string() }],
|
||
metadata: None,
|
||
};
|
||
|
||
let assistant_message = repo.add_message(assistant_message_request)?;
|
||
println!("✅ 助手消息保存成功,ID: {}", assistant_message.id);
|
||
println!("💾 会话消息保存完成,session_id: {}", session_id);
|
||
|
||
Ok(assistant_message.id)
|
||
}
|
||
|
||
/// 执行带有完整contents的RAG查询
|
||
async fn execute_rag_grounding_with_contents(
|
||
&mut self,
|
||
contents: Vec<ContentPart>,
|
||
rag_config: RagGroundingConfig,
|
||
request_system_prompt: Option<String>,
|
||
) -> Result<RagGroundingResponse> {
|
||
println!("🔑 开始获取访问令牌...");
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
println!("✅ 访问令牌获取成功");
|
||
|
||
// 创建客户端配置
|
||
let client_config = self.create_gemini_client(&access_token);
|
||
|
||
// 构建数据存储路径
|
||
let datastore_path = format!(
|
||
"projects/{}/locations/{}/collections/default_collection/dataStores/{}",
|
||
rag_config.project_id,
|
||
rag_config.location,
|
||
rag_config.data_store_id
|
||
);
|
||
|
||
// 构建工具配置 (Vertex AI Search) - 只使用支持的字段
|
||
let tools = vec![VertexAISearchTool {
|
||
retrieval: VertexAIRetrieval {
|
||
vertex_ai_search: VertexAISearchConfig {
|
||
datastore: datastore_path,
|
||
filter: rag_config.search_filter.clone(),
|
||
},
|
||
disable_attribution: Some(false), // 保留归属信息
|
||
},
|
||
}];
|
||
|
||
// 构建生成配置
|
||
let generation_config = GenerationConfig {
|
||
temperature: rag_config.temperature,
|
||
top_k: 32,
|
||
top_p: 1.0,
|
||
max_output_tokens: rag_config.max_output_tokens,
|
||
};
|
||
|
||
// 获取系统提示
|
||
let system_prompt = if let Some(system_prompt) = &request_system_prompt {
|
||
Some(system_prompt.clone())
|
||
} else {
|
||
rag_config.system_prompt.clone()
|
||
};
|
||
|
||
// 准备请求数据
|
||
let request_data = serde_json::json!({
|
||
"contents": contents,
|
||
"tools": tools,
|
||
"generationConfig": generation_config,
|
||
"toolConfig": {
|
||
"functionCallingConfig": {
|
||
"mode": "ANY",
|
||
"allowedFunctionNames": ["retrieval"]
|
||
}
|
||
},
|
||
"systemInstruction": system_prompt.map(|prompt| {
|
||
serde_json::json!({
|
||
"role": "system",
|
||
"parts": [{"text": prompt}]
|
||
})
|
||
})
|
||
});
|
||
|
||
// 发送请求到Cloudflare Gateway (使用beta API)
|
||
let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, rag_config.model_id);
|
||
println!("🌐 准备发送请求到: {}", generate_url);
|
||
println!("📊 请求数据: {}", serde_json::to_string(&request_data).unwrap_or_default());
|
||
|
||
// 重试机制
|
||
let mut last_error = None;
|
||
for attempt in 0..self.config.max_retries {
|
||
println!("🔄 尝试第 {} 次请求...", attempt + 1);
|
||
match self.send_rag_grounding_request(&generate_url, &client_config, &request_data).await {
|
||
Ok(response) => {
|
||
return Ok(RagGroundingResponse {
|
||
answer: response.answer,
|
||
grounding_metadata: response.grounding_metadata,
|
||
response_time_ms: 0, // 这里会在调用方设置
|
||
model_used: rag_config.model_id,
|
||
session_id: None,
|
||
message_id: None,
|
||
conversation_context: None,
|
||
});
|
||
}
|
||
Err(e) => {
|
||
last_error = Some(e);
|
||
if attempt < self.config.max_retries - 1 {
|
||
tokio::time::sleep(tokio::time::Duration::from_secs(self.config.retry_delay)).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("RAG Grounding查询失败,已重试{}次: {}", self.config.max_retries, last_error.unwrap()))
|
||
}
|
||
|
||
/// 将消息内容转换为Gemini API的Part格式
|
||
fn convert_message_content_to_parts(&self, content: &[MessageContent]) -> Result<Vec<Part>> {
|
||
let mut parts = Vec::new();
|
||
for item in content {
|
||
match item {
|
||
MessageContent::Text { text } => {
|
||
parts.push(Part::Text { text: text.clone() });
|
||
}
|
||
MessageContent::File { file_uri, mime_type, .. } => {
|
||
parts.push(Part::FileData {
|
||
file_data: FileData {
|
||
mime_type: mime_type.clone(),
|
||
file_uri: file_uri.clone(),
|
||
}
|
||
});
|
||
}
|
||
MessageContent::InlineData { data, mime_type, .. } => {
|
||
parts.push(Part::InlineData {
|
||
inline_data: InlineData {
|
||
mime_type: mime_type.clone(),
|
||
data: data.clone(),
|
||
}
|
||
});
|
||
}
|
||
}
|
||
}
|
||
Ok(parts)
|
||
}
|
||
|
||
/// 发送RAG Grounding请求
|
||
async fn send_rag_grounding_request(
|
||
&self,
|
||
url: &str,
|
||
client_config: &ClientConfig,
|
||
request_data: &serde_json::Value,
|
||
) -> Result<RagGroundingResponse> {
|
||
println!("📤 构建HTTP请求...");
|
||
let mut request_builder = self.client
|
||
.post(url)
|
||
.json(request_data);
|
||
|
||
// 添加请求头
|
||
for (key, value) in &client_config.headers {
|
||
request_builder = request_builder.header(key, value);
|
||
}
|
||
|
||
println!("🚀 发送HTTP请求...");
|
||
let response = request_builder
|
||
.send()
|
||
.await
|
||
.map_err(|e| anyhow!("发送RAG请求失败: {}", e))?;
|
||
|
||
println!("📨 收到HTTP响应,状态码: {}", response.status());
|
||
|
||
let status = response.status();
|
||
println!("📖 开始读取响应内容...");
|
||
let response_text = response.text().await
|
||
.map_err(|e| anyhow!("读取RAG响应失败: {}", e))?;
|
||
|
||
println!("📄 响应内容: {}", response_text);
|
||
|
||
if !status.is_success() {
|
||
println!("❌ HTTP请求失败,状态码: {}", status);
|
||
return Err(anyhow!("RAG请求失败 ({}): {}", status, response_text));
|
||
}
|
||
|
||
println!("🔍 开始解析响应内容...");
|
||
// 解析响应
|
||
self.parse_rag_grounding_response(&response_text)
|
||
}
|
||
|
||
/// 解析RAG Grounding响应
|
||
fn parse_rag_grounding_response(&self, response_text: &str) -> Result<RagGroundingResponse> {
|
||
// 使用容错JSON解析器
|
||
let mut parser = TolerantJsonParser::new(Some(ParserConfig {
|
||
max_text_length: 10_000_000, // 10MB限制
|
||
enable_comments: true,
|
||
enable_unquoted_keys: true,
|
||
enable_trailing_commas: true,
|
||
timeout_ms: 30000, // 30秒超时
|
||
recovery_strategies: vec![
|
||
RecoveryStrategy::StandardJson,
|
||
RecoveryStrategy::ManualFix,
|
||
RecoveryStrategy::RegexExtract,
|
||
RecoveryStrategy::PartialParse,
|
||
],
|
||
}))?;
|
||
|
||
let (response_json, _parse_stats) = parser.parse(response_text)
|
||
.map_err(|e| anyhow!("容错JSON解析失败: {}", e))?;
|
||
|
||
|
||
|
||
// 提取主要回答内容
|
||
let answer = response_json
|
||
.get("candidates")
|
||
.and_then(|c| c.as_array())
|
||
.and_then(|arr| arr.first())
|
||
.and_then(|candidate| candidate.get("content"))
|
||
.and_then(|content| content.get("parts"))
|
||
.and_then(|parts| parts.as_array())
|
||
.and_then(|arr| arr.first())
|
||
.and_then(|part| part.get("text"))
|
||
.and_then(|text| text.as_str())
|
||
.unwrap_or("无法解析响应内容")
|
||
.to_string();
|
||
|
||
// 提取Grounding元数据
|
||
let grounding_metadata = self.extract_grounding_metadata(&response_json);
|
||
|
||
Ok(RagGroundingResponse {
|
||
answer,
|
||
grounding_metadata,
|
||
response_time_ms: 0, // 将在调用方设置
|
||
model_used: "gemini-2.5-flash".to_string(),
|
||
session_id: None,
|
||
message_id: None,
|
||
conversation_context: None,
|
||
})
|
||
}
|
||
|
||
/// 提取Grounding元数据
|
||
fn extract_grounding_metadata(&self, response_json: &serde_json::Value) -> Option<GroundingMetadata> {
|
||
let grounding_metadata = response_json
|
||
.get("candidates")
|
||
.and_then(|c| c.as_array())
|
||
.and_then(|arr| arr.first())
|
||
.and_then(|candidate| candidate.get("groundingMetadata"))?;
|
||
|
||
// 打印grounding元数据的原始结构
|
||
/*
|
||
* "groundingMetadata": {
|
||
"retrievalQueries": [
|
||
"comfortable clothes for hot weather travel women",
|
||
"modest clothing tips Thailand tourist female",
|
||
"what to wear in Thailand female tourist summer",
|
||
"white woman tropical vacation outfits Thailand summer"
|
||
],
|
||
"groundingChunks": [
|
||
{
|
||
"retrievedContext": {
|
||
"uri": "gs://fashion_image_block/gallery_v2/models/Instagram_21129657_FREYA_KILLIN_20230828155546_1.jpg",
|
||
"title": "model",
|
||
"text": {}
|
||
}
|
||
},
|
||
*/
|
||
let sources = grounding_metadata
|
||
.get("groundingChunks")
|
||
.and_then(|sources| sources.as_array())
|
||
.map(|sources_array| {
|
||
|
||
sources_array
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(_index, chunk)| {
|
||
// 从 retrievedContext 中获取数据
|
||
let retrieved_context = chunk.get("retrievedContext")?;
|
||
|
||
let title = retrieved_context.get("title")?.as_str()?.to_string();
|
||
|
||
// 这个uri需要通过 convert_s3_to_cdn_url 转换
|
||
let uri = retrieved_context.get("uri")
|
||
.and_then(|u| u.as_str())
|
||
.map(|s| Self::convert_s3_to_cdn_url(s));
|
||
|
||
// 这个是个json
|
||
let content = retrieved_context.get("text").cloned();
|
||
|
||
|
||
|
||
Some(GroundingSource {
|
||
title,
|
||
uri,
|
||
content,
|
||
})
|
||
})
|
||
.collect()
|
||
})
|
||
.unwrap_or_default();
|
||
|
||
let search_queries = grounding_metadata
|
||
.get("retrievalQueries")
|
||
.and_then(|queries| queries.as_array())
|
||
.map(|queries_array| {
|
||
queries_array
|
||
.iter()
|
||
.filter_map(|q| q.as_str().map(|s| s.to_string()))
|
||
.collect()
|
||
})
|
||
.unwrap_or_default();
|
||
|
||
// 提取grounding supports
|
||
let grounding_supports = grounding_metadata
|
||
.get("groundingSupports")
|
||
.and_then(|supports| supports.as_array())
|
||
.map(|supports_array| {
|
||
supports_array
|
||
.iter()
|
||
.filter_map(|support| {
|
||
let grounding_chunk_indices = support
|
||
.get("groundingChunkIndices")
|
||
.and_then(|indices| indices.as_array())
|
||
.map(|indices_array| {
|
||
indices_array
|
||
.iter()
|
||
.filter_map(|idx| idx.as_u64().map(|i| i as usize))
|
||
.collect()
|
||
})
|
||
.unwrap_or_default();
|
||
|
||
let segment = support.get("segment").and_then(|seg| {
|
||
let start_index = seg.get("startIndex")?.as_u64()? as usize;
|
||
let end_index = seg.get("endIndex")?.as_u64()? as usize;
|
||
let text = seg.get("text")?.as_str()?.to_string();
|
||
|
||
Some(GroundingSegment {
|
||
start_index,
|
||
end_index,
|
||
text,
|
||
})
|
||
})?;
|
||
|
||
Some(GroundingSupport {
|
||
grounding_chunk_indices,
|
||
segment,
|
||
})
|
||
})
|
||
.collect()
|
||
});
|
||
|
||
Some(GroundingMetadata {
|
||
sources,
|
||
search_queries,
|
||
grounding_supports,
|
||
})
|
||
}
|
||
|
||
/// 使用自定义请求结构生成内容(支持多轮对话)
|
||
pub async fn generate_content_with_request(&mut self, request: GenerateContentRequest) -> Result<String> {
|
||
println!("🤖 开始多轮对话内容生成,包含 {} 条历史消息", request.contents.len());
|
||
|
||
// 获取访问令牌
|
||
let access_token = self.get_access_token().await?;
|
||
|
||
// 创建客户端配置
|
||
let client_config = self.create_gemini_client(&access_token);
|
||
|
||
// 发送请求到Cloudflare Gateway
|
||
let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, self.config.model_name);
|
||
|
||
// 重试机制
|
||
let mut last_error = None;
|
||
|
||
for attempt in 0..self.config.max_retries {
|
||
match self.send_generate_request(&generate_url, &client_config, &request).await {
|
||
Ok(result) => {
|
||
return self.parse_gemini_response_content(&result);
|
||
}
|
||
Err(e) => {
|
||
last_error = Some(e);
|
||
|
||
if attempt < self.config.max_retries - 1 {
|
||
tokio::time::sleep(tokio::time::Duration::from_secs(self.config.retry_delay)).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(anyhow!("多轮对话内容生成失败,已重试{}次: {}", self.config.max_retries, last_error.unwrap()))
|
||
}
|
||
|
||
fn convert_s3_to_cdn_url(s3_url: &str) -> String {
|
||
if s3_url.starts_with("s3://ap-northeast-2/modal-media-cache/") {
|
||
// 将 s3://ap-northeast-2/modal-media-cache/ 替换为 https://cdn.roasmax.cn/
|
||
s3_url.replace("s3://ap-northeast-2/modal-media-cache/", "https://cdn.roasmax.cn/")
|
||
} else if s3_url.starts_with("gs://fashion_image_block/") {
|
||
// 将 gs://fashion_image_block/ 替换为 https://cdn.roasmax.cn/fashion_image_block/
|
||
s3_url.replace("gs://", "https://storage.googleapis.com/")
|
||
} else if s3_url.starts_with("gs://") {
|
||
// 处理其他 gs:// 格式,转换为通用CDN格式
|
||
s3_url.replace("gs://", "https://storage.googleapis.com/")
|
||
} else if s3_url.starts_with("s3://") {
|
||
// 处理其他 s3:// 格式,转换为通用CDN格式
|
||
s3_url.replace("s3://", "https://cdn.roasmax.cn/")
|
||
} else {
|
||
// 如果不是预期的S3格式,返回原URL
|
||
s3_url.to_string()
|
||
}
|
||
}
|
||
}
|
||
|