chatrpg-api/src/logic/prompt.rs

use actix_web::{HttpResponse, web::Bytes};
use reqwest::Client;
use serde::Serialize;
use serde_json::{Value, json, from_str};
use futures_util::{StreamExt, stream};
use tokio::sync::oneshot;
use chrono::Utc;
use std::sync::{Arc, Mutex};
use crate::{
    http_error::HttpError,
    models::{
        game::Game,
        character::Character,
        session::Session,
        turn::Turn
    },
    DEEPINFRA_TOKEN
};

#[derive(Serialize)]
struct PromptMessage {
    pub role: PromptRole,
    pub content: String
}

#[derive(Serialize)]
#[serde(rename_all = "lowercase")]
enum PromptRole {
    User,
    System,
    Assistant
}

pub struct PromptResponse {
    pub message: oneshot::Receiver<(String, i32, i32)>,
    pub sse: HttpResponse
}

pub async fn prompt(
    game: Game,
    character: Character,
    sessions: Vec<Session>,
    turns: Vec<Turn>,
    new_input: Option<String>,
    new_game: bool
) -> Result<PromptResponse, HttpError> {
    let system_message = create_system_message(game, character, sessions, new_game);
    let mut prompt = create_turn_prompts(turns);
    prompt.insert(0, system_message);
    if let Some(p) = new_input {
        prompt.push(PromptMessage {
            role: PromptRole::User,
            content: p
        });
    }
    send_message(prompt).await
}

fn create_system_message(
    game: Game,
    character: Character,
    sessions: Vec<Session>,
    new_game: bool
) -> PromptMessage {
    let mut text = "".to_string();
    text += "You are acting as a game master for an RPG and running the game.\n";
    text += &format!("Here is the world context: {}\n", game.game_context);
    text += &format!(
        "Character name is {}, here is a description of the character: {}\n",
        character.name,
        character.description
    );

    if sessions.len() == 0 {
        text += "This is the first session in this game.\n";
    } else {
        for(i, session) in sessions.iter().enumerate() {
            match &session.summary {
                Some(s) => {text += &format!("Summary of session {}: {}\n", i, s)},
                None => ()
            }
        }
    }

    if new_game {
        text += "Start off a new session for the player";
    }

    PromptMessage {
        role: PromptRole::System,
        content: text
    }
}

fn create_turn_prompts(turns: Vec<Turn>) -> Vec<PromptMessage> {
    let mut prompts = Vec::new();

    for turn in turns {
        prompts.push(PromptMessage {
            role: PromptRole::User,
            content: turn.user_text
        });
        prompts.push(PromptMessage {
            role: PromptRole::Assistant,
            content: turn.llm_text
        });
    }

    prompts
}

async fn send_message(prompt: Vec<PromptMessage>) -> Result<PromptResponse, HttpError> {
    let client = Client::new();
    let key = DEEPINFRA_TOKEN.get().unwrap();

    let body = json!({
        "model": "Qwen/Qwen3-32B",
        "messages": prompt,
        "reasoning_effort": "none",
        "temperature": 0.7,
        "stream": true,
        "stream_options": {"include_usage": true},
        "max_tokens": 4096
    });

    let response = client
        .post("https://api.deepinfra.com/v1/chat/completions")
        .header("Authorization", format!("Bearer {}", key))
        .header("Content-Type", "application/json")
        .json(&body)
        .send()
        .await
        .map_err(|e| HttpError::InternalError(e.to_string()))?;

    let (tx, rx) = oneshot::channel::<(String, i32, i32)>();
    let full_text = Arc::new(Mutex::new(String::new()));
    let acc = full_text.clone();

    let stream = response.bytes_stream().map(move |chunk| -> Result<Bytes, HttpError> {
        match chunk {
            Ok(bytes) => {
                acc.lock().unwrap().push_str(&String::from_utf8_lossy(&bytes));
                Ok(bytes)
            },
            Err(e) => Err(HttpError::InternalError(e.to_string()))
        }
    }).chain(stream::once(async move {
        let raw = full_text.lock().unwrap().clone();

        let content = raw.lines()
            .filter_map(|l| l.strip_prefix("data: "))
            .filter(|d| *d != "[DONE]")
            .filter_map(|d| from_str::<Value>(d).ok())
            .filter_map(|v| v["choices"][0]["delta"]["content"].as_str().map(str::to_string))
            .collect::<String>();

        let (input_tokens, output_tokens) = raw.lines()
            .filter_map(|l| l.strip_prefix("data: "))
            .filter_map(|d| from_str::<Value>(d).ok())
            .find_map(|v| {
                let usage = v.get("usage")?;
                if usage.is_null() { return None; }
                Some(usage.clone())
            })
            .map(|u| (
                u["prompt_tokens"].as_i64().unwrap_or(0) as i32,
                u["completion_tokens"].as_i64().unwrap_or(0) as i32
            ))
            .unwrap_or((0, 0));

        let meta = json!({
            "tokens": input_tokens + (output_tokens * 2),
            "created_at": Utc::now(),
            "something": "else"
        });

        let sse_chunk = format!("event: meta\ndata: {}\n\n", meta.to_string());

        let _ = tx.send((content, input_tokens, output_tokens));
        Ok::<Bytes, HttpError>(Bytes::from(sse_chunk))
    }));

    Ok(PromptResponse {
        message: rx,
        sse: HttpResponse::Ok()
            .content_type("text/event-stream")
            .streaming(stream)
    })
}