itsamejms
119 lines
3.7 KiB
Rust
119 lines
3.7 KiB
Rust
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,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct LlmConfig {
|
|
pub api_url: String,
|
|
pub api_key: String,
|
|
pub model: String,
|
|
pub temperature: f32,
|
|
pub max_tokens: u32,
|
|
pub top_p: f32,
|
|
/// Ollama image-generation model (e.g. `x/flux2-klein:4b`). macOS-only via Ollama.
|
|
pub image_model: 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
|
|
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_model: "x/flux2-klein:4b".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: detect if we're talking to Ollama ───────────────
|
|
|
|
pub fn is_ollama(url: &str) -> bool {
|
|
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
|
|
} |