-**PUMA** aims to be a lightweight, high-performance inference engine for local AI. Play for fun.
+
+
+
+
+
-## How to Run
+**A lightweight, high-performance inference engine for local AI**
-### Build
+[](https://github.com/InftyAI/PUMA)
+[](https://github.com/InftyAI/PUMA/releases)
-Run `make build` to build the **puma** binary.
+
-### Run
+## ✨ Features
-Run `./puma help` to see all available commands.
+🔧 **Model Management** - Download, cache, and organize AI models from Hugging Face
-For example, you can run `./puma version` to see the binary version.
+🔍 **Advanced Filtering** - Search models with regex patterns and SQL-style queries
-## Supported Backends
+💻 **System Detection** - Automatic GPU detection and resource reporting
-Use [llama.cpp](https://github.com/ggerganov/llama.cpp) as the default backend for quick prototyping, will implement our own backend in the future.
+🚀 **OpenAI-Compatible API** - RESTful API with streaming support
+
+## Installation
+
+### Install with Cargo
+
+```bash
+cargo install puma
+```
+
+### Build from Source
+
+```bash
+# Clone the repository
+git clone https://github.com/InftyAI/PUMA.git
+cd PUMA
+
+# Build the binary
+make build
+
+# The binary will be available at ./puma
+./puma version
+```
+
+## Quick Start
+
+### CLI Usage
+
+```bash
+# Download a model
+puma pull inftyai/tiny-random-gpt2
+
+# List all models
+puma ls
+
+# Inspect model details
+puma inspect inftyai/tiny-random-gpt2
+
+# Check system info
+puma info
+
+# Remove a model
+puma rm inftyai/tiny-random-gpt2
+```
+
+### API Server
+
+```bash
+# Start the inference server with a model
+puma serve inftyai/tiny-random-gpt2
+
+# Server will start on http://0.0.0.0:8000
+# API endpoints:
+# POST /v1/chat/completions
+# POST /v1/completions
+# GET /v1/models
+# GET /v1/models/:model
+# GET /health
+```
+
+**Test the API:**
+
+```bash
+# Health check
+curl http://localhost:8000/health
+
+# Chat completion
+curl http://localhost:8000/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "inftyai/tiny-random-gpt2",
+ "messages": [{"role": "user", "content": "Hello!"}]
+ }'
+
+# Or use the test script
+./hack/scripts/test_api.sh
+```
+
+## Commands
+
+| Command | Status | Description |
+|---------|--------|-------------|
+| `pull ` | ✅ | Download model from provider |
+| `ls` | ✅ | List models (supports regex, label filters) |
+| `inspect ` | ✅ | Show detailed model information |
+| `rm ` | ✅ | Remove model and cache |
+| `info` | ✅ | Display system information |
+| `version` | ✅ | Show PUMA version |
+| `serve ` | ✅ | Start OpenAI-compatible API server with a model |
+| `ps` | 🚧 | List running models |
+| `run` | 🚧 | Start model inference |
+| `stop` | 🚧 | Stop running model |
+
+## Advanced Usage
+
+### Pattern Matching
+
+```bash
+# Substring match
+puma ls qwen
+
+# Prefix match
+puma ls "^inftyai/"
+
+# Alternation
+puma ls "llama-(2|3)"
+```
+
+### Label Filtering
+
+```bash
+# Single filter
+puma ls -l author=inftyai
+
+# Multiple filters (AND condition)
+puma ls -l author=inftyai,license=mit
+
+# Combine pattern + filter
+puma ls llama -l author=meta
+```
+
+**Available filters:** `author`, `task`, `license`, `provider`, `model_series`
+
+## API Server
+
+PUMA provides an OpenAI-compatible API server for model inference.
+
+### Starting the Server
+
+```bash
+# Start server with a model (default: 0.0.0.0:8000)
+puma serve inftyai/tiny-random-gpt2
+
+# Custom host and port
+puma serve inftyai/tiny-random-gpt2 --host 127.0.0.1 --port 3000
+
+# Model must be pulled first
+puma pull inftyai/tiny-random-gpt2
+```
+
+### API Endpoints
+
+#### Chat Completions (Recommended)
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "inftyai/tiny-random-gpt2",
+ "messages": [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Hello!"}
+ ],
+ "max_tokens": 100,
+ "temperature": 0.7
+ }'
+```
+
+#### Streaming (Server-Sent Events)
+```bash
+curl http://localhost:8000/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "inftyai/tiny-random-gpt2",
+ "messages": [{"role": "user", "content": "Tell me a story"}],
+ "stream": true
+ }'
+```
+
+#### List Models
+```bash
+# Returns the currently loaded model
+curl http://localhost:8000/v1/models
+```
+
+#### Health Check
+```bash
+curl http://localhost:8000/health
+# Returns: {"status":"ok"}
+```
+
+### OpenAI Python Client
+
+PUMA is compatible with the OpenAI Python SDK:
+
+```python
+from openai import OpenAI
+
+client = OpenAI(
+ base_url="http://localhost:8000/v1",
+ api_key="dummy" # Not required
+)
+
+response = client.chat.completions.create(
+ model="inftyai/tiny-random-gpt2",
+ messages=[
+ {"role": "user", "content": "Hello!"}
+ ]
+)
+
+print(response.choices[0].message.content)
+```
+
+### Inspect Output
+
+```bash
+$ puma inspect inftyai/tiny-random-gpt2
+
+name: inftyai/tiny-random-gpt2
+kind: Model
+spec:
+ author: inftyai
+ model_series: gpt2
+ task: text-generation
+ license: MIT
+ context_window: 2.05K
+ safetensors:
+ total: 7.00B
+ parameters:
+ f32: 7.00B
+ provider: huggingface
+ cache:
+ revision: abc123de
+ size: 1.24 GB
+ path: ~/.puma/cache/...
+status:
+ created: 2 hours ago
+ updated: 2 hours ago
+```
+
+## Model Management
+
+- **Database:** `~/.puma/models.db` (SQLite)
+- **Cache:** `~/.puma/cache/` (model files)
+
+Models are stored with lowercase names for case-insensitive matching.
+
+## Development
+
+```bash
+# Build
+make build
+
+# Run all tests
+make test
+
+# Test API manually
+./hack/scripts/test_api.sh
+```
+
+### Project Structure
+
+```
+puma/
+├── src/
+│ ├── api/ # OpenAI-compatible API
+│ ├── backend/ # Inference backends (Mock, MLX)
+│ ├── cli/ # Command implementations
+│ ├── downloader/ # HuggingFace download logic
+│ ├── registry/ # Model registry & metadata
+│ ├── storage/ # SQLite storage backend
+│ ├── system/ # System info detection
+│ └── utils/ # Formatting & helpers
+├── tests/ # Integration tests
+├── hack/ # Development scripts
+├── Cargo.toml # Rust dependencies
+└── Makefile # Build commands
+```
+
+## License
+
+Apache-2.0
+
+## Star History
+
+[](https://www.star-history.com/#inftyai/puma&Date)
diff --git a/docs/MLX_INTEGRATION.md b/docs/MLX_INTEGRATION.md
new file mode 100644
index 0000000..5efb505
--- /dev/null
+++ b/docs/MLX_INTEGRATION.md
@@ -0,0 +1,200 @@
+# MLX Integration Guide
+
+## Overview
+
+PUMA now supports Apple's MLX framework for high-performance inference on Apple Silicon. This integration uses [mlx-rs](https://github.com/oxiglade/mlx-rs) v0.25.3 as the binding layer.
+
+## Quick Start
+
+### Prerequisites
+
+- macOS 13.0+
+- Apple Silicon (M1/M2/M3/M4)
+- Xcode Command Line Tools: `xcode-select --install`
+
+### Building with MLX Support
+
+```bash
+# Build
+cargo build --release --features mlx
+
+# Run tests (on Apple Silicon)
+cargo test --features mlx
+
+# Run example
+cargo run --release --features mlx --example mlx_inference
+```
+
+### Starting Server with MLX
+
+```bash
+# Build with MLX feature
+cargo build --release --features mlx
+
+# Start server (will use MLX on Apple Silicon)
+./target/release/puma serve
+```
+
+The server will automatically detect and use MLX if:
+1. Running on Apple Silicon macOS
+2. Built with `--features mlx`
+
+Otherwise, it falls back to MockEngine.
+
+## Architecture
+
+```
+┌─────────────────────────────────────────┐
+│ PUMA CLI/API │
+└──────────────┬──────────────────────────┘
+ │
+ ┌───────┴────────┐
+ │ InferenceEngine│ (trait)
+ └───────┬────────┘
+ │
+ ┏━━━━━━━━━┻━━━━━━━━━┓
+ ┃ ┃
+┌───┴─────┐ ┌───────┴─────┐
+│MockEngine│ │ MlxEngine │
+└─────────┘ └───────┬──────┘
+ │
+ ┌──────┴──────┐
+ │ mlx-rs │ (v0.25.3)
+ └──────┬──────┘
+ │
+ ┌──────┴──────┐
+ │ mlx-sys │ (FFI)
+ └──────┬──────┘
+ │
+ ┌──────┴──────┐
+ │ MLX C++ │
+ └─────────────┘
+```
+
+## Current Implementation Status
+
+### ✅ Completed
+
+- Feature flag setup (`mlx` feature)
+- Platform detection (Apple Silicon only)
+- mlx-rs integration (v0.25.3)
+- Device initialization (GPU/CPU)
+- Basic InferenceEngine trait implementation
+- Conditional compilation for non-macOS platforms
+- Serve command auto-detection
+
+### 🚧 In Progress (Placeholder)
+
+- Model loading from PUMA cache
+- Tokenization (currently uses dummy tokens)
+- Token generation (placeholder implementation)
+- Streaming (basic structure in place)
+
+### 📋 TODO
+
+- [ ] Integrate proper tokenizer (tokenizers-rs or mlx-lm)
+- [ ] Implement actual model loading
+- [ ] Real token generation with mlx-rs
+- [ ] Model caching and memory management
+- [ ] Quantization support (4-bit, 8-bit)
+- [ ] Benchmark and optimize performance
+- [ ] Add more examples
+
+## Usage Example
+
+```rust
+use puma::backend::{MlxEngine, InferenceEngine};
+
+#[tokio::main]
+async fn main() -> Result<(), Box> {
+ // Create engine
+ let engine = MlxEngine::new()?;
+
+ // Generate text
+ let response = engine.generate(
+ "model-name",
+ "Hello, world!",
+ 100, // max_tokens
+ 0.7 // temperature
+ ).await?;
+
+ println!("Generated: {}", response.text);
+ Ok(())
+}
+```
+
+## API Endpoints
+
+Once MLX is enabled, the API server uses it automatically:
+
+```bash
+# Chat completion with MLX
+curl http://localhost:8000/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "your-model",
+ "messages": [{"role": "user", "content": "Hello!"}]
+ }'
+```
+
+## Performance Notes
+
+- MLX uses unified memory (CPU/GPU share memory)
+- Metal GPU acceleration enabled by default
+- Lazy evaluation for compute graphs
+- Multi-device support (CPU/GPU)
+
+## Troubleshooting
+
+### Build fails on non-macOS
+
+**Expected behavior.** MLX only works on macOS with Apple Silicon.
+
+### "MLX feature not enabled" error at runtime
+
+Build with the `mlx` feature:
+```bash
+cargo build --features mlx
+```
+
+### Server uses MockEngine instead of MLX
+
+Check:
+1. Running on Apple Silicon? `uname -m` should show `arm64`
+2. Built with MLX feature? `cargo build --features mlx`
+3. Check logs for MLX initialization
+
+### Compilation takes a long time
+
+MLX compilation (via mlx-rs) compiles the entire MLX framework. First build can take 10-20 minutes. Subsequent builds are cached.
+
+## Future: Custom MLX Bindings
+
+Current implementation uses mlx-rs. Future plan:
+
+1. **Phase 1** (Current): Use mlx-rs for rapid development
+2. **Phase 2**: Implement custom `puma-mlx-sys` FFI layer
+3. **Phase 3**: Build custom safe wrapper with PUMA-specific optimizations
+
+See `src/backend/mlx/README.md` for details.
+
+## References
+
+- [MLX Framework](https://github.com/ml-explore/mlx)
+- [mlx-rs Bindings](https://github.com/oxiglade/mlx-rs)
+- [mlx-rs Documentation](https://oxideai.github.io/mlx-rs/mlx_rs/)
+- [MLX Python Docs](https://ml-explore.github.io/mlx/build/html/index.html)
+
+## Contributing
+
+When contributing to MLX backend:
+
+1. Test on Apple Silicon macOS
+2. Ensure code compiles without `mlx` feature
+3. Use feature flags for MLX-specific code
+4. Add tests with `#[cfg(all(target_os = "macos", feature = "mlx"))]`
+5. Document any new MLX-specific features
+
+## License
+
+MLX integration in PUMA follows the same Apache-2.0 license as the main project.
diff --git a/examples/mlx_inference.rs b/examples/mlx_inference.rs
new file mode 100644
index 0000000..82b1b6c
--- /dev/null
+++ b/examples/mlx_inference.rs
@@ -0,0 +1,74 @@
+//! Example: Using MLX backend for inference
+//!
+//! This example demonstrates how to use PUMA with MLX backend on Apple Silicon.
+//!
+//! Build with: cargo build --release --features mlx
+//! Run with: cargo run --release --features mlx --example mlx_inference
+
+#[cfg(all(target_os = "macos", feature = "mlx"))]
+use puma::backend::MlxEngine;
+
+#[cfg(all(target_os = "macos", feature = "mlx"))]
+use puma::backend::InferenceEngine;
+
+#[cfg(all(target_os = "macos", feature = "mlx"))]
+#[tokio::main]
+async fn main() -> Result<(), Box> {
+ // Initialize tracing
+ tracing_subscriber::fmt()
+ .with_env_filter("info")
+ .init();
+
+ println!("🐆 PUMA MLX Inference Example\n");
+
+ // Create MLX engine
+ println!("Initializing MLX engine...");
+ let engine = MlxEngine::new()
+ .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+
+ println!("✓ MLX engine initialized\n");
+
+ // Generate text
+ println!("Generating text...");
+ let response = engine
+ .generate(
+ "test-model",
+ "Once upon a time",
+ 50, // max_tokens
+ 0.7, // temperature
+ )
+ .await?;
+
+ println!("\n📝 Generated Response:");
+ println!(" Text: {}", response.text);
+ println!(" Prompt tokens: {}", response.prompt_tokens);
+ println!(" Completion tokens: {}", response.completion_tokens);
+
+ // Test streaming
+ println!("\n🔄 Testing streaming generation...");
+ let mut stream = engine
+ .generate_stream(
+ "test-model",
+ "The quick brown fox",
+ 30,
+ 0.7,
+ )
+ .await?;
+
+ print!(" Tokens: ");
+ use tokio_stream::StreamExt;
+ while let Some(token) = stream.next().await {
+ print!("{}", token);
+ std::io::Write::flush(&mut std::io::stdout()).ok();
+ }
+ println!("\n\n✓ Example completed");
+
+ Ok(())
+}
+
+#[cfg(not(all(target_os = "macos", feature = "mlx")))]
+fn main() {
+ eprintln!("❌ This example requires macOS with Apple Silicon");
+ eprintln!("Build with: cargo run --features mlx --example mlx_inference");
+ std::process::exit(1);
+}
diff --git a/hack/README.md b/hack/README.md
new file mode 100644
index 0000000..ca9895f
--- /dev/null
+++ b/hack/README.md
@@ -0,0 +1,68 @@
+# Hack Directory
+
+Development and testing utilities for PUMA.
+
+## Structure
+
+```
+hack/
+└── scripts/ # Test and utility scripts
+ └── test_api.sh
+```
+
+## Scripts
+
+### `scripts/test_api.sh`
+
+Tests all PUMA API endpoints manually.
+
+**Usage:**
+```bash
+# Start PUMA server first
+./puma serve
+
+# In another terminal
+./hack/scripts/test_api.sh
+```
+
+**Tests:**
+- Health check
+- List models
+- Chat completion (non-streaming)
+- Chat completion (streaming)
+- Text completion
+
+**Requirements:**
+- Running PUMA server
+- `curl` and `jq` installed
+
+---
+
+
+## Adding New Scripts
+
+Place development and testing scripts in `hack/scripts/`:
+
+```bash
+# Create new script
+cat > hack/scripts/my_script.sh << 'EOF'
+#!/bin/bash
+# Your script here
+EOF
+
+# Make executable
+chmod +x hack/scripts/my_script.sh
+```
+
+---
+
+## Why "hack"?
+
+The `hack/` directory is a convention from Kubernetes and other projects for:
+- Development utilities
+- Test scripts
+- Build helpers
+- CI/CD scripts
+- One-off tools
+
+It keeps the root directory clean while providing a place for development tools.
diff --git a/hack/scripts/test_api.sh b/hack/scripts/test_api.sh
new file mode 100755
index 0000000..8caf5d4
--- /dev/null
+++ b/hack/scripts/test_api.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+
+echo "Testing PUMA OpenAI-Compatible API"
+echo "===================================="
+echo
+
+# Base URL
+BASE_URL="http://localhost:8000"
+
+echo "1. Health Check"
+curl -s "$BASE_URL/health"
+echo -e "\n"
+
+echo "2. List Models"
+curl -s "$BASE_URL/v1/models" | jq '.'
+echo
+
+echo "3. Chat Completion (Non-streaming)"
+curl -s "$BASE_URL/v1/chat/completions" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "test-model",
+ "messages": [
+ {"role": "user", "content": "Hello!"}
+ ],
+ "max_tokens": 50
+ }' | jq '.'
+echo
+
+echo "4. Chat Completion (Streaming)"
+curl -s -N "$BASE_URL/v1/chat/completions" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "test-model",
+ "messages": [
+ {"role": "user", "content": "Tell me a story"}
+ ],
+ "stream": true,
+ "max_tokens": 50
+ }'
+echo -e "\n"
+
+echo "5. Legacy Text Completion"
+curl -s "$BASE_URL/v1/completions" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "test-model",
+ "prompt": "Once upon a time",
+ "max_tokens": 50
+ }' | jq '.'
+echo
+
+echo "Done!"
diff --git a/site/images/logo-dark.svg b/site/images/logo-dark.svg
new file mode 100644
index 0000000..f11bffb
--- /dev/null
+++ b/site/images/logo-dark.svg
@@ -0,0 +1,102 @@
+
+
diff --git a/site/images/logo-light.svg b/site/images/logo-light.svg
new file mode 100644
index 0000000..fa61d5b
--- /dev/null
+++ b/site/images/logo-light.svg
@@ -0,0 +1,82 @@
+
+
diff --git a/src/api/chat.rs b/src/api/chat.rs
new file mode 100644
index 0000000..0dfcf4f
--- /dev/null
+++ b/src/api/chat.rs
@@ -0,0 +1,254 @@
+use axum::{
+ extract::State,
+ response::{
+ sse::{Event, KeepAlive, Sse},
+ IntoResponse, Response,
+ },
+ Json,
+};
+use futures::stream::StreamExt;
+use std::sync::Arc;
+use tokio_stream::wrappers::ReceiverStream;
+use uuid::Uuid;
+
+use crate::api::routes::AppState;
+use crate::api::types::{
+ ChatChoice, ChatChoiceDelta, ChatCompletionChunk, ChatCompletionRequest,
+ ChatCompletionResponse, ChatMessage, ChatMessageDelta, ErrorResponse, Usage,
+};
+use crate::backend::InferenceEngine;
+
+/// Main handler for chat completions
+pub async fn chat_completions(
+ State(state): State>,
+ Json(req): Json,
+) -> Response {
+ let engine = state.engine;
+ let registry = state.registry;
+
+ // Validate request
+ if req.messages.is_empty() {
+ return (
+ axum::http::StatusCode::BAD_REQUEST,
+ Json(ErrorResponse::new(
+ "messages cannot be empty".to_string(),
+ "invalid_request_error".to_string(),
+ )),
+ )
+ .into_response();
+ }
+
+ // Validate model exists
+ match registry.get_model(&req.model) {
+ Ok(None) => {
+ return (
+ axum::http::StatusCode::NOT_FOUND,
+ Json(ErrorResponse::new(
+ format!("Model '{}' not found", req.model),
+ "model_not_found".to_string(),
+ )),
+ )
+ .into_response();
+ }
+ Err(e) => {
+ return (
+ axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+ Json(ErrorResponse::new(
+ format!("Failed to check model: {}", e),
+ "internal_error".to_string(),
+ )),
+ )
+ .into_response();
+ }
+ Ok(Some(_)) => {
+ // Model exists, continue
+ }
+ }
+
+ if req.stream {
+ chat_completions_stream(engine, req).await.into_response()
+ } else {
+ match chat_completions_non_stream(engine, req).await {
+ Ok(response) => Json(response).into_response(),
+ Err(err) => (
+ axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+ Json(ErrorResponse::new(
+ err.to_string(),
+ "internal_error".to_string(),
+ )),
+ )
+ .into_response(),
+ }
+ }
+}
+
+/// Non-streaming chat completion
+async fn chat_completions_non_stream(
+ engine: Arc,
+ req: ChatCompletionRequest,
+) -> Result> {
+ let id = format!("chatcmpl-{}", Uuid::new_v4());
+ let created = chrono::Utc::now().timestamp();
+
+ // Convert messages to prompt
+ let prompt = format_chat_messages(&req.messages);
+
+ // Generate
+ let response = engine
+ .generate(
+ &req.model,
+ &prompt,
+ req.max_tokens.unwrap_or(100),
+ req.temperature.unwrap_or(0.7),
+ )
+ .await?;
+
+ Ok(ChatCompletionResponse {
+ id,
+ object: "chat.completion".to_string(),
+ created,
+ model: req.model,
+ choices: vec![ChatChoice {
+ index: 0,
+ message: ChatMessage {
+ role: "assistant".to_string(),
+ content: response.text,
+ },
+ finish_reason: "stop".to_string(),
+ }],
+ usage: Usage {
+ prompt_tokens: response.prompt_tokens,
+ completion_tokens: response.completion_tokens,
+ total_tokens: response.prompt_tokens + response.completion_tokens,
+ },
+ })
+}
+
+/// Streaming chat completion
+async fn chat_completions_stream(
+ engine: Arc,
+ req: ChatCompletionRequest,
+) -> Sse>> {
+ let id = format!("chatcmpl-{}", Uuid::new_v4());
+ let created = chrono::Utc::now().timestamp();
+ let model = req.model.clone();
+
+ let (tx, rx) = tokio::sync::mpsc::channel(100);
+
+ // Spawn task to generate tokens
+ tokio::spawn(async move {
+ let prompt = format_chat_messages(&req.messages);
+
+ // Send initial chunk with role
+ let initial_chunk = ChatCompletionChunk {
+ id: id.clone(),
+ object: "chat.completion.chunk".to_string(),
+ created,
+ model: model.clone(),
+ choices: vec![ChatChoiceDelta {
+ index: 0,
+ delta: ChatMessageDelta {
+ role: Some("assistant".to_string()),
+ content: None,
+ },
+ finish_reason: None,
+ }],
+ };
+
+ if tx
+ .send(Ok(
+ Event::default().data(serde_json::to_string(&initial_chunk).unwrap())
+ ))
+ .await
+ .is_err()
+ {
+ return;
+ }
+
+ // Stream tokens
+ match engine
+ .generate_stream(
+ &model,
+ &prompt,
+ req.max_tokens.unwrap_or(100),
+ req.temperature.unwrap_or(0.7),
+ )
+ .await
+ {
+ Ok(mut stream) => {
+ while let Some(token) = stream.next().await {
+ let chunk = ChatCompletionChunk {
+ id: id.clone(),
+ object: "chat.completion.chunk".to_string(),
+ created,
+ model: model.clone(),
+ choices: vec![ChatChoiceDelta {
+ index: 0,
+ delta: ChatMessageDelta {
+ role: None,
+ content: Some(token),
+ },
+ finish_reason: None,
+ }],
+ };
+
+ if tx
+ .send(Ok(
+ Event::default().data(serde_json::to_string(&chunk).unwrap())
+ ))
+ .await
+ .is_err()
+ {
+ break;
+ }
+ }
+ }
+ Err(e) => {
+ tracing::error!("Error generating stream: {}", e);
+ return;
+ }
+ }
+
+ // Send final chunk
+ let final_chunk = ChatCompletionChunk {
+ id: id.clone(),
+ object: "chat.completion.chunk".to_string(),
+ created,
+ model: model.clone(),
+ choices: vec![ChatChoiceDelta {
+ index: 0,
+ delta: ChatMessageDelta {
+ role: None,
+ content: None,
+ },
+ finish_reason: Some("stop".to_string()),
+ }],
+ };
+
+ let _ = tx
+ .send(Ok(
+ Event::default().data(serde_json::to_string(&final_chunk).unwrap())
+ ))
+ .await;
+ let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
+ });
+
+ Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default())
+}
+
+/// Format chat messages into a prompt
+fn format_chat_messages(messages: &[ChatMessage]) -> String {
+ messages
+ .iter()
+ .map(|m| {
+ if m.role == "system" {
+ format!("System: {}", m.content)
+ } else if m.role == "user" {
+ format!("User: {}", m.content)
+ } else {
+ format!("Assistant: {}", m.content)
+ }
+ })
+ .collect::>()
+ .join("\n")
+}
diff --git a/src/api/completions.rs b/src/api/completions.rs
new file mode 100644
index 0000000..497736b
--- /dev/null
+++ b/src/api/completions.rs
@@ -0,0 +1,113 @@
+use axum::{extract::State, response::IntoResponse, Json};
+use uuid::Uuid;
+
+use crate::api::routes::AppState;
+use crate::api::types::{
+ CompletionChoice, CompletionRequest, CompletionResponse, ErrorResponse, Usage,
+};
+use crate::backend::InferenceEngine;
+
+/// Handler for legacy text completions
+pub async fn completions(
+ State(state): State>,
+ Json(req): Json,
+) -> impl IntoResponse {
+ let engine = state.engine;
+ let registry = state.registry;
+
+ // Validate request
+ let prompt = req.prompt.to_string();
+ if prompt.is_empty() {
+ return (
+ axum::http::StatusCode::BAD_REQUEST,
+ Json(ErrorResponse::new(
+ "prompt cannot be empty".to_string(),
+ "invalid_request_error".to_string(),
+ )),
+ )
+ .into_response();
+ }
+
+ // TODO: Implement streaming support for /v1/completions
+ if req.stream {
+ return (
+ axum::http::StatusCode::BAD_REQUEST,
+ Json(ErrorResponse::new(
+ "Streaming not supported for /v1/completions endpoint".to_string(),
+ "invalid_request_error".to_string(),
+ )),
+ )
+ .into_response();
+ }
+
+ // Validate model exists
+ match registry.get_model(&req.model) {
+ Ok(None) => {
+ return (
+ axum::http::StatusCode::NOT_FOUND,
+ Json(ErrorResponse::new(
+ format!("Model '{}' not found", req.model),
+ "model_not_found".to_string(),
+ )),
+ )
+ .into_response();
+ }
+ Err(e) => {
+ return (
+ axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+ Json(ErrorResponse::new(
+ format!("Failed to check model: {}", e),
+ "internal_error".to_string(),
+ )),
+ )
+ .into_response();
+ }
+ Ok(Some(_)) => {
+ // Model exists, continue
+ }
+ }
+
+ let id = format!("cmpl-{}", Uuid::new_v4());
+ let created = chrono::Utc::now().timestamp();
+
+ // Generate
+ match engine
+ .generate(
+ &req.model,
+ &prompt,
+ req.max_tokens.unwrap_or(100),
+ req.temperature.unwrap_or(0.7),
+ )
+ .await
+ {
+ Ok(response) => {
+ let completion = CompletionResponse {
+ id,
+ object: "text_completion".to_string(),
+ created,
+ model: req.model,
+ choices: vec![CompletionChoice {
+ text: response.text,
+ index: 0,
+ logprobs: None,
+ finish_reason: "stop".to_string(),
+ }],
+ usage: Usage {
+ prompt_tokens: response.prompt_tokens,
+ completion_tokens: response.completion_tokens,
+ total_tokens: response.prompt_tokens + response.completion_tokens,
+ },
+ };
+
+ Json(completion).into_response()
+ }
+ Err(e) => (
+ axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+ Json(ErrorResponse::new(
+ e.to_string(),
+ "internal_error".to_string(),
+ )),
+ )
+ .into_response(),
+ }
+}
diff --git a/src/api/mod.rs b/src/api/mod.rs
new file mode 100644
index 0000000..5de16ff
--- /dev/null
+++ b/src/api/mod.rs
@@ -0,0 +1,8 @@
+pub mod chat;
+pub mod completions;
+pub mod models;
+pub mod routes;
+pub mod types;
+
+#[cfg(test)]
+mod tests;
diff --git a/src/api/models.rs b/src/api/models.rs
new file mode 100644
index 0000000..c8304e1
--- /dev/null
+++ b/src/api/models.rs
@@ -0,0 +1,80 @@
+use axum::{
+ extract::{Path, State},
+ response::IntoResponse,
+ Json,
+};
+
+use crate::api::routes::AppState;
+use crate::api::types::{ErrorResponse, Model, ModelList};
+use crate::backend::InferenceEngine;
+
+/// List all available models
+pub async fn list_models(
+ State(state): State>,
+) -> impl IntoResponse {
+ let registry = state.registry;
+ match registry.load_models(None) {
+ Ok(models) => {
+ let model_list = ModelList {
+ object: "list".to_string(),
+ data: models
+ .into_iter()
+ .map(|m| Model {
+ id: m.name.clone(),
+ object: "model".to_string(),
+ created: chrono::DateTime::parse_from_rfc3339(&m.created_at)
+ .map(|dt| dt.timestamp())
+ .unwrap_or(0),
+ owned_by: m.author.unwrap_or_else(|| "puma".to_string()),
+ })
+ .collect(),
+ };
+ Json(model_list).into_response()
+ }
+ Err(e) => (
+ axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+ Json(ErrorResponse::new(
+ format!("Failed to load models: {}", e),
+ "internal_error".to_string(),
+ )),
+ )
+ .into_response(),
+ }
+}
+
+/// Get a specific model by ID
+pub async fn get_model(
+ State(state): State>,
+ Path(model_id): Path,
+) -> impl IntoResponse {
+ let registry = state.registry;
+ match registry.get_model(&model_id) {
+ Ok(Some(model)) => {
+ let model_info = Model {
+ id: model.name.clone(),
+ object: "model".to_string(),
+ created: chrono::DateTime::parse_from_rfc3339(&model.created_at)
+ .map(|dt| dt.timestamp())
+ .unwrap_or(0),
+ owned_by: model.author.unwrap_or_else(|| "puma".to_string()),
+ };
+ Json(model_info).into_response()
+ }
+ Ok(None) => (
+ axum::http::StatusCode::NOT_FOUND,
+ Json(ErrorResponse::new(
+ format!("Model '{}' not found", model_id),
+ "model_not_found".to_string(),
+ )),
+ )
+ .into_response(),
+ Err(e) => (
+ axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+ Json(ErrorResponse::new(
+ format!("Failed to get model: {}", e),
+ "internal_error".to_string(),
+ )),
+ )
+ .into_response(),
+ }
+}
diff --git a/src/api/routes.rs b/src/api/routes.rs
new file mode 100644
index 0000000..dae6c0e
--- /dev/null
+++ b/src/api/routes.rs
@@ -0,0 +1,70 @@
+use axum::{
+ routing::{get, post},
+ Json, Router,
+};
+use serde::Serialize;
+use std::sync::Arc;
+use tower_http::{
+ cors::CorsLayer,
+ trace::{DefaultMakeSpan, DefaultOnRequest, DefaultOnResponse, TraceLayer},
+ LatencyUnit,
+};
+
+use crate::backend::InferenceEngine;
+use crate::registry::model_registry::ModelRegistry;
+
+use super::{chat, completions, models};
+
+/// Shared application state
+#[derive(Clone)]
+pub struct AppState {
+ pub engine: Arc,
+ pub registry: Arc,
+}
+
+/// Create the API router with all endpoints
+pub fn create_router(
+ engine: Arc,
+ registry: Arc,
+) -> Router {
+ let state = AppState { engine, registry };
+
+ Router::new()
+ // Chat completions (most important)
+ .route("/v1/chat/completions", post(chat::chat_completions::))
+ // Legacy completions
+ .route("/v1/completions", post(completions::completions::))
+ // Models
+ .route("/v1/models", get(models::list_models::))
+ .route("/v1/models/:model", get(models::get_model::))
+ // Health check
+ .route("/health", get(health_check))
+ // Pass state
+ .with_state(state)
+ // Enable request/response logging at INFO level
+ .layer(
+ TraceLayer::new_for_http()
+ .make_span_with(DefaultMakeSpan::new().level(tracing::Level::INFO))
+ .on_request(DefaultOnRequest::new().level(tracing::Level::INFO))
+ .on_response(
+ DefaultOnResponse::new()
+ .level(tracing::Level::INFO)
+ .latency_unit(LatencyUnit::Millis),
+ ),
+ )
+ // Enable CORS for browser clients
+ .layer(CorsLayer::permissive())
+}
+
+/// Health check response
+#[derive(Serialize)]
+struct HealthResponse {
+ status: String,
+}
+
+/// Health check endpoint
+async fn health_check() -> Json {
+ Json(HealthResponse {
+ status: "ok".to_string(),
+ })
+}
diff --git a/src/api/tests.rs b/src/api/tests.rs
new file mode 100644
index 0000000..fba05ca
--- /dev/null
+++ b/src/api/tests.rs
@@ -0,0 +1,437 @@
+//! API Integration Tests
+//!
+//! Tests the PUMA API endpoints using the Axum test utilities.
+//! These tests verify the entire request/response cycle through the router.
+
+use axum::{
+ body::Body,
+ http::{Request, StatusCode},
+};
+use serde_json::{json, Value};
+use std::sync::Arc;
+use tempfile::TempDir;
+use tower::util::ServiceExt; // for `oneshot` and `ready`
+
+use super::routes::create_router;
+use crate::backend::mock::MockEngine;
+use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry};
+
+/// Helper to create test app with a pre-registered test model
+/// Returns the router and the temp directory (which must be kept alive)
+fn create_test_app() -> (axum::Router, TempDir) {
+ let engine = Arc::new(MockEngine::new());
+ let temp_dir = TempDir::new().unwrap();
+ let registry = Arc::new(ModelRegistry::new(Some(temp_dir.path().to_path_buf())));
+
+ // Register a test model
+ let test_model = ModelInfo {
+ uuid: "test-uuid".to_string(),
+ name: "test-model".to_string(),
+ provider: "test".to_string(),
+ author: Some("test-author".to_string()),
+ task: Some("text-generation".to_string()),
+ model_series: Some("test-series".to_string()),
+ license: Some("MIT".to_string()),
+ created_at: chrono::Utc::now().to_rfc3339(),
+ updated_at: chrono::Utc::now().to_rfc3339(),
+ metadata: ModelMetadata {
+ cache: CacheInfo {
+ revision: "test-rev".to_string(),
+ size: 1000,
+ path: "/tmp/test-model".to_string(),
+ },
+ context_window: Some(2048),
+ safetensors: None,
+ },
+ };
+
+ registry
+ .register_model(test_model)
+ .expect("failed to register test model");
+
+ (create_router(engine, registry), temp_dir)
+}
+
+/// Helper to make a JSON request
+async fn make_json_request(
+ app: axum::Router,
+ method: &str,
+ uri: &str,
+ body: Option,
+) -> (StatusCode, Value) {
+ let mut request = Request::builder().uri(uri).method(method);
+
+ if body.is_some() {
+ request = request.header("content-type", "application/json");
+ }
+
+ let request = if let Some(body) = body {
+ request.body(Body::from(serde_json::to_vec(&body).unwrap()))
+ } else {
+ request.body(Body::empty())
+ }
+ .unwrap();
+
+ let response = app.oneshot(request).await.unwrap();
+
+ let status = response.status();
+ let body = axum::body::to_bytes(response.into_body(), usize::MAX)
+ .await
+ .unwrap();
+ let json: Value = serde_json::from_slice(&body).unwrap_or(json!({}));
+
+ // Debug output for failed requests
+ if !status.is_success() {
+ eprintln!("Request failed: {} {}", method, uri);
+ eprintln!("Status: {}", status);
+ eprintln!(
+ "Response: {}",
+ serde_json::to_string_pretty(&json).unwrap_or_default()
+ );
+ }
+
+ (status, json)
+}
+
+#[tokio::test]
+async fn test_health_check() {
+ let (app, _temp_dir) = create_test_app();
+ let (status, json) = make_json_request(app, "GET", "/health", None).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["status"], "ok");
+}
+
+#[tokio::test]
+async fn test_list_models() {
+ let (app, _temp_dir) = create_test_app();
+ let (status, json) = make_json_request(app, "GET", "/v1/models", None).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["object"], "list");
+ assert!(json["data"].is_array());
+}
+
+#[tokio::test]
+async fn test_chat_completion_non_streaming() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "messages": [
+ {"role": "user", "content": "Hello"}
+ ],
+ "max_tokens": 50,
+ "stream": false
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/chat/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["object"], "chat.completion");
+ assert!(json["id"].is_string());
+ assert_eq!(json["model"], "test-model");
+ assert!(json["choices"].is_array());
+ assert_eq!(json["choices"][0]["index"], 0);
+ assert_eq!(json["choices"][0]["message"]["role"], "assistant");
+ assert!(json["choices"][0]["message"]["content"].is_string());
+ assert_eq!(json["choices"][0]["finish_reason"], "stop");
+ assert!(json["usage"]["prompt_tokens"].is_number());
+ assert!(json["usage"]["completion_tokens"].is_number());
+ assert!(json["usage"]["total_tokens"].is_number());
+}
+
+#[tokio::test]
+async fn test_chat_completion_empty_messages() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "messages": [],
+ "stream": false
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/chat/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::BAD_REQUEST);
+ assert_eq!(json["error"]["type"], "invalid_request_error");
+ assert!(json["error"]["message"]
+ .as_str()
+ .unwrap()
+ .contains("messages cannot be empty"));
+}
+
+#[tokio::test]
+async fn test_text_completion() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "prompt": "Once upon a time",
+ "max_tokens": 50
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["object"], "text_completion");
+ assert!(json["id"].is_string());
+ assert_eq!(json["model"], "test-model");
+ assert!(json["choices"].is_array());
+ assert_eq!(json["choices"][0]["index"], 0);
+ assert!(json["choices"][0]["text"].is_string());
+ assert_eq!(json["choices"][0]["finish_reason"], "stop");
+ assert!(json["usage"]["prompt_tokens"].is_number());
+}
+
+#[tokio::test]
+async fn test_text_completion_empty_prompt() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "prompt": ""
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::BAD_REQUEST);
+ assert_eq!(json["error"]["type"], "invalid_request_error");
+ assert!(json["error"]["message"]
+ .as_str()
+ .unwrap()
+ .contains("prompt cannot be empty"));
+}
+
+#[tokio::test]
+async fn test_text_completion_streaming_not_supported() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "prompt": "Hello world",
+ "stream": true
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::BAD_REQUEST);
+ assert_eq!(json["error"]["type"], "invalid_request_error");
+ assert!(json["error"]["message"]
+ .as_str()
+ .unwrap()
+ .contains("Streaming not supported"));
+}
+
+#[tokio::test]
+async fn test_chat_completion_with_system_message() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "messages": [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Hello"}
+ ],
+ "max_tokens": 50
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/chat/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["object"], "chat.completion");
+ assert!(json["choices"][0]["message"]["content"].is_string());
+}
+
+#[tokio::test]
+async fn test_chat_completion_with_temperature() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "messages": [
+ {"role": "user", "content": "Hello"}
+ ],
+ "temperature": 0.5,
+ "max_tokens": 50
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/chat/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["object"], "chat.completion");
+}
+
+#[tokio::test]
+async fn test_chat_completion_default_values() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "messages": [
+ {"role": "user", "content": "Hello"}
+ ]
+ // No max_tokens, temperature, etc. - should use defaults
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/chat/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::OK);
+ assert_eq!(json["object"], "chat.completion");
+}
+
+#[tokio::test]
+async fn test_cors_headers() {
+ let (app, _temp_dir) = create_test_app();
+ let request = Request::builder()
+ .uri("/health")
+ .method("GET")
+ .header("Origin", "https://example.com")
+ .body(Body::empty())
+ .unwrap();
+
+ let response = app.oneshot(request).await.unwrap();
+
+ // CORS should be permissive
+ assert!(response
+ .headers()
+ .contains_key("access-control-allow-origin"));
+}
+
+#[tokio::test]
+async fn test_invalid_route() {
+ let (app, _temp_dir) = create_test_app();
+ let request = Request::builder()
+ .uri("/invalid/route")
+ .method("GET")
+ .body(Body::empty())
+ .unwrap();
+
+ let response = app.oneshot(request).await.unwrap();
+ assert_eq!(response.status(), StatusCode::NOT_FOUND);
+}
+
+#[tokio::test]
+async fn test_method_not_allowed() {
+ let (app, _temp_dir) = create_test_app();
+ // Try POST on GET-only endpoint
+ let request = Request::builder()
+ .uri("/health")
+ .method("POST")
+ .body(Body::empty())
+ .unwrap();
+
+ let response = app.oneshot(request).await.unwrap();
+ assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
+}
+
+#[tokio::test]
+async fn test_chat_completion_nonexistent_model() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "nonexistent-model",
+ "messages": [
+ {"role": "user", "content": "Hello"}
+ ],
+ "max_tokens": 50
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/chat/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::NOT_FOUND);
+ assert_eq!(json["error"]["type"], "model_not_found");
+ assert!(json["error"]["message"]
+ .as_str()
+ .unwrap()
+ .contains("nonexistent-model"));
+}
+
+#[tokio::test]
+async fn test_text_completion_nonexistent_model() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "nonexistent-model",
+ "prompt": "Hello world"
+ });
+
+ let (status, json) =
+ make_json_request(app, "POST", "/v1/completions", Some(request_body)).await;
+
+ assert_eq!(status, StatusCode::NOT_FOUND);
+ assert_eq!(json["error"]["type"], "model_not_found");
+ assert!(json["error"]["message"]
+ .as_str()
+ .unwrap()
+ .contains("nonexistent-model"));
+}
+
+#[tokio::test]
+async fn test_chat_completion_streaming() {
+ let (app, _temp_dir) = create_test_app();
+ let request_body = json!({
+ "model": "test-model",
+ "messages": [
+ {"role": "user", "content": "Hello"}
+ ],
+ "max_tokens": 50,
+ "stream": true
+ });
+
+ let request = Request::builder()
+ .uri("/v1/chat/completions")
+ .method("POST")
+ .header("content-type", "application/json")
+ .body(Body::from(serde_json::to_vec(&request_body).unwrap()))
+ .unwrap();
+
+ let response = app.oneshot(request).await.unwrap();
+
+ assert_eq!(response.status(), StatusCode::OK);
+ assert_eq!(
+ response.headers().get("content-type").unwrap(),
+ "text/event-stream"
+ );
+
+ // Read the streaming body
+ let body_bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
+ .await
+ .unwrap();
+ let body_text = String::from_utf8_lossy(&body_bytes);
+
+ // Parse SSE events (format: "data: {...}\n\n")
+ let events: Vec<&str> = body_text
+ .split("\n\n")
+ .filter(|line| line.starts_with("data: "))
+ .map(|line| line.strip_prefix("data: ").unwrap())
+ .collect();
+
+ assert!(!events.is_empty(), "Should have at least one event");
+
+ // Separate JSON events from [DONE]
+ let json_events: Vec<&str> = events.iter().filter(|&&e| e != "[DONE]").copied().collect();
+
+ assert!(
+ !json_events.is_empty(),
+ "Should have at least one JSON event"
+ );
+
+ // Check first event has role
+ let first_event: Value = serde_json::from_str(json_events[0]).unwrap();
+ assert_eq!(first_event["object"], "chat.completion.chunk");
+ assert_eq!(first_event["model"], "test-model");
+ assert_eq!(first_event["choices"][0]["delta"]["role"], "assistant");
+
+ // Check middle events have content
+ if json_events.len() > 2 {
+ let middle_event: Value = serde_json::from_str(json_events[1]).unwrap();
+ assert!(middle_event["choices"][0]["delta"]["content"].is_string());
+ }
+
+ // Check last JSON event has finish_reason
+ let last_json_event: Value = serde_json::from_str(json_events[json_events.len() - 1]).unwrap();
+ assert_eq!(last_json_event["choices"][0]["finish_reason"], "stop");
+
+ // Check stream ends with [DONE]
+ assert!(events.contains(&"[DONE]"), "Stream should end with [DONE]");
+}
diff --git a/src/api/types/mod.rs b/src/api/types/mod.rs
new file mode 100644
index 0000000..6415885
--- /dev/null
+++ b/src/api/types/mod.rs
@@ -0,0 +1,5 @@
+pub mod request;
+pub mod response;
+
+pub use request::*;
+pub use response::*;
diff --git a/src/api/types/request.rs b/src/api/types/request.rs
new file mode 100644
index 0000000..17e2556
--- /dev/null
+++ b/src/api/types/request.rs
@@ -0,0 +1,81 @@
+use serde::{Deserialize, Serialize};
+
+/// Chat completion request (OpenAI compatible)
+#[derive(Debug, Clone, Deserialize)]
+#[allow(dead_code)]
+pub struct ChatCompletionRequest {
+ pub model: String,
+ pub messages: Vec,
+ #[serde(default = "default_max_tokens")]
+ pub max_tokens: Option,
+ #[serde(default = "default_temperature")]
+ pub temperature: Option,
+ #[serde(default)]
+ pub top_p: Option,
+ #[serde(default)]
+ pub n: Option,
+ #[serde(default)]
+ pub stream: bool,
+ pub stop: Option,
+ #[serde(default)]
+ pub presence_penalty: Option,
+ #[serde(default)]
+ pub frequency_penalty: Option,
+}
+
+/// Chat message
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct ChatMessage {
+ pub role: String, // "system", "user", "assistant"
+ pub content: String,
+}
+
+/// Legacy text completion request
+#[derive(Debug, Clone, Deserialize)]
+#[allow(dead_code)]
+pub struct CompletionRequest {
+ pub model: String,
+ #[serde(alias = "prompt")]
+ pub prompt: StringOrArray,
+ #[serde(default = "default_max_tokens")]
+ pub max_tokens: Option,
+ #[serde(default = "default_temperature")]
+ pub temperature: Option,
+ #[serde(default)]
+ pub top_p: Option,
+ #[serde(default)]
+ pub n: Option,
+ #[serde(default)]
+ pub stream: bool,
+ pub stop: Option,
+ #[serde(default)]
+ pub presence_penalty: Option,
+ #[serde(default)]
+ pub frequency_penalty: Option,
+}
+
+/// Prompt can be string or array of strings
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub enum StringOrArray {
+ String(String),
+ Array(Vec),
+}
+
+impl std::fmt::Display for StringOrArray {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ StringOrArray::String(s) => write!(f, "{}", s),
+ StringOrArray::Array(arr) => write!(f, "{}", arr.join("\n")),
+ }
+ }
+}
+
+// Default values
+fn default_max_tokens() -> Option {
+ Some(100)
+}
+
+fn default_temperature() -> Option {
+ Some(0.7)
+}
diff --git a/src/api/types/response.rs b/src/api/types/response.rs
new file mode 100644
index 0000000..462ca03
--- /dev/null
+++ b/src/api/types/response.rs
@@ -0,0 +1,119 @@
+use serde::{Deserialize, Serialize};
+
+use super::request::ChatMessage;
+
+/// Chat completion response
+#[derive(Debug, Serialize)]
+pub struct ChatCompletionResponse {
+ pub id: String,
+ pub object: String, // "chat.completion"
+ pub created: i64,
+ pub model: String,
+ pub choices: Vec,
+ pub usage: Usage,
+}
+
+/// Choice in chat completion
+#[derive(Debug, Serialize)]
+pub struct ChatChoice {
+ pub index: usize,
+ pub message: ChatMessage,
+ pub finish_reason: String, // "stop", "length", "content_filter"
+}
+
+/// Streaming chat completion chunk
+#[derive(Debug, Serialize)]
+pub struct ChatCompletionChunk {
+ pub id: String,
+ pub object: String, // "chat.completion.chunk"
+ pub created: i64,
+ pub model: String,
+ pub choices: Vec,
+}
+
+/// Delta choice for streaming
+#[derive(Debug, Serialize)]
+pub struct ChatChoiceDelta {
+ pub index: usize,
+ pub delta: ChatMessageDelta,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub finish_reason: Option,
+}
+
+/// Delta message for streaming
+#[derive(Debug, Serialize)]
+pub struct ChatMessageDelta {
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub role: Option,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub content: Option,
+}
+
+/// Legacy completion response
+#[derive(Debug, Serialize)]
+pub struct CompletionResponse {
+ pub id: String,
+ pub object: String, // "text_completion"
+ pub created: i64,
+ pub model: String,
+ pub choices: Vec,
+ pub usage: Usage,
+}
+
+/// Choice in completion
+#[derive(Debug, Serialize)]
+pub struct CompletionChoice {
+ pub text: String,
+ pub index: usize,
+ pub logprobs: Option,
+ pub finish_reason: String,
+}
+
+/// Token usage statistics
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Usage {
+ pub prompt_tokens: usize,
+ pub completion_tokens: usize,
+ pub total_tokens: usize,
+}
+
+/// Model list response
+#[derive(Debug, Serialize)]
+pub struct ModelList {
+ pub object: String, // "list"
+ pub data: Vec,
+}
+
+/// Model information
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Model {
+ pub id: String,
+ pub object: String, // "model"
+ pub created: i64,
+ pub owned_by: String,
+}
+
+/// Error response
+#[derive(Debug, Serialize)]
+pub struct ErrorResponse {
+ pub error: ErrorDetail,
+}
+
+#[derive(Debug, Serialize)]
+pub struct ErrorDetail {
+ pub message: String,
+ pub r#type: String,
+ pub code: Option,
+}
+
+impl ErrorResponse {
+ pub fn new(message: String, error_type: String) -> Self {
+ Self {
+ error: ErrorDetail {
+ message,
+ r#type: error_type,
+ code: None,
+ },
+ }
+ }
+}
diff --git a/src/backend/engine.rs b/src/backend/engine.rs
new file mode 100644
index 0000000..2d464a5
--- /dev/null
+++ b/src/backend/engine.rs
@@ -0,0 +1,34 @@
+use std::io;
+use std::pin::Pin;
+use tokio_stream::Stream;
+
+/// Inference engine trait
+pub trait InferenceEngine: Send + Sync {
+ /// Generate text completion
+ fn generate(
+ &self,
+ model: &str,
+ prompt: &str,
+ max_tokens: usize,
+ temperature: f32,
+ ) -> impl std::future::Future
+