Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,364 @@
|
||||
//! 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"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user