70 lines
1.5 KiB
Rust
70 lines
1.5 KiB
Rust
use std::str::FromStr;
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
use crate::{client::GeminiClient, error::Result, prelude::TokenProvider};
|
|
|
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
|
pub enum Role {
|
|
User,
|
|
Model,
|
|
}
|
|
|
|
impl ToString for Role {
|
|
fn to_string(&self) -> String {
|
|
match self {
|
|
Role::User => "user".to_string(),
|
|
Role::Model => "model".to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl FromStr for Role {
|
|
type Err = ();
|
|
|
|
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
|
|
match s {
|
|
"user" => Ok(Role::User),
|
|
"model" => Ok(Role::Model),
|
|
_ => Err(()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
|
pub struct Message {
|
|
pub role: Role,
|
|
pub text: String,
|
|
}
|
|
|
|
impl Message {
|
|
pub fn new(role: Role, text: &str) -> Self {
|
|
Message {
|
|
role,
|
|
text: text.to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Serialize, Deserialize)]
|
|
pub struct Dialogue {
|
|
messages: Vec<Message>,
|
|
}
|
|
|
|
impl Dialogue {
|
|
pub fn new() -> Self {
|
|
Dialogue { messages: vec![] }
|
|
}
|
|
|
|
pub async fn do_turn<T: TokenProvider + Clone>(
|
|
&mut self,
|
|
gemini: &GeminiClient<T>,
|
|
message: &str,
|
|
) -> Result<Message> {
|
|
self.messages.push(Message::new(Role::User, message));
|
|
let response = gemini.prompt_conversation(&self.messages).await?;
|
|
self.messages.push(response.clone());
|
|
Ok(response)
|
|
}
|
|
}
|