mod bm25;
pub mod config;
mod discovery;
mod error;
mod eval;
mod explain;
mod introspection;
mod json_utils;
mod query_store;
mod types;
#[cfg(feature = "arrow")]
pub mod arrow;
pub use bm25::{Bm25Index, DocInfo, IndexOptions, SearchResult as Bm25SearchResult, TermInfo};
pub use config::{EngineBuilder, EngineConfig, EngineSection, FunctionsSection, QueriesSection};
pub use discovery::{
CategoryInfo, CategorySummary, DiscoveryRegistry, DiscoverySpec, ExampleSpec, IndexStats,
ParamSpec, RegistrationResult, ReturnSpec, ServerInfo, ServerSummary, ToolQueryResult,
ToolSpec,
};
pub use error::{EngineError, EvaluationErrorKind, Result};
pub use explain::{ExplainResult, ExplainStep, has_let_nodes};
pub use introspection::{FunctionDetail, SearchResult, SimilarFunctionsResult};
pub use json_utils::{FieldAnalysis, PathInfo, StatsResult};
pub use query_store::{QueryStore, StoredQuery};
pub use types::{
BatchEvaluateResult, BatchExpressionResult, EvalRequest, EvalResponse, ValidationResult,
};
use discovery::DiscoveryRegistry as DiscoveryRegistryInner;
use error::EngineError as EngineErrorInner;
use query_store::QueryStore as QueryStoreInner;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
pub use jpx_core::ast;
pub use jpx_core::query_library;
pub use jpx_core::{Category, Expression, FunctionInfo, FunctionRegistry, Runtime, compile, parse};
pub struct JpxEngine {
pub(crate) runtime: Runtime,
pub(crate) registry: FunctionRegistry,
pub(crate) discovery: Arc<RwLock<DiscoveryRegistryInner>>,
pub(crate) queries: Arc<RwLock<QueryStoreInner>>,
pub(crate) strict: bool,
}
impl JpxEngine {
pub fn new() -> Self {
Self::with_options(false)
}
pub fn with_options(strict: bool) -> Self {
let mut runtime = Runtime::new();
runtime.register_builtin_functions();
let mut registry = FunctionRegistry::new();
registry.register_all();
if !strict {
registry.apply(&mut runtime);
}
Self {
runtime,
registry,
discovery: Arc::new(RwLock::new(DiscoveryRegistryInner::new())),
queries: Arc::new(RwLock::new(QueryStoreInner::new())),
strict,
}
}
pub fn strict() -> Self {
Self::with_options(true)
}
pub fn from_config(config: EngineConfig) -> Result<Self> {
let strict = config.engine.strict;
let (runtime, registry) = config::build_runtime_from_config(&config.functions, strict);
let discovery = Arc::new(RwLock::new(DiscoveryRegistryInner::new()));
let queries = Arc::new(RwLock::new(QueryStoreInner::new()));
config::load_queries_into_store(&config.queries, &runtime, &queries)?;
Ok(Self {
runtime,
registry,
discovery,
queries,
strict,
})
}
pub fn builder() -> config::EngineBuilder {
config::EngineBuilder::new()
}
pub fn is_strict(&self) -> bool {
self.strict
}
pub fn runtime(&self) -> &Runtime {
&self.runtime
}
pub fn registry(&self) -> &FunctionRegistry {
&self.registry
}
pub fn discovery(&self) -> &Arc<RwLock<DiscoveryRegistryInner>> {
&self.discovery
}
pub fn queries(&self) -> &Arc<RwLock<QueryStoreInner>> {
&self.queries
}
pub fn define_query(
&self,
name: String,
expression: String,
description: Option<String>,
) -> Result<Option<StoredQuery>> {
let validation = self.validate(&expression);
if !validation.valid {
return Err(EngineErrorInner::InvalidExpression(
validation
.error
.unwrap_or_else(|| "Invalid expression".to_string()),
));
}
let query = StoredQuery {
name,
expression,
description,
};
self.queries
.write()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.define(query)
.pipe(Ok)
}
pub fn get_query(&self, name: &str) -> Result<Option<StoredQuery>> {
Ok(self
.queries
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.get(name)
.cloned())
}
pub fn delete_query(&self, name: &str) -> Result<Option<StoredQuery>> {
Ok(self
.queries
.write()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.delete(name))
}
pub fn list_queries(&self) -> Result<Vec<StoredQuery>> {
Ok(self
.queries
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.list()
.into_iter()
.cloned()
.collect())
}
pub fn run_query(&self, name: &str, input: &Value) -> Result<Value> {
let query = self
.get_query(name)?
.ok_or_else(|| EngineErrorInner::QueryNotFound(name.to_string()))?;
self.evaluate(&query.expression, input)
}
pub fn register_discovery(
&self,
spec: DiscoverySpec,
replace: bool,
) -> Result<RegistrationResult> {
Ok(self
.discovery
.write()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.register(spec, replace))
}
pub fn unregister_discovery(&self, server_name: &str) -> Result<bool> {
Ok(self
.discovery
.write()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.unregister(server_name))
}
pub fn query_tools(&self, query: &str, top_k: usize) -> Result<Vec<ToolQueryResult>> {
Ok(self
.discovery
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.query(query, top_k))
}
pub fn similar_tools(&self, tool_id: &str, top_k: usize) -> Result<Vec<ToolQueryResult>> {
Ok(self
.discovery
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.similar(tool_id, top_k))
}
pub fn list_discovery_servers(&self) -> Result<Vec<ServerSummary>> {
Ok(self
.discovery
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.list_servers())
}
pub fn list_discovery_categories(&self) -> Result<HashMap<String, CategorySummary>> {
Ok(self
.discovery
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.list_categories())
}
pub fn discovery_index_stats(&self) -> Result<Option<IndexStats>> {
Ok(self
.discovery
.read()
.map_err(|e| EngineErrorInner::Internal(e.to_string()))?
.index_stats())
}
pub fn get_discovery_schema(&self) -> Value {
DiscoveryRegistryInner::get_schema()
}
}
impl Default for JpxEngine {
fn default() -> Self {
Self::new()
}
}
trait Pipe: Sized {
fn pipe<T, F: FnOnce(Self) -> T>(self, f: F) -> T {
f(self)
}
}
impl<T> Pipe for T {}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_engine_creation() {
let engine = JpxEngine::new();
assert!(!engine.is_strict());
}
#[test]
fn test_engine_strict_mode() {
let engine = JpxEngine::strict();
assert!(engine.is_strict());
}
#[test]
fn test_engine_default() {
let engine = JpxEngine::default();
assert!(!engine.is_strict());
}
#[test]
fn test_query_store() {
let engine = JpxEngine::new();
engine
.define_query("count".to_string(), "length(@)".to_string(), None)
.unwrap();
let query = engine.get_query("count").unwrap().unwrap();
assert_eq!(query.expression, "length(@)");
let result = engine.run_query("count", &json!([1, 2, 3])).unwrap();
assert_eq!(result, json!(3));
let queries = engine.list_queries().unwrap();
assert_eq!(queries.len(), 1);
engine.delete_query("count").unwrap();
assert!(engine.get_query("count").unwrap().is_none());
}
#[test]
fn test_discovery() {
let engine = JpxEngine::new();
let spec: DiscoverySpec = serde_json::from_value(json!({
"server": {"name": "test-server", "version": "1.0.0"},
"tools": [
{"name": "test_tool", "description": "A test tool", "tags": ["test"]}
]
}))
.unwrap();
let result = engine.register_discovery(spec, false).unwrap();
assert!(result.ok);
assert_eq!(result.tools_indexed, 1);
let servers = engine.list_discovery_servers().unwrap();
assert_eq!(servers.len(), 1);
let tools = engine.query_tools("test", 10).unwrap();
assert!(!tools.is_empty());
assert!(engine.unregister_discovery("test-server").unwrap());
assert!(engine.list_discovery_servers().unwrap().is_empty());
}
#[test]
fn test_with_options_non_strict() {
let engine = JpxEngine::with_options(false);
assert!(!engine.is_strict());
}
#[test]
fn test_with_options_strict() {
let engine = JpxEngine::with_options(true);
assert!(engine.is_strict());
}
#[test]
fn test_from_config_default() {
let config = EngineConfig::default();
let engine = JpxEngine::from_config(config).unwrap();
assert!(!engine.is_strict());
}
#[test]
fn test_builder_default() {
let engine = JpxEngine::builder().build().unwrap();
assert!(!engine.is_strict());
}
#[test]
fn test_runtime_accessor() {
let engine = JpxEngine::new();
let runtime = engine.runtime();
let expr = runtime.compile("length(@)").unwrap();
let data = json!([1, 2, 3]);
let result = expr.search(&data).unwrap();
assert_eq!(result, json!(3));
}
#[test]
fn test_registry_accessor() {
let engine = JpxEngine::new();
let registry = engine.registry();
assert!(registry.get_function("upper").is_some());
assert!(registry.get_function("lower").is_some());
assert!(registry.is_enabled("upper"));
}
#[test]
fn test_discovery_accessor() {
let engine = JpxEngine::new();
let discovery = engine.discovery();
let guard = discovery.read().unwrap();
assert!(guard.list_servers().is_empty());
}
#[test]
fn test_queries_accessor() {
let engine = JpxEngine::new();
let queries = engine.queries();
let guard = queries.read().unwrap();
assert!(guard.is_empty());
}
#[test]
fn test_define_query_with_description() {
let engine = JpxEngine::new();
engine
.define_query(
"named".to_string(),
"length(@)".to_string(),
Some("Counts elements".to_string()),
)
.unwrap();
let query = engine.get_query("named").unwrap().unwrap();
assert_eq!(query.expression, "length(@)");
assert_eq!(query.description, Some("Counts elements".to_string()));
}
#[test]
fn test_define_query_invalid_expression() {
let engine = JpxEngine::new();
let result = engine.define_query("bad".to_string(), "invalid[".to_string(), None);
assert!(result.is_err());
match result.unwrap_err() {
EngineError::InvalidExpression(_) => {} other => panic!("Expected InvalidExpression, got {:?}", other),
}
}
#[test]
fn test_define_query_overwrite() {
let engine = JpxEngine::new();
let first = engine
.define_query("q".to_string(), "length(@)".to_string(), None)
.unwrap();
assert!(first.is_none());
let second = engine
.define_query("q".to_string(), "keys(@)".to_string(), None)
.unwrap();
assert!(second.is_some());
let old = second.unwrap();
assert_eq!(old.expression, "length(@)");
let current = engine.get_query("q").unwrap().unwrap();
assert_eq!(current.expression, "keys(@)");
}
#[test]
fn test_get_query_nonexistent() {
let engine = JpxEngine::new();
let result = engine.get_query("nonexistent").unwrap();
assert!(result.is_none());
}
#[test]
fn test_delete_query_nonexistent() {
let engine = JpxEngine::new();
let result = engine.delete_query("nonexistent").unwrap();
assert!(result.is_none());
}
#[test]
fn test_run_query_not_found() {
let engine = JpxEngine::new();
let result = engine.run_query("nonexistent", &json!({}));
assert!(result.is_err());
match result.unwrap_err() {
EngineError::QueryNotFound(name) => assert_eq!(name, "nonexistent"),
other => panic!("Expected QueryNotFound, got {:?}", other),
}
}
#[test]
fn test_list_queries_empty() {
let engine = JpxEngine::new();
let queries = engine.list_queries().unwrap();
assert!(queries.is_empty());
}
#[test]
fn test_list_queries_multiple() {
let engine = JpxEngine::new();
engine
.define_query("alpha".to_string(), "a".to_string(), None)
.unwrap();
engine
.define_query("beta".to_string(), "b".to_string(), None)
.unwrap();
engine
.define_query("gamma".to_string(), "c".to_string(), None)
.unwrap();
let queries = engine.list_queries().unwrap();
assert_eq!(queries.len(), 3);
}
#[test]
fn test_register_discovery_duplicate() {
let engine = JpxEngine::new();
let spec: DiscoverySpec = serde_json::from_value(json!({
"server": {"name": "dup-server", "version": "1.0.0"},
"tools": [
{"name": "tool_a", "description": "Tool A", "tags": ["test"]}
]
}))
.unwrap();
let first = engine.register_discovery(spec.clone(), false).unwrap();
assert!(first.ok);
let second = engine.register_discovery(spec, false).unwrap();
assert!(!second.ok);
assert!(second.warnings[0].contains("already registered"));
}
#[test]
fn test_unregister_nonexistent() {
let engine = JpxEngine::new();
let result = engine.unregister_discovery("nonexistent").unwrap();
assert!(!result);
}
#[test]
fn test_discovery_index_stats_empty() {
let engine = JpxEngine::new();
let stats = engine.discovery_index_stats().unwrap();
assert!(stats.is_none());
}
#[test]
fn test_get_discovery_schema() {
let engine = JpxEngine::new();
let schema = engine.get_discovery_schema();
assert!(schema.is_object());
assert!(schema.get("$schema").is_some());
assert_eq!(
schema.get("$schema").unwrap().as_str().unwrap(),
"http://json-schema.org/draft-07/schema#"
);
assert!(schema.get("title").is_some());
assert!(schema.get("properties").is_some());
}
}