365 lines
9.9 KiB
Rust
365 lines
9.9 KiB
Rust
//! Dialogue system for conversational AI
|
|
//!
|
|
//! Provides QA and multi-turn conversation support with context management.
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::VecDeque;
|
|
|
|
/// Dialogue context for multi-turn conversations
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct DialogueContext {
|
|
/// Conversation history
|
|
pub history: VecDeque<DialogueTurn>,
|
|
/// Maximum number of turns to keep in history
|
|
pub max_history: usize,
|
|
/// System prompt/context
|
|
pub system_prompt: Option<String>,
|
|
/// Session metadata
|
|
pub metadata: DialogueMetadata,
|
|
}
|
|
|
|
/// A single turn in a dialogue
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DialogueTurn {
|
|
/// Role of the speaker
|
|
pub role: Role,
|
|
/// Content of the message
|
|
pub content: String,
|
|
/// Timestamp of the turn
|
|
pub timestamp: chrono::DateTime<chrono::Utc>,
|
|
/// Optional metadata for this turn
|
|
pub metadata: Option<TurnMetadata>,
|
|
}
|
|
|
|
/// Speaker role in a conversation
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum Role {
|
|
/// User/human input
|
|
User,
|
|
/// AI assistant response
|
|
Assistant,
|
|
/// System message/prompt
|
|
System,
|
|
/// Tool/function call result
|
|
Tool,
|
|
}
|
|
|
|
/// Metadata for dialogue sessions
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct DialogueMetadata {
|
|
/// Session identifier
|
|
pub session_id: Option<String>,
|
|
/// User identifier
|
|
pub user_id: Option<String>,
|
|
/// Additional custom attributes
|
|
pub attributes: std::collections::HashMap<String, String>,
|
|
}
|
|
|
|
/// Metadata for individual turns
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct TurnMetadata {
|
|
/// Generation confidence score
|
|
pub confidence: Option<f32>,
|
|
/// Number of tokens in this turn
|
|
pub token_count: Option<usize>,
|
|
/// Processing time in milliseconds
|
|
pub latency_ms: Option<f64>,
|
|
}
|
|
|
|
impl DialogueContext {
|
|
/// Create a new dialogue context with specified history limit
|
|
pub fn new(max_history: usize) -> Self {
|
|
Self {
|
|
history: VecDeque::new(),
|
|
max_history,
|
|
system_prompt: None,
|
|
metadata: DialogueMetadata::default(),
|
|
}
|
|
}
|
|
|
|
/// Create a new dialogue context with a system prompt
|
|
pub fn with_system_prompt(max_history: usize, system_prompt: impl Into<String>) -> Self {
|
|
Self {
|
|
history: VecDeque::new(),
|
|
max_history,
|
|
system_prompt: Some(system_prompt.into()),
|
|
metadata: DialogueMetadata::default(),
|
|
}
|
|
}
|
|
|
|
/// Add a turn to the conversation
|
|
pub fn add_turn(&mut self, role: Role, content: impl Into<String>) {
|
|
let turn = DialogueTurn {
|
|
role,
|
|
content: content.into(),
|
|
timestamp: chrono::Utc::now(),
|
|
metadata: None,
|
|
};
|
|
self.history.push_back(turn);
|
|
|
|
// Trim history if needed
|
|
while self.history.len() > self.max_history {
|
|
self.history.pop_front();
|
|
}
|
|
}
|
|
|
|
/// Add a turn with metadata
|
|
pub fn add_turn_with_metadata(
|
|
&mut self,
|
|
role: Role,
|
|
content: impl Into<String>,
|
|
metadata: TurnMetadata,
|
|
) {
|
|
let turn = DialogueTurn {
|
|
role,
|
|
content: content.into(),
|
|
timestamp: chrono::Utc::now(),
|
|
metadata: Some(metadata),
|
|
};
|
|
self.history.push_back(turn);
|
|
|
|
while self.history.len() > self.max_history {
|
|
self.history.pop_front();
|
|
}
|
|
}
|
|
|
|
/// Add a user message
|
|
pub fn add_user_message(&mut self, content: impl Into<String>) {
|
|
self.add_turn(Role::User, content);
|
|
}
|
|
|
|
/// Add an assistant response
|
|
pub fn add_assistant_message(&mut self, content: impl Into<String>) {
|
|
self.add_turn(Role::Assistant, content);
|
|
}
|
|
|
|
/// Clear conversation history
|
|
pub fn clear(&mut self) {
|
|
self.history.clear();
|
|
}
|
|
|
|
/// Get the last N turns
|
|
pub fn last_n_turns(&self, n: usize) -> impl Iterator<Item = &DialogueTurn> {
|
|
self.history.iter().rev().take(n).rev()
|
|
}
|
|
|
|
/// Format the conversation for model input
|
|
pub fn format_for_model(&self) -> String {
|
|
let mut formatted = String::new();
|
|
|
|
if let Some(ref system) = self.system_prompt {
|
|
formatted.push_str(&format!("System: {system}\n\n"));
|
|
}
|
|
|
|
for turn in &self.history {
|
|
let role_str = match turn.role {
|
|
Role::User => "User",
|
|
Role::Assistant => "Assistant",
|
|
Role::System => "System",
|
|
Role::Tool => "Tool",
|
|
};
|
|
formatted.push_str(&format!("{}: {}\n", role_str, turn.content));
|
|
}
|
|
|
|
formatted
|
|
}
|
|
|
|
/// Get total token estimate (rough approximation)
|
|
pub fn estimate_tokens(&self) -> usize {
|
|
let mut total = 0;
|
|
|
|
if let Some(ref system) = self.system_prompt {
|
|
total += system.split_whitespace().count() * 4 / 3; // Rough token estimate
|
|
}
|
|
|
|
for turn in &self.history {
|
|
total += turn.content.split_whitespace().count() * 4 / 3;
|
|
}
|
|
|
|
total
|
|
}
|
|
|
|
/// Check if the conversation is empty
|
|
pub fn is_empty(&self) -> bool {
|
|
self.history.is_empty()
|
|
}
|
|
|
|
/// Get the number of turns
|
|
pub fn len(&self) -> usize {
|
|
self.history.len()
|
|
}
|
|
}
|
|
|
|
impl DialogueTurn {
|
|
/// Create a new dialogue turn
|
|
pub fn new(role: Role, content: impl Into<String>) -> Self {
|
|
Self {
|
|
role,
|
|
content: content.into(),
|
|
timestamp: chrono::Utc::now(),
|
|
metadata: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Question-Answering context and utilities
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct QAContext {
|
|
/// The question being asked
|
|
pub question: String,
|
|
/// Relevant context passages
|
|
pub context_passages: Vec<String>,
|
|
/// Expected answer type
|
|
pub answer_type: AnswerType,
|
|
}
|
|
|
|
/// Types of expected answers
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum AnswerType {
|
|
/// Short extractive answer
|
|
Extractive,
|
|
/// Longer abstractive answer
|
|
Abstractive,
|
|
/// Yes/No answer
|
|
Boolean,
|
|
/// Numeric answer
|
|
Numeric,
|
|
/// List/enumeration
|
|
List,
|
|
}
|
|
|
|
impl QAContext {
|
|
/// Create a new QA context
|
|
pub fn new(question: impl Into<String>) -> Self {
|
|
Self {
|
|
question: question.into(),
|
|
context_passages: Vec::new(),
|
|
answer_type: AnswerType::Abstractive,
|
|
}
|
|
}
|
|
|
|
/// Add a context passage
|
|
pub fn with_context(mut self, passage: impl Into<String>) -> Self {
|
|
self.context_passages.push(passage.into());
|
|
self
|
|
}
|
|
|
|
/// Set the expected answer type
|
|
pub fn with_answer_type(mut self, answer_type: AnswerType) -> Self {
|
|
self.answer_type = answer_type;
|
|
self
|
|
}
|
|
|
|
/// Format for model input
|
|
pub fn format_for_model(&self) -> String {
|
|
let mut formatted = String::new();
|
|
|
|
if !self.context_passages.is_empty() {
|
|
formatted.push_str("Context:\n");
|
|
for (i, passage) in self.context_passages.iter().enumerate() {
|
|
formatted.push_str(&format!("[{}] {}\n", i + 1, passage));
|
|
}
|
|
formatted.push('\n');
|
|
}
|
|
|
|
formatted.push_str(&format!("Question: {}\n", self.question));
|
|
formatted.push_str("Answer:");
|
|
|
|
formatted
|
|
}
|
|
}
|
|
|
|
/// Conversation manager for handling multiple dialogue sessions
|
|
pub struct ConversationTracker {
|
|
/// Active conversations by session ID
|
|
sessions: std::collections::HashMap<String, DialogueContext>,
|
|
/// Default max history for new sessions
|
|
default_max_history: usize,
|
|
}
|
|
|
|
impl ConversationTracker {
|
|
/// Create a new conversation tracker
|
|
pub fn new(default_max_history: usize) -> Self {
|
|
Self {
|
|
sessions: std::collections::HashMap::new(),
|
|
default_max_history,
|
|
}
|
|
}
|
|
|
|
/// Get or create a session
|
|
pub fn get_or_create(&mut self, session_id: &str) -> &mut DialogueContext {
|
|
self.sessions
|
|
.entry(session_id.to_string())
|
|
.or_insert_with(|| DialogueContext::new(self.default_max_history))
|
|
}
|
|
|
|
/// Get an existing session
|
|
pub fn get(&self, session_id: &str) -> Option<&DialogueContext> {
|
|
self.sessions.get(session_id)
|
|
}
|
|
|
|
/// Get a mutable reference to an existing session
|
|
pub fn get_mut(&mut self, session_id: &str) -> Option<&mut DialogueContext> {
|
|
self.sessions.get_mut(session_id)
|
|
}
|
|
|
|
/// Remove a session
|
|
pub fn remove(&mut self, session_id: &str) -> Option<DialogueContext> {
|
|
self.sessions.remove(session_id)
|
|
}
|
|
|
|
/// Get all active session IDs
|
|
pub fn session_ids(&self) -> impl Iterator<Item = &String> {
|
|
self.sessions.keys()
|
|
}
|
|
|
|
/// Get the number of active sessions
|
|
pub fn session_count(&self) -> usize {
|
|
self.sessions.len()
|
|
}
|
|
}
|
|
|
|
impl Default for ConversationTracker {
|
|
fn default() -> Self {
|
|
Self::new(20)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_dialogue_context() {
|
|
let mut ctx = DialogueContext::new(5);
|
|
|
|
ctx.add_user_message("Hello");
|
|
ctx.add_assistant_message("Hi there!");
|
|
|
|
assert_eq!(ctx.len(), 2);
|
|
assert!(!ctx.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_max_history() {
|
|
let mut ctx = DialogueContext::new(3);
|
|
|
|
for i in 0..5 {
|
|
ctx.add_user_message(format!("Message {}", i));
|
|
}
|
|
|
|
assert_eq!(ctx.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_qa_context() {
|
|
let qa = QAContext::new("What is Rust?")
|
|
.with_context("Rust is a systems programming language.")
|
|
.with_answer_type(AnswerType::Abstractive);
|
|
|
|
let formatted = qa.format_for_model();
|
|
assert!(formatted.contains("What is Rust?"));
|
|
assert!(formatted.contains("systems programming"));
|
|
}
|
|
}
|