refactor: Updated to the most recent Rust version with 2024 syntax

This commit is contained in:
2025-11-07 15:50:55 -07:00
parent 667c843fc0
commit 14549afd52
44 changed files with 377 additions and 371 deletions
Generated
+1 -1
View File
@@ -3089,7 +3089,7 @@ dependencies = [
[[package]] [[package]]
name = "loki-ai" name = "loki-ai"
version = "0.2.0" version = "0.0.1"
dependencies = [ dependencies = [
"ansi_colours", "ansi_colours",
"anyhow", "anyhow",
+7 -3
View File
@@ -1,13 +1,17 @@
[package] [package]
name = "loki-ai" name = "loki-ai"
version = "0.2.0" version = "0.0.1"
edition = "2021" edition = "2024"
authors = ["Alex Clarke <alex.j.tusa@gmail.com>"] authors = ["Alex Clarke <alex.j.tusa@gmail.com>"]
description = "An all-in-one, batteries included LLM CLI Tool" description = "An all-in-one, batteries included LLM CLI Tool"
keywords = ["chatgpt", "llm", "cli", "ai", "repl"]
homepage = "https://github.com/Dark-Alex-17/loki" homepage = "https://github.com/Dark-Alex-17/loki"
repository = "https://github.com/Dark-Alex-17/loki" repository = "https://github.com/Dark-Alex-17/loki"
categories = ["command-line-utilities"] categories = ["command-line-utilities"]
keywords = ["chatgpt", "llm", "cli", "ai", "repl"] readme = "README.md"
license = "MIT"
rust-version = "1.89.0"
exclude = [".github", "CONTRIBUTING.md"]
[dependencies] [dependencies]
anyhow = "1.0.69" anyhow = "1.0.69"
+3 -3
View File
@@ -1,6 +1,6 @@
use crate::client::{list_models, ModelType}; use crate::client::{ModelType, list_models};
use crate::config::{list_agents, Config}; use crate::config::{Config, list_agents};
use clap_complete::{generate, CompletionCandidate, Shell}; use clap_complete::{CompletionCandidate, Shell, generate};
use clap_complete_nushell::Nushell; use clap_complete_nushell::Nushell;
use std::ffi::OsStr; use std::ffi::OsStr;
use std::io; use std::io;
+4 -4
View File
@@ -1,15 +1,15 @@
mod completer; mod completer;
use crate::cli::completer::{ use crate::cli::completer::{
agent_completer, macro_completer, model_completer, rag_completer, role_completer, ShellCompletion, agent_completer, macro_completer, model_completer, rag_completer,
secrets_completer, session_completer, ShellCompletion, role_completer, secrets_completer, session_completer,
}; };
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use clap::ValueHint; use clap::ValueHint;
use clap::{crate_authors, crate_description, crate_name, crate_version, Parser}; use clap::{Parser, crate_authors, crate_description, crate_name, crate_version};
use clap_complete::ArgValueCompleter; use clap_complete::ArgValueCompleter;
use is_terminal::IsTerminal; use is_terminal::IsTerminal;
use std::io::{stdin, Read}; use std::io::{Read, stdin};
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)] #[command(author, version, about, long_about = None)]
+1 -1
View File
@@ -1,4 +1,4 @@
use anyhow::{anyhow, Result}; use anyhow::{Result, anyhow};
use chrono::Utc; use chrono::Utc;
use indexmap::IndexMap; use indexmap::IndexMap;
use parking_lot::RwLock; use parking_lot::RwLock;
+31 -29
View File
@@ -2,7 +2,7 @@ use super::*;
use crate::utils::{base64_decode, encode_uri, hex_encode, hmac_sha256, sha256, strip_think_tag}; use crate::utils::{base64_decode, encode_uri, hex_encode, hmac_sha256, sha256, strip_think_tag};
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use aws_smithy_eventstream::frame::{DecodedFrame, MessageFrameDecoder}; use aws_smithy_eventstream::frame::{DecodedFrame, MessageFrameDecoder};
use aws_smithy_eventstream::smithy::parse_response_headers; use aws_smithy_eventstream::smithy::parse_response_headers;
use bytes::BytesMut; use bytes::BytesMut;
@@ -11,7 +11,7 @@ use futures_util::StreamExt;
use indexmap::IndexMap; use indexmap::IndexMap;
use reqwest::{Client as ReqwestClient, Method, RequestBuilder}; use reqwest::{Client as ReqwestClient, Method, RequestBuilder};
use serde::Deserialize; use serde::Deserialize;
use serde_json::{json, Value}; use serde_json::{Value, json};
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub struct BedrockConfig { pub struct BedrockConfig {
@@ -222,29 +222,29 @@ async fn chat_completions_streaming(
debug!("stream-data: {smithy_type} {data}"); debug!("stream-data: {smithy_type} {data}");
match smithy_type { match smithy_type {
"contentBlockStart" => { "contentBlockStart" => {
if let Some(tool_use) = data["start"]["toolUse"].as_object() { if let Some(tool_use) = data["start"]["toolUse"].as_object()
if let (Some(id), Some(name)) = ( && let (Some(id), Some(name)) = (
json_str_from_map(tool_use, "toolUseId"), json_str_from_map(tool_use, "toolUseId"),
json_str_from_map(tool_use, "name"), json_str_from_map(tool_use, "name"),
) { )
if !function_name.is_empty() { {
if function_arguments.is_empty() { if !function_name.is_empty() {
function_arguments = String::from("{}"); if function_arguments.is_empty() {
} function_arguments = String::from("{}");
let arguments: Value = }
let arguments: Value =
function_arguments.parse().with_context(|| { function_arguments.parse().with_context(|| {
format!("Tool call '{function_name}' have non-JSON arguments '{function_arguments}'") format!("Tool call '{function_name}' have non-JSON arguments '{function_arguments}'")
})?; })?;
handler.tool_call(ToolCall::new( handler.tool_call(ToolCall::new(
function_name.clone(), function_name.clone(),
arguments, arguments,
Some(function_id.clone()), Some(function_id.clone()),
))?; ))?;
}
function_arguments.clear();
function_name = name.into();
function_id = id.into();
} }
function_arguments.clear();
function_name = name.into();
function_id = id.into();
} }
} }
"contentBlockDelta" => { "contentBlockDelta" => {
@@ -291,7 +291,9 @@ async fn chat_completions_streaming(
bail!("Invalid response data: {data} (smithy_type: {smithy_type})") bail!("Invalid response data: {data} (smithy_type: {smithy_type})")
} }
_ => { _ => {
bail!("Unrecognized message, message_type: {message_type}, smithy_type: {smithy_type}",); bail!(
"Unrecognized message, message_type: {message_type}, smithy_type: {smithy_type}",
);
} }
} }
} }
@@ -494,18 +496,18 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
if let Some(text) = json_str_from_map(reasoning_text, "text") { if let Some(text) = json_str_from_map(reasoning_text, "text") {
reasoning = Some(text.to_string()); reasoning = Some(text.to_string());
} }
} else if let Some(tool_use) = item["toolUse"].as_object() { } else if let Some(tool_use) = item["toolUse"].as_object()
if let (Some(id), Some(name), Some(input)) = ( && let (Some(id), Some(name), Some(input)) = (
json_str_from_map(tool_use, "toolUseId"), json_str_from_map(tool_use, "toolUseId"),
json_str_from_map(tool_use, "name"), json_str_from_map(tool_use, "name"),
tool_use.get("input"), tool_use.get("input"),
) { )
tool_calls.push(ToolCall::new( {
name.to_string(), tool_calls.push(ToolCall::new(
input.clone(), name.to_string(),
Some(id.to_string()), input.clone(),
)) Some(id.to_string()),
} ))
} }
} }
} }
+10 -10
View File
@@ -2,10 +2,10 @@ use super::openai::*;
use super::openai_compatible::*; use super::openai_compatible::*;
use super::*; use super::*;
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use reqwest::RequestBuilder; use reqwest::RequestBuilder;
use serde::Deserialize; use serde::Deserialize;
use serde_json::{json, Value}; use serde_json::{Value, json};
const API_BASE: &str = "https://api.cohere.ai/v2"; const API_BASE: &str = "https://api.cohere.ai/v2";
@@ -49,10 +49,10 @@ fn prepare_chat_completions(
let url = format!("{}/chat", api_base.trim_end_matches('/')); let url = format!("{}/chat", api_base.trim_end_matches('/'));
let mut body = openai_build_chat_completions_body(data, &self_.model); let mut body = openai_build_chat_completions_body(data, &self_.model);
if let Some(obj) = body.as_object_mut() { if let Some(obj) = body.as_object_mut()
if let Some(top_p) = obj.remove("top_p") { && let Some(top_p) = obj.remove("top_p")
obj.insert("p".to_string(), top_p); {
} obj.insert("p".to_string(), top_p);
} }
let mut request_data = RequestData::new(url, body); let mut request_data = RequestData::new(url, body);
@@ -218,10 +218,10 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let mut tool_calls = vec![]; let mut tool_calls = vec![];
if let Some(calls) = data["message"]["tool_calls"].as_array() { if let Some(calls) = data["message"]["tool_calls"].as_array() {
if text.is_empty() { if text.is_empty()
if let Some(tool_plain) = data["message"]["tool_plan"].as_str() { && let Some(tool_plain) = data["message"]["tool_plan"].as_str()
text = tool_plain.to_string(); {
} text = tool_plain.to_string();
} }
for call in calls { for call in calls {
if let (Some(name), Some(arguments), Some(id)) = ( if let (Some(name), Some(arguments), Some(id)) = (
+9 -9
View File
@@ -2,21 +2,21 @@ use super::*;
use crate::{ use crate::{
config::{Config, GlobalConfig, Input}, config::{Config, GlobalConfig, Input},
function::{eval_tool_calls, FunctionDeclaration, ToolCall, ToolResult}, function::{FunctionDeclaration, ToolCall, ToolResult, eval_tool_calls},
render::render_stream, render::render_stream,
utils::*, utils::*,
}; };
use crate::vault::Vault; use crate::vault::Vault;
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use fancy_regex::Regex; use fancy_regex::Regex;
use indexmap::IndexMap; use indexmap::IndexMap;
use inquire::{ use inquire::{
list_option::ListOption, required, validator::Validation, MultiSelect, Select, Text, MultiSelect, Select, Text, list_option::ListOption, required, validator::Validation,
}; };
use reqwest::{Client as ReqwestClient, RequestBuilder}; use reqwest::{Client as ReqwestClient, RequestBuilder};
use serde::Deserialize; use serde::Deserialize;
use serde_json::{json, Value}; use serde_json::{Value, json};
use std::sync::LazyLock; use std::sync::LazyLock;
use std::time::Duration; use std::time::Duration;
use tokio::sync::mpsc::unbounded_channel; use tokio::sync::mpsc::unbounded_channel;
@@ -180,11 +180,11 @@ pub trait Client: Sync + Send {
}; };
for (key, patch) in patch_map { for (key, patch) in patch_map {
let key = ESCAPE_SLASH_RE.replace_all(&key, r"\/"); let key = ESCAPE_SLASH_RE.replace_all(&key, r"\/");
if let Ok(regex) = Regex::new(&format!("^({key})$")) { if let Ok(regex) = Regex::new(&format!("^({key})$"))
if let Ok(true) = regex.is_match(self.model().name()) { && let Ok(true) = regex.is_match(self.model().name())
request_data.apply_patch(patch); {
return; request_data.apply_patch(patch);
} return;
} }
} }
} }
+4 -4
View File
@@ -119,10 +119,10 @@ impl MessageContent {
} }
for tool_result in tool_results { for tool_result in tool_results {
let mut parts = vec!["Call".to_string()]; let mut parts = vec!["Call".to_string()];
if let Some((agent_name, functions)) = agent_info { if let Some((agent_name, functions)) = agent_info
if functions.contains(&tool_result.call.name) { && functions.contains(&tool_result.call.name)
parts.push(agent_name.clone()) {
} parts.push(agent_name.clone())
} }
parts.push(tool_result.call.name.clone()); parts.push(tool_result.call.name.clone());
parts.push(tool_result.call.arguments.to_string()); parts.push(tool_result.call.arguments.to_string());
+6 -7
View File
@@ -1,13 +1,12 @@
use super::{ use super::{
list_all_models, list_client_names, ApiPatch, MessageContentToolCalls, RequestPatch, list_all_models, list_client_names,
message::{Message, MessageContent, MessageContentPart}, message::{Message, MessageContent, MessageContentPart},
ApiPatch, MessageContentToolCalls, RequestPatch,
}; };
use crate::config::Config; use crate::config::Config;
use crate::utils::{estimate_token_length, strip_think_tag}; use crate::utils::{estimate_token_length, strip_think_tag};
use anyhow::{bail, Result}; use anyhow::{Result, bail};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use std::fmt::Display; use std::fmt::Display;
@@ -275,10 +274,10 @@ impl Model {
pub fn guard_max_input_tokens(&self, messages: &[Message]) -> Result<()> { pub fn guard_max_input_tokens(&self, messages: &[Message]) -> Result<()> {
let total_tokens = self.total_tokens(messages) + BASIS_TOKENS; let total_tokens = self.total_tokens(messages) + BASIS_TOKENS;
if let Some(max_input_tokens) = self.data.max_input_tokens { if let Some(max_input_tokens) = self.data.max_input_tokens
if total_tokens >= max_input_tokens { && total_tokens >= max_input_tokens
bail!("Exceed max_input_tokens limit") {
} bail!("Exceed max_input_tokens limit")
} }
Ok(()) Ok(())
} }
+7 -7
View File
@@ -4,7 +4,7 @@ use super::*;
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use reqwest::RequestBuilder; use reqwest::RequestBuilder;
use serde::Deserialize; use serde::Deserialize;
use serde_json::{json, Value}; use serde_json::{Value, json};
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub struct OpenAICompatibleConfig { pub struct OpenAICompatibleConfig {
@@ -124,12 +124,12 @@ pub async fn generic_rerank(builder: RequestBuilder, _model: &Model) -> Result<R
if !status.is_success() { if !status.is_success() {
catch_error(&data, status.as_u16())?; catch_error(&data, status.as_u16())?;
} }
if data.get("results").is_none() && data.get("data").is_some() { if data.get("results").is_none()
if let Some(data_obj) = data.as_object_mut() { && data.get("data").is_some()
if let Some(value) = data_obj.remove("data") { && let Some(data_obj) = data.as_object_mut()
data_obj.insert("results".to_string(), value); && let Some(value) = data_obj.remove("data")
} {
} data_obj.insert("results".to_string(), value);
} }
let res_body: GenericRerankResBody = let res_body: GenericRerankResBody =
serde_json::from_value(data).context("Invalid rerank data")?; serde_json::from_value(data).context("Invalid rerank data")?;
+7 -7
View File
@@ -1,7 +1,7 @@
use super::{catch_error, ToolCall}; use super::{ToolCall, catch_error};
use crate::utils::AbortSignal; use crate::utils::AbortSignal;
use anyhow::{anyhow, bail, Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use futures_util::{Stream, StreamExt}; use futures_util::{Stream, StreamExt};
use reqwest::RequestBuilder; use reqwest::RequestBuilder;
use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt}; use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt};
@@ -214,11 +214,11 @@ impl JsonStreamParser {
} }
'}' => { '}' => {
self.balances.pop(); self.balances.pop();
if self.balances.is_empty() { if self.balances.is_empty()
if let Some(start) = self.start.take() { && let Some(start) = self.start.take()
let value: String = self.buffer[start..=i].iter().collect(); {
handle(&value)?; let value: String = self.buffer[start..=i].iter().collect();
} handle(&value)?;
} }
} }
']' => { ']' => {
+18 -18
View File
@@ -2,13 +2,13 @@ use super::*;
use crate::{ use crate::{
client::Model, client::Model,
function::{run_llm_function, Functions}, function::{Functions, run_llm_function},
}; };
use crate::vault::SECRET_RE; use crate::vault::SECRET_RE;
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use fancy_regex::Captures; use fancy_regex::Captures;
use inquire::{validator::Validation, Text}; use inquire::{Text, validator::Validation};
use rust_embed::Embed; use rust_embed::Embed;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::{ffi::OsStr, path::Path}; use std::{ffi::OsStr, path::Path};
@@ -530,23 +530,23 @@ impl AgentConfig {
if let Some(v) = read_env_value::<f64>(&with_prefix("top_p")) { if let Some(v) = read_env_value::<f64>(&with_prefix("top_p")) {
self.top_p = v; self.top_p = v;
} }
if let Ok(v) = env::var(with_prefix("global_tools")) { if let Ok(v) = env::var(with_prefix("global_tools"))
if let Ok(v) = serde_json::from_str(&v) { && let Ok(v) = serde_json::from_str(&v)
self.global_tools = v; {
} self.global_tools = v;
} }
if let Ok(v) = env::var(with_prefix("mcp_servers")) { if let Ok(v) = env::var(with_prefix("mcp_servers"))
if let Ok(v) = serde_json::from_str(&v) { && let Ok(v) = serde_json::from_str(&v)
self.mcp_servers = v; {
} self.mcp_servers = v;
} }
if let Some(v) = read_env_value::<String>(&with_prefix("agent_session")) { if let Some(v) = read_env_value::<String>(&with_prefix("agent_session")) {
self.agent_session = v; self.agent_session = v;
} }
if let Ok(v) = env::var(with_prefix("variables")) { if let Ok(v) = env::var(with_prefix("variables"))
if let Ok(v) = serde_json::from_str(&v) { && let Ok(v) = serde_json::from_str(&v)
self.variables = v; {
} self.variables = v;
} }
} }
@@ -619,10 +619,10 @@ pub fn list_agents() -> Vec<String> {
let mut agents = Vec::new(); let mut agents = Vec::new();
if let Ok(entries) = read_dir(agents_data_dir) { if let Ok(entries) = read_dir(agents_data_dir) {
for entry in entries.flatten() { for entry in entries.flatten() {
if entry.path().is_dir() { if entry.path().is_dir()
if let Some(name) = entry.file_name().to_str() { && let Some(name) = entry.file_name().to_str()
agents.push(name.to_string()); {
} agents.push(name.to_string());
} }
} }
} }
+4 -4
View File
@@ -1,13 +1,13 @@
use super::*; use super::*;
use crate::client::{ use crate::client::{
init_client, patch_messages, ChatCompletionsData, Client, ImageUrl, Message, MessageContent, ChatCompletionsData, Client, ImageUrl, Message, MessageContent, MessageContentPart,
MessageContentPart, MessageContentToolCalls, MessageRole, Model, MessageContentToolCalls, MessageRole, Model, init_client, patch_messages,
}; };
use crate::function::ToolResult; use crate::function::ToolResult;
use crate::utils::{base64_encode, is_loader_protocol, sha256, AbortSignal}; use crate::utils::{AbortSignal, base64_encode, is_loader_protocol, sha256};
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use indexmap::IndexSet; use indexmap::IndexSet;
use std::{collections::HashMap, fs::File, io::Read}; use std::{collections::HashMap, fs::File, io::Read};
use unicode_width::{UnicodeWidthChar, UnicodeWidthStr}; use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
+3 -3
View File
@@ -1,7 +1,7 @@
use crate::config::{ensure_parent_exists, Config, GlobalConfig, RoleLike}; use crate::config::{Config, GlobalConfig, RoleLike, ensure_parent_exists};
use crate::repl::{run_repl_command, split_args_text}; use crate::repl::{run_repl_command, split_args_text};
use crate::utils::{multiline_text, AbortSignal}; use crate::utils::{AbortSignal, multiline_text};
use anyhow::{anyhow, Result}; use anyhow::{Result, anyhow};
use indexmap::IndexMap; use indexmap::IndexMap;
use parking_lot::RwLock; use parking_lot::RwLock;
use rust_embed::Embed; use rust_embed::Embed;
+87 -89
View File
@@ -4,18 +4,18 @@ mod macros;
mod role; mod role;
mod session; mod session;
pub use self::agent::{complete_agent_variables, list_agents, Agent, AgentVariables}; pub use self::agent::{Agent, AgentVariables, complete_agent_variables, list_agents};
pub use self::input::Input; pub use self::input::Input;
pub use self::role::{ pub use self::role::{
Role, RoleLike, CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE, CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, Role, RoleLike, SHELL_ROLE,
}; };
use self::session::Session; use self::session::Session;
pub use macros::macro_execute; pub use macros::macro_execute;
use mem::take; use mem::take;
use crate::client::{ use crate::client::{
create_client_config, list_client_types, list_models, ClientConfig, MessageContentToolCalls, ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
Model, ModelType, ProviderModels, OPENAI_COMPATIBLE_PROVIDERS, ProviderModels, create_client_config, list_client_types, list_models,
}; };
use crate::function::{FunctionDeclaration, Functions, ToolResult}; use crate::function::{FunctionDeclaration, Functions, ToolResult};
use crate::rag::Rag; use crate::rag::Rag;
@@ -24,14 +24,14 @@ use crate::utils::*;
use crate::config::macros::Macro; use crate::config::macros::Macro;
use crate::mcp::{ use crate::mcp::{
McpRegistry, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_LIST_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_LIST_META_FUNCTION_NAME_PREFIX, McpRegistry,
}; };
use crate::vault::{create_vault_password_file, interpolate_secrets, GlobalVault, Vault}; use crate::vault::{GlobalVault, Vault, create_vault_password_file, interpolate_secrets};
use anyhow::{anyhow, bail, Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use fancy_regex::Regex; use fancy_regex::Regex;
use indexmap::IndexMap; use indexmap::IndexMap;
use indoc::formatdoc; use indoc::formatdoc;
use inquire::{list_option::ListOption, validator::Validation, Confirm, MultiSelect, Select, Text}; use inquire::{Confirm, MultiSelect, Select, Text, list_option::ListOption, validator::Validation};
use log::LevelFilter; use log::LevelFilter;
use parking_lot::RwLock; use parking_lot::RwLock;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
@@ -41,7 +41,7 @@ use std::sync::LazyLock;
use std::{ use std::{
env, env,
fs::{ fs::{
create_dir_all, read_dir, read_to_string, remove_dir_all, remove_file, File, OpenOptions, File, OpenOptions, create_dir_all, read_dir, read_to_string, remove_dir_all, remove_file,
}, },
io::Write, io::Write,
mem, mem,
@@ -50,7 +50,7 @@ use std::{
sync::{Arc, OnceLock}, sync::{Arc, OnceLock},
}; };
use syntect::highlighting::ThemeSet; use syntect::highlighting::ThemeSet;
use terminal_colorsaurus::{color_scheme, ColorScheme, QueryOptions}; use terminal_colorsaurus::{ColorScheme, QueryOptions, color_scheme};
use tokio::runtime::Handle; use tokio::runtime::Handle;
pub const TEMP_ROLE_NAME: &str = "temp"; pub const TEMP_ROLE_NAME: &str = "temp";
@@ -547,10 +547,10 @@ impl Config {
for entry in read_dir(Self::agent_data_dir(name))? { for entry in read_dir(Self::agent_data_dir(name))? {
let entry = entry?; let entry = entry?;
if let Some(file) = entry.file_name().to_str() { if let Some(file) = entry.file_name().to_str()
if allowed.contains(&file) { && allowed.contains(&file)
return Ok(entry.path()); {
} return Ok(entry.path());
} }
} }
@@ -783,24 +783,24 @@ impl Config {
} }
"enabled_mcp_servers" => { "enabled_mcp_servers" => {
let value: Option<String> = parse_value(value)?; let value: Option<String> = parse_value(value)?;
if let Some(servers) = value.as_ref() { if let Some(servers) = value.as_ref()
if let Some(registry) = &config.read().mcp_registry { && let Some(registry) = &config.read().mcp_registry
if registry.list_configured_servers().is_empty() { {
bail!( if registry.list_configured_servers().is_empty() {
"No MCP servers are configured. Please configure MCP servers first before setting 'enabled_mcp_servers'." bail!(
); "No MCP servers are configured. Please configure MCP servers first before setting 'enabled_mcp_servers'."
} );
}
if !servers.split(',').all(|s| { if !servers.split(',').all(|s| {
registry registry
.list_configured_servers() .list_configured_servers()
.contains(&s.trim().to_string()) .contains(&s.trim().to_string())
|| s == "all" || s == "all"
}) { }) {
bail!( bail!(
"Some of the specified MCP servers in 'enabled_mcp_servers' are configured. Please check your MCP server configuration." "Some of the specified MCP servers in 'enabled_mcp_servers' are configured. Please check your MCP server configuration."
); );
}
} }
} }
config.write().set_enabled_mcp_servers(value.clone()); config.write().set_enabled_mcp_servers(value.clone());
@@ -1399,18 +1399,16 @@ impl Config {
output, output,
continuous, continuous,
}) = &self.last_message }) = &self.last_message
&& (*continuous && !output.is_empty())
&& self.agent.is_some() == input.with_agent()
{ {
if (*continuous && !output.is_empty()) let ans = Confirm::new(
&& self.agent.is_some() == input.with_agent() "Start a session that incorporates the last question and answer?",
{ )
let ans = Confirm::new( .with_default(false)
"Start a session that incorporates the last question and answer?", .prompt()?;
) if ans {
.with_default(false) session.add_message(input, output)?;
.prompt()?;
if ans {
session.add_message(input, output)?;
}
} }
} }
} }
@@ -1520,11 +1518,11 @@ impl Config {
{ {
let mut config = config.write(); let mut config = config.write();
let compression_threshold = config.compression_threshold; let compression_threshold = config.compression_threshold;
if let Some(session) = config.session.as_mut() { if let Some(session) = config.session.as_mut()
if session.needs_compression(compression_threshold) { && session.needs_compression(compression_threshold)
session.set_compressing(true); {
needs_compression = true; session.set_compressing(true);
} needs_compression = true;
} }
}; };
if !needs_compression { if !needs_compression {
@@ -1587,11 +1585,11 @@ impl Config {
pub fn maybe_autoname_session(config: GlobalConfig) { pub fn maybe_autoname_session(config: GlobalConfig) {
let mut need_autoname = false; let mut need_autoname = false;
if let Some(session) = config.write().session.as_mut() { if let Some(session) = config.write().session.as_mut()
if session.need_autoname() { && session.need_autoname()
session.set_autonaming(true); {
need_autoname = true; session.set_autonaming(true);
} need_autoname = true;
} }
if !need_autoname { if !need_autoname {
return; return;
@@ -2439,15 +2437,15 @@ impl Config {
.unwrap_or_default() .unwrap_or_default()
.to_string(), .to_string(),
); );
if let Some(temperature) = role.temperature() { if let Some(temperature) = role.temperature()
if temperature != 0.0 { && temperature != 0.0
output.insert("temperature", temperature.to_string()); {
} output.insert("temperature", temperature.to_string());
} }
if let Some(top_p) = role.top_p() { if let Some(top_p) = role.top_p()
if top_p != 0.0 { && top_p != 0.0
output.insert("top_p", top_p.to_string()); {
} output.insert("top_p", top_p.to_string());
} }
if self.dry_run { if self.dry_run {
output.insert("dry_run", "true".to_string()); output.insert("dry_run", "true".to_string());
@@ -2458,10 +2456,10 @@ impl Config {
if self.save { if self.save {
output.insert("save", "true".to_string()); output.insert("save", "true".to_string());
} }
if let Some(wrap) = &self.wrap { if let Some(wrap) = &self.wrap
if wrap != "no" { && wrap != "no"
output.insert("wrap", wrap.clone()); {
} output.insert("wrap", wrap.clone());
} }
if !role.is_derived() { if !role.is_derived() {
output.insert("role", role.name().to_string()); output.insert("role", role.name().to_string());
@@ -2724,10 +2722,10 @@ impl Config {
if let Some(Some(v)) = read_env_bool(&get_env_name("save")) { if let Some(Some(v)) = read_env_bool(&get_env_name("save")) {
self.save = v; self.save = v;
} }
if let Ok(v) = env::var(get_env_name("keybindings")) { if let Ok(v) = env::var(get_env_name("keybindings"))
if v == "vi" { && v == "vi"
self.keybindings = v; {
} self.keybindings = v;
} }
if let Some(v) = read_env_value::<String>(&get_env_name("editor")) { if let Some(v) = read_env_value::<String>(&get_env_name("editor")) {
self.editor = v; self.editor = v;
@@ -2742,10 +2740,10 @@ impl Config {
if let Some(Some(v)) = read_env_bool(&get_env_name("function_calling_support")) { if let Some(Some(v)) = read_env_bool(&get_env_name("function_calling_support")) {
self.function_calling_support = v; self.function_calling_support = v;
} }
if let Ok(v) = env::var(get_env_name("mapping_tools")) { if let Ok(v) = env::var(get_env_name("mapping_tools"))
if let Ok(v) = serde_json::from_str(&v) { && let Ok(v) = serde_json::from_str(&v)
self.mapping_tools = v; {
} self.mapping_tools = v;
} }
if let Some(v) = read_env_value::<String>(&get_env_name("enabled_tools")) { if let Some(v) = read_env_value::<String>(&get_env_name("enabled_tools")) {
self.enabled_tools = v; self.enabled_tools = v;
@@ -2754,10 +2752,10 @@ impl Config {
if let Some(Some(v)) = read_env_bool(&get_env_name("mcp_server_support")) { if let Some(Some(v)) = read_env_bool(&get_env_name("mcp_server_support")) {
self.mcp_server_support = v; self.mcp_server_support = v;
} }
if let Ok(v) = env::var(get_env_name("mapping_mcp_servers")) { if let Ok(v) = env::var(get_env_name("mapping_mcp_servers"))
if let Ok(v) = serde_json::from_str(&v) { && let Ok(v) = serde_json::from_str(&v)
self.mapping_mcp_servers = v; {
} self.mapping_mcp_servers = v;
} }
if let Some(v) = read_env_value::<String>(&get_env_name("enabled_mcp_servers")) { if let Some(v) = read_env_value::<String>(&get_env_name("enabled_mcp_servers")) {
self.enabled_mcp_servers = v; self.enabled_mcp_servers = v;
@@ -2805,10 +2803,10 @@ impl Config {
self.rag_template = v; self.rag_template = v;
} }
if let Ok(v) = env::var(get_env_name("document_loaders")) { if let Ok(v) = env::var(get_env_name("document_loaders"))
if let Ok(v) = serde_json::from_str(&v) { && let Ok(v) = serde_json::from_str(&v)
self.document_loaders = v; {
} self.document_loaders = v;
} }
if let Some(Some(v)) = read_env_bool(&get_env_name("highlight")) { if let Some(Some(v)) = read_env_bool(&get_env_name("highlight")) {
@@ -2820,14 +2818,14 @@ impl Config {
if self.highlight && self.theme.is_none() { if self.highlight && self.theme.is_none() {
if let Some(v) = read_env_value::<String>(&get_env_name("theme")) { if let Some(v) = read_env_value::<String>(&get_env_name("theme")) {
self.theme = v; self.theme = v;
} else if *IS_STDOUT_TERMINAL { } else if *IS_STDOUT_TERMINAL
if let Ok(color_scheme) = color_scheme(QueryOptions::default()) { && let Ok(color_scheme) = color_scheme(QueryOptions::default())
let theme = match color_scheme { {
ColorScheme::Dark => "dark", let theme = match color_scheme {
ColorScheme::Light => "light", ColorScheme::Dark => "dark",
}; ColorScheme::Light => "light",
self.theme = Some(theme.into()); };
} self.theme = Some(theme.into());
} }
} }
if let Some(v) = read_env_value::<String>(&get_env_name("left_prompt")) { if let Some(v) = read_env_value::<String>(&get_env_name("left_prompt")) {
@@ -2934,7 +2932,7 @@ pub fn load_env_file() -> Result<()> {
continue; continue;
} }
if let Some((key, value)) = line.split_once('=') { if let Some((key, value)) = line.split_once('=') {
env::set_var(key.trim(), value.trim()); unsafe { env::set_var(key.trim(), value.trim()) };
} }
} }
Ok(()) Ok(())
+18 -21
View File
@@ -64,11 +64,11 @@ impl Role {
pub fn new(name: &str, content: &str) -> Self { pub fn new(name: &str, content: &str) -> Self {
let mut metadata = ""; let mut metadata = "";
let mut prompt = content.trim(); let mut prompt = content.trim();
if let Ok(Some(caps)) = RE_METADATA.captures(content) { if let Ok(Some(caps)) = RE_METADATA.captures(content)
if let (Some(metadata_value), Some(prompt_value)) = (caps.get(1), caps.get(2)) { && let (Some(metadata_value), Some(prompt_value)) = (caps.get(1), caps.get(2))
metadata = metadata_value.as_str().trim(); {
prompt = prompt_value.as_str().trim(); metadata = metadata_value.as_str().trim();
} prompt = prompt_value.as_str().trim();
} }
let mut prompt = prompt.to_string(); let mut prompt = prompt.to_string();
interpolate_variables(&mut prompt); interpolate_variables(&mut prompt);
@@ -77,23 +77,20 @@ impl Role {
prompt, prompt,
..Default::default() ..Default::default()
}; };
if !metadata.is_empty() { if !metadata.is_empty()
if let Ok(value) = serde_yaml::from_str::<Value>(metadata) { && let Ok(value) = serde_yaml::from_str::<Value>(metadata)
if let Some(value) = value.as_object() { && let Some(value) = value.as_object()
for (key, value) in value { {
match key.as_str() { for (key, value) in value {
"model" => role.model_id = value.as_str().map(|v| v.to_string()), match key.as_str() {
"temperature" => role.temperature = value.as_f64(), "model" => role.model_id = value.as_str().map(|v| v.to_string()),
"top_p" => role.top_p = value.as_f64(), "temperature" => role.temperature = value.as_f64(),
"enabled_tools" => { "top_p" => role.top_p = value.as_f64(),
role.enabled_tools = value.as_str().map(|v| v.to_string()) "enabled_tools" => role.enabled_tools = value.as_str().map(|v| v.to_string()),
} "enabled_mcp_servers" => {
"enabled_mcp_servers" => { role.enabled_mcp_servers = value.as_str().map(|v| v.to_string())
role.enabled_mcp_servers = value.as_str().map(|v| v.to_string())
}
_ => (),
}
} }
_ => (),
} }
} }
} }
+22 -20
View File
@@ -4,9 +4,9 @@ use super::*;
use crate::client::{Message, MessageContent, MessageRole}; use crate::client::{Message, MessageContent, MessageRole};
use crate::render::MarkdownRender; use crate::render::MarkdownRender;
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use fancy_regex::Regex; use fancy_regex::Regex;
use inquire::{validator::Validation, Confirm, Text}; use inquire::{Confirm, Text, validator::Validation};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::json; use serde_json::json;
use std::collections::HashMap; use std::collections::HashMap;
@@ -98,10 +98,10 @@ impl Session {
session.path = Some(path.display().to_string()); session.path = Some(path.display().to_string());
} }
if let Some(role_name) = &session.role_name { if let Some(role_name) = &session.role_name
if let Ok(role) = config.retrieve_role(role_name) { && let Ok(role) = config.retrieve_role(role_name)
session.role_prompt = role.prompt().to_string(); {
} session.role_prompt = role.prompt().to_string();
} }
session.update_tokens(); session.update_tokens();
@@ -473,23 +473,25 @@ impl Session {
pub fn guard_empty(&self) -> Result<()> { pub fn guard_empty(&self) -> Result<()> {
if !self.is_empty() { if !self.is_empty() {
bail!("Cannot perform this operation because the session has messages, please `.empty session` first."); bail!(
"Cannot perform this operation because the session has messages, please `.empty session` first."
);
} }
Ok(()) Ok(())
} }
pub fn add_message(&mut self, input: &Input, output: &str) -> Result<()> { pub fn add_message(&mut self, input: &Input, output: &str) -> Result<()> {
if input.continue_output().is_some() { if input.continue_output().is_some() {
if let Some(message) = self.messages.last_mut() { if let Some(message) = self.messages.last_mut()
if let MessageContent::Text(text) = &mut message.content { && let MessageContent::Text(text) = &mut message.content
*text = format!("{text}{output}"); {
} *text = format!("{text}{output}");
} }
} else if input.regenerate() { } else if input.regenerate() {
if let Some(message) = self.messages.last_mut() { if let Some(message) = self.messages.last_mut()
if let MessageContent::Text(text) = &mut message.content { && let MessageContent::Text(text) = &mut message.content
*text = output.to_string(); {
} *text = output.to_string();
} }
} else { } else {
if self.messages.is_empty() { if self.messages.is_empty() {
@@ -553,14 +555,14 @@ impl Session {
if len == 0 { if len == 0 {
messages = input.role().build_messages(input); messages = input.role().build_messages(input);
need_add_msg = false; need_add_msg = false;
} else if len == 1 && self.compressed_messages.len() >= 2 { } else if len == 1
if let Some(index) = self && self.compressed_messages.len() >= 2
&& let Some(index) = self
.compressed_messages .compressed_messages
.iter() .iter()
.rposition(|v| v.role == MessageRole::User) .rposition(|v| v.role == MessageRole::User)
{ {
messages.extend(self.compressed_messages[index..].to_vec()); messages.extend(self.compressed_messages[index..].to_vec());
}
} }
if need_add_msg { if need_add_msg {
messages.push(Message::new(MessageRole::User, input.message_content())); messages.push(Message::new(MessageRole::User, input.message_content()));
+2 -2
View File
@@ -6,12 +6,12 @@ use crate::{
use crate::config::ensure_parent_exists; use crate::config::ensure_parent_exists;
use crate::mcp::{MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_LIST_META_FUNCTION_NAME_PREFIX}; use crate::mcp::{MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_LIST_META_FUNCTION_NAME_PREFIX};
use crate::parsers::{bash, python}; use crate::parsers::{bash, python};
use anyhow::{anyhow, bail, Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use indexmap::IndexMap; use indexmap::IndexMap;
use indoc::formatdoc; use indoc::formatdoc;
use rust_embed::Embed; use rust_embed::Embed;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::{json, Value}; use serde_json::{Value, json};
use std::ffi::OsStr; use std::ffi::OsStr;
use std::fs::File; use std::fs::File;
use std::io::Write; use std::io::Write;
+5 -4
View File
@@ -15,11 +15,12 @@ mod vault;
extern crate log; extern crate log;
use crate::client::{ use crate::client::{
call_chat_completions, call_chat_completions_streaming, list_models, ModelType, ModelType, call_chat_completions, call_chat_completions_streaming, list_models,
}; };
use crate::config::{ use crate::config::{
ensure_parent_exists, list_agents, load_env_file, macro_execute, Agent, Config, GlobalConfig, Agent, CODE_ROLE, Config, EXPLAIN_SHELL_ROLE, GlobalConfig, Input, SHELL_ROLE,
Input, WorkingMode, CODE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE, TEMP_SESSION_NAME, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, load_env_file,
macro_execute,
}; };
use crate::render::render_error; use crate::render::render_error;
use crate::repl::Repl; use crate::repl::Repl;
@@ -27,7 +28,7 @@ use crate::utils::*;
use crate::cli::Cli; use crate::cli::Cli;
use crate::vault::Vault; use crate::vault::Vault;
use anyhow::{bail, Result}; use anyhow::{Result, bail};
use clap::{CommandFactory, Parser}; use clap::{CommandFactory, Parser};
use clap_complete::CompleteEnv; use clap_complete::CompleteEnv;
use inquire::Text; use inquire::Text;
+7 -5
View File
@@ -1,16 +1,16 @@
use crate::config::Config; use crate::config::Config;
use crate::utils::{abortable_run_with_spinner, AbortSignal}; use crate::utils::{AbortSignal, abortable_run_with_spinner};
use crate::vault::interpolate_secrets; use crate::vault::interpolate_secrets;
use anyhow::{anyhow, Context, Result}; use anyhow::{Context, Result, anyhow};
use futures_util::future::BoxFuture; use futures_util::future::BoxFuture;
use futures_util::{stream, StreamExt, TryStreamExt}; use futures_util::{StreamExt, TryStreamExt, stream};
use indoc::formatdoc; use indoc::formatdoc;
use rmcp::model::{CallToolRequestParam, CallToolResult}; use rmcp::model::{CallToolRequestParam, CallToolResult};
use rmcp::service::RunningService; use rmcp::service::RunningService;
use rmcp::transport::TokioChildProcess; use rmcp::transport::TokioChildProcess;
use rmcp::{RoleClient, ServiceExt}; use rmcp::{RoleClient, ServiceExt};
use serde::Deserialize; use serde::Deserialize;
use serde_json::{json, Value}; use serde_json::{Value, json};
use std::borrow::Cow; use std::borrow::Cow;
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use std::fs::OpenOptions; use std::fs::OpenOptions;
@@ -148,7 +148,9 @@ impl McpRegistry {
enabled_mcp_servers: Option<String>, enabled_mcp_servers: Option<String>,
) -> Result<()> { ) -> Result<()> {
if self.config.is_none() { if self.config.is_none() {
debug!("MCP config is not present; assuming MCP servers are disabled globally. Skipping MCP initialization"); debug!(
"MCP config is not present; assuming MCP servers are disabled globally. Skipping MCP initialization"
);
return Ok(()); return Ok(());
} }
+5 -5
View File
@@ -1,5 +1,5 @@
use crate::function::{FunctionDeclaration, JsonSchema}; use crate::function::{FunctionDeclaration, JsonSchema};
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use argc::{ChoiceValue, CommandValue, FlagOptionValue}; use argc::{ChoiceValue, CommandValue, FlagOptionValue};
use indexmap::IndexMap; use indexmap::IndexMap;
use std::fs::File; use std::fs::File;
@@ -95,10 +95,10 @@ fn with_description(mut schema: JsonSchema, describe: &str) -> JsonSchema {
} }
fn with_enum(mut schema: JsonSchema, choice: &Option<ChoiceValue>) -> JsonSchema { fn with_enum(mut schema: JsonSchema, choice: &Option<ChoiceValue>) -> JsonSchema {
if let Some(ChoiceValue::Values(values)) = choice { if let Some(ChoiceValue::Values(values)) = choice
if !values.is_empty() { && !values.is_empty()
schema.enum_value = Some(values.clone()); {
} schema.enum_value = Some(values.clone());
} }
schema schema
} }
+11 -12
View File
@@ -1,9 +1,9 @@
use crate::function::{FunctionDeclaration, JsonSchema}; use crate::function::{FunctionDeclaration, JsonSchema};
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use ast::{Stmt, StmtFunctionDef}; use ast::{Stmt, StmtFunctionDef};
use indexmap::IndexMap; use indexmap::IndexMap;
use rustpython_ast::{Constant, Expr, UnaryOp}; use rustpython_ast::{Constant, Expr, UnaryOp};
use rustpython_parser::{ast, Mode}; use rustpython_parser::{Mode, ast};
use serde_json::Value; use serde_json::Value;
use std::fs::File; use std::fs::File;
use std::io::Read; use std::io::Read;
@@ -104,12 +104,11 @@ fn python_to_function_declarations(
fn get_docstring_from_body(body: &[Stmt]) -> Option<String> { fn get_docstring_from_body(body: &[Stmt]) -> Option<String> {
let first = body.first()?; let first = body.first()?;
if let Stmt::Expr(expr_stmt) = first { if let Stmt::Expr(expr_stmt) = first
if let Expr::Constant(constant) = &*expr_stmt.value { && let Expr::Constant(constant) = &*expr_stmt.value
if let Constant::Str(s) = &constant.value { && let Constant::Str(s) = &constant.value
return Some(s.clone()); {
} return Some(s.clone());
}
} }
None None
} }
@@ -351,10 +350,10 @@ fn build_parameters_schema(params: &[Param], _description: &str) -> JsonSchema {
"str" "str"
}; };
if let Some(d) = &p.doc_desc { if let Some(d) = &p.doc_desc
if !d.is_empty() { && !d.is_empty()
schema.description = Some(d.clone()); {
} schema.description = Some(d.clone());
} }
apply_type_to_schema(ty, &mut schema); apply_type_to_schema(ty, &mut schema);
+12 -17
View File
@@ -7,11 +7,11 @@ use crate::utils::*;
mod serde_vectors; mod serde_vectors;
mod splitter; mod splitter;
use anyhow::{anyhow, bail, Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use bm25::{Language, SearchEngine, SearchEngineBuilder}; use bm25::{Language, SearchEngine, SearchEngineBuilder};
use hnsw_rs::prelude::*; use hnsw_rs::prelude::*;
use indexmap::{IndexMap, IndexSet}; use indexmap::{IndexMap, IndexSet};
use inquire::{required, validator::Validation, Confirm, Select, Text}; use inquire::{Confirm, Select, Text, required, validator::Validation};
use parking_lot::RwLock; use parking_lot::RwLock;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::json; use serde_json::json;
@@ -434,19 +434,18 @@ impl Rag {
} in loaded_documents } in loaded_documents
{ {
let hash = sha256(&contents); let hash = sha256(&contents);
if let Some(file_ids) = to_deleted.get_mut(&hash) { if let Some(file_ids) = to_deleted.get_mut(&hash)
if let Some((i, _)) = file_ids && let Some((i, _)) = file_ids
.iter() .iter()
.enumerate() .enumerate()
.find(|(_, v)| self.data.files[*v].path == path) .find(|(_, v)| self.data.files[*v].path == path)
{ {
if file_ids.len() == 1 { if file_ids.len() == 1 {
to_deleted.swap_remove(&hash); to_deleted.swap_remove(&hash);
} else { } else {
file_ids.remove(i); file_ids.remove(i);
}
continue;
} }
continue;
} }
let extension = metadata let extension = metadata
.swap_remove(EXTENSION_METADATA) .swap_remove(EXTENSION_METADATA)
@@ -679,7 +678,7 @@ impl Rag {
Err(e) => { Err(e) => {
return Err(e).with_context(|| { return Err(e).with_context(|| {
format!("Failed to create embedding after {retry_limit} attempts") format!("Failed to create embedding after {retry_limit} attempts")
})? })?;
} }
} }
}; };
@@ -926,11 +925,7 @@ fn add_documents() -> Result<Vec<String>> {
.split(';') .split(';')
.filter_map(|v| { .filter_map(|v| {
let v = v.trim().to_string(); let v = v.trim().to_string();
if v.is_empty() { if v.is_empty() { None } else { Some(v) }
None
} else {
Some(v)
}
}) })
.collect(); .collect();
Ok(paths) Ok(paths)
+2 -2
View File
@@ -1,7 +1,7 @@
use super::*; use super::*;
use base64::{engine::general_purpose::STANDARD, Engine}; use base64::{Engine, engine::general_purpose::STANDARD};
use serde::{de, Deserializer, Serializer}; use serde::{Deserializer, Serializer, de};
pub fn serialize<S>( pub fn serialize<S>(
vectors: &IndexMap<DocumentId, Vec<f32>>, vectors: &IndexMap<DocumentId, Vec<f32>>,
+6 -10
View File
@@ -120,10 +120,10 @@ impl RecursiveCharacterTextSplitter {
} }
}; };
if prev_chunk.is_some() { if prev_chunk.is_some()
if let Some(chunk_overlap_header) = chunk_overlap_header { && let Some(chunk_overlap_header) = chunk_overlap_header
page_content += chunk_overlap_header; {
} page_content += chunk_overlap_header;
} }
let metadata = metadatas[i].clone(); let metadata = metadatas[i].clone();
@@ -240,11 +240,7 @@ impl RecursiveCharacterTextSplitter {
fn join_docs(&self, docs: &[String], separator: &str) -> Option<String> { fn join_docs(&self, docs: &[String], separator: &str) -> Option<String> {
let text = docs.join(separator).trim().to_string(); let text = docs.join(separator).trim().to_string();
if text.is_empty() { if text.is_empty() { None } else { Some(text) }
None
} else {
Some(text)
}
} }
} }
@@ -309,7 +305,7 @@ mod tests {
use super::*; use super::*;
use indexmap::IndexMap; use indexmap::IndexMap;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use serde_json::{json, Value}; use serde_json::{Value, json};
fn build_metadata(source: &str) -> Value { fn build_metadata(source: &str) -> Value {
json!({ "source": source }) json!({ "source": source })
+5 -5
View File
@@ -1,7 +1,7 @@
use crate::utils::decode_bin; use crate::utils::decode_bin;
use ansi_colours::AsRGB; use ansi_colours::AsRGB;
use anyhow::{anyhow, Context, Result}; use anyhow::{Context, Result, anyhow};
use crossterm::style::{Color, Stylize}; use crossterm::style::{Color, Stylize};
use crossterm::terminal; use crossterm::terminal;
use std::collections::HashMap; use std::collections::HashMap;
@@ -122,10 +122,10 @@ impl MarkdownRender {
line_type = LineType::Normal; line_type = LineType::Normal;
} }
LineType::CodeBegin => { LineType::CodeBegin => {
if code_syntax.is_none() { if code_syntax.is_none()
if let Some(syntax) = self.syntax_set.find_syntax_by_first_line(line) { && let Some(syntax) = self.syntax_set.find_syntax_by_first_line(line)
code_syntax = Some(syntax.clone()); {
} code_syntax = Some(syntax.clone());
} }
line_type = LineType::CodeInner; line_type = LineType::CodeInner;
is_code = true; is_code = true;
+1 -1
View File
@@ -4,7 +4,7 @@ mod stream;
pub use self::markdown::{MarkdownRender, RenderOptions}; pub use self::markdown::{MarkdownRender, RenderOptions};
use self::stream::{markdown_stream, raw_stream}; use self::stream::{markdown_stream, raw_stream};
use crate::utils::{error_text, pretty_error, AbortSignal, IS_STDOUT_TERMINAL}; use crate::utils::{AbortSignal, IS_STDOUT_TERMINAL, error_text, pretty_error};
use crate::{client::SseEvent, config::GlobalConfig}; use crate::{client::SseEvent, config::GlobalConfig};
use anyhow::Result; use anyhow::Result;
+2 -2
View File
@@ -1,6 +1,6 @@
use super::{MarkdownRender, SseEvent}; use super::{MarkdownRender, SseEvent};
use crate::utils::{poll_abort_signal, spawn_spinner, AbortSignal}; use crate::utils::{AbortSignal, poll_abort_signal, spawn_spinner};
use anyhow::Result; use anyhow::Result;
use crossterm::{ use crossterm::{
@@ -8,7 +8,7 @@ use crossterm::{
terminal::{self, disable_raw_mode, enable_raw_mode}, terminal::{self, disable_raw_mode, enable_raw_mode},
}; };
use std::{ use std::{
io::{stdout, Stdout, Write}, io::{Stdout, Write, stdout},
time::Duration, time::Duration,
}; };
use textwrap::core::display_width; use textwrap::core::display_width;
+1 -1
View File
@@ -1,4 +1,4 @@
use super::{ReplCommand, REPL_COMMANDS}; use super::{REPL_COMMANDS, ReplCommand};
use crate::{config::GlobalConfig, utils::fuzzy_filter}; use crate::{config::GlobalConfig, utils::fuzzy_filter};
+10 -10
View File
@@ -8,23 +8,23 @@ use self::prompt::ReplPrompt;
use crate::client::{call_chat_completions, call_chat_completions_streaming}; use crate::client::{call_chat_completions, call_chat_completions_streaming};
use crate::config::{ use crate::config::{
macro_execute, AgentVariables, AssertState, Config, GlobalConfig, Input, LastMessage, AgentVariables, AssertState, Config, GlobalConfig, Input, LastMessage, StateFlags,
StateFlags, macro_execute,
}; };
use crate::render::render_error; use crate::render::render_error;
use crate::utils::{ use crate::utils::{
abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file, AbortSignal, AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file,
}; };
use crate::mcp::McpRegistry; use crate::mcp::McpRegistry;
use anyhow::{bail, Context, Result}; use anyhow::{Context, Result, bail};
use crossterm::cursor::SetCursorStyle; use crossterm::cursor::SetCursorStyle;
use fancy_regex::Regex; use fancy_regex::Regex;
use reedline::CursorConfig; use reedline::CursorConfig;
use reedline::{ use reedline::{
default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
ColumnarMenu, EditCommand, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline, ColumnarMenu, EditCommand, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
ReedlineEvent, ReedlineMenu, ValidationResult, Validator, Vi, ReedlineEvent, ReedlineMenu, ValidationResult, Validator, Vi, default_emacs_keybindings,
default_vi_insert_keybindings, default_vi_normal_keybindings,
}; };
use reedline::{MenuBuilder, Signal}; use reedline::{MenuBuilder, Signal};
use std::sync::LazyLock; use std::sync::LazyLock;
@@ -382,10 +382,10 @@ pub async fn run_repl_command(
abort_signal: AbortSignal, abort_signal: AbortSignal,
mut line: &str, mut line: &str,
) -> Result<bool> { ) -> Result<bool> {
if let Ok(Some(captures)) = MULTILINE_RE.captures(line) { if let Ok(Some(captures)) = MULTILINE_RE.captures(line)
if let Some(text_match) = captures.get(1) { && let Some(text_match) = captures.get(1)
line = text_match.as_str(); {
} line = text_match.as_str();
} }
match parse_command(line) { match parse_command(line) {
Some((cmd, args)) => match cmd { Some((cmd, args)) => match cmd {
+13 -13
View File
@@ -2,8 +2,8 @@ use anyhow::Result;
use crossterm::event::{self, Event, KeyCode, KeyModifiers}; use crossterm::event::{self, Event, KeyCode, KeyModifiers};
use std::{ use std::{
sync::{ sync::{
atomic::{AtomicBool, Ordering},
Arc, Arc,
atomic::{AtomicBool, Ordering},
}, },
time::Duration, time::Duration,
}; };
@@ -69,19 +69,19 @@ pub async fn wait_abort_signal(abort_signal: &AbortSignal) {
} }
pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result<bool> { pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result<bool> {
if event::poll(Duration::from_millis(25))? { if event::poll(Duration::from_millis(25))?
if let Event::Key(key) = event::read()? { && let Event::Key(key) = event::read()?
match key.code { {
KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => { match key.code {
abort_signal.set_ctrlc(); KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => {
return Ok(true); abort_signal.set_ctrlc();
} return Ok(true);
KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => {
abort_signal.set_ctrld();
return Ok(true);
}
_ => {}
} }
KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => {
abort_signal.set_ctrld();
return Ok(true);
}
_ => {}
} }
} }
Ok(false) Ok(false)
+1 -1
View File
@@ -3,7 +3,7 @@ use anyhow::Context;
#[cfg(not(any(target_os = "android", target_os = "emscripten")))] #[cfg(not(any(target_os = "android", target_os = "emscripten")))]
mod internal { mod internal {
use arboard::Clipboard; use arboard::Clipboard;
use base64::{engine::general_purpose::STANDARD, Engine as _}; use base64::{Engine as _, engine::general_purpose::STANDARD};
use std::sync::{LazyLock, Mutex}; use std::sync::{LazyLock, Mutex};
static CLIPBOARD: LazyLock<Mutex<Option<Clipboard>>> = static CLIPBOARD: LazyLock<Mutex<Option<Clipboard>>> =
+1 -1
View File
@@ -10,7 +10,7 @@ use std::{
process::Command, process::Command,
}; };
use anyhow::{anyhow, bail, Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use dirs::home_dir; use dirs::home_dir;
use std::sync::LazyLock; use std::sync::LazyLock;
+1 -1
View File
@@ -1,4 +1,4 @@
use base64::{engine::general_purpose::STANDARD, Engine}; use base64::{Engine, engine::general_purpose::STANDARD};
use hmac::{Hmac, Mac}; use hmac::{Hmac, Mac};
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
+1 -1
View File
@@ -1,6 +1,6 @@
use std::{cell::RefCell, rc::Rc}; use std::{cell::RefCell, rc::Rc};
use html_to_markdown::{markdown, TagHandler}; use html_to_markdown::{TagHandler, markdown};
pub fn html_to_md(html: &str) -> String { pub fn html_to_md(html: &str) -> String {
let mut handlers: Vec<TagHandler> = vec![ let mut handlers: Vec<TagHandler> = vec![
+1 -1
View File
@@ -1,7 +1,7 @@
use anyhow::Result; use anyhow::Result;
use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyModifiers}; use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyModifiers};
use crossterm::terminal::{disable_raw_mode, enable_raw_mode}; use crossterm::terminal::{disable_raw_mode, enable_raw_mode};
use std::io::{stdout, Write}; use std::io::{Write, stdout};
/// Reads a single character from stdin without requiring Enter /// Reads a single character from stdin without requiring Enter
/// Returns the character if it's one of the valid options, or the default if Enter is pressed /// Returns the character if it's one of the valid options, or the default if Enter is pressed
+1 -1
View File
@@ -1,6 +1,6 @@
use super::*; use super::*;
use anyhow::{anyhow, Context, Result}; use anyhow::{Context, Result, anyhow};
use indexmap::IndexMap; use indexmap::IndexMap;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::collections::HashMap; use std::collections::HashMap;
+1 -1
View File
@@ -29,7 +29,7 @@ pub use self::variables::*;
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use fancy_regex::Regex; use fancy_regex::Regex;
use fuzzy_matcher::{skim::SkimMatcherV2, FuzzyMatcher}; use fuzzy_matcher::{FuzzyMatcher, skim::SkimMatcherV2};
use is_terminal::IsTerminal; use is_terminal::IsTerminal;
use std::borrow::Cow; use std::borrow::Cow;
use std::sync::LazyLock; use std::sync::LazyLock;
+19 -20
View File
@@ -1,7 +1,7 @@
use std::fs; use std::fs;
use std::path::{Component, Path, PathBuf}; use std::path::{Component, Path, PathBuf};
use anyhow::{bail, Result}; use anyhow::{Result, bail};
use fancy_regex::Regex; use fancy_regex::Regex;
use indexmap::IndexSet; use indexmap::IndexSet;
use path_absolutize::Absolutize; use path_absolutize::Absolutize;
@@ -97,11 +97,12 @@ pub fn to_absolute_path(path: &str) -> Result<String> {
pub fn resolve_home_dir(path: &str) -> String { pub fn resolve_home_dir(path: &str) -> String {
let mut path = path.to_string(); let mut path = path.to_string();
if path.starts_with("~/") || path.starts_with("~\\") { if (path.starts_with("~/") || path.starts_with("~\\"))
if let Some(home_dir) = dirs::home_dir() { && let Some(home_dir) = dirs::home_dir()
path.replace_range(..1, &home_dir.display().to_string()); {
} path.replace_range(..1, &home_dir.display().to_string());
} }
path path
} }
@@ -241,22 +242,20 @@ fn add_file(files: &mut IndexSet<String>, suffixes: Option<&Vec<String>>, path:
fn is_valid_extension(suffixes: Option<&Vec<String>>, path: &Path) -> bool { fn is_valid_extension(suffixes: Option<&Vec<String>>, path: &Path) -> bool {
let filename_regex = Regex::new(r"^.+\.*").unwrap(); let filename_regex = Regex::new(r"^.+\.*").unwrap();
if let Some(suffixes) = suffixes { if let Some(suffixes) = suffixes
if !suffixes.is_empty() { && !suffixes.is_empty()
if let Ok(Some(_)) = filename_regex.find(&suffixes.join(",")) { {
let file_name = path if let Ok(Some(_)) = filename_regex.find(&suffixes.join(",")) {
.file_name() let file_name = path
.and_then(|v| v.to_str()) .file_name()
.expect("invalid filename") .and_then(|v| v.to_str())
.to_string(); .expect("invalid filename")
return suffixes.contains(&file_name); .to_string();
} else if let Some(extension) = return suffixes.contains(&file_name);
path.extension().map(|v| v.to_string_lossy().to_string()) } else if let Some(extension) = path.extension().map(|v| v.to_string_lossy().to_string()) {
{ return suffixes.contains(&extension);
return suffixes.contains(&extension);
}
return false;
} }
return false;
} }
true true
} }
+2 -2
View File
@@ -1,8 +1,8 @@
use super::*; use super::*;
use anyhow::{anyhow, bail, Context, Result}; use anyhow::{Context, Result, anyhow, bail};
use fancy_regex::Regex; use fancy_regex::Regex;
use futures_util::{stream, StreamExt}; use futures_util::{StreamExt, stream};
use http::header::CONTENT_TYPE; use http::header::CONTENT_TYPE;
use reqwest::Url; use reqwest::Url;
use scraper::{Html, Selector}; use scraper::{Html, Selector};
+3 -3
View File
@@ -1,10 +1,10 @@
use super::{poll_abort_signal, wait_abort_signal, AbortSignal, IS_STDOUT_TERMINAL}; use super::{AbortSignal, IS_STDOUT_TERMINAL, poll_abort_signal, wait_abort_signal};
use anyhow::{bail, Result}; use anyhow::{Result, bail};
use crossterm::{cursor, queue, style, terminal}; use crossterm::{cursor, queue, style, terminal};
use std::{ use std::{
future::Future, future::Future,
io::{stdout, Write}, io::{Write, stdout},
time::Duration, time::Duration,
}; };
use tokio::{ use tokio::{
+2 -2
View File
@@ -9,9 +9,9 @@ use crate::config::Config;
use crate::vault::utils::ensure_password_file_initialized; use crate::vault::utils::ensure_password_file_initialized;
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use fancy_regex::Regex; use fancy_regex::Regex;
use gman::providers::local::LocalProvider;
use gman::providers::SecretProvider; use gman::providers::SecretProvider;
use inquire::{required, Password, PasswordDisplayMode}; use gman::providers::local::LocalProvider;
use inquire::{Password, PasswordDisplayMode, required};
use std::sync::{Arc, LazyLock}; use std::sync::{Arc, LazyLock};
use tokio::runtime::Handle; use tokio::runtime::Handle;
+20 -8
View File
@@ -1,11 +1,11 @@
use crate::config::ensure_parent_exists; use crate::config::ensure_parent_exists;
use crate::vault::{Vault, SECRET_RE}; use crate::vault::{SECRET_RE, Vault};
use anyhow::anyhow;
use anyhow::Result; use anyhow::Result;
use anyhow::anyhow;
use gman::providers::local::LocalProvider; use gman::providers::local::LocalProvider;
use indoc::formatdoc; use indoc::formatdoc;
use inquire::validator::Validation; use inquire::validator::Validation;
use inquire::{min_length, required, Confirm, Password, PasswordDisplayMode, Text}; use inquire::{Confirm, Password, PasswordDisplayMode, Text, min_length, required};
use std::borrow::Cow; use std::borrow::Cow;
use std::path::PathBuf; use std::path::PathBuf;
@@ -21,11 +21,16 @@ pub fn ensure_password_file_initialized(local_provider: &mut LocalProvider) -> R
if !file_contents.trim().is_empty() { if !file_contents.trim().is_empty() {
Ok(()) Ok(())
} else { } else {
Err(anyhow!("The configured password file '{}' is empty. Please populate it with a password and try again.", vault_password_file.display())) Err(anyhow!(
"The configured password file '{}' is empty. Please populate it with a password and try again.",
vault_password_file.display()
))
} }
} }
} else { } else {
Err(anyhow!("A password file is required to utilize the Loki vault. Please configure a password file in your config file and try again.")) Err(anyhow!(
"A password file is required to utilize the Loki vault. Please configure a password file in your config file and try again."
))
} }
} }
@@ -40,7 +45,9 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
{ {
let file_contents = std::fs::read_to_string(&vault_password_file)?; let file_contents = std::fs::read_to_string(&vault_password_file)?;
if !file_contents.trim().is_empty() { if !file_contents.trim().is_empty() {
debug!("create_vault_password_file was called but the password file already exists and is non-empty"); debug!(
"create_vault_password_file was called but the password file already exists and is non-empty"
);
return Ok(()); return Ok(());
} }
} }
@@ -56,7 +63,10 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
.prompt()?; .prompt()?;
if !ans { if !ans {
return Err(anyhow!("The configured password file '{}' is empty. Please populate it with a password and try again.", vault_password_file.display())); return Err(anyhow!(
"The configured password file '{}' is empty. Please populate it with a password and try again.",
vault_password_file.display()
));
} }
let password = Password::new("Enter a password to encrypt all vault secrets:") let password = Password::new("Enter a password to encrypt all vault secrets:")
@@ -85,7 +95,9 @@ pub fn create_vault_password_file(vault: &mut Vault) -> Result<()> {
.prompt()?; .prompt()?;
if !ans { if !ans {
return Err(anyhow!("A password file is required to utilize the Loki vault. Please configure a password file in your config file and try again.")); return Err(anyhow!(
"A password file is required to utilize the Loki vault. Please configure a password file in your config file and try again."
));
} }
let password_file: PathBuf = Text::new("Enter the path to the password file to create:") let password_file: PathBuf = Text::new("Enter the path to the password file to create:")