dm-pal/src-tauri/src/llm/mod.rs

use serde::{Deserialize, Serialize};
use std::sync::Mutex;

// ─── LLM Event (for streaming) ───────────────────────────────

#[derive(Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "data", rename_all = "camelCase")]
pub enum LlmEvent {
    Token(String),
    Done(String),   // full text
    Error(String),
}

// ─── App State ────────────────────────────────────────────────

pub struct AppState {
    pub config: Mutex<LlmConfig>,
    pub rag: crate::rag::RagStore,
    pub gen: crate::generations::GenerationStore,
    /// Where lore.db, generations.db, and cached images live. Configurable via
    /// the `dataDir` key in the `dm-pal-prefs.json` store; defaults to
    /// `$APPDATA/dm-toolkit` when unset. Set once at startup.
    pub data_dir: std::path::PathBuf,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmConfig {
    /// API dialect: "ollama" (native /api/*) or "openai" (/v1/*).
    /// Empty = sniff the URL (legacy configs + defaults). Provider presets
    /// set it explicitly — a custom-port Ollama would otherwise be sniffed
    /// as OpenAI and fail confusingly. serde default: old stored configs.
    #[serde(default)]
    pub provider: String,
    pub api_url: String,
    pub api_key: String,
    pub model: String,
    pub temperature: f32,
    pub max_tokens: u32,
    pub top_p: f32,
    /// Base URL of a stable-diffusion.cpp `sd-server` (A1111-compatible API),
    /// e.g. `http://localhost:1234`. Cross-platform — no longer Ollama/macOS-only.
    pub image_api_url: String,
    /// Ollama embedding model for lore RAG (e.g. `nomic-embed-text`).
    pub embed_model: String,
}

impl Default for LlmConfig {
    fn default() -> Self {
        Self {
            // Default to Ollama local server; also works with LM Studio, llama.cpp server, or OpenAI
            provider: String::new(),
            api_url: "http://localhost:11434".to_string(),
            api_key: String::new(),
            model: "llama3.2".to_string(),
            temperature: 0.7,
            max_tokens: 512,
            top_p: 0.9,
            image_api_url: "http://localhost:1234".to_string(),
            embed_model: "nomic-embed-text".to_string(),
        }
    }
}

// ─── Generation Request ───────────────────────────────────────

#[derive(Debug, Deserialize)]
pub struct GenerateRequest {
    pub prompt: String,
    pub system: Option<String>,
    pub temperature: Option<f32>,
    pub max_tokens: Option<u32>,
    /// Optional RAG query: when set, the top lore chunks for this query are
    /// retrieved and prepended to the system prompt so generation stays
    /// consistent with the user's world bible.
    pub rag_query: Option<String>,
}

// ─── OpenAI-Compatible Chat Response ─────────────────────────

#[derive(Debug, Deserialize)]
pub struct ChatResponse {
    pub choices: Vec<ChatChoice>,
}

#[derive(Debug, Deserialize)]
pub struct ChatChoice {
    pub message: ChatMessage,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: String,
    pub content: String,
}

// ─── Ollama Generate Response ─────────────────────────────────

#[derive(Debug, Deserialize)]
pub struct OllamaGenerateResponse {
    pub response: String,
    pub done: bool,
}

#[derive(Debug, Deserialize)]
pub struct OllamaChatResponse {
    pub message: ChatMessage,
    pub done: bool,
}

// ─── Helper: are we talking Ollama? ───────────────────────

/// True when we should speak Ollama's native /api/* protocol. An explicit
/// provider on the config wins; empty falls back to URL sniffing so legacy
/// and default configs keep working.
pub fn is_ollama(provider: &str, url: &str) -> bool {
    if provider.eq_ignore_ascii_case("ollama") {
        return true;
    }
    if provider.eq_ignore_ascii_case("openai") {
        return false;
    }
    url.contains("localhost:11434") || url.contains("127.0.0.1:11434")
}

// ─── Build system + user messages from a prompt ──────────────

pub fn build_messages(system: Option<&str>, prompt: &str) -> Vec<ChatMessage> {
    let mut messages = Vec::new();
    if let Some(sys) = system {
        messages.push(ChatMessage {
            role: "system".to_string(),
            content: sys.to_string(),
        });
    }
    messages.push(ChatMessage {
        role: "user".to_string(),
        content: prompt.to_string(),
    });
    messages
}