Skip to main content

bamboo_agent/server/services/
gemini_model_mapping_service.rs

1use crate::core::model_mapping::GeminiModelMapping;
2
3/// Resolve a Gemini model name to the actual backend model.
4///
5/// This no longer performs any file I/O. The mapping should be sourced from the
6/// unified `Config` (`config.gemini_model_mapping`) and passed in.
7pub fn resolve_model(mapping: &GeminiModelMapping, gemini_model: &str) -> ModelResolution {
8    log::info!(
9        "Resolving Gemini model '{}', available mappings: {:?}",
10        gemini_model,
11        mapping.mappings
12    );
13
14    // Match by keyword in model name (case-insensitive)
15    let model_lower = gemini_model.to_lowercase();
16
17    // Extract model type from Gemini model names like:
18    // - gemini-pro, gemini-ultra, gemini-1.5-pro, gemini-1.5-flash, etc.
19    let model_type = if model_lower.contains("ultra") {
20        "ultra"
21    } else if model_lower.contains("1.5") && model_lower.contains("flash") {
22        "flash-1.5"
23    } else if model_lower.contains("1.5") && model_lower.contains("pro") {
24        "pro-1.5"
25    } else if model_lower.contains("flash") {
26        "flash"
27    } else if model_lower.contains("pro") {
28        "pro"
29    } else {
30        log::warn!(
31            "No Gemini model mapping found for '{}', falling back to default model",
32            gemini_model
33        );
34        return ModelResolution {
35            mapped_model: String::new(),
36            response_model: gemini_model.to_string(),
37        };
38    };
39
40    if let Some(mapped) = mapping
41        .mappings
42        .get(model_type)
43        .filter(|value| !value.trim().is_empty())
44    {
45        log::info!(
46            "Model '{}' (type: {}) mapped to '{}'",
47            gemini_model,
48            model_type,
49            mapped
50        );
51        return ModelResolution {
52            mapped_model: mapped.to_string(),
53            response_model: gemini_model.to_string(),
54        };
55    }
56
57    log::warn!(
58        "No mapping configured for model type '{}', falling back to default model",
59        model_type
60    );
61
62    ModelResolution {
63        mapped_model: String::new(),
64        response_model: gemini_model.to_string(),
65    }
66}
67
68#[derive(Clone)]
69pub struct ModelResolution {
70    pub mapped_model: String,
71    pub response_model: String,
72}