feat(core): conversation history with bounded length
split llm.rs into a module with separate client and history submodules. ConversationHistory holds an optional system prompt plus a FIFO of user/ assistant turns capped at max_turns; oldest turns evict on overflow, system prompt is preserved across truncation and clear().
This commit is contained in:
parent
23c57b855b
commit
40f2f8b5f9
3 changed files with 123 additions and 0 deletions
188
crates/jarvis-core/src/llm/client.rs
Normal file
188
crates/jarvis-core/src/llm/client.rs
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use std::env;
|
||||
use thiserror::Error;
|
||||
|
||||
pub const DEFAULT_BASE_URL: &str = "https://api.groq.com/openai/v1";
|
||||
pub const DEFAULT_MODEL: &str = "llama-3.3-70b-versatile";
|
||||
pub const ENV_TOKEN: &str = "GROQ_TOKEN";
|
||||
pub const ENV_BASE_URL: &str = "GROQ_BASE_URL";
|
||||
pub const ENV_MODEL: &str = "GROQ_MODEL";
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ConfigError {
|
||||
#[error("missing environment variable: {0}")]
|
||||
MissingEnv(&'static str),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum LlmError {
|
||||
#[error("HTTP request failed: {0}")]
|
||||
Http(#[from] reqwest::Error),
|
||||
#[error("failed to parse response JSON: {0}")]
|
||||
Deserialize(String),
|
||||
#[error("API returned error (status {status}): {body}")]
|
||||
Api { status: u16, body: String },
|
||||
#[error("response contained no choices")]
|
||||
EmptyResponse,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
pub fn user(content: impl Into<String>) -> Self {
|
||||
Self { role: "user".to_string(), content: content.into() }
|
||||
}
|
||||
|
||||
pub fn system(content: impl Into<String>) -> Self {
|
||||
Self { role: "system".to_string(), content: content.into() }
|
||||
}
|
||||
|
||||
pub fn assistant(content: impl Into<String>) -> Self {
|
||||
Self { role: "assistant".to_string(), content: content.into() }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ChatRequest<'a> {
|
||||
model: &'a str,
|
||||
messages: &'a [ChatMessage],
|
||||
max_tokens: u32,
|
||||
temperature: f32,
|
||||
top_p: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatResponse {
|
||||
choices: Vec<Choice>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct Choice {
|
||||
message: ResponseMessage,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ResponseMessage {
|
||||
content: String,
|
||||
}
|
||||
|
||||
pub struct LlmClient {
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
model: String,
|
||||
http: reqwest::blocking::Client,
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
pub fn new(
|
||||
base_url: impl Into<String>,
|
||||
api_key: impl Into<String>,
|
||||
model: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
base_url: base_url.into(),
|
||||
api_key: api_key.into(),
|
||||
model: model.into(),
|
||||
http: reqwest::blocking::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_env() -> Result<Self, ConfigError> {
|
||||
let api_key = env::var(ENV_TOKEN).map_err(|_| ConfigError::MissingEnv(ENV_TOKEN))?;
|
||||
let base_url = env::var(ENV_BASE_URL).unwrap_or_else(|_| DEFAULT_BASE_URL.to_string());
|
||||
let model = env::var(ENV_MODEL).unwrap_or_else(|_| DEFAULT_MODEL.to_string());
|
||||
Ok(Self::new(base_url, api_key, model))
|
||||
}
|
||||
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
pub fn complete(&self, messages: &[ChatMessage], max_tokens: u32) -> Result<String, LlmError> {
|
||||
let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
|
||||
let body = ChatRequest {
|
||||
model: &self.model,
|
||||
messages,
|
||||
max_tokens,
|
||||
temperature: 0.7,
|
||||
top_p: 1.0,
|
||||
};
|
||||
|
||||
let resp = self
|
||||
.http
|
||||
.post(&url)
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&body)
|
||||
.send()?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp.text()?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(LlmError::Api { status: status.as_u16(), body: text });
|
||||
}
|
||||
|
||||
let parsed: ChatResponse =
|
||||
serde_json::from_str(&text).map_err(|e| LlmError::Deserialize(e.to_string()))?;
|
||||
parsed
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.map(|c| c.message.content)
|
||||
.ok_or(LlmError::EmptyResponse)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_valid_chat_response() {
|
||||
let raw = r#"{
|
||||
"id": "chatcmpl-xyz",
|
||||
"object": "chat.completion",
|
||||
"model": "llama-3.3-70b-versatile",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": { "role": "assistant", "content": "Hello there!" },
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
],
|
||||
"usage": { "prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8 }
|
||||
}"#;
|
||||
|
||||
let parsed: ChatResponse = serde_json::from_str(raw).expect("valid JSON should parse");
|
||||
assert_eq!(parsed.choices.len(), 1);
|
||||
assert_eq!(parsed.choices[0].message.content, "Hello there!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_fails_when_token_missing() {
|
||||
let prev = env::var(ENV_TOKEN).ok();
|
||||
unsafe { env::remove_var(ENV_TOKEN); }
|
||||
let result = LlmClient::from_env();
|
||||
let err = match result {
|
||||
Ok(_) => panic!("expected ConfigError when {ENV_TOKEN} is unset"),
|
||||
Err(e) => e,
|
||||
};
|
||||
match err {
|
||||
ConfigError::MissingEnv(name) => assert_eq!(name, ENV_TOKEN),
|
||||
}
|
||||
if let Some(v) = prev {
|
||||
unsafe { env::set_var(ENV_TOKEN, v); }
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_message_helpers_set_role() {
|
||||
assert_eq!(ChatMessage::user("hi").role, "user");
|
||||
assert_eq!(ChatMessage::system("ctx").role, "system");
|
||||
assert_eq!(ChatMessage::assistant("ok").role, "assistant");
|
||||
}
|
||||
}
|
||||
115
crates/jarvis-core/src/llm/history.rs
Normal file
115
crates/jarvis-core/src/llm/history.rs
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
use super::client::ChatMessage;
|
||||
|
||||
pub struct ConversationHistory {
|
||||
system: Option<ChatMessage>,
|
||||
turns: Vec<ChatMessage>,
|
||||
max_turns: usize,
|
||||
}
|
||||
|
||||
impl ConversationHistory {
|
||||
pub fn new(system_prompt: impl Into<String>, max_turns: usize) -> Self {
|
||||
Self {
|
||||
system: Some(ChatMessage::system(system_prompt)),
|
||||
turns: Vec::new(),
|
||||
max_turns: max_turns.max(1),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn without_system(max_turns: usize) -> Self {
|
||||
Self {
|
||||
system: None,
|
||||
turns: Vec::new(),
|
||||
max_turns: max_turns.max(1),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push_user(&mut self, content: impl Into<String>) {
|
||||
self.turns.push(ChatMessage::user(content));
|
||||
self.truncate();
|
||||
}
|
||||
|
||||
pub fn push_assistant(&mut self, content: impl Into<String>) {
|
||||
self.turns.push(ChatMessage::assistant(content));
|
||||
self.truncate();
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> Vec<ChatMessage> {
|
||||
let mut out = Vec::with_capacity(self.turns.len() + 1);
|
||||
if let Some(s) = &self.system {
|
||||
out.push(s.clone());
|
||||
}
|
||||
out.extend(self.turns.iter().cloned());
|
||||
out
|
||||
}
|
||||
|
||||
pub fn turns(&self) -> &[ChatMessage] {
|
||||
&self.turns
|
||||
}
|
||||
|
||||
pub fn clear(&mut self) {
|
||||
self.turns.clear();
|
||||
}
|
||||
|
||||
fn truncate(&mut self) {
|
||||
if self.turns.len() > self.max_turns {
|
||||
let drop = self.turns.len() - self.max_turns;
|
||||
self.turns.drain(0..drop);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn evicts_oldest_turns_but_keeps_system_prompt() {
|
||||
let mut h = ConversationHistory::new("you are jarvis", 4);
|
||||
h.push_user("u1");
|
||||
h.push_assistant("a1");
|
||||
h.push_user("u2");
|
||||
h.push_assistant("a2");
|
||||
h.push_user("u3");
|
||||
|
||||
let snap = h.snapshot();
|
||||
assert_eq!(snap[0].role, "system");
|
||||
assert_eq!(snap[0].content, "you are jarvis");
|
||||
assert_eq!(snap.len(), 1 + 4);
|
||||
|
||||
let contents: Vec<&str> = snap.iter().skip(1).map(|m| m.content.as_str()).collect();
|
||||
assert_eq!(contents, vec!["a1", "u2", "a2", "u3"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_with_no_system_returns_only_turns() {
|
||||
let mut h = ConversationHistory::without_system(3);
|
||||
h.push_user("hi");
|
||||
h.push_assistant("hello");
|
||||
let snap = h.snapshot();
|
||||
assert_eq!(snap.len(), 2);
|
||||
assert_eq!(snap[0].role, "user");
|
||||
assert_eq!(snap[1].role, "assistant");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cap_of_zero_is_clamped_to_one() {
|
||||
let mut h = ConversationHistory::new("sys", 0);
|
||||
h.push_user("a");
|
||||
h.push_user("b");
|
||||
let snap = h.snapshot();
|
||||
assert_eq!(snap.len(), 2);
|
||||
assert_eq!(snap[0].role, "system");
|
||||
assert_eq!(snap[1].content, "b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_removes_turns_but_not_system() {
|
||||
let mut h = ConversationHistory::new("sys", 4);
|
||||
h.push_user("u");
|
||||
h.push_assistant("a");
|
||||
h.clear();
|
||||
let snap = h.snapshot();
|
||||
assert_eq!(snap.len(), 1);
|
||||
assert_eq!(snap[0].role, "system");
|
||||
}
|
||||
}
|
||||
8
crates/jarvis-core/src/llm/mod.rs
Normal file
8
crates/jarvis-core/src/llm/mod.rs
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
mod client;
|
||||
mod history;
|
||||
|
||||
pub use client::{
|
||||
ChatMessage, ConfigError, LlmClient, LlmError,
|
||||
DEFAULT_BASE_URL, DEFAULT_MODEL, ENV_BASE_URL, ENV_MODEL, ENV_TOKEN,
|
||||
};
|
||||
pub use history::ConversationHistory;
|
||||
Loading…
Add table
Add a link
Reference in a new issue