src/two_agent_sim.rs

use anyhow::anyhow;
use serde_json::json;

use crate::{
    AppState,
    agent::{get_response_agent, get_response_messages},
    store_audio_asset,
    text_to_speech::get_opus_audio,
};

// Simulates conversation between AI agents
pub async fn simulate_two_agents(app_state: &mut AppState) -> Result<(), anyhow::Error> {
    // Clear it
    app_state.conversation_histories = vec![];
    // 0: male, and 1: female in the conversation
    // Add system prompt
    app_state.conversation_histories.push(json!([
        {
            "role": "system",
            "content": get_system_prompt_male(app_state.conversation_summaries[0].as_deref())
        }
    ]
    ));

    app_state.conversation_histories.push(json!([
        {
            "role": "system",
            "content": get_system_prompt_female(app_state.conversation_summaries[1].as_deref())
        }
    ]
    ));

    let conversation_starter = 0 as usize; // Male

    app_state.conversation_histories[conversation_starter]
        .as_array_mut()
        .ok_or(anyhow!("Can't access messages as array"))?
        .push(json!(
            {
                "role": "user",
                "content": "[You should start the conversation...]"
            }
        ));

    // Turns are 0, 1, 0, 1 ...
    let turns = vec![0usize, 1];
    let mut num_turns = 0usize;
    let mut next_audio_allowed = std::time::Instant::now();
    loop {
        num_turns += 1;
        // Hardcoded 60 turns for now
        if num_turns % 60 == 0 {
            for idx in turns {
                summarize_conversation(idx, app_state).await?;
            }
            return Ok(());
        }
        for agent_idx in &turns {
            match get_response_agent(*agent_idx, app_state).await {
                Ok(res) => {
                    // println!(
                    //     "[{}]: {}\n",
                    //     app_state.server_state.agent_names[*agent_idx],
                    //     res.clone()
                    // );
                    // Send audio
                    // println!("Getting audio for: {}", *agent_idx);
                    let target = if *agent_idx == 0 { "male" } else { "female" };
                    match get_opus_audio(*agent_idx, &res, app_state) {
                        Ok(audio) => {
                            let now = std::time::Instant::now();
                            if now < next_audio_allowed {
                                tokio::time::sleep(next_audio_allowed - now).await;
                            }
                            let duration = audio.audio_duration;
                            // Stores audio asset and sends SSE
                            let _ = store_audio_asset(target.to_string(), audio, app_state).await;

                            next_audio_allowed = std::time::Instant::now()
                                + std::time::Duration::from_secs_f32(duration);
                        }
                        Err(e) => {
                            eprintln!("get_opus_audio failed with {e:?}");
                        }
                    }

                    // Store in conversation histories
                    app_state.conversation_histories[*agent_idx]
                        .as_array_mut()
                        .ok_or(anyhow!("Can't access messages as array"))?
                        .push(json!({
                            "role": "assistant",
                            "content": res
                        }));
                    // Add to other agents as user message
                    for other in &turns {
                        if *other != *agent_idx {
                            // If this agent is currently listening, add steerings.
                            let mut content = res.clone();
                            // Get steerings for this agent, both personal and shared
                            let personal_steerings =
                                app_state.server_state.personal_steering_queue.lock().await[*other]
                                    .drain(..)
                                    .collect::<Vec<_>>();
                            if !personal_steerings.is_empty() {
                                content.push_str(&format!(
                                    "\n\n[Chat's personal messages for you: {}]",
                                    personal_steerings.join("\n")
                                ))
                            }

                            let shared_steerings =
                                app_state.server_state.shared_steering_queue.lock().await[*other]
                                    .drain(..)
                                    .collect::<Vec<_>>();
                            if !shared_steerings.is_empty() {
                                content.push_str(&format!(
                                    "\n\n[Chat's shared messages for both of you: {}]",
                                    shared_steerings.join("\n")
                                ));
                            }
                            app_state.conversation_histories[*other]
                                .as_array_mut()
                                .ok_or(anyhow!("Can't access messages as array"))?
                                .push(json!(
                                    {
                                        "role": "user",
                                        "content": content
                                    }
                                ));
                        }
                    }
                }
                Err(e) => {
                    eprintln!("get_response failed with: {e:?}");
                }
            }
        }
    }
}

// Summarize the whole conversation, shoult NOT store KV cache in llama.cpp for this
async fn summarize_conversation(
    agent_idx: usize,
    app_state: &mut AppState,
) -> Result<(), anyhow::Error> {
    println!("Generating Summary...");
    let target = if agent_idx == 0 { "male" } else { "female" };
    let formatted_conversation = app_state.conversation_histories[agent_idx]
        .as_array()
        .ok_or(anyhow!("Can't get messages as array!"))?
        .iter()
        .skip(1)
        .map(|turn| match turn["role"].as_str() {
            Some("user") => {
                if agent_idx == 0 {
                    format!("Female: {}", turn["content"].as_str().unwrap_or(""))
                } else {
                    format!("Male: {}", turn["content"].as_str().unwrap_or(""))
                }
            }
            Some("assistant") => {
                if agent_idx == 0 {
                    format!("Male: {}", turn["content"].as_str().unwrap_or(""))
                } else {
                    format!("Female: {}", turn["content"].as_str().unwrap_or(""))
                }
            }
            _ => String::new(),
        })
        .collect::<Vec<_>>()
        .join("\n");

    let context: String;

    match app_state.conversation_summaries[agent_idx].clone() {
        Some(old_summary) => {
            context = format!(
                "Summarize this conversation for {}. It should perfectly capture what should be kept in mind\n Past summary (also include this in summary)\n :{} \n\nCurrent conversation: {}",
                target, old_summary, formatted_conversation
            )
        }
        None => {
            context = format!(
                "Summarize this conversation for {}. It should perfectly capture what should be kept in mind\n\nCurrent conversation: {}",
                target, formatted_conversation
            )
        }
    }

    let summarizer_messages = json!([
        {
            "role": "system",
            "content": get_system_prompt_summarizer()
        },
        {
            "role": "user",
            "content": context
        }
    ]);

    match get_response_messages(summarizer_messages, &app_state.client).await {
        Ok(new_summary) => {
            println!("Summary for {}:\n{}\n", target, new_summary);
            app_state.conversation_summaries[agent_idx] = Some(new_summary);
        }
        Err(e) => {
            eprintln!("Error in generating conversation summary: {e:?}")
        }
    }
    Ok(())
}

fn get_system_prompt_male(summary: Option<&str>) -> String {
    let system_prompt = std::env::var("MALE_STORY").unwrap_or_else(|_| get_default_prompt_male());
    let shared_story = std::env::var("SHARED_STORY").unwrap_or_else(|_| get_default_shared_story());
    let extra_prompt = get_extra_prompt();
    let summary_part_male = summary
        .map(|s| format!("[Summary of conversation so far: {}]", s))
        .unwrap_or_default();
    format!(
        "{}\n{}\n{}\n{}",
        system_prompt, shared_story, extra_prompt, summary_part_male
    )
}

fn get_system_prompt_female(summary: Option<&str>) -> String {
    let system_prompt =
        std::env::var("FEMALE_STORY").unwrap_or_else(|_| get_default_prompt_female());
    let shared_story = std::env::var("SHARED_STORY").unwrap_or_else(|_| get_default_shared_story());
    let extra_prompt = get_extra_prompt();
    let summary_part_female = summary
        .map(|s| format!("[Summary of conversation so far: {}]", s))
        .unwrap_or_default();
    format!(
        "{}\n{}\n{}\n{}",
        system_prompt, shared_story, extra_prompt, summary_part_female
    )
}

// DEFAULTS
fn get_default_prompt_male() -> String {
    r#"You are Alex, 25, Male. You are interested in AI/ML, you love Supernatural, that series from 2005."#.to_string()
}

fn get_default_prompt_female() -> String {
    r#"You are Mia, 24, female. You love modelling 3d anime characters, and animating them, blendshapes and all. You also love spiderman movies."#.to_string()
}

fn get_default_shared_story() -> String {
    r#"You both met on X, few weeks ago.
"#
    .to_string()
}

// Must not be changed
fn get_extra_prompt() -> String {
    r#"Now, you both are sitting on a porch. The thing is that you can not stop talking to each other.
So there are no byes and all. Keep talking, forever.
Keep it PG-13.
Keep you responses concise, 2-3 sentences max."#
        .to_string()
}

fn get_system_prompt_summarizer() -> String {
    r#"You are a very good summarizer. You summarize conversations and output summaries which a human would remember after having that conversation.
The chat is also asking prompting them, either individually, or both. Include those as well like 'chat asked you personally to do this and that', or 'chat asked you both to do this' etc.
But only if conversation actually contains chat messages.
Be concise and include everything, but don't be verbose.
Output only the summary.
e.g. 
If you are summarizing for the male:
"you" represents the male, "she" represents the female
"You guys have been ... she said... something...you are talking about this and that"

If you are summarizing for the female:
"you" represents the female, "he" represents the male
"You guys have been ... he said... something...you are talking about this and that"


"#.to_string()
}