Spaces:
Sleeping
Sleeping
File size: 2,644 Bytes
2887ce2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
use crate::chat::constant::AVAILABLE_MODELS;
use serde::{Deserialize, Serialize};
#[derive(Clone)]
pub enum UsageCheck {
None,
Default,
All,
Custom(Vec<&'static str>),
}
impl Default for UsageCheck {
fn default() -> Self {
Self::Default
}
}
impl Serialize for UsageCheck {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
use serde::ser::SerializeStruct;
let mut state = serializer.serialize_struct("UsageCheck", 1)?;
match self {
UsageCheck::None => {
state.serialize_field("type", "none")?;
}
UsageCheck::Default => {
state.serialize_field("type", "default")?;
}
UsageCheck::All => {
state.serialize_field("type", "all")?;
}
UsageCheck::Custom(models) => {
state.serialize_field("type", "list")?;
state.serialize_field("content", &models.join(","))?;
}
}
state.end()
}
}
impl<'de> Deserialize<'de> for UsageCheck {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(tag = "type", content = "content")]
enum UsageCheckHelper {
#[serde(rename = "none")]
None,
#[serde(rename = "default")]
Default,
#[serde(rename = "all")]
All,
#[serde(rename = "list")]
Custom(String),
}
let helper = UsageCheckHelper::deserialize(deserializer)?;
Ok(match helper {
UsageCheckHelper::None => UsageCheck::None,
UsageCheckHelper::Default => UsageCheck::Default,
UsageCheckHelper::All => UsageCheck::All,
UsageCheckHelper::Custom(list) => {
if list.is_empty() {
return Ok(UsageCheck::None);
}
let models: Vec<&'static str> = list
.split(',')
.filter_map(|model| {
let model = model.trim();
AVAILABLE_MODELS
.iter()
.find(|m| m.id == model)
.map(|m| m.id)
})
.collect();
if models.is_empty() {
UsageCheck::None
} else {
UsageCheck::Custom(models)
}
}
})
}
}
|