Skip to content

Commit

Permalink
feat: jina and serper API checks (#133)
Browse files Browse the repository at this point in the history
* added `jina` and `serper` checks

* tiny typo fixes, update workflows
  • Loading branch information
erhant authored Oct 14, 2024
1 parent e0386e4 commit 09c0fa7
Show file tree
Hide file tree
Showing 8 changed files with 243 additions and 31 deletions.
49 changes: 24 additions & 25 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ default-members = ["compute"]

[workspace.package]
edition = "2021"
version = "0.2.14"
version = "0.2.15"
license = "Apache-2.0"
readme = "README.md"

Expand Down
1 change: 0 additions & 1 deletion workflows/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ ollama-workflows = { git = "https://github.com/andthattoo/ollama-workflows" }
# async stuff
tokio-util.workspace = true
tokio.workspace = true
async-trait.workspace = true

# serialize & deserialize
serde.workspace = true
Expand Down
92 changes: 92 additions & 0 deletions workflows/src/apis/jina.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
use eyre::{eyre, Context, Result};
use reqwest::Client;
use std::env;

/// Makes a request for `example.com`.
const JINA_EXAMPLE_ENDPOINT: &str = "https://r.jina.ai/https://example.com";
const ENV_VAR_NAME: &str = "JINA_API_KEY";

/// Jina-specific configurations.
#[derive(Debug, Clone, Default)]
pub struct JinaConfig {
/// API key, if available.
api_key: Option<String>,
}

impl JinaConfig {
/// Looks at the environment variables for Jina API key.
pub fn new() -> Self {
Self {
api_key: env::var(ENV_VAR_NAME).ok(),
}
}

/// Sets the API key for Jina.
pub fn with_api_key(mut self, api_key: String) -> Self {
self.api_key = Some(api_key);
self
}

/// Checks API KEY, and if it exists tries a dummy request.
/// Fails if the provided API KEY is not authorized enough for the dummy request.
///
/// Equivalent cURL is as follows:
///
/// ```sh
/// curl 'https://r.jina.ai/https://example.com' \
/// -H "Authorization: Bearer jina_key"
/// ```
pub async fn check_optional(&self) -> Result<()> {
// check API key
let Some(api_key) = &self.api_key else {
log::debug!("Jina API key not found, skipping Jina check");
return Ok(());
};
log::info!("Jina API key found, checking Jina service");

// make a dummy request models
let client = Client::new();
let request = client
.get(JINA_EXAMPLE_ENDPOINT)
.header("Authorization", format!("Bearer {}", api_key))
.build()
.wrap_err("failed to build request")?;

let response = client
.execute(request)
.await
.wrap_err("failed to send request")?;

// parse response
if response.status().is_client_error() {
return Err(eyre!("Failed to make Jina request",))
.wrap_err(response.text().await.unwrap_or_default());
}

log::info!("Jina check succesful!");

Ok(())
}
}

#[cfg(test)]
mod tests {
use super::*;

#[tokio::test]
#[ignore = "requires Jina API key"]
async fn test_jina_check() {
let _ = dotenvy::dotenv();
assert!(env::var(ENV_VAR_NAME).is_ok());
let res = JinaConfig::new().check_optional().await;
assert!(res.is_ok(), "should pass with api key");

env::set_var(ENV_VAR_NAME, "i-dont-work");
let res = JinaConfig::new().check_optional().await;
assert!(res.is_err(), "should fail with bad api key");

env::remove_var(ENV_VAR_NAME);
let res = JinaConfig::new().check_optional().await;
assert!(res.is_ok(), "should pass without api key");
}
}
5 changes: 5 additions & 0 deletions workflows/src/apis/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
mod jina;
pub use jina::JinaConfig;

mod serper;
pub use serper::SerperConfig;
98 changes: 98 additions & 0 deletions workflows/src/apis/serper.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
use eyre::{eyre, Context, Result};
use reqwest::Client;
use std::env;

/// Makes a search request.
const SERPER_EXAMPLE_ENDPOINT: &str = "https://google.serper.dev/search";
const ENV_VAR_NAME: &str = "SERPER_API_KEY";

/// Serper-specific configurations.
#[derive(Debug, Clone, Default)]
pub struct SerperConfig {
/// API key, if available.
api_key: Option<String>,
}

impl SerperConfig {
/// Looks at the environment variables for Serper API key.
pub fn new() -> Self {
Self {
api_key: env::var(ENV_VAR_NAME).ok(),
}
}

/// Sets the API key for Serper.
pub fn with_api_key(mut self, api_key: String) -> Self {
self.api_key = Some(api_key);
self
}

/// Check if Serper API KEY exists and if it does, tries a dummy request.
/// Fails if the provided API KEY is not authorized enough for the dummy request.
///
/// Equivalent cURL is as follows:
///
/// ```sh
/// curl -X POST 'https://google.serper.dev/search' \
/// -H 'X-API-KEY: API_KEY' \
/// -H 'Content-Type: application/json' \
/// -d '{
/// "q": "Your search query here"
/// }'
/// ```
pub async fn check_optional(&self) -> Result<()> {
// check API key
let Some(api_key) = &self.api_key else {
log::debug!("Serper API key not found, skipping Serper check");
return Ok(());
};
log::info!("Serper API key found, checking Serper service");

// make a dummy request
let client = Client::new();
let request = client
.post(SERPER_EXAMPLE_ENDPOINT)
.header("X-API-KEY", api_key)
.header("Content-Type", "application/json")
.body("{\"q\": \"Your search query here\"}")
.build()
.wrap_err("failed to build request")?;

let response = client
.execute(request)
.await
.wrap_err("failed to send request")?;

// parse response
if response.status().is_client_error() {
return Err(eyre!("Failed to make Serper request",))
.wrap_err(response.text().await.unwrap_or_default());
}

log::info!("Serper check succesful!");

Ok(())
}
}

#[cfg(test)]
mod tests {
use super::*;

#[tokio::test]
#[ignore = "requires Serper API key"]
async fn test_serper_check() {
let _ = dotenvy::dotenv();
assert!(env::var(ENV_VAR_NAME).is_ok());
let res = SerperConfig::new().check_optional().await;
assert!(res.is_ok(), "should pass with api key");

env::set_var(ENV_VAR_NAME, "i-dont-work");
let res = SerperConfig::new().check_optional().await;
assert!(res.is_err(), "should fail with bad api key");

env::remove_var(ENV_VAR_NAME);
let res = SerperConfig::new().check_optional().await;
assert!(res.is_ok(), "should pass without api key");
}
}
Loading

0 comments on commit 09c0fa7

Please sign in to comment.