Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 2 additions & 10 deletions src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use serde_json::{Map, Value};
use tokio::time::sleep;

use crate::error::KagiError;
use crate::http;
use crate::parser::parse_assistant_thread_list;
#[cfg(test)]
use crate::types::ApiMeta;
Expand All @@ -31,11 +32,6 @@ use crate::types::{
TranslateWarning, TranslationSuggestionsResponse, WordInsightsResponse,
};

const USER_AGENT: &str = concat!(
"kagi-cli/",
env!("CARGO_PKG_VERSION"),
" (+https://github.com/Microck/kagi-cli)"
);
const KAGI_SUMMARIZE_URL: &str = "https://kagi.com/api/v0/summarize";
const KAGI_SUBSCRIBER_SUMMARIZE_URL: &str = "https://kagi.com/mother/summary_labs";
const KAGI_NEWS_LATEST_URL: &str = "https://news.kagi.com/api/batches/latest";
Expand Down Expand Up @@ -2179,11 +2175,7 @@ where
}

fn build_client() -> Result<Client, KagiError> {
Client::builder()
.user_agent(USER_AGENT)
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}")))
http::client_30s()
}

fn map_transport_error(error: reqwest::Error) -> KagiError {
Expand Down
41 changes: 41 additions & 0 deletions src/http.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
use std::sync::OnceLock;
use std::time::Duration;

use reqwest::Client;

use crate::error::KagiError;

const USER_AGENT: &str = concat!(
"kagi-cli/",
env!("CARGO_PKG_VERSION"),
" (+https://github.com/Microck/kagi-cli)"
);

static CLIENT_20S: OnceLock<Result<Client, String>> = OnceLock::new();
static CLIENT_30S: OnceLock<Result<Client, String>> = OnceLock::new();

pub fn client_20s() -> Result<Client, KagiError> {
cached_client(&CLIENT_20S, Duration::from_secs(20))
}

pub fn client_30s() -> Result<Client, KagiError> {
cached_client(&CLIENT_30S, Duration::from_secs(30))
}

fn cached_client(
slot: &OnceLock<Result<Client, String>>,
timeout: Duration,
) -> Result<Client, KagiError> {
let result = slot.get_or_init(|| {
Client::builder()
.user_agent(USER_AGENT)
.timeout(timeout)
.build()
.map_err(|error| format!("failed to build HTTP client: {error}"))
});

result
.as_ref()
.cloned()
.map_err(|error| KagiError::Network(error.clone()))
}
54 changes: 22 additions & 32 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ mod auth;
mod auth_wizard;
mod cli;
mod error;
mod http;
mod parser;
mod quick;
mod search;
Expand Down Expand Up @@ -53,7 +54,7 @@ struct SearchRequestOptions {
no_personalized: bool,
}

#[tokio::main]
#[tokio::main(flavor = "current_thread")]
async fn main() {
if let Err(error) = run().await {
eprintln!("{error}");
Expand Down Expand Up @@ -746,10 +747,9 @@ async fn run_batch_search(
let semaphore_clone = Arc::clone(&semaphore);
let credentials_clone = credentials.clone();
let options_clone = options.clone();
let format_clone = format.clone();
let query_clone = query.clone();

let handle: tokio::task::JoinHandle<Result<(String, String), KagiError>> =
let handle: tokio::task::JoinHandle<Result<(String, SearchResponse), KagiError>> =
tokio::spawn(async move {
let _permit = semaphore_clone.acquire().await;
rate_limiter_clone.acquire().await?;
Expand All @@ -758,19 +758,7 @@ async fn run_batch_search(

let response = execute_search_request(&request, credentials_clone).await?;

let output = match format_clone.as_str() {
"pretty" => format_pretty_response(&response, use_color),
"compact" => serde_json::to_string(&response).map_err(|error| {
KagiError::Parse(format!("failed to serialize search response: {error}"))
})?,
"markdown" => format_markdown_response(&response),
"csv" => format_csv_response(&response),
_ => serde_json::to_string_pretty(&response).map_err(|error| {
KagiError::Parse(format!("failed to serialize search response: {error}"))
})?,
};

Ok((query, output))
Ok((query, response))
});

handles.push(handle);
Expand Down Expand Up @@ -804,32 +792,34 @@ async fn run_batch_search(
if format == "json" || format == "compact" {
// For machine-readable formats, create a proper JSON envelope
let queries: Vec<String> = results.iter().map(|(query, _)| query.clone()).collect();
let mut results_json = serde_json::json!({
let results_payload = results
.into_iter()
.map(|(_, response)| serde_json::to_value(response))
.collect::<Result<Vec<_>, _>>()
.map_err(|error| {
KagiError::Parse(format!("failed to serialize batch search response: {error}"))
})?;
let results_json = serde_json::json!({
"queries": queries,
"results": []
"results": results_payload
});

let results_array = results_json["results"].as_array_mut().unwrap();

for (query, output) in results {
// Parse the individual JSON output and add to array
let parsed: serde_json::Value = serde_json::from_str(&output).map_err(|e| {
KagiError::Parse(format!(
"failed to parse batch result for '{}': {}",
query, e
))
})?;
results_array.push(parsed);
}

if format == "compact" {
println!("{}", serde_json::to_string(&results_json)?);
} else {
println!("{}", serde_json::to_string_pretty(&results_json)?);
}
} else {
// For human-readable formats, output with headers
for (query, output) in results {
for (query, response) in results {
let output = match format.as_str() {
"pretty" => format_pretty_response(&response, use_color),
"markdown" => format_markdown_response(&response),
"csv" => format_csv_response(&response),
_ => serde_json::to_string_pretty(&response).map_err(|error| {
KagiError::Parse(format!("failed to serialize search response: {error}"))
})?,
};
println!("=== Results for: {} ===", query);
println!("{}", output);
println!();
Expand Down
8 changes: 2 additions & 6 deletions src/quick.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@ use scraper::Html;
use serde::Deserialize;

use crate::error::KagiError;
use crate::http;
use crate::search::{SearchRequest, validate_lens_value};
use crate::types::{
QuickMessage, QuickMeta, QuickReferenceCollection, QuickReferenceItem, QuickResponse,
};

const USER_AGENT: &str = "kagi-cli/0.1.0 (+https://github.com/)";
const KAGI_QUICK_ANSWER_URL: &str = "https://kagi.com/mother/context";

pub async fn execute_quick(
Expand Down Expand Up @@ -452,11 +452,7 @@ fn format_client_error_suffix(body: &str) -> String {
}

fn build_client() -> Result<Client, KagiError> {
Client::builder()
.user_agent(USER_AGENT)
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}")))
http::client_30s()
}

fn map_transport_error(error: reqwest::Error) -> KagiError {
Expand Down
8 changes: 2 additions & 6 deletions src/search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@ use reqwest::{Client, StatusCode, header};
use serde::Deserialize;

use crate::error::KagiError;
use crate::http;
use crate::parser::parse_search_results;
use crate::types::{SearchResponse, SearchResult};

const KAGI_SEARCH_URL: &str = "https://kagi.com/html/search";
const KAGI_API_SEARCH_URL: &str = "https://kagi.com/api/v0/search";
const USER_AGENT: &str = "kagi-cli/0.1.0 (+https://github.com/)";
const UNAUTHENTICATED_MARKERS: [&str; 3] = [
"<title>Kagi Search - A Premium Search Engine</title>",
"Welcome to Kagi",
Expand Down Expand Up @@ -406,11 +406,7 @@ fn looks_unauthenticated(body: &str) -> bool {
}

fn build_client() -> Result<Client, KagiError> {
Client::builder()
.user_agent(USER_AGENT)
.timeout(std::time::Duration::from_secs(20))
.build()
.map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}")))
http::client_20s()
}

fn map_transport_error(error: reqwest::Error) -> KagiError {
Expand Down
Loading