unfunny/baibot-withmcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for prompt variables (bot name, date/time, model id)

Browse Source 
Fixes https://github.com/etkecc/baibot/issues/10

This also includes them in the default prompts (for newly-created agents),
so that people can get a better experience out of the box.
This commit is contained in:
Slavi Pantaleev
2024-09-21 14:05:19 +00:00
parent 0ee663ee92
commit 2a5a2d6a4d
26 changed files with 218 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/agent/mod.rs
View File

@@ -19,3 +19,7 @@ pub use instantiation::Result as AgentInstantiationResult;
pub use provider::{AgentProvider, AgentProviderInfo, ControllerTrait};
pub use purpose::AgentPurpose;
pub(super) fn default_prompt() -> &'static str {
"You are a brief, but helpful bot called {{ baibot_name }} powered by the {{ baibot_model_id }} model. The date/time now is: {{ baibot_now_utc }}."
}

4
src/agent/provider/anthropic/config.rs
View File

@@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize};
use anthropic_rs::models::claude::ClaudeModel;
use crate::agent::provider::ConfigTrait;
use crate::agent::{default_prompt, provider::ConfigTrait};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
@@ -58,7 +58,7 @@ impl Default for TextGenerationConfig {
fn default() -> Self {
Self {
model_id: default_text_model_id(),
prompt: Some("You are a brief, but helpful bot.".to_owned()),
prompt: Some(default_prompt().to_owned()),
temperature: super::super::default_temperature(),
max_response_tokens: 8192,
max_context_tokens: 204_800,

36
src/agent/provider/anthropic/controller.rs
View File

@@ -95,11 +95,12 @@ impl ControllerTrait for Controller {
));
};
let prompt_text = params
.prompt_override
.unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
.trim()
.to_owned();
let prompt_text = params.prompt_variables.format(
params
.prompt_override
.unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
.trim(),
);
let prompt_message = if prompt_text.is_empty() {
None
@@ -225,20 +226,25 @@ impl ControllerTrait for Controller {
}
}
fn text_generation_prompt(&self) -> Option<String> {
let Some(text_generation_config) = &self.config.text_generation else {
return None;
};
fn text_generation_model_id(&self) -> Option<String> {
self.config
.text_generation
.as_ref()
.map(|config| config.model_id.to_owned())
}
text_generation_config.prompt.clone()
fn text_generation_prompt(&self) -> Option<String> {
self.config
.text_generation
.as_ref()
.and_then(|config| config.prompt.clone())
}
fn text_generation_temperature(&self) -> Option<f32> {
let Some(text_generation_config) = &self.config.text_generation else {
return None;
};
Some(text_generation_config.temperature)
self.config
.text_generation
.as_ref()
.map(|config| config.temperature)
}
fn text_to_speech_voice(&self) -> Option<String> {

10
src/agent/provider/controller.rs
View File

@@ -13,6 +13,8 @@ pub trait ControllerTrait {
fn ping(&self) -> impl std::future::Future<Output = anyhow::Result<PingResult>> + Send;
fn text_generation_model_id(&self) -> Option<String>;
fn text_generation_prompt(&self) -> Option<String>;
fn text_generation_temperature(&self) -> Option<f32>;
@@ -63,6 +65,14 @@ impl ControllerTrait for ControllerType {
}
}
fn text_generation_model_id(&self) -> Option<String> {
match &self {
ControllerType::OpenAI(controller) => controller.text_generation_model_id(),
ControllerType::OpenAICompat(controller) => controller.text_generation_model_id(),
ControllerType::Anthropic(controller) => controller.text_generation_model_id(),
}
}
fn text_generation_prompt(&self) -> Option<String> {
match &self {
ControllerType::OpenAI(controller) => controller.text_generation_prompt(),

4
src/agent/provider/entity/mod.rs
View File

@@ -9,5 +9,7 @@ pub use agent_provider::{AgentProvider, AgentProviderInfo};
pub use image_generation::{ImageGenerationParams, ImageGenerationResult};
pub use ping::PingResult;
pub use speech_to_text::{SpeechToTextParams, SpeechToTextResult};
pub use text_generation::{TextGenerationParams, TextGenerationResult};
pub use text_generation::{
TextGenerationParams, TextGenerationPromptVariables, TextGenerationResult,
};
pub use text_to_speech::{TextToSpeechParams, TextToSpeechResult};

5
src/agent/provider/entity/text_generation.rs → src/agent/provider/entity/text_generation/mod.rs
View File

@@ -1,8 +1,13 @@
mod prompt_variables;
pub use prompt_variables::TextGenerationPromptVariables;
#[derive(Default)]
pub struct TextGenerationParams {
pub context_management_enabled: bool,
pub prompt_override: Option<String>,
pub temperature_override: Option<f32>,
pub prompt_variables: TextGenerationPromptVariables,
}
pub struct TextGenerationResult {

75
src/agent/provider/entity/text_generation/prompt_variables.rs Normal file
View File

@@ -0,0 +1,75 @@
use chrono::{DateTime, Utc};
use std::collections::HashMap;
pub struct TextGenerationPromptVariables {
map: HashMap<String, String>,
}
impl Default for TextGenerationPromptVariables {
fn default() -> Self {
Self::new("unnamed", "unknown-model", Utc::now())
}
}
impl TextGenerationPromptVariables {
pub fn new(bot_name: &str, model_id: &str, utc_time: DateTime<Utc>) -> Self {
let mut map = HashMap::new();
map.insert("baibot_name".to_string(), bot_name.to_string());
map.insert("baibot_model_id".to_string(), model_id.to_string());
map.insert("baibot_now_utc".to_string(), format_utc_time(utc_time));
Self { map }
}
pub fn format(&self, text: &str) -> String {
let mut formatted_text = text.to_string();
for (key, value) in &self.map {
let placeholder = format!("{{{{ {} }}}}", key);
formatted_text = formatted_text.replace(&placeholder, value);
}
formatted_text
}
}
fn format_utc_time(time: DateTime<Utc>) -> String {
time.format("%Y-%m-%d (%A), %H:%M:%S UTC").to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{TimeZone, Timelike};
#[test]
fn test_new() {
// Intentionally injecting some sub-seconds to ensure formatting would ignore them.
let now_utc = Utc
.with_ymd_and_hms(2024, 9, 20, 18, 34, 15)
.unwrap()
.with_nanosecond(250000000)
.unwrap();
let variables = TextGenerationPromptVariables::new("baibot", "gpt-4o", now_utc);
assert_eq!(
variables.map.get("baibot_name"),
Some(&"baibot".to_string())
);
assert_eq!(
variables.map.get("baibot_model_id"),
Some(&"gpt-4o".to_string())
);
assert_eq!(
variables.map.get("baibot_now_utc"),
Some(&format_utc_time(now_utc))
);
let prompt = "Hello, I'm {{ baibot_name }} using {{ baibot_model_id }}. The date/time now is {{ baibot_now_utc }}.";
let expected = "Hello, I'm baibot using gpt-4o. The date/time now is 2024-09-20 (Friday), 18:34:15 UTC.";
assert_eq!(variables.format(prompt), expected);
}
}

2
src/agent/provider/mod.rs
View File

@@ -21,5 +21,5 @@ pub use config::ConfigTrait;
pub use entity::{
AgentProvider, AgentProviderInfo, ImageGenerationParams, PingResult, SpeechToTextParams,
SpeechToTextResult, TextGenerationParams, TextToSpeechParams,
SpeechToTextResult, TextGenerationParams, TextGenerationPromptVariables, TextToSpeechParams,
};

4
src/agent/provider/openai/config.rs
View File

@@ -1,6 +1,6 @@
use serde::{Deserialize, Serialize};
use crate::agent::provider::ConfigTrait;
use crate::agent::{default_prompt, provider::ConfigTrait};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
@@ -66,7 +66,7 @@ impl Default for TextGenerationConfig {
fn default() -> Self {
Self {
model_id: default_text_model_id(),
prompt: Some("You are a brief, but helpful bot.".to_owned()),
prompt: Some(default_prompt().to_owned()),
temperature: super::super::default_temperature(),
max_response_tokens: 16_384,
max_context_tokens: 128_000,

36
src/agent/provider/openai/controller.rs
View File

@@ -86,11 +86,12 @@ impl ControllerTrait for Controller {
));
};
let prompt_text = params
.prompt_override
.unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
.trim()
.to_owned();
let prompt_text = params.prompt_variables.format(
params
.prompt_override
.unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
.trim(),
);
let prompt_message = if prompt_text.is_empty() {
None
@@ -391,20 +392,25 @@ impl ControllerTrait for Controller {
}
}
fn text_generation_prompt(&self) -> Option<String> {
let Some(text_generation_config) = &self.config.text_generation else {
return None;
};
fn text_generation_model_id(&self) -> Option<String> {
self.config
.text_generation
.as_ref()
.map(|config| config.model_id.to_owned())
}
text_generation_config.prompt.clone()
fn text_generation_prompt(&self) -> Option<String> {
self.config
.text_generation
.as_ref()
.and_then(|config| config.prompt.clone())
}
fn text_generation_temperature(&self) -> Option<f32> {
let Some(text_generation_config) = &self.config.text_generation else {
return None;
};
Some(text_generation_config.temperature)
self.config
.text_generation
.as_ref()
.map(|config| config.temperature)
}
fn text_to_speech_voice(&self) -> Option<String> {

3
src/agent/provider/openai_compat/config.rs
View File

@@ -1,5 +1,6 @@
use serde::{Deserialize, Serialize};
use crate::agent::default_prompt;
use crate::agent::provider::openai::{
ImageGenerationConfig as OpenAIImageGenerationConfig,
SpeechToTextConfig as OpenAISpeechToTextConfig,
@@ -75,7 +76,7 @@ impl Default for TextGenerationConfig {
fn default() -> Self {
Self {
model_id: default_text_model_id(),
prompt: Some("You are a brief, but helpful bot.".to_owned()),
prompt: Some(default_prompt().to_owned()),
temperature: super::super::default_temperature(),
max_response_tokens: 4096,
max_context_tokens: 128_000,

36
src/agent/provider/openai_compat/controller.rs
View File

@@ -84,11 +84,12 @@ impl ControllerTrait for Controller {
));
};
let prompt_text = params
.prompt_override
.unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
.trim()
.to_owned();
let prompt_text = params.prompt_variables.format(
params
.prompt_override
.unwrap_or(self.text_generation_prompt().unwrap_or("".to_owned()))
.trim(),
);
let prompt_message = if prompt_text.is_empty() {
None
@@ -409,20 +410,25 @@ impl ControllerTrait for Controller {
}
}
fn text_generation_prompt(&self) -> Option<String> {
let Some(text_generation_config) = &self.config.text_generation else {
return None;
};
fn text_generation_model_id(&self) -> Option<String> {
self.config
.text_generation
.as_ref()
.map(|config| config.model_id.to_owned())
}
text_generation_config.prompt.clone()
fn text_generation_prompt(&self) -> Option<String> {
self.config
.text_generation
.as_ref()
.and_then(|config| config.prompt.clone())
}
fn text_generation_temperature(&self) -> Option<f32> {
let Some(text_generation_config) = &self.config.text_generation else {
return None;
};
Some(text_generation_config.temperature)
self.config
.text_generation
.as_ref()
.map(|config| config.temperature)
}
fn text_to_speech_voice(&self) -> Option<String> {

19
src/controller/chat_completion/mod.rs
View File

@@ -4,7 +4,9 @@ use mxlink::{MatrixLink, MessageResponseType};
use tracing::Instrument;
use crate::agent::provider::{SpeechToTextParams, TextGenerationParams};
use crate::agent::provider::{
SpeechToTextParams, TextGenerationParams, TextGenerationPromptVariables,
};
use crate::agent::AgentInstance;
use crate::agent::AgentPurpose;
use crate::agent::ControllerTrait;
@@ -405,6 +407,16 @@ async fn handle_stage_text_generation(
let start_time = std::time::Instant::now();
let controller = agent.controller();
let prompt_variables = TextGenerationPromptVariables::new(
bot.name(),
&controller
.text_generation_model_id()
.unwrap_or("unknown-model".to_owned()),
chrono::Utc::now(),
);
let params = TextGenerationParams {
context_management_enabled: message_context
.room_config_context()
@@ -417,10 +429,11 @@ async fn handle_stage_text_generation(
temperature_override: message_context
.room_config_context()
.text_generation_temperature_override(),
prompt_variables,
};
let result = agent
.controller()
let result = controller
.generate_text(conversation, params)
.instrument(span)
.await;