Skip to main content

oxibonsai_runtime/
lib.rs

1//! # OxiBonsai Runtime
2//!
3//! High-level inference engine, sampling, tokenizer bridge, and
4//! OpenAI-compatible HTTP server for OxiBonsai.
5//!
6//! This crate ties together [`oxibonsai_core`], [`oxibonsai_kernels`],
7//! and [`oxibonsai_model`] into a production-ready inference runtime:
8//!
9//! - **[`InferenceEngine`]** — orchestrates prefill + decode with metrics
10//! - **[`Sampler`]** — temperature, top-k, top-p, repetition penalty
11//! - **[`SamplingPreset`]** — named parameter sets (Greedy, Balanced, Creative, ...)
12//! - **[`SamplerBuilder`] / [`ConfigBuilder`] / [`EngineBuilder`]** — ergonomic setup
13//! - **[`server`]** — Axum-based `/v1/chat/completions` server (feature-gated)
14//! - **[`InferenceMetrics`]** — Prometheus-compatible counters, gauges, histograms
15//! - **[`CircuitBreaker`]** — resilience pattern for cascading-failure protection
16//! - **[`HealthReport`]** — structured health checks for ops monitoring
17//!
18//! ## Quick Start
19//!
20//! ```rust,no_run
21//! use oxibonsai_core::config::Qwen3Config;
22//! use oxibonsai_runtime::engine::InferenceEngine;
23//! use oxibonsai_runtime::presets::SamplingPreset;
24//!
25//! let config = Qwen3Config::tiny_test();
26//! let params = SamplingPreset::Balanced.params();
27//! let mut engine = InferenceEngine::new(config, params, 42);
28//!
29//! let tokens = engine.generate(&[151644, 872], 10)
30//!     .expect("generation should succeed");
31//! ```
32
33pub mod adaptive_lookahead;
34pub mod adaptive_sampling;
35#[cfg(feature = "server")]
36pub mod admin;
37#[cfg(feature = "server")]
38pub mod api_extensions;
39#[cfg(feature = "server")]
40pub mod api_types;
41#[cfg(not(target_arch = "wasm32"))]
42pub mod async_engine;
43pub mod auto_tuner;
44pub mod batch_engine;
45pub mod beam_search;
46pub mod builders;
47pub mod circuit_breaker;
48#[cfg(feature = "server")]
49pub mod completions;
50pub mod config;
51pub mod constrained_decoding;
52pub mod context_manager;
53pub mod continuous_batch;
54pub mod convenience;
55pub mod dedup;
56#[cfg(feature = "server")]
57pub mod distributed;
58pub mod embedding_index;
59#[cfg(feature = "server")]
60pub mod embeddings;
61pub mod engine;
62pub mod engine_pool;
63pub mod error;
64pub mod grammar;
65pub mod health;
66pub mod hot_reload;
67pub mod json_schema;
68pub mod kv_cache_policy;
69pub mod memory;
70pub mod metrics;
71pub mod middleware;
72pub mod model_cache;
73pub mod multi_model;
74pub mod native_tokenizer;
75pub mod nbest;
76pub mod ngram_cache;
77pub mod pipeline;
78pub mod prefix_cache_engine;
79pub mod presets;
80pub mod profiler;
81pub mod quality_metrics;
82#[cfg(feature = "rag")]
83pub mod rag_server;
84pub mod rate_limiter;
85pub mod recovery;
86pub mod request_id;
87pub mod request_metrics;
88pub mod request_queue;
89pub mod sampling;
90pub mod sampling_advanced;
91pub mod semantic_cache;
92#[cfg(feature = "server")]
93pub mod server;
94pub mod speculative;
95pub mod stream_metrics;
96pub mod streaming;
97pub mod token_budget;
98pub mod token_healing;
99pub mod tokenizer_bridge;
100#[cfg(feature = "server")]
101pub mod tool_calling;
102pub mod tracing_setup;
103pub mod wasm_api;
104#[cfg(feature = "server")]
105pub mod web_ui;
106
107pub use adaptive_lookahead::{AdaptiveLookahead, AdaptiveLookaheadConfig, AdaptiveLookaheadError};
108pub use adaptive_sampling::{
109    AdaptiveSamplerChain, AdaptiveStrategy, EntropyCooling, GenerationState, RepetitionAdaptation,
110    ScheduledDecay,
111};
112pub use auto_tuner::{
113    AutoTuner, CpuArch, CpuFeatures, KernelBenchmark, KvCacheType, MemoryBudget, SimdTier,
114    TuningRecommendation,
115};
116pub use builders::{ConfigBuilder, EngineBuilder, SamplerBuilder};
117pub use circuit_breaker::{CircuitBreaker, CircuitBreakerConfig, CircuitState};
118pub use config::OxiBonsaiConfig;
119pub use constrained_decoding::{
120    AllowListConstraint, ConstrainedSampler, ConstrainedSamplerBuilder, ConstraintError,
121    JsonConstraint, JsonParseState, LengthConstraint, NoConstraint, RegexConstraint,
122    SequenceConstraint, TokenConstraint,
123};
124pub use convenience::{GenerationResult, MemoryEstimate, ModelFileInfo, TokenStats};
125pub use dedup::{DedupCache, DedupStats, RequestKey};
126#[cfg(feature = "server")]
127pub use distributed::{
128    ConsistentHashRing, CoordinatorConfig, DistributedCoordinator, NodeInfo, NodeRegistry,
129};
130pub use engine::InferenceEngine;
131pub use error::{RuntimeError, RuntimeResult};
132pub use grammar::{
133    compile_json_schema, compile_json_schema_str, compile_regex, parse_bnf, parse_gbnf,
134    BnfParseError, EarleyRecognizer, GbnfParseError, Grammar, GrammarConstraint,
135    JsonSchemaCompileError, RegexCompileError, Rule, Symbol,
136};
137pub use health::{HealthReport, HealthStatus};
138pub use hot_reload::{HotReloadCoordinator, ModelVersion, ReloadLog};
139pub use json_schema::{
140    parse_schema, schema_example, schema_template, validate_against_schema, SchemaError,
141    SchemaState, SchemaType,
142};
143pub use kv_cache_policy::{KvCacheLevel, KvCachePolicy, KvCachePolicyConfig, KvCachePolicyError};
144pub use memory::{get_rss_bytes, MemoryProfiler, MemorySnapshot};
145pub use metrics::InferenceMetrics;
146pub use multi_model::{
147    AdapterRef, AdapterStack, EndpointStatus, ModelEndpoint, ModelId, ModelListEntry,
148    ModelRegistry, ModelRouter, RoutingError,
149};
150pub use native_tokenizer::{NativeTokenizerBridge, NativeTokenizerError};
151pub use nbest::{Hypothesis, NBestDecoder, NBestList};
152pub use presets::SamplingPreset;
153pub use profiler::{flop_counter, AggregateStats, ProfileEvent, ProfileTrace, Profiler};
154pub use quality_metrics::{
155    extract_ngrams, perplexity_from_logprobs, repetition_penalty_rate, self_bleu, token_entropy,
156    BatchQualityAnalyzer, BleuScore, DiversityMetrics, GenerationQualityReport, RepetitionMetrics,
157};
158pub use recovery::{ErrorClass, RecoveryStrategy};
159pub use request_id::RequestId;
160pub use request_metrics::{
161    AggregateRateSnapshot, RequestRateAggregator, RequestRateSnapshot, RequestRateTracker,
162};
163pub use sampling::Sampler;
164pub use sampling_advanced::{
165    EtaSampler, LcgRng, MinPSampler, MirostatV1Sampler, MirostatV2Sampler, SamplerChain,
166    SamplerStep, TypicalSampler,
167};
168pub use stream_metrics::{RequestStreamMetrics, StreamMetricsSnapshot, StreamingMetricsAggregator};
169pub use token_budget::{
170    BudgetConfig, BudgetError, BudgetPolicy, GlobalTokenBudget, RequestBudget, TokenCostEstimate,
171};
172pub use tokenizer_bridge::TokenizerBridge;
173#[cfg(feature = "server")]
174pub use tool_calling::{
175    build_tool_constraint, make_tool_call, new_tool_call_id, select_tool, validate_tool_arguments,
176    ToolCallError, ToolRegistry,
177};
178pub use tracing_setup::{init_tracing, TracingConfig};
179#[cfg(feature = "server")]
180pub use web_ui::create_ui_router;