diff --git a/.claude/agents/docs-writer.md b/.claude/agents/docs-writer.md index b5c4ee3f..cf749c3a 100644 --- a/.claude/agents/docs-writer.md +++ b/.claude/agents/docs-writer.md @@ -65,7 +65,7 @@ Available types: `note`, `abstract`, `info`, `tip`, `success`, `question`, `warn === "From Source" ```bash - git clone https://github.com/jonsaadfalcon/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis uv sync --extra dev ``` diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 09b55a30..8fd6f9d2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -48,3 +48,32 @@ jobs: - name: Run tests run: uv run pytest tests/ -v --tb=short + + rust: + runs-on: ubuntu-latest + defaults: + run: + working-directory: rust + steps: + - uses: actions/checkout@v4 + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + with: + components: clippy + + - name: Cargo cache + uses: actions/cache@v4 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + rust/target + key: rust-${{ runner.os }}-${{ hashFiles('rust/Cargo.lock') }} + restore-keys: rust-${{ runner.os }}- + + - name: Clippy + run: cargo clippy --workspace --all-targets -- -D warnings + + - name: Test + run: cargo test --workspace diff --git a/.gitignore b/.gitignore index 71989b63..82ab9ef4 100644 --- a/.gitignore +++ b/.gitignore @@ -23,6 +23,7 @@ env/ # Testing .pytest_cache/ +.ruff_cache/ .coverage htmlcov/ @@ -42,7 +43,12 @@ Thumbs.db # Project *.sqlite *.db +*.jsonl +*.npz results/ +logs/ +traces/ +coding_task_* get-pip.py # MkDocs build output @@ -70,5 +76,9 @@ docs/plans/ # Tauri auto-generated schemas **/src-tauri/gen/schemas/ -# Rust build artifacts -rust/target/ \ No newline at end of file +# NFS lock artifacts +.nfs* +**/.nfs* + +# Research output +research_mining_* diff --git a/CLAUDE.md b/CLAUDE.md index 9c25ed0f..6d15d4b5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,246 +2,96 @@ This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. -## Project Status +## Project Overview -OpenJarvis is a research framework for studying on-device AI systems. Phase 25 complete. Five composable pillars: Intelligence, Engine, Agents, Tools (with storage + MCP), and Learning — with trace-driven learning as a cross-cutting concern. Speech subsystem (STT) with pluggable backends. ~3405 tests pass (~44 skipped for optional deps). Python SDK (`Jarvis` class), composition layer (`SystemBuilder`/`JarvisSystem`), eval framework (20 real benchmarks), composable recipes, agent templates, bundled skills, operator recipes, trace-driven learning pipeline, Docker deployment, Tauri desktop app, 41+ tools, 20+ CLI commands, 40+ API endpoints all ready. +OpenJarvis is a modular AI assistant backend / research framework for on-device AI systems. It is organized around **five composable "pillars"** that are wired together via a config-driven `JarvisSystem` composition layer. ## Build & Development Commands +**Package manager:** [uv](https://github.com/astral-sh/uv) (lock file `uv.lock` is tracked for reproducibility) + ```bash -uv sync --extra dev # Install deps + dev tools -uv run pytest tests/ -v # Run ~3295 tests (~44 skipped if optional deps missing) -uv run ruff check src/ tests/ # Lint -uv run jarvis --version # 1.0.0 -uv run jarvis ask "Hello" # Query via discovered engine (direct mode) -uv run jarvis ask --agent simple "Hello" # SimpleAgent route -uv run jarvis ask --agent orchestrator "Hello" # OrchestratorAgent route -uv run jarvis ask --agent orchestrator --tools calculator,think "What is 2+2?" -uv run jarvis ask --agent native_react --tools calculator "What is 2+2?" # NativeReActAgent -uv run jarvis ask --agent react "Hello" # Alias for native_react -uv run jarvis ask --agent native_openhands "Hello" # NativeOpenHandsAgent (CodeAct) -uv run jarvis ask --agent openhands "Hello" # Real OpenHands SDK (requires openhands-sdk) -uv run jarvis ask --router heuristic "Hello" # Explicit heuristic policy -uv run jarvis ask --no-context "Hello" # Query without memory context injection -uv run jarvis model list # List models from running engines -uv run jarvis model info qwen3:8b # Show model details -uv run jarvis memory index ./docs/ # Index documents into memory -uv run jarvis memory search "topic" # Search memory for relevant chunks -uv run jarvis memory stats # Show memory backend statistics -uv run jarvis telemetry stats # Show aggregated telemetry stats -uv run jarvis telemetry export --format json # Export records as JSON -uv run jarvis telemetry export --format csv # Export records as CSV -uv run jarvis telemetry clear --yes # Delete all telemetry records -uv run jarvis channel list # List available messaging channels -uv run jarvis channel send slack "Hello" # Send a message to a channel -uv run jarvis channel status # Show channel bridge connection status -uv run jarvis scheduler create "Check weather" --type cron --value "0 9 * * *" -uv run jarvis scheduler list # List scheduled tasks -uv run jarvis scheduler start # Start scheduler daemon (foreground) -uv run jarvis bench run # Run all benchmarks against engine -uv run jarvis bench run -b energy -w 5 -n 20 --json # Energy benchmark with warmup -uv run jarvis serve --port 8000 # OpenAI-compatible API server (requires openjarvis[server]) -uv run jarvis doctor # Run diagnostic checks (config, engines, models, deps) -uv run jarvis doctor --json # Machine-readable diagnostics -uv run jarvis start # Start server as background daemon -uv run jarvis stop # Stop background daemon -uv run jarvis restart # Restart background daemon -uv run jarvis status # Show daemon status (PID, uptime) -uv run jarvis chat # Interactive REPL (/quit, /clear, /model, /help, /history) -uv run jarvis chat --agent orchestrator --tools calculator # REPL with agent -uv run jarvis agent list # List registered agents -uv run jarvis agent info native_react # Show agent details -uv run jarvis workflow list # List available workflows -uv run jarvis workflow run my_workflow # Execute a workflow -uv run jarvis skill list # List installed skills -uv run jarvis skill install path/to/skill.toml # Install a skill -uv run jarvis vault set MY_KEY # Store encrypted credential -uv run jarvis vault get MY_KEY # Retrieve credential -uv run jarvis vault list # List stored keys -uv run jarvis add github # Quick-add MCP server (github, slack, postgres, etc.) -uv run jarvis eval list # List 15 benchmarks and backends -uv run jarvis eval run -b supergpqa -m "qwen3:8b" # Single benchmark run -uv run jarvis eval run -c src/openjarvis/evals/configs/suite.toml # Suite mode (models x benchmarks) -uv run jarvis eval compare result1.jsonl result2.jsonl # Compare runs side-by-side -uv run jarvis eval report result.jsonl # Detailed report with per-subject breakdown -uv run jarvis --help # Show all subcommands -uv run jarvis init --force # Detect hardware, write ~/.openjarvis/config.toml -# Eval framework (direct module invocation) -source .env # Load API keys before running evals -uv run python -m openjarvis.evals run -c src/openjarvis/evals/configs/glm-4.7-flash-openhands.toml -v # Run eval suite from TOML config -uv run python -m openjarvis.evals run -b supergpqa -m "qwen3:8b" -n 50 # Run single benchmark -uv run python -m openjarvis.evals summarize results/supergpqa_qwen3-8b.jsonl # Summarize results +# Install core + dev dependencies +uv sync --extra dev + +# Lint (ruff, rules: E/F/I/W, target Python 3.10) +uv run ruff check src/ tests/ + +# Run full test suite +uv run pytest tests/ -v --tb=short + +# Run a single test file / single test +uv run pytest tests/agents/test_native_react.py -v +uv run pytest tests/agents/test_native_react.py::test_function_name -v + +# Run tests by marker (live, cloud, nvidia, amd, apple, slow) +uv run pytest -m "not live and not cloud" tests/ + +# CLI entry point +uv run jarvis --help + +# Run evals +uv run python -m openjarvis.evals --config src/openjarvis/evals/configs/.toml + +# Docs (MkDocs Material) +uv sync --extra docs +uv run mkdocs serve ``` -### Config File Conventions - -- **Runtime config (source of truth):** `configs/openjarvis/config.toml` — Pillar-aligned OpenJarvis config. Copied to `~/.openjarvis/config.toml` at runtime (which is where `load_config()` reads from). -- **Eval suite configs:** `src/openjarvis/evals/configs/*.toml` — TOML configs defining models x benchmarks matrices. -- **API keys:** `.env` file in project root (gitignored). Source with `source .env` before running evals or cloud operations. -- **Never save configs to `~/.openjarvis/` directly** — always maintain the canonical copy in `configs/openjarvis/` and copy/symlink to `~/.openjarvis/`. - -### Python SDK - -```python -from openjarvis import Jarvis - -j = Jarvis() # Uses default config + auto-detected engine -j = Jarvis(model="qwen3:8b") # Override model -j = Jarvis(engine_key="ollama") # Override engine - -response = j.ask("Hello") # Returns string -full = j.ask_full("Hello") # Returns dict with content, usage, model, engine -response = j.ask("Hello", agent="orchestrator", tools=["calculator"]) - -j.memory.index("./docs/") # Index documents -results = j.memory.search("topic") # Search memory -j.memory.stats() # Backend stats - -j.list_models() # Available models -j.list_engines() # Registered engines -j.close() # Release resources +**Rust workspace** (in `rust/`): +```bash +cd rust +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace ``` -- **Package manager:** `uv` with `hatchling` build backend -- **Config:** `pyproject.toml` with extras for optional backends (e.g., `openjarvis[inference-vllm]`, `openjarvis[inference-mlx]`, `openjarvis[memory-colbert]`, `openjarvis[server]`, `openjarvis[openclaw]`, `openjarvis[energy-amd]`, `openjarvis[energy-apple]`, `openjarvis[energy-all]`, `openjarvis[security-signing]`, `openjarvis[sandbox-wasm]`, `openjarvis[dashboard]`, `openjarvis[browser]`, `openjarvis[media]`, `openjarvis[pdf]`, `openjarvis[channel-line]`, `openjarvis[channel-viber]`, `openjarvis[channel-reddit]`, `openjarvis[channel-mastodon]`, `openjarvis[channel-xmpp]`, `openjarvis[channel-rocketchat]`, `openjarvis[channel-zulip]`, `openjarvis[channel-twitch]`, `openjarvis[channel-nostr]`, `openjarvis[speech]`, `openjarvis[speech-deepgram]`) -- **CLI entry point:** `jarvis` (Click-based) — subcommands: `init`, `ask`, `serve`, `start`, `stop`, `restart`, `status`, `chat`, `model`, `memory`, `telemetry`, `bench`, `eval`, `channel`, `scheduler`, `doctor`, `agent`, `workflow`, `skill`, `vault`, `add` -- **Python:** 3.10+ required -- **Node.js:** 22+ required only for OpenClaw agent +**Frontend / Desktop** (Tauri + Vite, in `frontend/` and `desktop/`): +```bash +cd frontend && npm install && npm run dev +cd desktop && npm install && npm run tauri dev +``` -## Architecture +## Architecture: The Five Pillars -OpenJarvis is a research framework for on-device AI organized around **five composable pillars**, each with a clear ABC interface and a decorator-based registry for runtime discovery. +All pillars are wired together by `JarvisSystem` (`src/openjarvis/system.py`) which is constructed from `configs/openjarvis/config.toml` (or `~/.openjarvis/config.toml`). -### Five Pillars +### 1. Intelligence (`src/openjarvis/intelligence/`) +Model selection, provider routing. Config section: `[intelligence]`. -1. **Intelligence** (`src/openjarvis/intelligence/`) — Model definition, catalog, and generation defaults. `ModelRegistry` maps model keys to `ModelSpec`. `IntelligenceConfig` holds model identity (default/fallback model, model_path, checkpoint_path, quantization, preferred_engine, provider) and generation defaults (temperature, max_tokens, top_p, top_k, repetition_penalty, stop_sequences). Model catalog maintains `BUILTIN_MODELS` with auto-discovery via `merge_discovered_models()`. -2. **Engine** (`src/openjarvis/engine/`) — The inference runtime. Backends: vLLM, SGLang, Ollama, llama.cpp, MLX, LM Studio. All implement `InferenceEngine` ABC with `generate()`, `stream()`, `list_models()`, `health()`. Engines extract and pass through `tool_calls` in OpenAI format. -3. **Agents** (`src/openjarvis/agents/`) — Pluggable logic for queries, tool/API calls, memory. Hierarchy: `BaseAgent` ABC (helpers: `_emit_turn_start/end`, `_build_messages`, `_generate`, `_max_turns_result`, `_strip_think_tags`, `_check_continuation`) → `ToolUsingAgent` (adds `tools`, `ToolExecutor`, `max_turns`). Agents: `SimpleAgent` (single-turn), `OrchestratorAgent` (multi-turn tool loop), `NativeReActAgent` (Thought-Action-Observation, key `"native_react"`, alias `"react"`), `NativeOpenHandsAgent` (CodeAct, key `"native_openhands"`), `RLMAgent` (recursive LM), `OpenHandsAgent` (real `openhands-sdk`, key `"openhands"`, requires Python 3.12+), `OpenClawAgent` (HTTP/subprocess transport), `ClaudeCodeAgent` (Claude Agent SDK via Node.js, key `"claude_code"`), `SandboxedAgent` (Docker wrapper, key `"sandboxed"`), `MonitorOperativeAgent` (long-horizon monitoring with 4 configurable strategies: memory extraction, observation compression, retrieval, task decomposition, key `"monitor_operative"`). `accepts_tools` class attribute for CLI/SDK auto-detection. Agents call `engine.generate()` directly — telemetry handled by `InstrumentedEngine` wrapper. -4. **Tools** (`src/openjarvis/tools/`) — All tools managed via MCP (Model Context Protocol). - - **API tools**: `CalculatorTool`, `ThinkTool`, `FileReadTool`, `FileWriteTool`, `WebSearchTool`, `CodeInterpreterTool`, `LLMTool`, `ShellExecTool`, `ApplyPatchTool`, `HttpRequestTool`, `DatabaseQueryTool`, `PDFExtractTool`, `ImageGenerateTool`, `AudioTranscribeTool` — all implement `BaseTool` ABC - - **Git tools** (`git_tool.py`): `GitStatusTool`, `GitDiffTool`, `GitCommitTool`, `GitLogTool` - - **Browser tools** (`browser.py`): `BrowserNavigateTool`, `BrowserClickTool`, `BrowserTypeTool`, `BrowserScreenshotTool`, `BrowserExtractTool` (Playwright, optional `[browser]`) - - **Agent tools** (`agent_tools.py`): `AgentSpawnTool`, `AgentSendTool`, `AgentListTool`, `AgentKillTool` - - **Storage tools** (`storage_tools.py`): `MemoryStoreTool`, `MemoryRetrieveTool`, `MemorySearchTool`, `MemoryIndexTool` - - **Storage backends** (`tools/storage/`): SQLite/FTS5 (default), FAISS, ColBERTv2, BM25, Hybrid (RRF fusion), KnowledgeGraph. All implement `MemoryBackend` ABC. Canonical import: `from openjarvis.tools.storage.sqlite import SQLiteMemory`. - - **Scheduler tools** (`scheduler/tools.py`): 5 MCP tools for task scheduling - - **Knowledge graph tools** (`knowledge_tools.py`): `KGAddEntityTool`, `KGAddRelationTool`, `KGQueryTool`, `KGNeighborsTool` - - **MCP adapter** (`mcp_adapter.py`): `MCPToolAdapter` wraps external MCP tools as native `BaseTool`; `MCPToolProvider` discovers from server - - **MCP server** (`mcp/server.py`): Exposes all built-in tools via JSON-RPC `tools/list` + `tools/call` (MCP spec 2025-11-25) - - **MCP templates** (`tools/templates/`): `ToolTemplate` dynamically constructs tools from TOML specs. 10 builtin templates. `discover_templates()` auto-discovers. - - **`ToolExecutor`**: dispatch with RBAC check + taint check, `timeout_seconds` on `ToolSpec` (default 30s via `ThreadPoolExecutor`), event bus integration - - All registered via `@ToolRegistry.register("name")` decorator -5. **Learning** (`src/openjarvis/learning/`) — Structured learning with nested per-pillar sub-policies. `LearningConfig` sections: `routing` (heuristic/learned/grpo/bandit), `intelligence` (none/sft), `agent` (none/agent_advisor/icl_updater), `metrics` (accuracy/latency/cost/efficiency weights). Policies: `SFTRouterPolicy` (query→model from traces), `AgentAdvisorPolicy` (LM-guided), `ICLUpdaterPolicy` (in-context with example DB, versioning, rollback, quality gates), `GRPORouterPolicy` (softmax sampling, group relative advantage, per-query-class weights), `BanditRouterPolicy` (Thompson Sampling / UCB1, per-arm stats). `SkillDiscovery` mines tool subsequences from traces to auto-generate skill manifests. Router policies: `HeuristicRouter`, `TraceDrivenPolicy`. Orchestrator training subpackage provides SFT and GRPO pipelines. **Trace-driven learning pipeline**: `TrainingDataMiner` (extracts SFT pairs from traces with quality filters), `LoRATrainer` (LoRA fine-tuning with configurable rank/alpha, requires torch), `AgentConfigEvolver` (LM-guided agent config recommendations from trace patterns), `LearningOrchestrator` (wired into `SystemBuilder`, orchestrates mine→train→evolve cycle on schedule). +### 2. Agent (`src/openjarvis/agents/`) +Multi-turn reasoning and tool use. Agents register via `@AgentRegistry.register()` decorator. Key agents: `simple`, `native_react`, `native_openhands`, `orchestrator`, `monitor_operative`, `claude_code`, `rlm`. Config section: `[agent]`. -### Speech Subsystem +### 3. Tools (`src/openjarvis/tools/`) +Built-in tools (code_interpreter, web_search, file_read, shell_exec, calculator, think, browser, etc.) plus MCP adapter. Tool storage backends in `tools/storage/`. Config section: `[tools]`. -- **Speech** (`src/openjarvis/speech/`) — Speech-to-text with pluggable backends. `SpeechBackend` ABC (`transcribe()`, `health()`, `supported_formats()`). `TranscriptionResult` + `Segment` dataclasses. Backends: `FasterWhisperBackend` (local, CTranslate2, key `"faster-whisper"`), `OpenAIWhisperBackend` (cloud, `whisper-1`, key `"openai"`), `DeepgramSpeechBackend` (cloud, `nova-2`, key `"deepgram"`). Auto-discovery with local-first priority. `SpeechConfig` in `JarvisConfig`. `SpeechRegistry` for backend registration. Wired into `SystemBuilder.speech()`, `create_app()`, `jarvis serve`. API: `POST /v1/speech/transcribe` (multipart), `GET /v1/speech/health`. Frontend: `useSpeech` hook + `MicButton` component. Tauri: `transcribe_audio` + `speech_health` commands. Optional deps: `openjarvis[speech]` (faster-whisper), `openjarvis[speech-deepgram]` (deepgram-sdk). +### 4. Engine (`src/openjarvis/engine/`) +Inference runtime abstraction. All engines implement `InferenceEngine` (defined in `engine/_stubs.py`) and use OpenAI-compatible chat completions. Supported: vLLM, Ollama, llama.cpp, SGLang, MLX, cloud (OpenAI/Anthropic/Google), LiteLLM, Apple FM, Exo, Nexa. Discovery in `engine/_discovery.py`. Config section: `[engine]`. -### Cross-cutting Systems +### 5. Learning (`src/openjarvis/learning/`) +Improvement methodologies: router policies (heuristic, bandit, trace-based), SFT/GRPO training, ICL updater, agent evolution, skill discovery. Orchestrated by `LearningOrchestrator`. Config section: `[learning]`. -- **Traces** (`src/openjarvis/traces/`) — Full interaction recording. `Trace` captures `TraceStep`s (route, retrieve, generate, tool_call, respond) with timing. `TraceStore` (SQLite), `TraceCollector` (auto-wraps agents), `TraceAnalyzer` (stats for learning). -- **Telemetry** (`src/openjarvis/telemetry/`) — `InstrumentedEngine` wraps any engine, publishing events to SQLite via `TelemetryStore`. `TelemetryAggregator` for read-only queries. `EnergyMonitor` ABC with vendor-specific implementations: `NvidiaEnergyMonitor` (hw counters/polling), `AmdEnergyMonitor` (amdsmi), `AppleEnergyMonitor` (zeus-ml), `RaplEnergyMonitor` (sysfs). `EnergyBatch` for batch-level energy-per-token. `SteadyStateDetector` for thermal equilibrium (CV-based). -- **Security** (`src/openjarvis/security/`) — `SecretScanner` + `PIIScanner` (implement `BaseScanner` ABC). `GuardrailsEngine` wraps engines with input/output scanning (WARN/REDACT/BLOCK modes). `AuditLogger` with Merkle hash chain (SHA-256 tamper-evidence). `CapabilityPolicy` RBAC (10 capabilities with glob matching, enforced in `ToolExecutor`). `TaintLabel`/`TaintSet` information flow control with `SINK_POLICY`. Ed25519 signing via `cryptography` (optional `[security-signing]`). `file_policy.py` for sensitive file detection. `InjectionScanner` (11 regex patterns: prompt override, identity override, code/shell injection, exfiltration, jailbreak, delimiter injection). `check_ssrf()` SSRF protection (RFC 1918, loopback, link-local, cloud metadata blocking). `RateLimiter` with `TokenBucket` (thread-safe, per-key). `run_sandboxed()` subprocess isolation (`os.setsid`, process group kill, env clearing). Security HTTP middleware (7 headers: CSP, HSTS, X-Frame-Options, etc.). +### Supporting Systems +- **Core** (`core/`): `RegistryBase` pattern (decorator-based registration), types (`Message`, `Conversation`, `ToolCall`), config loader with hardware detection, `EventBus`. +- **Evals** (`evals/`): Benchmark framework with TOML configs. Datasets: SuperGPQA, GPQA, MMLU-Pro, MATH-500, GAIA, SWE-bench, FRAMES, SimpleQA, TerminalBench, PaperArena, etc. Run via `python -m openjarvis.evals`. +- **Channels** (`channels/`): Chat platform integrations (Telegram, Discord, Slack, WhatsApp, Signal, IRC, Matrix, etc.). +- **Telemetry** (`telemetry/`): GPU monitoring, energy measurement (NVIDIA/AMD/Apple/RAPL), latency instrumentation, vLLM metrics. +- **Traces** (`traces/`): Execution trace recording for analysis. +- **MCP** (`mcp/`): Model Context Protocol server. +- **Security** (`security/`): PII scanning, capability policies. +- **Server** (`server/`): FastAPI REST API. +- **SDK** (`sdk.py`): High-level `Jarvis` and `JarvisSystem` classes, `MemoryHandle` for memory operations. +- **Rust** (`rust/`): Parallel Rust implementation with PyO3 bindings. Workspace crates mirror Python pillars (core, engine, agents, tools, learning, telemetry, traces, security, mcp, python). -### Composition & Infrastructure +## Key Patterns -- **Composition Layer** (`system.py`) — `SystemBuilder` fluent builder → `JarvisSystem` with `ask()`, `close()`. Wires engine, model, agent, tools, telemetry, traces, workflow, sessions, capability policy. -- **SDK** (`sdk.py`) — `Jarvis` class: high-level sync API with `ask()`/`ask_full()`, `MemoryHandle`, lazy init, telemetry. Also exports `JarvisSystem`/`SystemBuilder`. -- **Benchmarks** (`bench/`) — `LatencyBenchmark`, `ThroughputBenchmark`, `EnergyBenchmark`. All registered via `BenchmarkRegistry`. CLI: `jarvis bench run`. -- **Eval Framework** (`src/openjarvis/evals/`) — 20 real benchmark datasets: SuperGPQA, GPQA, MMLU-Pro, MATH-500, Natural Reasoning, HLE, SimpleQA, WildChat, IPW, GAIA, FRAMES, SWE-bench, SWEfficiency, TerminalBench, TerminalBench Native, LogHub, AMA-Bench, LifelongAgentBench, WebChoreArena, WorkArena++. Scorer types: MCQ letter extraction, LLM-judge, exact match, structural validation. `EvalRunner` with parallel execution and episode mode (sequential processing with shared agent state). `EnvironmentProvider` ABC for Docker/ServiceNow environments. CLI: `jarvis eval list|run|compare|report`. -- **Recipes** (`src/openjarvis/recipes/`) — Composable TOML configs that wire all 5 pillars. `Recipe` dataclass with `to_builder_kwargs()`. `load_recipe()`, `discover_recipes()`, `resolve_recipe()`. 3 built-in recipes in `data/`: coding_assistant, research_assistant, general_assistant. Operator recipes in `data/operators/`: researcher (4h cycle), correspondent (5min interval), sentinel (2h cycle), monitor (2h cycle, causality graph + hybrid retrieval). -- **Agent Templates** (`src/openjarvis/templates/`) — Pre-configured TOML manifests with system prompts, tool sets, behavioral parameters. `AgentTemplate` dataclass, `load_template()`, `discover_templates()`. 15 built-in templates in `data/` (code-reviewer, debugger, architect, deep-researcher, fact-checker, summarizer, etc.). -- **Bundled Skills** (`src/openjarvis/skills/data/`) — 20 ready-to-use TOML skill manifests. Categories: file management (organizer, deduplicator, backup), research (web-summarize, topic-research, knowledge-extract), code quality (lint, test-gen, security-scan, dependency-audit), productivity (email-draft, meeting-notes, daily-digest), document processing (compare, translate, data-analyze). -- **OpenClaw** (`agents/openclaw*.py`) — `OpenClawAgent` with `HttpTransport`/`SubprocessTransport`, JSON-line protocol, `ProviderPlugin`, `MemorySearchManager`. -- **API Server** (`server/`) — OpenAI-compatible via `jarvis serve` (FastAPI + uvicorn). Endpoints: `POST /v1/chat/completions`, `GET /v1/models`, `GET /health`, channel endpoints. SSE streaming. -- **Channels** (`channels/`) — `BaseChannel` ABC. `OpenClawChannelBridge` (WebSocket/HTTP to OpenClaw gateway). `WhatsAppBaileysChannel` (Baileys protocol, Node.js bridge, QR auth). Phase 21 channels: `LINEChannel`, `ViberChannel`, `MessengerChannel`, `RedditChannel`, `MastodonChannel`, `XMPPChannel`, `RocketChatChannel`, `ZulipChannel`, `TwitchChannel`, `NostrChannel`. All follow `BaseChannel` ABC with env var fallbacks, `@ChannelRegistry.register()`, `EventBus` integration. -- **Sandbox** (`sandbox/`) — `ContainerRunner` (Docker/Podman lifecycle, mount validation). `WasmRunner` (wasmtime-py, fuel/memory limits, optional `[sandbox-wasm]`). `SandboxedAgent` transparent wrapper. `create_sandbox_runner()` factory. `MountAllowlist` with path traversal prevention. -- **Scheduler** (`scheduler/`) — `TaskScheduler` with cron/interval/once scheduling, SQLite persistence, 5 MCP tools, event bus. CLI: `jarvis scheduler create|list|pause|resume|cancel|logs|start`. -- **Agent Hardening** (`agents/loop_guard.py`) — `LoopGuard`: SHA-256 hash tracking (identical calls), ping-pong detection (A-B-A-B patterns), poll-tool budget, context overflow recovery. `BaseAgent._check_continuation()` auto-resumes on `finish_reason=length`. -- **Workflow Engine** (`workflow/`) — DAG-based `WorkflowGraph` (cycle detection, topological sort, parallel stages via `ThreadPoolExecutor`). `WorkflowBuilder` fluent API. `WorkflowEngine` executes against `JarvisSystem`. TOML loader. Node types: agent, tool, condition, parallel, loop, transform. -- **Skills** (`skills/`) — `SkillManifest`/`SkillExecutor` (sequential tool steps with template rendering). Ed25519 signature verification. `SkillTool` adapter wraps skills as invocable tools. TOML loader. -- **Knowledge Graph** (`tools/storage/knowledge_graph.py`) — `KnowledgeGraphMemory(MemoryBackend)`: SQLite entity-relation store. `add_entity()`, `add_relation()`, `neighbors()`, `query_pattern()`. Registered as `"knowledge_graph"`. -- **Sessions** (`sessions/`) — `SessionStore` (SQLite): cross-channel persistent sessions. `SessionIdentity` canonical user across channels. `consolidate()` summarizes old messages, `decay()` removes expired. -- **A2A Protocol** (`a2a/`) — Google Agent-to-Agent spec (JSON-RPC 2.0). `A2AServer` (tasks/send, tasks/get, tasks/cancel, `/.well-known/agent.json`). `A2AClient`. `A2AAgentTool` adapter. -- **TUI Dashboard** (`cli/dashboard.py`) — `textual`-based terminal dashboard (optional `[dashboard]`). Panels: system status, event stream, telemetry, agent activity, sessions. -- **Desktop App** (`desktop/`) — Tauri 2.0 native desktop application. 5 dashboard panels: EnergyDashboard (real-time power monitoring with recharts), TraceDebugger (timeline inspection with step-type color coding), LearningCurve (policy visualization, GRPO/bandit stats), MemoryBrowser (search + stats), AdminPanel (health, agents, server control). Tauri commands proxy to OpenJarvis REST API. Plugins: notification, shell, global-shortcut, autostart, updater, single-instance. CI: `.github/workflows/desktop.yml` (Linux/macOS/Windows). -- **Vault** (`cli/vault_cmd.py`) — Fernet-encrypted credential store at `~/.openjarvis/vault.enc` with auto-generated key (`0o600` permissions). -- **MCP Quick-Add** (`cli/add_cmd.py`) — `jarvis add ` with 8 templates (github, filesystem, slack, postgres, brave-search, memory, puppeteer, google-maps). Saves JSON config to `~/.openjarvis/mcp/`. +- **Registry pattern**: Components (agents, engines, memory backends, tools, channels, etc.) self-register via `@XRegistry.register("key")` decorators. Tests auto-clear all registries via `conftest.py` fixture. +- **Optional dependencies**: Heavy deps are extras in `pyproject.toml` (e.g., `inference-cloud`, `memory-faiss`, `channel-telegram`). Import failures are caught with try/except so the core stays lightweight. +- **OpenAI-compatible**: All engines expose an OpenAI-format chat completions interface. `messages_to_dicts()` in `engine/_base.py` handles conversion. +- **Config-driven**: TOML configs control everything. `load_config()` detects hardware, fills defaults, then overlays user overrides. -### Core Module (`src/openjarvis/core/`) +## Testing Conventions -- `registry.py` — `RegistryBase[T]` generic base. Subclasses: `ModelRegistry`, `EngineRegistry`, `MemoryRegistry`, `AgentRegistry`, `ToolRegistry`, `RouterPolicyRegistry`, `BenchmarkRegistry`, `ChannelRegistry`, `LearningRegistry`, `SkillRegistry`. -- `types.py` — `Message`, `Conversation`, `ModelSpec`, `ToolResult`, `TelemetryRecord`, `StepType`, `TraceStep`, `Trace`, `RoutingContext`. -- `config.py` — `JarvisConfig` dataclass hierarchy with TOML loader. Config classes for each pillar/subsystem. TOML sections: `[engine]` (+ nested `[engine.ollama]`, `[engine.vllm]`, `[engine.sglang]`, `[engine.llamacpp]`, `[engine.mlx]`, `[engine.lmstudio]`), `[intelligence]`, `[agent]`, `[tools.storage]`, `[tools.mcp]`, `[tools.browser]`, `[learning]` (+ nested routing/intelligence/agent/metrics), `[server]`, `[telemetry]`, `[traces]`, `[channel]`, `[security]` (+ `[security.capabilities]`, `ssrf_protection`, `rate_limit_*`), `[sandbox]`, `[scheduler]`, `[workflow]`, `[sessions]`, `[a2a]`. Backward-compat: `engine.ollama_host` → `engine.ollama.host`, `agent.default_tools` → `agent.tools`, TOML migration for cross-section moves. -- `events.py` — Pub/sub event bus (synchronous dispatch). ~30 EventType values covering inference, tools, memory, agents, telemetry, traces, channels, security, scheduler, workflow, skills, sessions, A2A. - -### Docker & Deployment - -- `deploy/docker/Dockerfile` — Multi-stage: Python 3.12-slim, `.[server]`, entrypoint `jarvis serve` -- `deploy/docker/Dockerfile.gpu` — NVIDIA CUDA 12.4 variant -- `deploy/docker/Dockerfile.gpu.rocm` — AMD ROCm 6.2 variant -- `deploy/docker/docker-compose.yml` — `jarvis` (8000) + `ollama` (11434). ROCm override: `docker-compose.gpu.rocm.yml` -- `deploy/systemd/openjarvis.service`, `deploy/launchd/com.openjarvis.plist` - -### Query Flow - -User query → Security scanning (input) → Intelligence resolves model → Agentic Logic (tools/memory) → Memory retrieval → Context injection → Engine generates → Security scanning (output) → Trace recorded → Telemetry recorded → Learning policies update. - -### API Surface - -OpenAI-compatible server via `jarvis serve`: -- **Core**: `POST /v1/chat/completions`, `GET /v1/models`, `GET /health` -- **Channels**: `GET /v1/channels`, `POST /v1/channels/send`, `GET /v1/channels/status` -- **Agents**: `GET /v1/agents`, `POST /v1/agents`, `DELETE /v1/agents/{id}`, `POST /v1/agents/{id}/message` -- **Memory**: `POST /v1/memory/store`, `POST /v1/memory/search`, `GET /v1/memory/stats` -- **Traces**: `GET /v1/traces`, `GET /v1/traces/{id}` -- **Telemetry**: `GET /v1/telemetry/stats`, `GET /v1/telemetry/energy` -- **Learning**: `GET /v1/learning/stats`, `GET /v1/learning/policy` -- **Skills**: `GET /v1/skills`, `POST /v1/skills`, `DELETE /v1/skills/{name}` -- **Speech**: `POST /v1/speech/transcribe`, `GET /v1/speech/health` -- **Sessions**: `GET /v1/sessions`, `GET /v1/sessions/{id}` -- **Budget**: `GET /v1/budget`, `PUT /v1/budget/limits` -- **Metrics**: `GET /metrics` (Prometheus-compatible) -- **WebSocket**: `WS /v1/chat/stream` (JSON chunked streaming) -- SSE streaming on `/v1/chat/completions` with `stream=true` - -## Key Design Patterns - -- **Registry pattern:** All extensible components use `@XRegistry.register("name")` decorator for registration and runtime discovery. -- **ABC interfaces:** Each pillar defines an ABC. Implement the ABC + register via decorator to add a new backend. -- **Offline-first:** Cloud APIs are optional. All core functionality works without network. -- **Hardware-aware:** Auto-detect GPU vendor/model/VRAM via `nvidia-smi`, `rocm-smi`, `system_profiler`, `/proc/cpuinfo`. Recommend engine accordingly. -- **Telemetry opt-in:** `InstrumentedEngine` wraps inference transparently. Agents unaware of telemetry. -- **Backward-compat shims:** `intelligence/` re-exports from `learning/`, `agents/react.py` re-exports as `ReActAgent`, registry alias `"react"` → `NativeReActAgent`. Old config keys continue to work. -- **`ensure_registered()` pattern:** Benchmark and learning modules use lazy registration to survive registry clearing in tests. - -## Development Phases - -| Version | Phase | Delivers | -|---------|-------|----------| -| v0.1 | 0 | Scaffolding, registries, core types, config, CLI skeleton | -| v0.2 | 1 | Intelligence + Inference — `jarvis ask` end-to-end | -| v0.3 | 2 | Memory backends, document indexing, context injection | -| v0.4 | 3 | Agents, tool system, OpenAI-compatible API server | -| v0.5 | 4 | Learning, telemetry aggregation, `--router` CLI | -| v1.0 | 5 | SDK, OpenClaw infra, benchmarks, Docker | -| v1.1 | 6 | Trace system, trace-driven learning, pluggable agents | -| v1.2 | 7 | 5-pillar restructuring, composition layer, MCP, structured learning | -| v1.3 | 8 | Intelligence = "The Model", routing → Learning, engine selection | -| v1.4 | 9 | Pillar-aligned config, nested configs, TOML migration | -| v1.5 | 10 | Agent restructuring, BaseAgent/ToolUsingAgent, `accepts_tools`, OpenHands SDK | -| v1.6 | 11 | NanoClaw subsumption: ClaudeCodeAgent, WhatsApp Baileys, Docker sandbox, TaskScheduler | -| v1.7 | 12 | EnergyMonitor ABC (NVIDIA/AMD/Apple/RAPL), EnergyBatch, SteadyStateDetector | -| v1.8 | 13 | `jarvis doctor`/`init`, MLX engine, AMD multi-GPU, PWA, ROCm Docker | -| v1.9 | 14 | Agent hardening: LoopGuard, RBAC CapabilityPolicy, taint tracking, Merkle audit, Ed25519 | -| v2.0 | 15 | WorkflowEngine (DAG), SkillSystem, KnowledgeGraphMemory, SessionStore | -| v2.1 | 16 | A2A protocol, MCP templates, WasmRunner, TUI dashboard | -| v2.2 | 17 | Production tool parity: FileWrite, ApplyPatch, ShellExec, Git, HTTP, DB, Browser, Agent, Media, PDF tools. SSRF protection, injection scanner, rate limiter, subprocess sandbox, security middleware | -| v2.3 | 18 | CLI expansion (20 commands): daemon, chat REPL, agent, workflow, skill, vault, add. API expansion (40+ endpoints): agents, memory, traces, telemetry, learning, skills, sessions, budget, metrics, WebSocket streaming | -| v2.4 | 19 | Learning productionization: GRPO (softmax/advantage), BanditRouter (Thompson/UCB1), SkillDiscovery (trace mining), ICL updates (versioning/rollback/quality gates) | -| v2.5 | 20 | Tauri 2.0 desktop app: energy dashboard, trace debugger, learning curve visualization, memory browser, admin panel. CI for Linux/macOS/Windows | -| v2.6 | 21 | 10 new channels: LINE, Viber, Messenger, Reddit, Mastodon, XMPP, Rocket.Chat, Zulip, Twitch, Nostr | -| v2.7 | 22 | Operators: persistent, scheduled autonomous agents with recipe + schedule + channel output | -| v2.8 | 23 | Differentiated functionalities: trace-driven learning pipeline (TrainingDataMiner, LoRATrainer, AgentConfigEvolver, LearningOrchestrator), 15 real IPW benchmarks, composable recipes, 15 agent templates, 20 bundled skills, 3 operator recipes | -| v2.9 | 24 | Speech subsystem: SpeechBackend ABC, SpeechRegistry, 3 backends (FasterWhisper local, OpenAI cloud, Deepgram cloud), auto-discovery, API endpoints, frontend MicButton + useSpeech hook, Tauri commands, SystemBuilder wiring | -| v3.0 | 25 | Operator benchmarks: LogHub, AMA-Bench, LifelongAgentBench, WebChoreArena, WorkArena++. MonitorOperativeAgent with 4 configurable strategies (memory extraction, observation compression, retrieval, task decomposition). EvalRunner episode mode. EnvironmentProvider ABC. browser_axtree tool | +- Tests mirror `src/` structure under `tests/`. +- Markers: `live` (needs running engine), `cloud` (needs API keys), `nvidia`/`amd`/`apple` (GPU-specific), `slow`. +- `conftest.py` provides hardware fixtures (`hardware_nvidia`, `hardware_apple`, etc.) and `mock_engine` factory. +- E501 line length is relaxed for `evals/datasets/*.py` and `evals/scorers/*.py`. diff --git a/NOTES.md b/NOTES.md deleted file mode 100644 index b507a1f4..00000000 --- a/NOTES.md +++ /dev/null @@ -1,492 +0,0 @@ -# OpenJarvis Development Notes - -Living document tracking implementation progress, testing state, lessons learned, dead ends, and practices for ongoing development. Updated across sessions. - ---- - -## Current State (2026-02-21) - -- **Version:** 1.0.0 (trace system added, targeting v1.1) -- **All 6 roadmap phases complete** (Phase 0 through Phase 5) + Phase 6 trace system in progress -- **Tests:** 576 passed, 8 skipped, 0 failures -- **Lint:** ruff clean (`select = ["E", "F", "I", "W"]`) -- **Source files:** 76 Python files in `src/openjarvis/` -- **Test files:** 78 Python files in `tests/` -- **Python:** 3.13 (compatible with 3.10+) -- **Package manager:** `uv` with `hatchling` build backend - -### 8 Skipped Tests (Optional Dependencies) - -| Test | Missing Dep | Install Extra | -|------|-------------|---------------| -| `tests/memory/test_bm25.py` | `rank_bm25` | `openjarvis[memory-bm25]` | -| `tests/memory/test_colbert.py` | `colbert` | `openjarvis[memory-colbert]` | -| `tests/memory/test_embeddings.py` | `sentence_transformers` | `openjarvis[memory-faiss]` | -| `tests/memory/test_faiss.py` | `faiss` | `openjarvis[memory-faiss]` | -| `tests/server/test_models_pydantic.py` | `pydantic` | `openjarvis[server]` | -| `tests/server/test_routes.py` | `fastapi` | `openjarvis[server]` | -| `tests/test_integration.py:165` | `fastapi` | `openjarvis[server]` | -| `tests/test_integration.py:190` | `fastapi` | `openjarvis[server]` | - ---- - -## Phase Completion Log - -| Phase | Version | Deliverables | Test Count (cumulative) | -|-------|---------|-------------|------------------------| -| Phase 0 | v0.1 | Scaffolding, registries, core types, config, CLI skeleton, event bus | ~60 | -| Phase 1 | v0.2 | Intelligence + Inference — `jarvis ask` end-to-end, heuristic router, engine discovery, basic telemetry | ~160 | -| Phase 2 | v0.3 | Memory — SQLite/FAISS/ColBERT/BM25/Hybrid backends, document ingest pipeline, context injection, `jarvis memory` CLI | ~270 | -| Phase 3 | v0.4 | Agents (Simple/Orchestrator/Custom/OpenClaw stub), tool system (Calculator/Think/Retrieval/LLM/FileRead), OpenAI-compatible API server, `jarvis serve` | ~360 | -| Phase 4 | v0.5 | Learning — HeuristicRouter, HeuristicRewardFunction, GRPORouterPolicy stub, TelemetryAggregator, `jarvis telemetry` CLI, `--router` CLI option | ~432 | -| Phase 5 | v1.0 | SDK (`Jarvis` class), OpenClaw infrastructure (protocol/transport/plugin), benchmarks (`jarvis bench`), Docker, docs | ~520 | -| Phase 6 | v1.1 | Trace system (TraceStore, TraceCollector, TraceAnalyzer), trace-driven learning (TraceDrivenPolicy) | ~576 | - ---- - -## Architecture Quick Reference - -### Directory Layout - -``` -src/openjarvis/ -├── __init__.py # __version__ = "1.0.0", exports Jarvis, MemoryHandle -├── sdk.py # Python SDK: Jarvis class + MemoryHandle -├── core/ -│ ├── registry.py # RegistryBase[T] + 7 typed registries -│ ├── types.py # Message, Conversation, ModelSpec, ToolResult, TelemetryRecord -│ ├── config.py # JarvisConfig dataclass hierarchy, TOML loader -│ └── events.py # EventBus pub/sub (synchronous) -├── intelligence/ # ModelRegistry, HeuristicRouter, model catalog -├── traces/ # TraceStore, TraceCollector, TraceAnalyzer -├── learning/ # RouterPolicyRegistry, HeuristicRouter, TraceDrivenPolicy, GRPO stub -├── memory/ # SQLite/FAISS/ColBERT/BM25/Hybrid backends, chunking, ingest -├── agents/ # Simple/Orchestrator/Custom/OpenClaw agents + protocol/transport -├── engine/ # Ollama/vLLM/llama.cpp/Cloud engine wrappers -├── tools/ # Calculator/Think/Retrieval/LLM/FileRead tools -├── bench/ # Latency/Throughput benchmarks, BenchmarkSuite -├── telemetry/ # TelemetryStore, TelemetryAggregator, instrumented_generate -├── server/ # FastAPI OpenAI-compatible API server -└── cli/ # Click CLI: init, ask, serve, model, memory, telemetry, bench -``` - -### 7 Registries - -All use `RegistryBase[T]` with `@XRegistry.register("name")` or `register_value()`: - -1. `ModelRegistry` — `ModelSpec` objects -2. `EngineRegistry` — `InferenceEngine` implementations -3. `MemoryRegistry` — `MemoryBackend` implementations -4. `AgentRegistry` — `BaseAgent` implementations -5. `ToolRegistry` — `BaseTool` implementations -6. `RouterPolicyRegistry` — `RouterPolicy` implementations -7. `BenchmarkRegistry` — `BaseBenchmark` implementations - ---- - -## Patterns and Practices - -### The `ensure_registered()` Pattern - -**Problem:** The `_clean_registries` autouse fixture in `tests/conftest.py` calls `.clear()` on every registry before each test. Module-level `@XRegistry.register("name")` decorators only fire once at import time (Python caches modules in `sys.modules`). After registry clearing, the decorations never re-fire, leaving registries empty for subsequent tests. - -**Solution:** Use lazy registration via `ensure_registered()`: - -```python -# src/openjarvis/bench/latency.py -_registered = False - -def ensure_registered() -> None: - global _registered - if _registered: - return - from openjarvis.core.registry import BenchmarkRegistry - if not BenchmarkRegistry.contains("latency"): - BenchmarkRegistry.register_value("latency", LatencyBenchmark) - _registered = True -``` - -Then in `__init__.py`: -```python -def ensure_registered() -> None: - from openjarvis.bench.latency import ensure_registered as _reg_latency - _reg_latency() -``` - -And in test files, use an autouse fixture: -```python -@pytest.fixture(autouse=True) -def _register_latency(): - from openjarvis.bench import ensure_registered - ensure_registered() -``` - -**Where this pattern is used:** `bench/latency.py`, `bench/throughput.py`, `learning/heuristic_policy.py`, `learning/grpo_policy.py`, `learning/heuristic_reward.py` - -**Where this pattern is NOT needed:** Agents, engines, memory backends, and tools use `@register` decorators that work fine because their test files explicitly import and re-register as needed, or the test module import triggers registration. - -### Test Infrastructure - -- **`tests/conftest.py`** — `_clean_registries` autouse fixture clears all 7 registries + clears `EventBus` default listeners before each test. Critical for test isolation. -- **Mock engine pattern** — Almost every test that touches the engine layer uses a `MagicMock()` with `.engine_id`, `.health()`, `.list_models()`, `.generate()` stubbed: - ```python - def _make_engine(content="Hello"): - engine = MagicMock() - engine.engine_id = "mock" - engine.health.return_value = True - engine.list_models.return_value = ["test-model"] - engine.generate.return_value = { - "content": content, - "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, - "model": "test-model", - "finish_reason": "stop", - } - return engine - ``` -- **CLI tests** use Click's `CliRunner` with `patch("openjarvis.cli.X.get_engine", ...)` to mock the engine layer. -- **Memory tests** use `tmp_path` fixture for SQLite DB paths and test files. -- **Optional dep tests** use `pytest.importorskip("module_name")` at module level. - -### Config Defaults - -`JarvisConfig()` with no arguments produces sane defaults: -- Engine: auto-discover (Ollama, vLLM, llama.cpp, cloud in priority order) -- Memory: `sqlite` backend, `~/.openjarvis/memory.db` -- Agent: `simple` (no Node.js dependency) -- Intelligence: `qwen3:8b` default, `qwen3:0.6b` fallback -- Telemetry: enabled, `~/.openjarvis/telemetry.db` -- Learning: `heuristic` default policy - -### File Naming Conventions - -- ABCs and shared dataclasses: `_stubs.py` (e.g., `agents/_stubs.py`, `bench/_stubs.py`, `tools/_stubs.py`) -- Internal helpers: `_discovery.py`, `_base.py` (underscore prefix) -- CLI commands: `*_cmd.py` (e.g., `bench_cmd.py`, `telemetry_cmd.py`, `memory_cmd.py`) -- Test files mirror source: `tests/agents/test_openclaw.py` tests `src/openjarvis/agents/openclaw.py` - -### Import Structure - -- Package `__init__.py` files import submodules to trigger registration -- Try/except around optional dependency imports: - ```python - try: - from openjarvis.engine.ollama import OllamaEngine # noqa: F401 - except ImportError: - pass - ``` -- Top-level `openjarvis/__init__.py` exports: `Jarvis`, `MemoryHandle`, `__version__` - ---- - -## Dead Ends and Gotchas - -### 1. `@register` Decorator vs. `ensure_registered()` - -**Dead end:** Initially used `@BenchmarkRegistry.register("latency")` class decorator in `bench/latency.py`. This caused ~10 test failures because: -- Registry cleared between tests by `conftest.py` -- Module already in `sys.modules`, so `import openjarvis.bench` is a no-op on second import -- Registry stays empty after clearing - -**Fix:** Switched to `ensure_registered()` pattern (see above). This is the pattern already used by `learning/` modules. - -**Rule of thumb:** If a module is imported at package init time AND its registry gets cleared in tests, use `ensure_registered()`. If registration only happens in test fixtures or explicit calls, `@register` is fine. - -### 2. Chunk Attribute Names - -`memory/chunking.py` `Chunk` dataclass uses `content` (not `text`). `ChunkConfig` uses `chunk_overlap` (not `overlap`). Easy to get wrong because these aren't obvious from the field names alone. Always read `_stubs.py` or the actual dataclass before using. - -### 3. Test Content Size for Chunking - -`ChunkConfig.min_chunk_size=50` tokens by default. A test string like `"This is test content."` produces 0 chunks. Use at least ~100 words: -```python -words = " ".join(f"word{i}" for i in range(100)) -``` - -### 4. Version String Locations - -Version is defined in **three places** that must stay in sync: -1. `src/openjarvis/__init__.py` — `__version__ = "1.0.0"` -2. `pyproject.toml` — `version = "1.0.0"` -3. `src/openjarvis/server/app.py` — FastAPI `version="1.0.0"` constructor arg - -Tests that check version: `tests/cli/test_cli.py::test_version_flag` - -### 5. Server Import Guards - -The `server/` module requires `fastapi`, `uvicorn`, `pydantic`. These are behind the `[server]` optional extra. All test files that touch server code use `pytest.importorskip("fastapi")`. The server `__init__.py` wraps imports in try/except. - -### 6. `patch()` Targets for Engine Mocking - -When mocking `get_engine` in CLI tests, the patch target must be the *importing module*, not the source module: -```python -# CORRECT — patches where it's imported -patch("openjarvis.cli.bench_cmd.get_engine", return_value=("mock", engine)) - -# WRONG — patches the source, doesn't affect the already-imported reference -patch("openjarvis.engine._discovery.get_engine", return_value=("mock", engine)) -``` - -Same for SDK tests: `patch("openjarvis.sdk.get_engine", ...)`. - -### 7. EventBus Clearing - -`EventBus()` creates a new instance each time, but `EventBus._default_listeners` is a class variable. The `conftest.py` fixture resets it. If tests subscribe to events, subscriptions won't persist across tests. - -### 8. Module Shadowing in CLI Package - -In `cli/__init__.py`, `from openjarvis.cli.ask import ask` imports the Click command. This shadows the module name. When you try `mock.patch("openjarvis.cli.ask.get_engine")`, Python resolves `openjarvis.cli.ask` as the Click command (via attribute lookup on the package), not the module. - -**Fix:** Use `importlib.import_module("openjarvis.cli.ask")` to get the actual module object, then `mock.patch.object(module, "get_engine")`. - ---- - -## Post-v1.0: Unimplemented Ideas from VISION.md - -These are mentioned in `VISION.md` but not in the roadmap phases. They represent future work: - -### Learning / Router -- [ ] Learned router via GRPO (Group Relative Policy Optimization) — `GRPORouterPolicy` is a stub -- [ ] Preference learning from user feedback -- [ ] Continual fine-tuning on accumulated trajectories -- [ ] Multi-objective optimization: quality vs. latency vs. energy vs. cost - -### Memory -- [ ] ConversationMemory — sliding window with automatic summarization of older turns -- [ ] Personal Notes — user-created persistent notes and preferences -- [ ] Episodic Memory — records of past interactions, tool uses, and outcomes -- [ ] Vector DB adapters (Qdrant, ChromaDB) for users with existing infrastructure - -### Tools -- [ ] WebSearch tool (Tavily, SearXNG, DuckDuckGo) -- [ ] CodeInterpreter tool (sandboxed Python execution) -- [ ] FileWrite tool (safe file writing with path validation) -- [ ] MCP (Model Context Protocol) compatibility - -### Engines -- [ ] SGLang engine backend (structured generation, constrained decoding) -- [ ] MLX engine backend (Apple Silicon native, Metal acceleration) -- [ ] Complete vLLM integration (tensor parallelism config, multi-GPU) - -### OpenClaw -- [ ] Full OpenClaw gateway integration (WebSocket, `:18789`) -- [ ] OpenClaw skill composition -- [ ] Context compaction in OpenClaw agent -- [ ] `openjarvis-openclaw` as separate plugin package (currently inline) - -### Infrastructure -- [ ] Documentation site (MkDocs or similar) -- [ ] Getting started guide -- [ ] Plugin development guide -- [ ] API reference docs -- [ ] CI/CD pipeline -- [ ] PyPI publishing - ---- - -## Testing Recipes - -### Run all tests -```bash -uv sync --extra dev -uv run pytest tests/ -v --tb=short -``` - -### Run a specific module's tests -```bash -uv run pytest tests/bench/ -v -uv run pytest tests/sdk/ -v -uv run pytest tests/agents/test_openclaw.py -v -``` - -### Run with optional deps (server) -```bash -uv sync --extra dev --extra server -uv run pytest tests/server/ -v # No longer skipped -``` - -### Lint -```bash -uv run ruff check src/ tests/ -uv run ruff check src/ tests/ --fix # Auto-fix -``` - -### Quick smoke test -```bash -uv run jarvis --version # 1.0.0 -uv run jarvis --help # All subcommands -python -c "from openjarvis import Jarvis; print(Jarvis)" -``` - ---- - -## Adding New Components - -### New Benchmark - -1. Create `src/openjarvis/bench/my_benchmark.py`: - ```python - from openjarvis.bench._stubs import BaseBenchmark, BenchmarkResult - - class MyBenchmark(BaseBenchmark): - @property - def name(self) -> str: return "my-bench" - @property - def description(self) -> str: return "Description" - def run(self, engine, model, *, num_samples=10) -> BenchmarkResult: ... - - _registered = False - def ensure_registered(): - global _registered - if _registered: return - from openjarvis.core.registry import BenchmarkRegistry - if not BenchmarkRegistry.contains("my-bench"): - BenchmarkRegistry.register_value("my-bench", MyBenchmark) - _registered = True - ``` -2. Import in `bench/__init__.py` `ensure_registered()` -3. Add test file `tests/bench/test_my_benchmark.py` with autouse fixture calling `ensure_registered()` - -### New Tool - -1. Create `src/openjarvis/tools/my_tool.py`: - ```python - from openjarvis.core.registry import ToolRegistry - from openjarvis.tools._stubs import BaseTool, ToolSpec - - @ToolRegistry.register("my-tool") - class MyTool(BaseTool): - @property - def spec(self) -> ToolSpec: ... - def execute(self, input: str, **params) -> str: ... - ``` -2. Import in `tools/__init__.py` -3. Add test file `tests/tools/test_my_tool.py` - -### New Memory Backend - -1. Create `src/openjarvis/memory/my_backend.py`: - ```python - from openjarvis.core.registry import MemoryRegistry - from openjarvis.memory._stubs import MemoryBackend, RetrievalResult - - @MemoryRegistry.register("my-backend") - class MyBackend(MemoryBackend): - def store(self, content, *, source="", metadata=None) -> str: ... - def retrieve(self, query, top_k=5) -> list[RetrievalResult]: ... - def delete(self, doc_id) -> bool: ... - def clear(self) -> None: ... - ``` -2. Import in `memory/__init__.py` with try/except for optional deps -3. Add test file with `pytest.importorskip()` if using optional deps -4. Add optional dep group in `pyproject.toml` if needed - -### New Agent - -1. Create `src/openjarvis/agents/my_agent.py`: - ```python - from openjarvis.agents._stubs import AgentResult, BaseAgent - from openjarvis.core.registry import AgentRegistry - - @AgentRegistry.register("my-agent") - class MyAgent(BaseAgent): - agent_id = "my-agent" - def __init__(self, engine, model, *, bus=None, **kwargs): ... - def run(self, input, context=None, **kwargs) -> AgentResult: ... - ``` -2. Import in `agents/__init__.py` -3. Add test file `tests/agents/test_my_agent.py` - -### New Engine - -1. Create `src/openjarvis/engine/my_engine.py`: - ```python - from openjarvis.core.registry import EngineRegistry - from openjarvis.engine._stubs import InferenceEngine - - @EngineRegistry.register("my-engine") - class MyEngine(InferenceEngine): - engine_id = "my-engine" - def generate(self, messages, *, model, **kwargs) -> dict: ... - def stream(self, messages, *, model, **kwargs): ... - def list_models(self) -> list[str]: ... - def health(self) -> bool: ... - ``` -2. Import in `engine/__init__.py` with try/except -3. Add to `_discovery.py` engine priority list if auto-discoverable - ---- - -## Session Log - -### Session 1 (2026-02-16) — Phase 5 Implementation - -**Scope:** Full Phase 5 (v1.0) — SDK, OpenClaw, Benchmarks, Docker, Docs - -**Work completed:** -- Step 1: Added `BenchmarkRegistry` to `core/registry.py`, updated `conftest.py` -- Step 2: Created `bench/` package — `_stubs.py`, `latency.py`, `throughput.py`, `__init__.py`; CLI `bench_cmd.py` -- Step 3: Created `sdk.py` — `Jarvis` class + `MemoryHandle`; updated `__init__.py` exports -- Step 4: Created OpenClaw infra — `openclaw_protocol.py`, `openclaw_transport.py`, `openclaw_plugin.py`; rewrote `openclaw.py` from stub -- Step 5: Created `Dockerfile`, `Dockerfile.gpu`, `docker-compose.yml`, `deploy/systemd/openjarvis.service`, `deploy/launchd/com.openjarvis.plist` -- Step 6: Version bump to 1.0.0, updated `README.md`, `CLAUDE.md` - -**Bugs fixed during implementation:** -1. Ruff lint: 17 issues (E501, I001, F401, F841) — all fixed -2. Registry clearing broke `@register` decorators — switched to `ensure_registered()` for bench modules -3. `ChunkConfig(overlap=...)` should be `ChunkConfig(chunk_overlap=...)` — fixed -4. `chunk.text` should be `chunk.content` — fixed -5. Test content too short for chunking (0 chunks produced) — used 100 words - -**Final: 520 passed, 8 skipped, 0 failures, ruff clean** - -### Session 2 (2026-02-17) — Test Fixes + Live vLLM Testing - -**Scope:** Fix broken tests, set up live vLLM inference testing - -**Work completed:** -- Fixed 6 failed + 13 errored tests in `tests/cli/test_ask_router.py` and `tests/cli/test_ask_agent.py` - - **Root cause:** `from openjarvis.cli.ask import ask` in `cli/__init__.py` shadows the `ask` module with the Click command object. When `mock.patch("openjarvis.cli.ask.get_engine")` resolves, it tries to patch an attribute on the Click command, not the module. - - **Fix:** Use `importlib.import_module("openjarvis.cli.ask")` + `mock.patch.object(_ask_mod, "get_engine")` instead of string-based patching. -- Added tool fallback in `_openai_compat.py`: if server returns 400 when tools are sent (e.g., vLLM without `--enable-auto-tool-choice`), retry without tools. -- Verified live vLLM testing: existing vLLM server on port 8003 with `Qwen/Qwen3-8B` -- Tested: `jarvis ask`, `jarvis bench run`, `jarvis model list`, `jarvis memory index/search`, `jarvis telemetry stats`, SDK `Jarvis.ask()` and `ask_full()` - -**Gotcha discovered:** -8. **Module shadowing with `from X import Y`** — If a package's `__init__.py` does `from openjarvis.cli.ask import ask`, then `openjarvis.cli.ask` in `sys.modules` is the *module*, but accessing it via attribute lookup on `openjarvis.cli` gives the imported *object* (the Click command). Use `importlib.import_module()` for reliable module access when patching. - -**Live vLLM setup notes:** -- vLLM 0.15.1 running on Lambda cluster (8x A100-SXM4-80GB) -- Config: `~/.openjarvis/config.toml` with `vllm_host = "http://localhost:8003"` and `default_model = "Qwen/Qwen3-8B"` -- Tool calling requires `--enable-auto-tool-choice --tool-call-parser hermes` flags on vLLM server -- Without tool support, orchestrator falls back to reasoning-only mode - -**Final: 520 passed, 8 skipped, 0 failures, ruff clean** - -### Session 3 (2026-02-21) — Trace System & Research Direction - -**Scope:** Design new research direction (abstractions for local AI), implement trace system - -**Design decisions made:** -- OpenJarvis repositioned as a research framework for studying on-device AI -- Four core abstractions: Intelligence, Engine, Agentic Logic, Memory -- Learning is a cross-cutting concern driven by interaction traces -- Agentic Logic should be pluggable — users bring their own architecture (ReAct, OpenHands-style, etc.) -- Trace collection is the bridge between static and learned agents -- Evolve existing codebase rather than full redesign -- Name stays as OpenJarvis -- Learning focus: telemetry-driven routing/tool policies (lightweight, always-on) -- Agent-model coupling: loose (any agent, any model) - -**Work completed:** -- Added `StepType` enum, `TraceStep`, `Trace` dataclasses to `core/types.py` -- Added `TRACE_STEP`, `TRACE_COMPLETE` event types to `core/events.py` -- Created `traces/` package: - - `store.py` — `TraceStore`: SQLite-backed, save/get/list with filters, event bus subscription - - `collector.py` — `TraceCollector`: wraps any `BaseAgent`, subscribes to EventBus, records steps automatically - - `analyzer.py` — `TraceAnalyzer`: per-route stats, per-tool stats, summaries, export, query-type filtering -- Created `learning/trace_policy.py` — `TraceDrivenPolicy`: learns routing from trace outcomes, batch/online updates, registered as `"learned"` policy -- Registered `TraceDrivenPolicy` in `learning/__init__.py` -- Added 56 new tests across 4 test files in `tests/traces/` and `tests/learning/test_trace_policy.py` -- Updated all markdown documentation (README, VISION, ROADMAP, NOTES, CLAUDE) - -**Final: 576 passed, 8 skipped, 0 failures, ruff clean** diff --git a/ROADMAP.md b/ROADMAP.md deleted file mode 100644 index 1578045d..00000000 --- a/ROADMAP.md +++ /dev/null @@ -1,374 +0,0 @@ -# OpenJarvis Roadmap - -Phased development plan for OpenJarvis. Phases are ordered to maximize early usability: foundation first, then intelligence + inference (so you can ask questions), then memory (so it remembers), then agents (so it can act), then learning (so it improves). - ---- - -## Phase 0 — Foundation (~2-3 weeks) - -**Goal:** Repository scaffolding, core abstractions, and CLI skeleton. Nothing runs yet, but all interfaces are defined. - -**Version milestone:** v0.1 - -### Repository structure - -``` -OpenJarvis/ -├── pyproject.toml # uv/hatchling, all deps + extras -├── src/ -│ └── openjarvis/ -│ ├── __init__.py -│ ├── core/ -│ │ ├── registry.py # RegistryBase[T] + typed registries -│ │ ├── types.py # Message, Conversation, ModelSpec, ToolResult, TelemetryRecord -│ │ ├── config.py # JarvisConfig dataclass hierarchy, TOML loader -│ │ └── events.py # Event bus: pub/sub for inter-pillar telemetry -│ ├── intelligence/ # Phase 1 -│ ├── memory/ # Phase 2 -│ ├── agents/ # Phase 3 -│ ├── engine/ # Phase 1 -│ ├── learning/ # Phase 4 -│ └── cli/ # CLI entry points -├── tests/ -├── VISION.md -├── ROADMAP.md -└── README.md -``` - -### Deliverables - -- [ ] **Registry system** — `RegistryBase[T]` adapted from IPW's `registry.py`. Typed subclasses: - - `ModelRegistry` — model specs and metadata - - `EngineRegistry` — inference engine implementations - - `MemoryRegistry` — memory backend implementations - - `AgentRegistry` — agent implementations - - `ToolRegistry` — tools with `ToolSpec` metadata (category, cost, latency, capabilities) - -- [ ] **Core types** (`core/types.py`): - - `Message` — role + content + metadata (tool calls, images, etc.) - - `Conversation` — ordered list of messages with sliding window support - - `ModelSpec` — model ID, parameter count, quantization, context length, hardware compatibility - - `ToolResult` — tool name + output + usage + cost - - `TelemetryRecord` — timestamp, model, tokens, latency, energy (optional), cost - -- [ ] **Config system** (`core/config.py`): - - `JarvisConfig` dataclass hierarchy: `EngineConfig`, `IntelligenceConfig`, `MemoryConfig`, `AgentConfig` - - TOML config file at `~/.openjarvis/config.toml` - - Hardware auto-detection: GPU vendor/model/VRAM/platform → populate defaults - -- [ ] **Event bus** (`core/events.py`): - - Simple pub/sub for inter-pillar communication - - Telemetry events flow without tight coupling between pillars - - Synchronous dispatch (async optional later) - -- [ ] **CLI skeleton** (Click-based): - - `jarvis init` — create `~/.openjarvis/config.toml` with auto-detected defaults - - `jarvis ask` — placeholder (wired in Phase 1) - - `jarvis serve` — placeholder (wired in Phase 3) - - `jarvis model list|pull|info` — placeholder (wired in Phase 1) - - `jarvis memory index|search|stats` — placeholder (wired in Phase 2) - ---- - -## Phase 1 — Intelligence + Inference Engine (~3-4 weeks) - -**Goal:** You can ask OpenJarvis a question and get an answer. Models run locally or via cloud APIs. Basic telemetry records every call. - -**Version milestone:** v0.2 — first usable version - -### Inference Engine - -- [ ] **`InferenceEngine` ABC:** - ```python - class InferenceEngine(ABC): - def generate(self, model: str, messages: list[Message], **params) -> Response: ... - def stream(self, model: str, messages: list[Message], **params) -> Iterator[ResponseChunk]: ... - def list_models(self) -> list[ModelSpec]: ... - def health(self) -> bool: ... - ``` - -- [ ] **Engine implementations:** - - `OllamaEngine` — wraps Ollama HTTP API (`/api/chat`, `/api/tags`). Apple Silicon + NVIDIA. - - `VLLMEngine` — wraps vLLM OpenAI-compatible API. Multi-GPU, tensor parallelism. - - `LlamaCppEngine` — wraps `llama-cpp-python` or llama.cpp server. Maximum compatibility. - - `CloudEngine` — unified wrapper for OpenAI, Anthropic, and Google APIs. Key-based routing. - -- [ ] **Model management:** - - Auto-discovery from running engines (poll `/api/tags`, `/v1/models`) - - `ModelSpec` with hardware compatibility matrix (min VRAM, supported engines, quantization options) - - `jarvis model list` shows all available models across engines - - `jarvis model info ` shows spec, hardware requirements, estimated performance - -### Intelligence - -- [ ] **Hardware profiles:** - - Auto-detect: `nvidia-smi`, `rocm-smi`, `system_profiler` (macOS), `/proc/cpuinfo` - - Map GPU to capabilities: VRAM, compute capability, FP8/FP4 support, unified memory - - Recommend engine: Apple Silicon → Ollama/MLX, NVIDIA datacenter → vLLM, AMD → vLLM+ROCm, CPU → llama.cpp - -- [ ] **Heuristic Router V0:** - - Rule-based routing: short queries (< 50 tokens) → small model, complex (reasoning keywords, multi-step) → large model, code patterns → code specialist - - Fallback chains: if preferred model unavailable, try next in chain - - Configurable via `~/.openjarvis/config.toml` - -- [ ] **Basic telemetry:** - - Wrap every `generate()` / `stream()` call with timing + token counting - - Record to SQLite: model, prompt tokens, completion tokens, latency, cost estimate - - `TelemetryRecord` stored via event bus, accumulated for future learning phase - -### Wire-up - -- [ ] **`jarvis ask "What is X?"` works end-to-end:** - 1. Parse query → detect complexity → route to model - 2. Generate response via selected engine - 3. Record telemetry - 4. Print response (with optional `--json` output) - ---- - -## Phase 2 — Memory / Storage (~3-4 weeks) - -**Goal:** OpenJarvis remembers conversations, can index your documents, and injects relevant context into prompts. - -**Version milestone:** v0.3 - -### Memory backends - -- [ ] **`MemoryBackend` ABC:** - ```python - class MemoryBackend(ABC): - def store(self, content: str, metadata: dict) -> str: ... # Returns doc ID - def retrieve(self, query: str, k: int = 10) -> list[Result]: ... - def delete(self, doc_id: str) -> bool: ... - def clear(self) -> None: ... - ``` - -- [ ] **Memory subtypes:** - - `ConversationMemory` — sliding window (configurable size) + automatic summarization of older turns via LLM call - - `KnowledgeBase` — indexed document collection with multi-backend search - -- [ ] **Backend implementations:** - - - **`SQLiteMemory`** — FTS5 full-text search, zero-config default. Always available, no extra dependencies. - - - **`FAISSMemory`** — Dense neural retrieval. Encodes documents with `sentence-transformers`, builds FAISS index (IVF or flat depending on collection size). GPU-accelerated when available. - - - **`ColBERTMemory`** — ColBERTv2 late interaction retrieval. Best retrieval quality. - - Package: `colbert-ai[torch,faiss-gpu]` - - Indexing: `Indexer(checkpoint="colbertv2.0", config=ColBERTConfig(nbits=2))` → `indexer.index(name, collection)` - - Search: `Searcher(index=name)` → `searcher.search(query, k=10)` returns `(passage_ids, ranks, scores)` - - Token-level MaxSim matching: each query token attends to each document token, max-pooled per query token, summed - - 2-bit residual compression keeps indexes compact (~50x smaller than full embeddings) - - Millisecond query latency, substantially better than single-vector methods on complex queries - - - **`BM25Memory`** — Keyword search baseline using `rank-bm25`. No GPU, no embeddings. Fast and effective for keyword-heavy queries. - - - **`HybridMemory`** — Combines BM25 with a dense backend (FAISS or ColBERT) using Reciprocal Rank Fusion (RRF): - ``` - RRF_score(d) = sum(1 / (k + rank_i(d))) for each retriever i - ``` - Configurable `k` parameter (default 60). Best overall retrieval when you don't know the query type. - -### Document pipeline - -- [ ] **Indexing pipeline:** PDF / Markdown / plain text / code → chunking (configurable size + overlap) → embedding (if using dense/ColBERT backend) → index -- [ ] **Context injection:** auto-retrieve top-k relevant chunks before each LLM call, inject into prompt with source attribution (`[Source: filename:line]`) - -### CLI - -- [ ] `jarvis memory index ` — index a file or directory -- [ ] `jarvis memory search ` — search across all memory backends -- [ ] `jarvis memory stats` — show index sizes, document counts, backend status - ---- - -## Phase 3 — Agentic Logic (~3-4 weeks) - -**Goal:** OpenJarvis can use tools, reason over multiple turns, and serve an OpenAI-compatible API. The default agent is OpenClaw's Pi. - -**Version milestone:** v0.4 - -### Agent framework - -- [ ] **`BaseAgent` ABC:** - ```python - class BaseAgent(ABC): - def run(self, input: str, context: AgentContext) -> AgentResult: ... - ``` - `AgentContext` carries: conversation history, memory handle, tool registry, telemetry recorder, model router. - `AgentResult` contains: response text, tool calls made, tokens used, telemetry data. - -- [ ] **Agent implementations:** - - - **`OpenClawAgent`** (default) — wraps OpenClaw's Pi agent runtime (`@mariozechner/pi-coding-agent` v0.52.12+). Two modes: - 1. **HTTP mode:** OpenClaw gateway running locally on `:18789`. Communicate via WebSocket. Best for persistent sessions. - 2. **Subprocess mode:** invoke `node` with `runEmbeddedPiAgent()` call, JSON over stdin/stdout. No gateway needed. - - Capabilities: multi-turn reasoning, tool calling, streaming, skill composition, context compaction - - Requires Node.js 22+ - - - **`SimpleAgent`** — single-turn: parse query → call model → return response. No tool calling, no multi-turn. Works without Node.js. Good for testing and simple Q&A. - - - **`OrchestratorAgent`** — multi-turn with model selection per step. Adapted from IPW's executor pattern. Each reasoning step can route to a different model (e.g., fast model for planning, large model for synthesis). - - - **`CustomAgent`** — template class for user-defined agents. Subclass `BaseAgent`, implement `run()`, register with `@AgentRegistry.register("my-agent")`. - -### Tool system - -- [ ] **`BaseTool` ABC:** - ```python - class BaseTool(ABC): - name: str - spec: ToolSpec # category, cost_estimate, latency_estimate, capabilities - def execute(self, input: str, **params) -> ToolResult: ... - ``` - -- [ ] **Built-in tools:** - - `Calculator` — evaluate math expressions - - `WebSearch` — search the web (Tavily, SearXNG, or DuckDuckGo) - - `CodeInterpreter` — execute Python in sandboxed environment - - `FileRead` / `FileWrite` — local file operations - - `Think` — internal reasoning scratchpad (zero-cost tool for chain-of-thought) - - `Retrieval` — wired to memory backends, returns relevant documents - - `LLMTool` — call another LLM as a tool (for model composition) - -- [ ] **`ToolRegistry`** with discovery: - - `ToolSpec` metadata: category, estimated latency, estimated cost, required API keys, capabilities list - - Auto-discover available tools based on installed packages and environment - -### API server - -- [ ] **OpenAI-compatible API server:** - - `POST /v1/chat/completions` — standard chat completion with tool use - - `GET /v1/models` — list available models - - Streaming via Server-Sent Events (SSE) - - `jarvis serve --port 8000 --agent openclaw` - -### CLI - -- [ ] `jarvis serve --port 8000 --agent ` — start API server -- [ ] `jarvis ask` now supports `--agent ` flag - ---- - -## Phase 4 — Learning Approach (placeholder) - -**Goal:** Stub interfaces for the learned router. No ML training in this phase — just the contracts and telemetry plumbing so everything is ready when we build it. - -**Version milestone:** v0.5 - -### Stubs - -- [ ] **`RouterPolicy` ABC:** - ```python - class RouterPolicy(ABC): - def select_model(self, query: str, context: RoutingContext) -> ModelSpec: ... - ``` - The heuristic router from Phase 1 implements this as the default. - -- [ ] **`RewardFunction` ABC:** - ```python - class RewardFunction(ABC): - def compute(self, trajectory: Trajectory) -> float: ... - ``` - Placeholder implementations: `QualityReward` (LLM-judge), `LatencyReward` (inverse latency), `EnergyReward` (inverse energy), `CostReward` (inverse cost), `CompositeReward` (weighted combination). - -- [ ] **`TelemetryAggregator`:** - - Reads `TelemetryRecord` entries from SQLite (accumulated since Phase 1) - - Computes per-model statistics: average latency, token throughput, cost, quality (when graded) - - Exports training-ready datasets for the future GRPO pipeline - -- [ ] **Design document:** `docs/learning-pipeline.md` describing the planned GRPO training pipeline: - - Trajectory generation from Phases 1-3 telemetry - - Reward model training - - Policy optimization with GRPO - - Online evaluation and rollout strategy - ---- - -## Phase 5 — Integration & Polish (~3-4 weeks) - -**Goal:** Production-ready packaging, OpenClaw integration, benchmarking, SDK, and documentation. - -**Version milestone:** v1.0 - -### OpenClaw integration - -- [ ] **`openjarvis-openclaw` plugin package:** - - `register()` hook implementing OpenClaw's plugin API - - `registerProvider()` — wraps OpenJarvis as an OpenClaw `ProviderPlugin` (routes through OpenJarvis intelligence + engine) - - `registerTool()` — exposes OpenJarvis tools to OpenClaw - - `MemorySearchManager` — implements OpenClaw's `search()` / `sync()` / `status()` interface, backed by OpenJarvis memory - -### Deployment - -- [ ] **Dockerfile** — multi-stage build with optional GPU support -- [ ] **docker-compose.yml** — OpenJarvis + Ollama/vLLM + optional gateway -- [ ] **Service files** — systemd (Linux) and launchd (macOS) for running as a system service - -### Python SDK - -- [ ] **Programmatic API:** - ```python - from openjarvis import Jarvis - - j = Jarvis() # Auto-loads config - response = await j.ask("Explain transformers") # Uses router + engine - await j.memory.index("~/papers/") # Index documents - results = await j.memory.search("attention mechanism") - ``` - -### Benchmarking - -- [ ] **`jarvis bench` CLI:** - - `BaseBenchmark` / `DatasetBenchmark` ABCs (adapted from IPW's `BenchmarkSuite`) - - Run benchmarks across models, measure accuracy + latency + energy - - Output JSONL results + summary JSON - -### Documentation - -- [ ] Documentation site (MkDocs or similar) -- [ ] Getting started guide -- [ ] Plugin development guide -- [ ] API reference - ---- - -## Phase 6 — Trace System & Learning (~ongoing) - -**Goal:** Full interaction-level trace recording, trace-driven learning, and pluggable agentic architectures. The foundation for studying local AI systems. - -**Version milestone:** v1.1 - -### Trace System (complete) - -- [x] **`Trace` and `TraceStep` types** — full interaction recording with step types: route, retrieve, generate, tool_call, respond -- [x] **`TraceStore`** — SQLite-backed append-only store with filtering (by agent, model, outcome, time range) -- [x] **`TraceCollector`** — wraps any `BaseAgent`, subscribes to EventBus, records steps automatically -- [x] **`TraceAnalyzer`** — read-only query layer: per-route stats, per-tool stats, summaries, query-type filtering, export -- [x] **`TraceDrivenPolicy`** — learns routing from trace outcomes, batch and online updates, registered as `"learned"` policy -- [x] **Event bus integration** — `TRACE_STEP` and `TRACE_COMPLETE` event types - -### Next Steps - -- [ ] Wire `TraceCollector` into SDK/CLI for automatic trace collection -- [ ] `jarvis trace` CLI subcommand (list, inspect, export traces) -- [ ] User feedback mechanisms (thumbs up/down, quality scores) -- [ ] Hierarchical memory (episodic/semantic/procedural layers) -- [ ] Pluggable agentic architectures (ReAct, tree-of-thought, custom loops) -- [ ] Prompt optimization from traces (DSPy-style compilation for local models) -- [ ] Model weight updates from traces (LoRA/QLoRA finetuning) -- [ ] GAIA benchmark evaluation with local models - ---- - -## Version Summary - -| Version | Phase | What you get | -|---------|-------|-------------| -| **v0.1** | Phase 0 | Scaffolding, registries, config, CLI skeleton | -| **v0.2** | Phase 1 | `jarvis ask` works — local & cloud inference with telemetry | -| **v0.3** | Phase 2 | Memory — index docs, conversation history, context injection | -| **v0.4** | Phase 3 | Agents + tools + OpenAI-compatible API server | -| **v0.5** | Phase 4 | Learning stubs — router policy interface, telemetry aggregation | -| **v1.0** | Phase 5 | Production — SDK, OpenClaw plugin, Docker, benchmarks, docs | -| **v1.1** | Phase 6 | Trace system, trace-driven learning, pluggable agent architectures | diff --git a/VISION.md b/VISION.md deleted file mode 100644 index c01930cc..00000000 --- a/VISION.md +++ /dev/null @@ -1,311 +0,0 @@ -# OpenJarvis - -**Programming abstractions for on-device AI.** - -OpenJarvis defines the abstractions needed to study and build AI systems that run entirely on local hardware. Instead of locking you into one model, one memory system, or one inference engine, OpenJarvis lets you compose your own stack across four core abstractions — then swap any piece without touching the rest. - -Built for researchers studying local AI systems and developers who want full control over their AI stack. Every interaction generates a trace; the system learns from its own usage to improve over time. - -## Why Local AI Needs New Abstractions - -Cloud AI treats intelligence as a **service** — you send a request, get a response, pay per token. Local AI treats intelligence as a **resource** — it lives on your machine, it's always available, it has fixed capabilities, it can be modified, and it accumulates state over time. This inversion changes everything: - -- **Fixed resource budget** — You have 8-24GB VRAM, period. Scheduling and allocation are first-class problems. -- **Persistent state is free, compute is expensive** — The opposite of cloud. You can store everything forever but can only run one model at a time. -- **You own the weights** — Fine-tuning, RL, prompt compilation are all possible on your data, your hardware, with immediate feedback loops. -- **Hardware heterogeneity** — Apple Silicon, NVIDIA consumer GPUs, AMD, CPU-only — each with different optimal strategies. -- **Every interaction is a learning signal** — Traces accumulate locally, enabling the system to learn routing, tool selection, and memory strategies from personal usage patterns. - ---- - -## The Five Pillars - -OpenJarvis is organized around five composable pillars. Each pillar defines a clear interface; implementations are discovered at runtime via a decorator-based registry system. - -``` -┌─────────────────────────────────────────────────────────────────────┐ -│ OpenJarvis │ -│ │ -│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ -│ │ Intelligence │ │ Learning │ │ Memory / │ │ -│ │ (Models) │ │ Approach │ │ Storage │ │ -│ │ │ │ (Router) │ │ │ │ -│ └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ │ -│ │ │ │ │ -│ ▼ ▼ ▼ │ -│ ┌──────────────────────────────────────────────┐ │ -│ │ Agentic Logic │ │ -│ │ (Orchestration, Tools, Reasoning) │ │ -│ └──────────────────────┬───────────────────────┘ │ -│ │ │ -│ ▼ │ -│ ┌──────────────────────────────────────────────┐ │ -│ │ Inference Engine │ │ -│ │ (vLLM, Ollama, llama.cpp, SGLang, MLX) │ │ -│ └──────────────────────────────────────────────┘ │ -└─────────────────────────────────────────────────────────────────────┘ -``` - ---- - -### 1. Intelligence (Model Layer) - -**What it does:** Manages available language models — local and cloud — and routes queries to the best model for the task. - -**What's pluggable:** Models, model providers, routing heuristics. - -**Supported at launch:** - -| Category | Models | -|----------|--------| -| Open-source (local) | Qwen3 8B, Qwen3 32B, GPT OSS 120B, Kimi-K2.5, MiniMax-M2.5 | -| Cloud APIs | Claude (Anthropic), GPT-4o / GPT-5 (OpenAI), Gemini (Google) | - -**Key components:** - -- **`ModelRegistry`** — decorator-based registry mapping model keys to `ModelSpec` objects (parameter count, quantization, hardware compatibility, context length) -- **Heuristic Router (V0)** — rule-based routing: short queries → small model, complex reasoning → large model, code → code specialist, fallback chains for unavailable models -- **Auto-discovery** — detects models from running inference engines (Ollama, vLLM) and available API keys - ---- - -### 2. Learning Approach (Router Policy) - -**What it does:** Determines *which* model handles a given query. Static policies use rules; learned policies update from interaction traces. - -**What's pluggable:** Routing policy, reward functions, training pipeline, trace analyzers. - -**Implemented:** -- **Heuristic routing** — rule-based routing based on query characteristics (length, complexity keywords, domain detection), fallback chains -- **Trace-driven routing** — learns from accumulated interaction traces which model/agent/tool combinations produce the best outcomes for different query types. Registered as `"learned"` policy. -- **Trace system** — every interaction generates a `Trace` recording the full sequence of steps (route, retrieve, generate, tool_call, respond) with timing, inputs, outputs, and outcomes. Stored in SQLite via `TraceStore`. -- **Trace analysis** — `TraceAnalyzer` computes per-route stats, per-tool stats, success rates, and query-type distributions from stored traces. - -**Future:** -- Learned router via GRPO (Group Relative Policy Optimization) -- Preference learning from user feedback -- Continual fine-tuning on accumulated trajectories -- Multi-objective optimization: quality vs. latency vs. energy vs. cost - ---- - -### 3. Memory / Storage - -**What it does:** Provides persistent, searchable memory across conversations, documents, and personal notes. Memory is automatically injected into prompts with source attribution. - -**What's pluggable:** Storage backends, retrieval strategies, embedding models, chunking strategies. - -**Memory types:** -- **Conversation Memory** — sliding window with automatic summarization of older turns -- **Knowledge Base** — indexed documents (PDF, Markdown, code, text) with multi-backend search -- **Personal Notes** — user-created persistent notes and preferences -- **Episodic Memory** — records of past interactions, tool uses, and outcomes - -**Backend implementations:** - -| Backend | Type | Description | -|---------|------|-------------| -| **SQLite** (default) | Keyword + FTS | FTS5 full-text search. Zero dependencies, zero config. Always available. | -| **FAISS** | Dense retrieval | Neural semantic search via `sentence-transformers` + FAISS indexes. | -| **ColBERTv2** | Late interaction | Token-level MaxSim matching with 2-bit residual compression. Best retrieval quality. Uses `colbert-ai` package with `Indexer` for offline indexing and `Searcher` for millisecond-latency queries. | -| **BM25** | Sparse retrieval | Classic keyword search baseline. Fast, no GPU needed. | -| **Hybrid** | Fusion | BM25 + dense (or ColBERT) with Reciprocal Rank Fusion (RRF). Best of both worlds. | -| **Vector DB adapters** | Dense retrieval | Qdrant, ChromaDB connectors for users with existing vector infrastructure. | - -**ColBERTv2 details:** -- Late interaction model: queries and documents are encoded independently, then matched at the token level via MaxSim -- 2-bit residual compression keeps indexes compact while preserving quality -- Offline indexing via `Indexer(checkpoint="colbertv2.0", config=ColBERTConfig(nbits=2))` -- Millisecond query latency via `Searcher(index=name).search(query, k=10)` -- Substantially better retrieval quality than single-vector dense methods on complex queries - ---- - -### 4. Agentic Logic - -**What it does:** Orchestrates multi-turn reasoning, tool calling, and task execution. The agent layer sits between the user and the model, managing context, tools, and conversation flow. - -**What's pluggable:** Agent implementations, tools, tool registries, execution strategies. - -**Agent implementations:** - -| Agent | Description | -|-------|-------------| -| **`OpenClawAgent`** (default) | Wraps OpenClaw's Pi agent runtime. Multi-turn reasoning, tool calling, streaming responses, skill composition, context compaction. Two modes: **HTTP** (WebSocket to OpenClaw gateway on `:18789`) or **subprocess** (invoke `node` with `runEmbeddedPiAgent()`, JSON over stdin/stdout). Requires Node.js 22+. | -| **`SimpleAgent`** | Single-turn: query → model → response. No tool calling. Works without Node.js. Good for quick answers and testing. | -| **`OrchestratorAgent`** | Multi-turn with per-step model selection. Adapted from IPW's executor pattern. Routes each reasoning step to the optimal model. | -| **`CustomAgent`** | Template for user-defined agent logic. Subclass `BaseAgent`, implement `run()`, register with `AgentRegistry`. | - -**Tool system:** -- `BaseTool` ABC with `ToolSpec` metadata (category, cost estimate, latency estimate, capabilities) -- `ToolRegistry` — runtime-discoverable tool catalog -- Built-in tools: Calculator, WebSearch, CodeInterpreter, FileRead/Write, Think, Retrieval (wired to memory backends), LLM-as-tool -- MCP (Model Context Protocol) compatible - -**API server:** -- OpenAI-compatible `/v1/chat/completions` and `/v1/models` endpoints -- Streaming via Server-Sent Events (SSE) -- Drop-in replacement for any OpenAI-compatible client - ---- - -### 5. Inference Engine - -**What it does:** Manages the actual LLM inference runtime — loading models, generating tokens, managing GPU memory. - -**What's pluggable:** Engine backends, hardware profiles, quantization strategies. - -**Supported engines:** - -| Engine | Best for | GPU | CPU | -|--------|----------|-----|-----| -| **vLLM** | High-throughput server, multi-GPU, production | NVIDIA, AMD | — | -| **SGLang** | Structured generation, constrained decoding | NVIDIA, AMD | — | -| **Ollama** | Easy setup, Apple Silicon, single-model | NVIDIA, Apple | Yes | -| **llama.cpp** | Maximum hardware compatibility, GGUF models | NVIDIA, AMD, Apple | Yes | -| **MLX** | Apple Silicon native, Metal acceleration | Apple | Apple | - -**Hardware auto-detection:** -- Detects GPU vendor (NVIDIA/AMD/Apple), model, VRAM, compute capability -- Recommends the best engine for detected hardware -- Apple Silicon → Ollama or MLX; NVIDIA datacenter → vLLM; AMD → vLLM with ROCm; CPU-only → llama.cpp - ---- - -## Query Flow - -``` -User query - │ - ▼ -┌──────────┐ ┌──────────┐ ┌──────────┐ -│ Agentic │────▶│ Memory │────▶│ Context │ -│ Logic │ │ Retrieve │ │ Inject │ -└────┬─────┘ └──────────┘ └────┬─────┘ - │ │ - ▼ ▼ -┌──────────┐ ┌──────────┐ ┌──────────┐ -│ Learning │────▶│ Model │────▶│ Inference│ -│ (Router) │ │ Select │ │ Engine │ -└──────────┘ └──────────┘ └────┬─────┘ - │ - ▼ - ┌──────────┐ - │ Response │ - │ + Telem. │ - └──────────┘ -``` - -1. **Agentic Logic** receives the user query, determines if tools or memory are needed -2. **Memory** retrieves relevant context (conversation history, documents, notes) -3. **Context Injection** assembles the full prompt with retrieved content and source attribution -4. **Learning/Router** selects the best model for this query based on routing policy (heuristic or trace-driven) -5. **Inference Engine** runs the selected model and streams the response -6. **Trace** records the full interaction sequence: every routing decision, memory retrieval, tool call, and generation step with timing and outcomes -7. **Learning** periodically updates routing policies from accumulated traces - ---- - -## User Scenarios - -### Developer on M4 Max MacBook Pro (128 GB unified memory) - -```toml -# ~/.openjarvis/config.toml -[engine] -backend = "ollama" # Native Apple Silicon support - -[intelligence] -default_model = "qwen3-32b" # Fits in 128 GB unified memory -fallback = "qwen3-8b" - -[memory] -backend = "sqlite" # Zero-config, always works -retrieval = "hybrid" # BM25 + FAISS for local docs - -[agent] -type = "openclaw" # Full agent capabilities -mode = "subprocess" # No separate gateway needed -``` - -Day-to-day: codes with `jarvis ask`, indexes project docs with `jarvis memory index`, runs a local OpenAI-compatible server with `jarvis serve` for editor integration. - -### Researcher on DGX Spark (2x B200, 384 GB GPU memory) - -```toml -[engine] -backend = "vllm" -tensor_parallel = 2 - -[intelligence] -default_model = "qwen3-235b-a22b" -router = "heuristic" # Route small queries to 8B, large to 235B - -[memory] -backend = "colbert" # Best retrieval quality for papers -knowledge_base = "~/papers/" - -[agent] -type = "orchestrator" # Multi-model orchestration -``` - -Running benchmarks with `jarvis bench`, profiling energy per query, comparing model efficiency across hardware configurations. - -### Privacy-Focused Offline Setup - -```toml -[engine] -backend = "llamacpp" # No server needed -network = "offline" - -[intelligence] -default_model = "qwen3-8b-q4" # Quantized to fit available RAM - -[memory] -backend = "sqlite" # Everything local -retrieval = "bm25" # No neural models needed - -[agent] -type = "simple" # No external dependencies -``` - -Fully air-gapped. No cloud APIs, no network calls, no telemetry export. All data stays on the machine. - ---- - -## Comparison - -| Feature | OpenJarvis | Ollama | LangChain | OpenClaw | vLLM | -|---------|-----------|--------|-----------|----------|------| -| **Focus** | Composable AI backend | Model runner | LLM app framework | AI coding assistant | Inference server | -| **Model management** | Multi-engine, auto-detect | Single engine | Bring your own | Cloud-first | Single engine | -| **Memory** | Multi-backend retrieval | None | Vector store wrappers | Conversation only | None | -| **Agents** | Pluggable (Pi, custom) | None | Chain-based | Pi agent (built-in) | None | -| **Inference** | vLLM/SGLang/Ollama/llama.cpp/MLX | Ollama only | External | External | vLLM only | -| **Hardware-aware** | Auto-detect + recommend | Manual | No | No | Manual | -| **Telemetry** | Energy, latency, cost | None | Callbacks | Basic | Metrics | -| **Offline** | Full support | Full support | Partial | No | Full support | -| **API** | OpenAI-compatible | OpenAI-compatible | Custom | Custom | OpenAI-compatible | -| **Language** | Python | Go | Python | TypeScript | Python | - -OpenJarvis is **not** a replacement for these tools — it *composes* them. Ollama and vLLM are inference engine options. OpenClaw's Pi agent is the default agentic logic. LangChain-style chains can be implemented as custom agents. - ---- - -## Design Principles - -1. **Pluggable everything** — every component is registered and discoverable at runtime. Swap models, engines, memory backends, and agents without code changes. - -2. **Registry-driven** — `RegistryBase[T]` pattern (adapted from IPW) provides type-safe, decorator-based registration for all extensible components: `ModelRegistry`, `EngineRegistry`, `MemoryRegistry`, `AgentRegistry`, `ToolRegistry`. - -3. **Offline-first** — works without network access. Cloud APIs are optional enhancements, never requirements. - -4. **Telemetry-native** — every inference call records timing, token counts, and (when hardware supports it) energy consumption. Data lands in SQLite for analysis. - -5. **Hardware-aware** — auto-detects GPU vendor, model, VRAM, and platform. Recommends the best engine and model configuration for your hardware. - -6. **Python-first** — core is pure Python (3.10+). Node.js required only for OpenClaw agent integration. No Java, no JVM, no heavy runtimes. - -7. **OpenAI-compatible API** — `jarvis serve` exposes `/v1/chat/completions` and `/v1/models`. Any client that speaks OpenAI protocol works out of the box. - -8. **Standalone** — OpenJarvis is a self-contained backend. OpenClaw is one possible frontend; so is `curl`, a Python SDK call, or any OpenAI-compatible client. diff --git a/desktop/src-tauri/tauri.conf.json b/desktop/src-tauri/tauri.conf.json index 7798f682..be84736e 100644 --- a/desktop/src-tauri/tauri.conf.json +++ b/desktop/src-tauri/tauri.conf.json @@ -58,7 +58,7 @@ "updater": { "pubkey": "dW50cnVzdGVkIGNvbW1lbnQ6IG1pbmlzaWduIHB1YmxpYyBrZXk6IDFFNzUzMzhEOEY2MjNEMDMKUldRRFBXS1BqVE4xSG8vK0lkUWN4WnZQYVIrbmc4RmpoOGlJWTBLTE15RlIya3JvQisvdUR3a0QK", "endpoints": [ - "https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/latest.json" + "https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/latest.json" ] } } diff --git a/docs/api/agents.md b/docs/api/agents.md deleted file mode 100644 index 4a894266..00000000 --- a/docs/api/agents.md +++ /dev/null @@ -1,92 +0,0 @@ -# Agents Module - -The agents module implements the agentic logic pillar. All agents implement -the `BaseAgent` ABC with a `run()` method. Agents handle queries by -coordinating tool calls, memory retrieval, and inference engine interactions. -The module also includes the OpenClaw infrastructure for interoperating with -external agent frameworks via HTTP or subprocess transport. - -## Abstract Base Classes and Context - -### BaseAgent - -::: openjarvis.agents._stubs.BaseAgent - options: - show_source: true - members_order: source - -### ToolUsingAgent - -::: openjarvis.agents._stubs.ToolUsingAgent - options: - show_source: true - members_order: source - -### AgentContext - -::: openjarvis.agents._stubs.AgentContext - options: - show_source: true - members_order: source - -### AgentResult - -::: openjarvis.agents._stubs.AgentResult - options: - show_source: true - members_order: source - ---- - -## Agent Implementations - -### SimpleAgent - -::: openjarvis.agents.simple.SimpleAgent - options: - show_source: true - members_order: source - -### OrchestratorAgent - -::: openjarvis.agents.orchestrator.OrchestratorAgent - options: - show_source: true - members_order: source - -### NativeReActAgent - -::: openjarvis.agents.native_react.NativeReActAgent - options: - show_source: true - members_order: source - -### NativeOpenHandsAgent - -::: openjarvis.agents.native_openhands.NativeOpenHandsAgent - options: - show_source: true - members_order: source - -### RLMAgent - -::: openjarvis.agents.rlm.RLMAgent - options: - show_source: true - members_order: source - -### OpenHandsAgent - -::: openjarvis.agents.openhands.OpenHandsAgent - options: - show_source: true - members_order: source - - - -!!! note "OpenClaw Infrastructure" - The OpenClaw protocol, transport, and plugin modules (`openclaw_protocol.py`, - `openclaw_transport.py`, `openclaw_plugin.py`, `openclaw.py`) are part of the - OpenClaw agent infrastructure and require the `openjarvis[openclaw]` extra. - See the [architecture documentation](../architecture/agents.md#openclaw-infrastructure) - for protocol and transport details. diff --git a/docs/api/bench.md b/docs/api/bench.md deleted file mode 100644 index f92012f0..00000000 --- a/docs/api/bench.md +++ /dev/null @@ -1,48 +0,0 @@ -# Benchmarks Module - -The benchmarks module provides a framework for measuring inference engine -performance. All benchmarks implement the `BaseBenchmark` ABC and are -registered via `BenchmarkRegistry`. The `BenchmarkSuite` runner executes -a collection of benchmarks and aggregates results into JSONL or summary -format. - -## Abstract Base Class and Runner - -### BaseBenchmark - -::: openjarvis.bench._stubs.BaseBenchmark - options: - show_source: true - members_order: source - -### BenchmarkResult - -::: openjarvis.bench._stubs.BenchmarkResult - options: - show_source: true - members_order: source - -### BenchmarkSuite - -::: openjarvis.bench._stubs.BenchmarkSuite - options: - show_source: true - members_order: source - ---- - -## Benchmark Implementations - -### LatencyBenchmark - -::: openjarvis.bench.latency.LatencyBenchmark - options: - show_source: true - members_order: source - -### ThroughputBenchmark - -::: openjarvis.bench.throughput.ThroughputBenchmark - options: - show_source: true - members_order: source diff --git a/docs/api/channels.md b/docs/api/channels.md deleted file mode 100644 index 6bd3cb28..00000000 --- a/docs/api/channels.md +++ /dev/null @@ -1,128 +0,0 @@ -# API Reference: Channels - -The `openjarvis.channels` package provides the channel messaging abstraction and the OpenClaw gateway bridge. All public classes and types are documented below. - -For usage examples, CLI commands, and configuration, see the [Channels user guide](../user-guide/channels.md). For the architectural design and listener loop internals, see [Channels architecture](../architecture/channels.md). - ---- - -## Types and Enums - -### ChannelStatus - -Connection status values for a channel. Returned by `BaseChannel.status()`. - -::: openjarvis.channels._stubs.ChannelStatus - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### ChannelMessage - -Dataclass representing a message received from or sent to a channel. All fields correspond to the JSON payload exchanged with the gateway. - -::: openjarvis.channels._stubs.ChannelMessage - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### ChannelHandler - -Type alias for message handler callbacks: - -```python -ChannelHandler = Callable[[ChannelMessage], Optional[str]] -``` - -Handlers are called synchronously from the listener thread when a message arrives. The optional `str` return value is reserved for future auto-reply routing and has no effect in the current implementation. - -!!! warning "Thread safety" - Handlers run on the listener thread, not the caller thread. Protect shared state with locks or `queue.Queue`. - -::: openjarvis.channels._stubs.ChannelHandler - options: - show_source: true - show_root_heading: true - heading_level: 4 - ---- - -## BaseChannel - -Abstract base class for all channel implementations. Subclasses must implement all six abstract methods and register via `@ChannelRegistry.register("name")`. - -```python title="custom_channel.py" -from openjarvis.channels._stubs import BaseChannel, ChannelMessage, ChannelStatus -from openjarvis.core.registry import ChannelRegistry - - -@ChannelRegistry.register("my-channel") -class MyChannel(BaseChannel): - channel_id = "my-channel" - - def connect(self) -> None: ... - def disconnect(self) -> None: ... - def send(self, channel, content, *, conversation_id="", metadata=None) -> bool: ... - def status(self) -> ChannelStatus: ... - def list_channels(self) -> list[str]: ... - def on_message(self, handler) -> None: ... -``` - -::: openjarvis.channels._stubs.BaseChannel - options: - show_source: true - show_root_heading: true - heading_level: 3 - ---- - -## OpenClawChannelBridge - -`OpenClawChannelBridge` connects to the OpenClaw gateway over WebSocket, with automatic HTTP fallback when the `websockets` package is not installed or a WebSocket send fails. It is registered as `"openclaw"` in `ChannelRegistry`. - -```python title="bridge_example.py" -from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge -from openjarvis.channels._stubs import ChannelMessage -from openjarvis.core.events import EventBus - -bus = EventBus() -bridge = OpenClawChannelBridge( - gateway_url="ws://127.0.0.1:18789/ws", - reconnect_interval=5.0, - bus=bus, -) - - -def on_message(msg: ChannelMessage) -> None: - print(f"[{msg.channel}] {msg.sender}: {msg.content}") - - -bridge.on_message(on_message) -bridge.connect() - -bridge.send("notifications", "Hello!") -channels = bridge.list_channels() - -bridge.disconnect() -``` - -### Constructor Parameters - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `gateway_url` | `str` | `ws://127.0.0.1:18789/ws` | WebSocket URL of the OpenClaw gateway | -| `reconnect_interval` | `float` | `5.0` | Seconds to wait before reconnecting after a disconnect | -| `bus` | `EventBus` | `None` | Event bus for publishing `CHANNEL_MESSAGE_RECEIVED` and `CHANNEL_MESSAGE_SENT` events | - -### Events Published - -| Event | When | -|-------|------| -| `CHANNEL_MESSAGE_RECEIVED` | Message received from gateway WebSocket | -| `CHANNEL_MESSAGE_SENT` | Message successfully delivered via WebSocket or HTTP | - -!!! note "Optional dependency" - `OpenClawChannelBridge` requires the `openjarvis[openclaw]` extra. - See the [Channels architecture](../architecture/channels.md) for design details. diff --git a/docs/api/core.md b/docs/api/core.md deleted file mode 100644 index b1c23f02..00000000 --- a/docs/api/core.md +++ /dev/null @@ -1,339 +0,0 @@ -# Core Module - -The core module contains the foundational abstractions shared across all -OpenJarvis pillars: the decorator-based registry system for runtime discovery, -canonical data types, configuration loading with hardware detection, and the -pub/sub event bus for inter-pillar communication. - -## Registry System - -The registry pattern is the primary extension mechanism. Each pillar has its -own typed registry subclass. New implementations are added by decorating a -class with `@XRegistry.register("name")`. - -### RegistryBase - -::: openjarvis.core.registry.RegistryBase - options: - show_source: true - members_order: source - -### ModelRegistry - -::: openjarvis.core.registry.ModelRegistry - options: - show_source: true - members_order: source - -### EngineRegistry - -::: openjarvis.core.registry.EngineRegistry - options: - show_source: true - members_order: source - -### MemoryRegistry - -::: openjarvis.core.registry.MemoryRegistry - options: - show_source: true - members_order: source - -### AgentRegistry - -::: openjarvis.core.registry.AgentRegistry - options: - show_source: true - members_order: source - -### ToolRegistry - -::: openjarvis.core.registry.ToolRegistry - options: - show_source: true - members_order: source - -### RouterPolicyRegistry - -::: openjarvis.core.registry.RouterPolicyRegistry - options: - show_source: true - members_order: source - -### BenchmarkRegistry - -::: openjarvis.core.registry.BenchmarkRegistry - options: - show_source: true - members_order: source - -### ChannelRegistry - -::: openjarvis.core.registry.ChannelRegistry - options: - show_source: true - members_order: source - -### LearningRegistry - -::: openjarvis.core.registry.LearningRegistry - options: - show_source: true - members_order: source - ---- - -## Types - -Canonical data types shared across all OpenJarvis pillars, including message -structures, model specifications, telemetry records, and trace objects. - -### Role - -::: openjarvis.core.types.Role - options: - show_source: true - members_order: source - -### Quantization - -::: openjarvis.core.types.Quantization - options: - show_source: true - members_order: source - -### StepType - -::: openjarvis.core.types.StepType - options: - show_source: true - members_order: source - -### ToolCall - -::: openjarvis.core.types.ToolCall - options: - show_source: true - members_order: source - -### Message - -::: openjarvis.core.types.Message - options: - show_source: true - members_order: source - -### Conversation - -::: openjarvis.core.types.Conversation - options: - show_source: true - members_order: source - -### ModelSpec - -::: openjarvis.core.types.ModelSpec - options: - show_source: true - members_order: source - -### ToolResult - -::: openjarvis.core.types.ToolResult - options: - show_source: true - members_order: source - -### TelemetryRecord - -::: openjarvis.core.types.TelemetryRecord - options: - show_source: true - members_order: source - -### TraceStep - -::: openjarvis.core.types.TraceStep - options: - show_source: true - members_order: source - -### Trace - -::: openjarvis.core.types.Trace - options: - show_source: true - members_order: source - -### RoutingContext - -::: openjarvis.core.types.RoutingContext - options: - show_source: true - members_order: source - ---- - -## Configuration - -Configuration loading, hardware detection, and engine recommendation. User -configuration lives at `~/.openjarvis/config.toml`. The `load_config()` -function detects hardware, fills sensible defaults, then overlays any user -overrides from the TOML file. - -### JarvisConfig - -::: openjarvis.core.config.JarvisConfig - options: - show_source: true - members_order: source - -### EngineConfig - -::: openjarvis.core.config.EngineConfig - options: - show_source: true - members_order: source - -### IntelligenceConfig - -::: openjarvis.core.config.IntelligenceConfig - options: - show_source: true - members_order: source - -### LearningConfig - -::: openjarvis.core.config.LearningConfig - options: - show_source: true - members_order: source - -### MemoryConfig - -::: openjarvis.core.config.MemoryConfig - options: - show_source: true - members_order: source - -### AgentConfig - -::: openjarvis.core.config.AgentConfig - options: - show_source: true - members_order: source - -### ServerConfig - -::: openjarvis.core.config.ServerConfig - options: - show_source: true - members_order: source - -### TelemetryConfig - -::: openjarvis.core.config.TelemetryConfig - options: - show_source: true - members_order: source - -### TracesConfig - -::: openjarvis.core.config.TracesConfig - options: - show_source: true - members_order: source - -### StorageConfig - -::: openjarvis.core.config.StorageConfig - options: - show_source: true - members_order: source - -### MCPConfig - -::: openjarvis.core.config.MCPConfig - options: - show_source: true - members_order: source - -### ToolsConfig - -::: openjarvis.core.config.ToolsConfig - options: - show_source: true - members_order: source - -### HardwareInfo - -::: openjarvis.core.config.HardwareInfo - options: - show_source: true - members_order: source - -### GpuInfo - -::: openjarvis.core.config.GpuInfo - options: - show_source: true - members_order: source - -### detect_hardware - -::: openjarvis.core.config.detect_hardware - options: - show_source: true - -### load_config - -::: openjarvis.core.config.load_config - options: - show_source: true - -### recommend_engine - -::: openjarvis.core.config.recommend_engine - options: - show_source: true - ---- - -## Event Bus - -Thread-safe pub/sub event bus for inter-pillar telemetry. Any pillar can emit -events (e.g. `INFERENCE_END`) and any other pillar can react without direct -coupling. - -### EventType - -::: openjarvis.core.events.EventType - options: - show_source: true - members_order: source - -### Event - -::: openjarvis.core.events.Event - options: - show_source: true - members_order: source - -### EventBus - -::: openjarvis.core.events.EventBus - options: - show_source: true - members_order: source - -### get_event_bus - -::: openjarvis.core.events.get_event_bus - options: - show_source: true - -### reset_event_bus - -::: openjarvis.core.events.reset_event_bus - options: - show_source: true diff --git a/docs/api/engine.md b/docs/api/engine.md deleted file mode 100644 index dd7b17ee..00000000 --- a/docs/api/engine.md +++ /dev/null @@ -1,96 +0,0 @@ -# Engine Module - -The engine module implements the inference runtime pillar. All backends -implement the `InferenceEngine` ABC with `generate()`, `stream()`, -`list_models()`, and `health()` methods. The discovery subsystem probes -running engines and selects the best available backend based on -configuration and health checks. - -## Abstract Base Class - -### InferenceEngine - -::: openjarvis.engine._stubs.InferenceEngine - options: - show_source: true - members_order: source - -### EngineConnectionError - -::: openjarvis.engine._base.EngineConnectionError - options: - show_source: true - -### messages_to_dicts - -::: openjarvis.engine._base.messages_to_dicts - options: - show_source: true - ---- - -## Engine Implementations - -### OllamaEngine - -::: openjarvis.engine.ollama.OllamaEngine - options: - show_source: true - members_order: source - -### OpenAI-Compatible Engines (vLLM, SGLang, llama.cpp, MLX, LM Studio) - -These engines are thin subclasses of `_OpenAICompatibleEngine`, created -via data-driven registration in `openai_compat_engines.py`. They share -the same implementation and differ only in `engine_id` and default host. - -| Engine | Registry Key | Default Host | -|--------|-------------|--------------| -| `VLLMEngine` | `vllm` | `http://localhost:8000` | -| `SGLangEngine` | `sglang` | `http://localhost:30000` | -| `LlamaCppEngine` | `llamacpp` | `http://localhost:8080` | -| `MLXEngine` | `mlx` | `http://localhost:8080` | -| `LMStudioEngine` | `lmstudio` | `http://localhost:1234` | - -::: openjarvis.engine._openai_compat._OpenAICompatibleEngine - options: - show_source: true - members_order: source - -### CloudEngine - -::: openjarvis.engine.cloud.CloudEngine - options: - show_source: true - members_order: source - -### estimate_cost - -::: openjarvis.engine.cloud.estimate_cost - options: - show_source: true - ---- - -## Engine Discovery - -Functions for probing running engines, aggregating available models, -and selecting the best engine for a given configuration. - -### get_engine - -::: openjarvis.engine._discovery.get_engine - options: - show_source: true - -### discover_engines - -::: openjarvis.engine._discovery.discover_engines - options: - show_source: true - -### discover_models - -::: openjarvis.engine._discovery.discover_models - options: - show_source: true diff --git a/docs/api/evals.md b/docs/api/evals.md deleted file mode 100644 index dfa1b47b..00000000 --- a/docs/api/evals.md +++ /dev/null @@ -1,1213 +0,0 @@ -# Evaluation Framework - -The `openjarvis-evals` package provides a structured harness for measuring model quality -across research benchmarks. It is a separate package from the main `openjarvis` library, -installed from the `evals/` directory, and exposes a CLI (`openjarvis-eval`) plus a -Python API for programmatic use. - -The framework is organized around four ABCs — `InferenceBackend`, `DatasetProvider`, -`Scorer`, and the concrete `EvalRunner` — wired together by `RunConfig`. A TOML-based -suite configuration system expands a models-by-benchmarks matrix into individual -`RunConfig` objects so an entire comparison table can be launched from a single file. - -!!! note "Installation" - The evaluation framework is a separate package. Install it from the repository root: - - ```bash - cd evals/ - uv pip install -e ".[dev]" - # or - pip install -e ".[dev]" - ``` - - The package requires Python 3.10+, `openjarvis>=1.0.0`, and `datasets>=2.14`. - ---- - -## Core Types (`evals.core.types`) - -These dataclasses are the shared vocabulary for every component in the framework. - -### EvalRecord - -A single evaluation sample loaded from a dataset. - -```python -from evals.core.types import EvalRecord -``` - -| Field | Type | Default | Description | -|-------|------|---------|-------------| -| `record_id` | `str` | — | Unique identifier for this sample | -| `problem` | `str` | — | The prompt or question presented to the model | -| `reference` | `str` | — | Ground-truth answer used for scoring | -| `category` | `str` | — | Task category: `"chat"`, `"reasoning"`, `"rag"`, or `"agentic"` | -| `subject` | `str` | `""` | Subject area or sub-topic within the benchmark | -| `metadata` | `Dict[str, Any]` | `{}` | Benchmark-specific extra fields (options, difficulty, file paths, etc.) | - -```python -record = EvalRecord( - record_id="supergpqa-0", - problem="What is the capital of France?\nOptions:\nA. Berlin\nB. Paris\nC. Madrid", - reference="B", - category="reasoning", - subject="geography", - metadata={"difficulty": "easy", "options": ["Berlin", "Paris", "Madrid"]}, -) -``` - ---- - -### EvalResult - -The result of running inference on a single `EvalRecord`. - -```python -from evals.core.types import EvalResult -``` - -| Field | Type | Default | Description | -|-------|------|---------|-------------| -| `record_id` | `str` | — | Matches the source `EvalRecord.record_id` | -| `model_answer` | `str` | — | Raw text output from the model | -| `is_correct` | `Optional[bool]` | `None` | Scoring verdict; `None` if scoring could not be determined | -| `score` | `Optional[float]` | `None` | Numeric score (typically `1.0` / `0.0`); may be `None` if `is_correct` is `None` | -| `latency_seconds` | `float` | `0.0` | Wall-clock generation time | -| `prompt_tokens` | `int` | `0` | Input token count from usage metadata | -| `completion_tokens` | `int` | `0` | Output token count from usage metadata | -| `cost_usd` | `float` | `0.0` | Estimated inference cost in USD | -| `error` | `Optional[str]` | `None` | Exception message if inference or scoring failed | -| `scoring_metadata` | `Dict[str, Any]` | `{}` | Scorer-specific details (extracted letter, judge output, match type, etc.) | -| `ttft` | `float` | `0.0` | Time to first token (seconds) | -| `energy_joules` | `float` | `0.0` | GPU energy consumed (joules) | -| `power_watts` | `float` | `0.0` | Average GPU power draw (watts) | -| `gpu_utilization_pct` | `float` | `0.0` | Average GPU utilization (%) | -| `throughput_tok_per_sec` | `float` | `0.0` | Output token throughput (tokens/sec) | -| `mfu_pct` | `float` | `0.0` | Model FLOPs Utilization (%) | -| `mbu_pct` | `float` | `0.0` | Memory Bandwidth Utilization (%) | -| `ipw` | `float` | `0.0` | Intelligence Per Watt: `accuracy / power_watts` | -| `ipj` | `float` | `0.0` | Intelligence Per Joule: `accuracy / energy_joules` | - -!!! tip "Distinguishing errors from wrong answers" - A non-`None` `error` field means inference itself failed. When `error` is `None` but - `is_correct` is `None`, scoring was attempted but the scorer could not determine a - verdict (for example, the judge returned an unparseable response). - ---- - -### RunConfig - -Configuration for a single evaluation run (one model on one benchmark). - -```python -from evals.core.types import RunConfig -``` - -| Field | Type | Default | Description | -|-------|------|---------|-------------| -| `benchmark` | `str` | — | Benchmark name: `"supergpqa"`, `"gaia"`, `"frames"`, or `"wildchat"` | -| `backend` | `str` | — | Backend identifier: `"jarvis-direct"` or `"jarvis-agent"` | -| `model` | `str` | — | Model identifier passed to the backend (e.g., `"qwen3:8b"`, `"gpt-4o"`) | -| `max_samples` | `Optional[int]` | `None` | Limit the dataset to this many records; `None` uses the full dataset | -| `max_workers` | `int` | `4` | Number of parallel threads for inference | -| `temperature` | `float` | `0.0` | Sampling temperature | -| `max_tokens` | `int` | `2048` | Maximum output tokens per sample | -| `judge_model` | `str` | `"gpt-5-mini-2025-08-07"` | Model identifier used by the LLM judge scorer | -| `engine_key` | `Optional[str]` | `None` | Override the OpenJarvis engine (`"ollama"`, `"vllm"`, `"cloud"`, etc.) | -| `agent_name` | `Optional[str]` | `None` | Agent name for `jarvis-agent` backend; defaults to `"orchestrator"` | -| `tools` | `List[str]` | `[]` | Tool names enabled for the agent (e.g., `["calculator", "file_read"]`) | -| `output_path` | `Optional[str]` | `None` | JSONL output file path; auto-generated from benchmark and model name if `None` | -| `seed` | `int` | `42` | Random seed for dataset shuffling | -| `dataset_split` | `Optional[str]` | `None` | Override the dataset split (e.g., `"validation"`, `"test"`) | -| `telemetry` | `bool` | `False` | Enable GPU telemetry capture | -| `gpu_metrics` | `bool` | `False` | Enable GPU metric polling via `pynvml` | -| `metadata` | `Dict[str, Any]` | `{}` | Model hardware metadata for efficiency calculations (populated by `expand_suite()`) | - -```python -config = RunConfig( - benchmark="supergpqa", - backend="jarvis-direct", - model="qwen3:8b", - max_samples=100, - max_workers=8, - engine_key="ollama", - output_path="results/supergpqa_qwen3-8b.jsonl", -) -``` - ---- - -### RunSummary - -Aggregate statistics produced by `EvalRunner.run()` at the end of a completed run. - -```python -from evals.core.types import RunSummary -``` - -| Field | Type | Default | Description | -|-------|------|---------|-------------| -| `benchmark` | `str` | — | Benchmark name | -| `category` | `str` | — | Task category (inferred from records; falls back to `benchmark` name) | -| `backend` | `str` | — | Backend used | -| `model` | `str` | — | Model identifier | -| `total_samples` | `int` | — | Total records processed (including errors) | -| `scored_samples` | `int` | — | Records where `is_correct` is not `None` | -| `correct` | `int` | — | Records where `is_correct` is `True` | -| `accuracy` | `float` | — | `correct / scored_samples`; rounded to 4 decimal places | -| `errors` | `int` | — | Records where inference or scoring raised an exception | -| `mean_latency_seconds` | `float` | — | Mean wall-clock latency across all successful inferences | -| `total_cost_usd` | `float` | — | Sum of `cost_usd` across all records | -| `per_subject` | `Dict[str, Dict[str, float]]` | `{}` | Per-subject breakdown: `{subject: {accuracy, total, scored, correct}}` | -| `started_at` | `float` | `0.0` | Unix timestamp at run start | -| `ended_at` | `float` | `0.0` | Unix timestamp at run end | -| `accuracy_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for per-sample accuracy (binary 0/1) | -| `latency_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for inference latency | -| `ttft_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for time-to-first-token | -| `energy_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for GPU energy (joules) | -| `power_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for GPU power (watts) | -| `gpu_utilization_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for GPU utilization (%) | -| `throughput_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for token throughput | -| `mfu_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for Model FLOPs Utilization (%) | -| `mbu_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for Memory Bandwidth Utilization (%) | -| `ipw_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for Intelligence Per Watt | -| `ipj_stats` | `Optional[MetricStats]` | `None` | Descriptive statistics for Intelligence Per Joule | -| `total_energy_joules` | `float` | `0.0` | Total GPU energy consumed across all samples | - -The runner also writes a `.summary.json` file alongside the JSONL output, containing -the serialized `RunSummary`. - ---- - -### MetricStats - -Descriptive statistics for a single metric across samples. - -```python -from evals.core.types import MetricStats -``` - -| Field | Type | Default | Description | -|-------|------|---------|-------------| -| `mean` | `float` | `0.0` | Arithmetic mean | -| `median` | `float` | `0.0` | Median value | -| `min` | `float` | `0.0` | Minimum value | -| `max` | `float` | `0.0` | Maximum value | -| `std` | `float` | `0.0` | Standard deviation (0.0 for single-element lists) | - -`MetricStats` is computed by `_metric_stats()` in the runner and serialized to -JSON by `_metric_stats_to_dict()`. Fields in `RunSummary` like `accuracy_stats`, -`energy_stats`, `mfu_stats`, etc. are `Optional[MetricStats]` — they are `None` -when no positive values were observed for that metric. - ---- - -## Suite Config Types (`evals.core.types`) - -These dataclasses map directly to sections in a TOML eval suite config file. -They are populated by `load_eval_config()` and consumed by `expand_suite()`. - -### MetaConfig - -```python -@dataclass -class MetaConfig: - name: str = "" - description: str = "" -``` - -Maps to the `[meta]` TOML section. Both fields are optional and used only for -display output in the CLI. - ---- - -### DefaultsConfig - -```python -@dataclass -class DefaultsConfig: - temperature: float = 0.0 - max_tokens: int = 2048 -``` - -Maps to `[defaults]`. These values are the lowest-priority settings in the merge -precedence: `benchmark-level > model-level > [defaults] > built-in defaults`. - ---- - -### JudgeConfig - -```python -@dataclass -class JudgeConfig: - model: str = "gpt-4o" - provider: Optional[str] = None - temperature: float = 0.0 - max_tokens: int = 1024 -``` - -Maps to `[judge]`. The judge model is used by LLM-as-judge scorers (GAIA, FRAMES, -WildChat, SuperGPQA). The `provider` field is reserved for future routing; currently -the judge backend is always constructed with `engine_key="cloud"`. - ---- - -### ExecutionConfig - -```python -@dataclass -class ExecutionConfig: - max_workers: int = 4 - output_dir: str = "results/" - seed: int = 42 - telemetry: bool = False - gpu_metrics: bool = False -``` - -Maps to `[run]`. `output_dir` is the base directory for all JSONL output files; -individual filenames are auto-generated as `{benchmark}_{model-slug}.jsonl`. -When `telemetry` is enabled, the runner captures GPU energy, power, utilization, -and throughput per sample via `InstrumentedEngine`. When `gpu_metrics` is enabled, -`GpuMonitor` polls GPU sensors via `pynvml` during inference. - ---- - -### ModelConfig - -```python -@dataclass -class ModelConfig: - name: str = "" - engine: Optional[str] = None - provider: Optional[str] = None - temperature: Optional[float] = None - max_tokens: Optional[int] = None - param_count_b: float = 0.0 - active_params_b: Optional[float] = None - gpu_peak_tflops: float = 0.0 - gpu_peak_bandwidth_gb_s: float = 0.0 - num_gpus: int = 1 -``` - -Maps to each `[[models]]` entry. `name` is required. `temperature` and `max_tokens` -override `[defaults]` for every benchmark this model runs against, unless a -benchmark-level override also exists. The hardware parameters (`param_count_b`, -`active_params_b`, `gpu_peak_tflops`, `gpu_peak_bandwidth_gb_s`, `num_gpus`) are -used to compute MFU (Model FLOPs Utilization) and MBU (Memory Bandwidth Utilization) -per sample. These are flowed into `RunConfig.metadata` by `expand_suite()`. - ---- - -### BenchmarkConfig - -```python -@dataclass -class BenchmarkConfig: - name: str = "" - backend: str = "jarvis-direct" - max_samples: Optional[int] = None - split: Optional[str] = None - agent: Optional[str] = None - tools: List[str] = field(default_factory=list) - judge_model: Optional[str] = None - temperature: Optional[float] = None - max_tokens: Optional[int] = None -``` - -Maps to each `[[benchmarks]]` entry. `name` is required. `backend` must be one of -`"jarvis-direct"` or `"jarvis-agent"`. `judge_model` overrides `[judge].model` for -this benchmark only. - ---- - -### EvalSuiteConfig - -The top-level config object returned by `load_eval_config()`. - -```python -@dataclass -class EvalSuiteConfig: - meta: MetaConfig - defaults: DefaultsConfig - judge: JudgeConfig - run: ExecutionConfig - models: List[ModelConfig] - benchmarks: List[BenchmarkConfig] -``` - -`expand_suite(suite)` iterates over `models x benchmarks` to produce one `RunConfig` -per pair, applying the merge precedence rules documented in `DefaultsConfig`. - ---- - -## Config Module (`evals.core.config`) - -```python -from evals.core.config import load_eval_config, expand_suite, EvalConfigError -``` - -### EvalConfigError - -```python -class EvalConfigError(Exception): ... -``` - -Raised by `load_eval_config()` for structural validation failures: missing required -fields, invalid backend names, or empty `[[models]]` / `[[benchmarks]]` lists. - ---- - -### load_eval_config - -```python -def load_eval_config(path: str | Path) -> EvalSuiteConfig -``` - -Load and validate an eval suite configuration from a TOML file. - -Uses the standard library `tomllib` on Python 3.11+ and the `tomli` backport on -Python 3.10. - -**Parameters:** - -| Parameter | Type | Description | -|-----------|------|-------------| -| `path` | `str \| Path` | Path to the TOML config file | - -**Returns:** `EvalSuiteConfig` - -**Raises:** - -- `EvalConfigError` — structural validation failures (missing `name`, invalid backend, no models/benchmarks defined) -- `FileNotFoundError` — if the config file does not exist - -```python -from evals.core.config import load_eval_config - -suite = load_eval_config("evals/configs/full-suite.toml") -print(f"{len(suite.models)} models, {len(suite.benchmarks)} benchmarks") -``` - ---- - -### expand_suite - -```python -def expand_suite(suite: EvalSuiteConfig) -> List[RunConfig] -``` - -Expand an `EvalSuiteConfig` into a flat list of `RunConfig` objects, one per -model-benchmark pair, with all override layers merged. - -**Merge precedence (highest wins):** - -1. Benchmark-level (`BenchmarkConfig.temperature`, `.max_tokens`, `.judge_model`) -2. Model-level (`ModelConfig.temperature`, `.max_tokens`) -3. Suite defaults (`DefaultsConfig`) -4. Built-in dataclass defaults - -Output paths are auto-generated as `{output_dir}/{benchmark}_{model-slug}.jsonl`, -where `model-slug` replaces `/` and `:` with `-`. - -**Parameters:** - -| Parameter | Type | Description | -|-----------|------|-------------| -| `suite` | `EvalSuiteConfig` | Parsed suite configuration | - -**Returns:** `List[RunConfig]` — one entry per model-benchmark combination. - -```python -from evals.core.config import load_eval_config, expand_suite - -suite = load_eval_config("evals/configs/full-suite.toml") -run_configs = expand_suite(suite) # e.g., 3 models x 4 benchmarks = 12 RunConfigs -for rc in run_configs: - print(f"{rc.benchmark} / {rc.model} -> {rc.output_path}") -``` - ---- - -## Abstract Base Classes - -### InferenceBackend (`evals.core.backend`) - -```python -from evals.core.backend import InferenceBackend -``` - -Base class for all inference backends. A backend wraps an engine or agent and -provides a uniform text-in / text-out interface for the runner. - -```python -class InferenceBackend(ABC): - backend_id: str -``` - -**Class attribute:** - -| Attribute | Type | Description | -|-----------|------|-------------| -| `backend_id` | `str` | Registry identifier (e.g., `"jarvis-direct"`, `"jarvis-agent"`) | - -**Abstract methods:** - -#### generate - -```python -@abstractmethod -def generate( - self, - prompt: str, - *, - model: str, - system: str = "", - temperature: float = 0.0, - max_tokens: int = 2048, -) -> str -``` - -Generate a response and return the text content only. - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `prompt` | `str` | — | User message or formatted problem text | -| `model` | `str` | — | Model identifier | -| `system` | `str` | `""` | Optional system prompt | -| `temperature` | `float` | `0.0` | Sampling temperature | -| `max_tokens` | `int` | `2048` | Maximum output tokens | - -**Returns:** `str` — model text output. - ---- - -#### generate_full - -```python -@abstractmethod -def generate_full( - self, - prompt: str, - *, - model: str, - system: str = "", - temperature: float = 0.0, - max_tokens: int = 2048, -) -> Dict[str, Any] -``` - -Generate a response and return full details including usage and cost metadata. - -**Returns:** `dict` with keys: - -| Key | Type | Description | -|-----|------|-------------| -| `content` | `str` | Model text output | -| `usage` | `dict` | Token usage (`prompt_tokens`, `completion_tokens`) | -| `model` | `str` | Model identifier used | -| `latency_seconds` | `float` | Wall-clock generation time | -| `cost_usd` | `float` | Estimated inference cost | - ---- - -#### close - -```python -def close(self) -> None -``` - -Release resources held by the backend (connections, engine handles, etc.). -The default implementation is a no-op; subclasses override as needed. - ---- - -### DatasetProvider (`evals.core.dataset`) - -```python -from evals.core.dataset import DatasetProvider -``` - -Base class for all evaluation dataset providers. Datasets are loaded lazily via -`load()` and then consumed record-by-record through `iter_records()`. - -```python -class DatasetProvider(ABC): - dataset_id: str - dataset_name: str -``` - -**Class attributes:** - -| Attribute | Type | Description | -|-----------|------|-------------| -| `dataset_id` | `str` | Short identifier matching the CLI benchmark name | -| `dataset_name` | `str` | Human-readable display name | - -**Abstract methods:** - -#### load - -```python -@abstractmethod -def load( - self, - *, - max_samples: Optional[int] = None, - split: Optional[str] = None, - seed: Optional[int] = None, -) -> None -``` - -Load the dataset, optionally downloading from HuggingFace Hub. Must be called -before `iter_records()`. - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `max_samples` | `Optional[int]` | `None` | Truncate to this many records after shuffling | -| `split` | `Optional[str]` | `None` | Dataset split override (e.g., `"test"`, `"validation"`) | -| `seed` | `Optional[int]` | `None` | Shuffle seed; `None` preserves original order | - ---- - -#### iter_records - -```python -@abstractmethod -def iter_records(self) -> Iterable[EvalRecord] -``` - -Iterate over the loaded `EvalRecord` objects. Raises if called before `load()`. - ---- - -#### size - -```python -@abstractmethod -def size(self) -> int -``` - -Return the count of loaded records. - ---- - -### Scorer (`evals.core.scorer`) - -```python -from evals.core.scorer import Scorer, LLMJudgeScorer -``` - -Base class for all scorers. A scorer compares a model's answer to the reference -in an `EvalRecord` and returns a correctness verdict with optional metadata. - -```python -class Scorer(ABC): - scorer_id: str -``` - -**Class attribute:** - -| Attribute | Type | Description | -|-----------|------|-------------| -| `scorer_id` | `str` | Short identifier matching the benchmark name | - -**Abstract method:** - -#### score - -```python -@abstractmethod -def score( - self, - record: EvalRecord, - model_answer: str, -) -> Tuple[Optional[bool], Dict[str, Any]] -``` - -Score a model answer against the reference. - -| Parameter | Type | Description | -|-----------|------|-------------| -| `record` | `EvalRecord` | The source sample including `reference` and `metadata` | -| `model_answer` | `str` | Raw text output from the model | - -**Returns:** `(is_correct, metadata)` tuple where: - -- `is_correct` is `True`, `False`, or `None` (if scoring could not be determined) -- `metadata` is a `dict` of scorer-specific details stored in `EvalResult.scoring_metadata` - ---- - -### LLMJudgeScorer - -```python -class LLMJudgeScorer(Scorer): - def __init__(self, judge_backend: InferenceBackend, judge_model: str) -> None -``` - -Convenience base class for scorers that call an LLM to evaluate answers. -Exposes `_ask_judge()` to subclasses. - -```python -def _ask_judge( - self, - prompt: str, - *, - system: str = "", - temperature: float = 0.0, - max_tokens: int = 1024, -) -> str -``` - -Send a prompt to the judge LLM and return the response text. Delegates to -`judge_backend.generate()`. - ---- - -## EvalRunner (`evals.core.runner`) - -```python -from evals.core.runner import EvalRunner -``` - -The `EvalRunner` wires together a `RunConfig`, `DatasetProvider`, `InferenceBackend`, -and `Scorer` and executes the benchmark. Inference is parallelized using a -`ThreadPoolExecutor`. Results are written to JSONL incrementally so progress is -not lost if the run is interrupted. - -### Constructor - -```python -class EvalRunner: - def __init__( - self, - config: RunConfig, - dataset: DatasetProvider, - backend: InferenceBackend, - scorer: Scorer, - ) -> None -``` - -| Parameter | Type | Description | -|-----------|------|-------------| -| `config` | `RunConfig` | Run parameters (model, workers, output path, etc.) | -| `dataset` | `DatasetProvider` | Dataset to evaluate against | -| `backend` | `InferenceBackend` | Inference backend for generation | -| `scorer` | `Scorer` | Scorer for comparing model answers to references | - -### run - -```python -def run(self) -> RunSummary -``` - -Execute the full evaluation and return aggregate statistics. - -The method: - -1. Calls `dataset.load()` with the `RunConfig` sampling parameters -2. Submits all records to a `ThreadPoolExecutor` with `config.max_workers` threads -3. For each record, calls `backend.generate_full()` then `scorer.score()` -4. Writes each `EvalResult` to a JSONL file as it completes -5. Writes a `.summary.json` alongside the JSONL at the end - -**Returns:** `RunSummary` - -```python title="programmatic_eval.py" -from evals.core.types import RunConfig -from evals.core.runner import EvalRunner -from evals.datasets.supergpqa import SuperGPQADataset -from evals.backends.jarvis_direct import JarvisDirectBackend -from evals.scorers.supergpqa_mcq import SuperGPQAScorer - -config = RunConfig( - benchmark="supergpqa", - backend="jarvis-direct", - model="qwen3:8b", - max_samples=50, - engine_key="ollama", -) - -dataset = SuperGPQADataset() -backend = JarvisDirectBackend(engine_key="ollama") -judge_backend = JarvisDirectBackend(engine_key="cloud") -scorer = SuperGPQAScorer(judge_backend=judge_backend, judge_model="gpt-4o") - -runner = EvalRunner(config, dataset, backend, scorer) -summary = runner.run() - -print(f"Accuracy: {summary.accuracy:.4f} ({summary.correct}/{summary.scored_samples})") -print(f"Mean latency: {summary.mean_latency_seconds:.2f}s") -print(f"Total cost: ${summary.total_cost_usd:.4f}") - -backend.close() -judge_backend.close() -``` - ---- - -## Backends - -### JarvisDirectBackend (`evals.backends.jarvis_direct`) - -```python -from evals.backends.jarvis_direct import JarvisDirectBackend -``` - -Engine-level inference via `SystemBuilder`. Routes directly to the configured -`InferenceEngine` without an agent loop, making it the fastest backend and -appropriate for benchmarks that do not require tool use. - -```python -class JarvisDirectBackend(InferenceBackend): - backend_id = "jarvis-direct" - - def __init__(self, engine_key: Optional[str] = None) -> None -``` - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `engine_key` | `Optional[str]` | `None` | OpenJarvis engine identifier. `None` uses the auto-discovered engine from `~/.openjarvis/config.toml` | - -Telemetry and traces are disabled for eval runs. The backend calls -`SystemBuilder().engine(engine_key).telemetry(False).traces(False).build()`. - -**Compatible benchmarks:** `supergpqa`, `frames`, `wildchat` (any benchmark that -does not require multi-step tool calling). - -=== "Local model" - - ```python - backend = JarvisDirectBackend(engine_key="ollama") - text = backend.generate("What is 2+2?", model="qwen3:8b") - ``` - -=== "Cloud model" - - ```python - backend = JarvisDirectBackend(engine_key="cloud") - text = backend.generate("What is 2+2?", model="gpt-4o") - ``` - ---- - -### JarvisAgentBackend (`evals.backends.jarvis_agent`) - -```python -from evals.backends.jarvis_agent import JarvisAgentBackend -``` - -Agent-level inference via `JarvisSystem.ask()`. Wraps the full OpenJarvis agent -harness, enabling multi-turn tool-calling loops for agentic benchmarks. - -```python -class JarvisAgentBackend(InferenceBackend): - backend_id = "jarvis-agent" - - def __init__( - self, - engine_key: Optional[str] = None, - agent_name: str = "orchestrator", - tools: Optional[List[str]] = None, - ) -> None -``` - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `engine_key` | `Optional[str]` | `None` | OpenJarvis engine identifier | -| `agent_name` | `str` | `"orchestrator"` | Agent to use (`"orchestrator"`, `"react"`, etc.) | -| `tools` | `Optional[List[str]]` | `None` | Tool names to enable (e.g., `["calculator", "file_read"]`) | - -The `generate_full()` return dict includes two additional keys beyond the standard -`InferenceBackend` contract: - -| Key | Type | Description | -|-----|------|-------------| -| `turns` | `int` | Number of agent turns completed | -| `tool_results` | `list` | Tool call results from the agent loop | - -**Compatible benchmarks:** `gaia` (requires file reading and multi-step reasoning). - -```python -backend = JarvisAgentBackend( - engine_key="ollama", - agent_name="orchestrator", - tools=["file_read", "calculator"], -) -result = backend.generate_full( - "How many pages is the attached PDF?", - model="qwen3:8b", -) -print(result["content"]) -print(f"Completed in {result['turns']} turn(s)") -backend.close() -``` - ---- - -## Dataset Providers - -### SuperGPQADataset (`evals.datasets.supergpqa`) - -```python -from evals.datasets.supergpqa import SuperGPQADataset -``` - -Loads the SuperGPQA multiple-choice benchmark from HuggingFace (`m-a-p/SuperGPQA`). -Records have `category="reasoning"` and `subject` set to the discipline subfield. - -```python -class SuperGPQADataset(DatasetProvider): - dataset_id = "supergpqa" - dataset_name = "SuperGPQA" -``` - -- **Default split:** `"train"` -- **HuggingFace path:** `m-a-p/SuperGPQA` -- Each problem is formatted with lettered options (A, B, C, ...) and the instruction - "Respond with the correct letter only." -- `record.reference` is the correct answer letter (e.g., `"B"`). - ---- - -### GAIADataset (`evals.datasets.gaia`) - -```python -from evals.datasets.gaia import GAIADataset -``` - -Loads the GAIA agentic benchmark from HuggingFace (`gaia-benchmark/GAIA`). -Records have `category="agentic"` and `subject` set to `level_1`, `level_2`, or -`level_3`. - -```python -class GAIADataset(DatasetProvider): - dataset_id = "gaia" - dataset_name = "GAIA" - - def __init__(self, cache_dir: Optional[str] = None) -> None -``` - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `cache_dir` | `Optional[str]` | `~/.cache/gaia_benchmark` | Local directory for HuggingFace snapshot download | - -- **Default split:** `"validation"` -- **Default subset:** `"2023_all"` -- Downloads the full dataset snapshot including associated files (PDFs, images, CSVs) - referenced in questions. File paths are embedded in the problem prompt. - -!!! warning "Dataset access" - GAIA requires accepting the HuggingFace dataset terms of service and being logged - in with `huggingface-cli login` before the snapshot download can proceed. - ---- - -### FRAMESDataset (`evals.datasets.frames`) - -```python -from evals.datasets.frames import FRAMESDataset -``` - -Loads the FRAMES multi-hop factual retrieval benchmark from HuggingFace -(`google/frames-benchmark`). Records have `category="rag"` and `subject` set to -the reasoning type(s) (e.g., `"multi-hop, temporal"`). - -```python -class FRAMESDataset(DatasetProvider): - dataset_id = "frames" - dataset_name = "FRAMES" -``` - -- **Default split:** `"test"` -- Wikipedia article links referenced in each question are included in the problem prompt. - ---- - -### WildChatDataset (`evals.datasets.wildchat`) - -```python -from evals.datasets.wildchat import WildChatDataset -``` - -Loads the WildChat-1M dataset (`allenai/WildChat-1M`) and filters to English -single-turn conversations for chat quality evaluation. Records have -`category="chat"` and `subject="conversation"`. - -```python -class WildChatDataset(DatasetProvider): - dataset_id = "wildchat" - dataset_name = "WildChat" -``` - -- **Default split:** `"train"` -- Filters by `language == "english"` and exactly two turns (one user + one assistant). -- `record.problem` is the user message; `record.reference` is the original assistant - response used as the quality baseline by the judge scorer. - ---- - -## Scorers - -### SuperGPQAScorer (`evals.scorers.supergpqa_mcq`) - -```python -from evals.scorers.supergpqa_mcq import SuperGPQAScorer -``` - -LLM-based letter extraction followed by exact match against the reference letter. -The judge LLM extracts the final answer letter from potentially verbose model -responses, then compares it to `record.reference`. - -```python -class SuperGPQAScorer(LLMJudgeScorer): - scorer_id = "supergpqa" -``` - -**Scoring metadata keys:** - -| Key | Description | -|-----|-------------| -| `reference_letter` | Correct answer letter from the dataset | -| `candidate_letter` | Letter extracted by the judge LLM | -| `valid_letters` | Valid answer letters for this question (e.g., `"ABCD"`) | -| `reason` | Set to `"missing_reference_letter"` or `"no_choice_letter_extracted"` on failure | - ---- - -### GAIAScorer (`evals.scorers.gaia_exact`) - -```python -from evals.scorers.gaia_exact import GAIAScorer, exact_match -``` - -Normalized exact match with an LLM fallback for semantic comparison. Tries exact -match first (no API call); falls back to the judge LLM only when exact match fails. - -```python -class GAIAScorer(LLMJudgeScorer): - scorer_id = "gaia" -``` - -**Normalization rules for exact match:** - -- Numbers: strips `$`, `%`, `,` then compares as `float` -- Lists (comma- or semicolon-separated): splits and compares element-by-element -- Strings: lowercases, strips whitespace and punctuation - -**Scoring metadata keys:** - -| Key | Description | -|-----|-------------| -| `match_type` | `"exact"` or `"llm_fallback"` | -| `raw_judge_output` | Full LLM judge response (llm_fallback only) | -| `extracted_answer` | Answer extracted by the judge (llm_fallback only) | - -The `exact_match` helper function is also exported and can be used independently: - -```python -from evals.scorers.gaia_exact import exact_match - -assert exact_match("$1,000", "1000") is True -assert exact_match("paris", "Paris") is True -assert exact_match("3, 5", "3,5") is True -``` - ---- - -### FRAMESScorer (`evals.scorers.frames_judge`) - -```python -from evals.scorers.frames_judge import FRAMESScorer -``` - -LLM-as-judge scorer for FRAMES multi-hop factual retrieval. Uses a structured -grading rubric that focuses on semantic equivalence, ignoring formatting and -capitalization differences. - -```python -class FRAMESScorer(LLMJudgeScorer): - scorer_id = "frames" -``` - -**Scoring metadata keys:** - -| Key | Description | -|-----|-------------| -| `raw_judge_output` | Full LLM judge response | -| `extracted_answer` | Answer extracted by the judge | - ---- - -### WildChatScorer (`evals.scorers.wildchat_judge`) - -```python -from evals.scorers.wildchat_judge import WildChatScorer -``` - -Dual-comparison LLM-as-judge for chat quality. Runs two comparisons — once with -the model answer as Assistant A and once as Assistant B — to reduce position bias. -The model answer is considered correct if it wins or ties in either comparison. - -```python -class WildChatScorer(LLMJudgeScorer): - scorer_id = "wildchat" -``` - -The judge uses a five-point verdict scale: `[[A>>B]]`, `[[A>B]]`, `[[A=B]]`, -`[[B>A]]`, `[[B>>A]]`. A tie (`A=B`) is counted as correct. - -**Scoring metadata keys:** - -| Key | Description | -|-----|-------------| -| `generated_as_a` | `{verdict, response}` from the first comparison pass | -| `generated_as_b` | `{verdict, response}` from the second comparison pass | - ---- - -## CLI Reference - -The evaluation framework ships a `openjarvis-eval` CLI built with Click. - -### openjarvis-eval run - -Run a single benchmark or a full suite from a TOML config. - -```bash title="Single run" -openjarvis-eval run \ - --benchmark supergpqa \ - --model qwen3:8b \ - --engine ollama \ - --max-samples 100 \ - --max-workers 8 \ - --output results/supergpqa_qwen3-8b.jsonl -``` - -```bash title="Suite run from TOML" -openjarvis-eval run --config evals/configs/full-suite.toml -``` - -| Option | Short | Default | Description | -|--------|-------|---------|-------------| -| `--config` | `-c` | `None` | TOML suite config file; enables suite mode | -| `--benchmark` | `-b` | — | Benchmark name (required in single-run mode) | -| `--backend` | | `jarvis-direct` | `jarvis-direct` or `jarvis-agent` | -| `--model` | `-m` | — | Model identifier (required in single-run mode) | -| `--engine` | `-e` | `None` | Engine key override | -| `--agent` | | `orchestrator` | Agent name for `jarvis-agent` backend | -| `--tools` | | `""` | Comma-separated tool names | -| `--max-samples` | `-n` | `None` | Sample limit | -| `--max-workers` | `-w` | `4` | Parallel threads | -| `--judge-model` | | `gpt-4o` | LLM judge model | -| `--output` | `-o` | auto | JSONL output path | -| `--seed` | | `42` | Shuffle seed | -| `--split` | | `None` | Dataset split override | -| `--temperature` | | `0.0` | Sampling temperature | -| `--max-tokens` | | `2048` | Maximum output tokens | -| `--verbose` | `-v` | `False` | Enable debug logging | - -### openjarvis-eval run-all - -Run all four benchmarks against a single model. - -```bash -openjarvis-eval run-all \ - --model qwen3:8b \ - --engine ollama \ - --max-samples 50 \ - --output-dir results/ -``` - -### openjarvis-eval summarize - -Recompute summary statistics from an existing JSONL output file. - -```bash -openjarvis-eval summarize results/supergpqa_qwen3-8b.jsonl -``` - -### openjarvis-eval list - -List all available benchmarks and backends. - -```bash -openjarvis-eval list -``` - ---- - -## TOML Suite Config Format - -A suite config drives a full `models x benchmarks` comparison matrix with a single -command. All sections except `[[models]]` and `[[benchmarks]]` are optional. - -```toml title="evals/configs/full-suite.toml" -[meta] -name = "full-suite-v1" -description = "Evaluate all benchmarks against production models" - -[defaults] -temperature = 0.0 -max_tokens = 2048 - -[judge] -model = "gpt-4o" -temperature = 0.0 -max_tokens = 1024 - -[run] -max_workers = 4 -output_dir = "results/" -seed = 42 - -# One [[models]] entry per model to evaluate -[[models]] -name = "qwen3:8b" -engine = "ollama" -temperature = 0.3 - -[[models]] -name = "gpt-4o" -provider = "openai" - -# One [[benchmarks]] entry per benchmark -[[benchmarks]] -name = "supergpqa" -backend = "jarvis-direct" -max_samples = 200 - -[[benchmarks]] -name = "gaia" -backend = "jarvis-agent" -agent = "orchestrator" -tools = ["file_read", "calculator"] -max_samples = 50 -judge_model = "claude-sonnet-4-20250514" # override judge for this benchmark - -[[benchmarks]] -name = "frames" -backend = "jarvis-direct" -max_samples = 100 - -[[benchmarks]] -name = "wildchat" -backend = "jarvis-direct" -max_samples = 150 -temperature = 0.7 -``` - -```bash -openjarvis-eval run --config evals/configs/full-suite.toml -# Suite: full-suite-v1 -# 2 model(s) x 4 benchmark(s) = 8 run(s) -``` - ---- - -## See Also - -- [Benchmarks Module](bench.md) — `openjarvis.bench` performance benchmarks (latency, throughput) for the inference engine, separate from the eval framework -- [Telemetry & Traces](telemetry.md) — `openjarvis.telemetry` and `openjarvis.traces` for production monitoring -- [Python SDK](../user-guide/python-sdk.md) — `Jarvis` class used internally by eval backends -- [Agents](../user-guide/agents.md) — Agent implementations invoked by `JarvisAgentBackend` diff --git a/docs/api/index.md b/docs/api/index.md deleted file mode 100644 index 8db37f6c..00000000 --- a/docs/api/index.md +++ /dev/null @@ -1,26 +0,0 @@ -# API Reference - -This section contains the auto-generated API reference for the OpenJarvis Python -package. Documentation is extracted directly from the source code docstrings -using [mkdocstrings](https://mkdocstrings.github.io/). - -Each page documents a major module of the framework, including abstract base -classes, concrete implementations, and utility functions. - -## Modules - -| Module | Description | -|--------|-------------| -| [SDK](sdk.md) | High-level `Jarvis` class and `MemoryHandle` proxy | -| [Core](core.md) | Registries, types, configuration, and event bus | -| [Engine](engine.md) | Inference engine ABC and backends (Ollama, vLLM, llama.cpp, SGLang, Cloud) | -| [Agents](agents.md) | Agent ABC and implementations (Simple, Orchestrator, OpenClaw, Custom) | -| [Memory](memory.md) | Memory backend ABC and implementations (SQLite, FAISS, ColBERT, BM25, Hybrid) | -| [Tools](tools.md) | Tool ABC, executor, and built-in tools (Calculator, Think, Retrieval, LLM, FileRead) | -| [Intelligence](intelligence.md) | Heuristic router and model catalog | -| [Learning](learning.md) | Router policies, reward functions, and trace-driven learning | -| [Traces](traces.md) | Trace storage, collection, and analysis | -| [Telemetry](telemetry.md) | Telemetry storage, aggregation, and instrumented wrappers | -| [Benchmarks](bench.md) | Benchmark ABC, suite runner, latency and throughput benchmarks | -| [Evals](evals.md) | Evaluation framework: backends, dataset providers, scorers, and suite runner | -| [Server](server.md) | FastAPI application, OpenAI-compatible routes, and Pydantic models | diff --git a/docs/api/intelligence.md b/docs/api/intelligence.md deleted file mode 100644 index 3a927417..00000000 --- a/docs/api/intelligence.md +++ /dev/null @@ -1,47 +0,0 @@ -# Intelligence Module - -The intelligence module handles model management and query routing. It -provides the `HeuristicRouter` which selects models based on query -characteristics (code detection, math detection, query length, urgency), -and a model catalog of well-known local and cloud models with their -specifications. - -## Heuristic Router - -### HeuristicRouter - -::: openjarvis.intelligence.router.HeuristicRouter - options: - show_source: true - members_order: source - -### build_routing_context - -::: openjarvis.intelligence.router.build_routing_context - options: - show_source: true - ---- - -## Model Catalog - -Built-in catalog of well-known model specifications, including local models -(Qwen, Llama, Mistral, DeepSeek) and cloud models (OpenAI, Anthropic, Google). - -### BUILTIN_MODELS - -::: openjarvis.intelligence.model_catalog.BUILTIN_MODELS - options: - show_source: true - -### register_builtin_models - -::: openjarvis.intelligence.model_catalog.register_builtin_models - options: - show_source: true - -### merge_discovered_models - -::: openjarvis.intelligence.model_catalog.merge_discovered_models - options: - show_source: true diff --git a/docs/api/learning.md b/docs/api/learning.md deleted file mode 100644 index d2e6f1e9..00000000 --- a/docs/api/learning.md +++ /dev/null @@ -1,98 +0,0 @@ -# Learning Module - -The learning module implements learning policies that improve routing, agent, -and tool decisions based on historical interaction outcomes. The module provides -a `LearningPolicy` ABC taxonomy with specialized sub-ABCs for intelligence -(model routing), agent behavior, and tool selection. It also includes reward -functions for scoring inference results. - -## Abstract Base Classes - -### RouterPolicy - -::: openjarvis.intelligence._stubs.RouterPolicy - options: - show_source: true - members_order: source - -### QueryAnalyzer - -::: openjarvis.intelligence._stubs.QueryAnalyzer - options: - show_source: true - members_order: source - -### RoutingContext - -`RoutingContext` is defined in [`core/types.py`](core.md#openjarvis.core.types.RoutingContext). - -### RewardFunction - -::: openjarvis.learning._stubs.RewardFunction - options: - show_source: true - members_order: source - -### LearningPolicy Taxonomy - -The learning system defines a hierarchy of learning policy ABCs: - -- **`LearningPolicy`** -- base ABC for all learning policies -- **`IntelligenceLearningPolicy`** -- specialization for model routing decisions -- **`AgentLearningPolicy`** -- specialization for agent behavior advice (ICL examples, tool-use strategies) - ---- - -## Policy Implementations - -### TraceDrivenPolicy - -::: openjarvis.learning.trace_policy.TraceDrivenPolicy - options: - show_source: true - members_order: source - -### classify_query - -::: openjarvis.learning.trace_policy.classify_query - options: - show_source: true - -### SFTRouterPolicy - -::: openjarvis.learning.sft_policy.SFTRouterPolicy - options: - show_source: true - members_order: source - -### AgentAdvisorPolicy - -::: openjarvis.learning.agent_advisor.AgentAdvisorPolicy - options: - show_source: true - members_order: source - -### ICLUpdaterPolicy - -::: openjarvis.learning.icl_updater.ICLUpdaterPolicy - options: - show_source: true - members_order: source - -### GRPORouterPolicy - -::: openjarvis.learning.grpo_policy.GRPORouterPolicy - options: - show_source: true - members_order: source - ---- - -## Reward Functions - -### HeuristicRewardFunction - -::: openjarvis.learning.heuristic_reward.HeuristicRewardFunction - options: - show_source: true - members_order: source diff --git a/docs/api/memory.md b/docs/api/memory.md deleted file mode 100644 index 15126dde..00000000 --- a/docs/api/memory.md +++ /dev/null @@ -1,183 +0,0 @@ -# Memory Module - -The memory module implements persistent searchable storage for document -retrieval. All backends implement the `MemoryBackend` ABC with `store()`, -`retrieve()`, `delete()`, and `clear()` methods. The module also includes -a document ingestion pipeline (chunking, file reading) and context injection -for augmenting prompts with retrieved knowledge. - -!!! note "Canonical import location" - Memory backends have moved to `openjarvis.tools.storage.*`. The old - `openjarvis.memory.*` imports still work via backward-compatibility shims. - -## Abstract Base Class - -### MemoryBackend - -::: openjarvis.tools.storage._stubs.MemoryBackend - options: - show_source: true - members_order: source - -### RetrievalResult - -::: openjarvis.tools.storage._stubs.RetrievalResult - options: - show_source: true - members_order: source - ---- - -## Backend Implementations - -### SQLiteMemory - -::: openjarvis.tools.storage.sqlite.SQLiteMemory - options: - show_source: true - members_order: source - -### FAISSMemory - -::: openjarvis.tools.storage.faiss_backend.FAISSMemory - options: - show_source: true - members_order: source - -### ColBERTMemory - -::: openjarvis.tools.storage.colbert_backend.ColBERTMemory - options: - show_source: true - members_order: source - -### BM25Memory - -::: openjarvis.tools.storage.bm25.BM25Memory - options: - show_source: true - members_order: source - -### HybridMemory - -::: openjarvis.tools.storage.hybrid.HybridMemory - options: - show_source: true - members_order: source - -### reciprocal_rank_fusion - -::: openjarvis.tools.storage.hybrid.reciprocal_rank_fusion - options: - show_source: true - ---- - -## Document Chunking - -Splits text into fixed-size chunks with configurable overlap, respecting -paragraph boundaries. - -### ChunkConfig - -::: openjarvis.tools.storage.chunking.ChunkConfig - options: - show_source: true - members_order: source - -### Chunk - -::: openjarvis.tools.storage.chunking.Chunk - options: - show_source: true - members_order: source - -### chunk_text - -::: openjarvis.tools.storage.chunking.chunk_text - options: - show_source: true - ---- - -## Document Ingestion - -File reading, type detection, and directory walking for the ingestion -pipeline. - -### DocumentMeta - -::: openjarvis.tools.storage.ingest.DocumentMeta - options: - show_source: true - members_order: source - -### detect_file_type - -::: openjarvis.tools.storage.ingest.detect_file_type - options: - show_source: true - -### read_document - -::: openjarvis.tools.storage.ingest.read_document - options: - show_source: true - -### ingest_path - -::: openjarvis.tools.storage.ingest.ingest_path - options: - show_source: true - ---- - -## Context Injection - -Retrieves relevant memory and injects it into prompts as system messages -with source attribution. - -### ContextConfig - -::: openjarvis.tools.storage.context.ContextConfig - options: - show_source: true - members_order: source - -### inject_context - -::: openjarvis.tools.storage.context.inject_context - options: - show_source: true - -### format_context - -::: openjarvis.tools.storage.context.format_context - options: - show_source: true - -### build_context_message - -::: openjarvis.tools.storage.context.build_context_message - options: - show_source: true - ---- - -## Embeddings - -Abstraction layer for text embedding models used by dense retrieval backends. - -### Embedder - -::: openjarvis.tools.storage.embeddings.Embedder - options: - show_source: true - members_order: source - -### SentenceTransformerEmbedder - -::: openjarvis.tools.storage.embeddings.SentenceTransformerEmbedder - options: - show_source: true - members_order: source diff --git a/docs/api/sdk.md b/docs/api/sdk.md deleted file mode 100644 index 6c9a63e8..00000000 --- a/docs/api/sdk.md +++ /dev/null @@ -1,20 +0,0 @@ -# SDK - -The SDK module provides the high-level Python interface for OpenJarvis. The -`Jarvis` class wraps engine discovery, agent dispatch, memory operations, and -telemetry into a single convenient entry point. `MemoryHandle` acts as a lazy -proxy for memory backend operations. - -## Jarvis - -::: openjarvis.sdk.Jarvis - options: - show_source: true - members_order: source - -## MemoryHandle - -::: openjarvis.sdk.MemoryHandle - options: - show_source: true - members_order: source diff --git a/docs/api/security.md b/docs/api/security.md deleted file mode 100644 index a4cf024d..00000000 --- a/docs/api/security.md +++ /dev/null @@ -1,210 +0,0 @@ -# API Reference: Security - -The `openjarvis.security` package provides text scanning, content guardrails, file path filtering, and audit logging. All public components are documented below. - -For usage examples and configuration, see the [Security user guide](../user-guide/security.md). For the architectural design, see [Security architecture](../architecture/security.md). - ---- - -## Types - -Core data types shared across the security subsystem. - -### ThreatLevel - -Severity classification for individual scan findings. Ordered from least to most severe: `LOW` < `MEDIUM` < `HIGH` < `CRITICAL`. - -::: openjarvis.security.types.ThreatLevel - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### RedactionMode - -Controls the action taken by `GuardrailsEngine` when findings are detected. - -::: openjarvis.security.types.RedactionMode - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### SecurityEventType - -Categories of security events recorded by `AuditLogger`. - -::: openjarvis.security.types.SecurityEventType - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### ScanFinding - -A single match produced by a scanner. Includes the pattern name, matched text, position, threat level, and a human-readable description. - -::: openjarvis.security.types.ScanFinding - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### ScanResult - -Aggregated result from one or more scanner passes. The `clean` property returns `True` when no findings were detected; `highest_threat` returns the most severe `ThreatLevel` found. - -::: openjarvis.security.types.ScanResult - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### SecurityEvent - -A recorded security event, as persisted by `AuditLogger`. - -::: openjarvis.security.types.SecurityEvent - options: - show_source: true - show_root_heading: true - heading_level: 4 - ---- - -## BaseScanner - -`BaseScanner` is the abstract base class for all scanner implementations. Implement both `scan()` and `redact()` to create a custom scanner. - -::: openjarvis.security._stubs.BaseScanner - options: - show_source: true - show_root_heading: true - heading_level: 3 - ---- - -## SecretScanner - -Detects API keys, tokens, passwords, and credentials in text using regex patterns. See the [pattern reference table](../user-guide/security.md#pattern-reference) in the user guide for the full list of patterns and their threat levels. - -::: openjarvis.security.scanner.SecretScanner - options: - show_source: true - show_root_heading: true - heading_level: 3 - ---- - -## PIIScanner - -Detects personally identifiable information including email addresses, Social Security Numbers, credit card numbers, phone numbers, and public IP addresses. - -::: openjarvis.security.scanner.PIIScanner - options: - show_source: true - show_root_heading: true - heading_level: 3 - ---- - -## GuardrailsEngine - -`GuardrailsEngine` wraps any `InferenceEngine` with security scanning on both input and output. It implements the full `InferenceEngine` interface, so it can be used anywhere an engine is expected. - -!!! note "Registration" - `GuardrailsEngine` is **not** registered in `EngineRegistry`. Instantiate it directly by wrapping an existing engine instance. - -::: openjarvis.security.guardrails.GuardrailsEngine - options: - show_source: true - show_root_heading: true - heading_level: 3 - -### SecurityBlockError - -Raised by `GuardrailsEngine` when `mode=RedactionMode.BLOCK` and findings are detected during a scan. Catch this exception to handle blocked requests gracefully. - -```python -from openjarvis.security.guardrails import GuardrailsEngine, SecurityBlockError -from openjarvis.security.types import RedactionMode - -guarded = GuardrailsEngine(engine, mode=RedactionMode.BLOCK) - -try: - response = guarded.generate(messages, model="qwen3:8b") -except SecurityBlockError as exc: - # exc.args[0] describes the direction and finding count - print(f"Request blocked: {exc}") -``` - -::: openjarvis.security.guardrails.SecurityBlockError - options: - show_source: true - show_root_heading: true - heading_level: 4 - ---- - -## File Policy - -Functions and constants for filtering sensitive file paths. Used internally by `FileReadTool` and the memory ingest path. - -### DEFAULT_SENSITIVE_PATTERNS - -The default set of glob patterns used to identify sensitive files. This is a `frozenset[str]` exported from `openjarvis.security.file_policy`. - -See the [sensitive file patterns table](../user-guide/security.md#sensitive-file-patterns) in the user guide for the complete list. - -::: openjarvis.security.file_policy.DEFAULT_SENSITIVE_PATTERNS - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### is_sensitive_file - -::: openjarvis.security.file_policy.is_sensitive_file - options: - show_source: true - show_root_heading: true - heading_level: 4 - -### filter_sensitive_paths - -::: openjarvis.security.file_policy.filter_sensitive_paths - options: - show_source: true - show_root_heading: true - heading_level: 4 - ---- - -## AuditLogger - -Append-only SQLite-backed storage for security events. Subscribes to `SECURITY_SCAN`, `SECURITY_ALERT`, and `SECURITY_BLOCK` events on the `EventBus` when a bus is provided. - -The default database path is `~/.openjarvis/audit.db`, overridable via `security.audit_log_path` in `config.toml`. - -```python title="audit_logger_example.py" -from openjarvis.core.events import EventBus -from openjarvis.security.audit import AuditLogger -from openjarvis.security.guardrails import GuardrailsEngine -from openjarvis.security.types import RedactionMode - -bus = EventBus() -audit = AuditLogger(bus=bus) - -guarded = GuardrailsEngine(engine, mode=RedactionMode.WARN, bus=bus) - -# Security events from guarded engine are now persisted automatically -events = audit.query(limit=10) -print(f"Logged {audit.count()} events") -audit.close() -``` - -::: openjarvis.security.audit.AuditLogger - options: - show_source: true - show_root_heading: true - heading_level: 3 diff --git a/docs/api/server.md b/docs/api/server.md deleted file mode 100644 index 68b20213..00000000 --- a/docs/api/server.md +++ /dev/null @@ -1,139 +0,0 @@ -# Server Module - -The server module provides an OpenAI-compatible API server built on FastAPI -and uvicorn. It exposes `POST /v1/chat/completions` (with streaming via SSE), -`GET /v1/models`, and `GET /health` endpoints. Pydantic models ensure -request/response validation matching the OpenAI API format. - -## Application Factory - -### create_app - -::: openjarvis.server.app.create_app - options: - show_source: true - ---- - -## Route Handlers - -### chat_completions - -::: openjarvis.server.routes.chat_completions - options: - show_source: true - -### list_models - -::: openjarvis.server.routes.list_models - options: - show_source: true - -### health - -::: openjarvis.server.routes.health - options: - show_source: true - ---- - -## Pydantic Request/Response Models - -### ChatMessage - -::: openjarvis.server.models.ChatMessage - options: - show_source: true - members_order: source - -### ChatCompletionRequest - -::: openjarvis.server.models.ChatCompletionRequest - options: - show_source: true - members_order: source - -### ChatCompletionResponse - -::: openjarvis.server.models.ChatCompletionResponse - options: - show_source: true - members_order: source - -### Choice - -::: openjarvis.server.models.Choice - options: - show_source: true - members_order: source - -### ChoiceMessage - -::: openjarvis.server.models.ChoiceMessage - options: - show_source: true - members_order: source - -### UsageInfo - -::: openjarvis.server.models.UsageInfo - options: - show_source: true - members_order: source - -### ChatCompletionChunk - -::: openjarvis.server.models.ChatCompletionChunk - options: - show_source: true - members_order: source - -### DeltaMessage - -::: openjarvis.server.models.DeltaMessage - options: - show_source: true - members_order: source - -### StreamChoice - -::: openjarvis.server.models.StreamChoice - options: - show_source: true - members_order: source - -### ModelObject - -::: openjarvis.server.models.ModelObject - options: - show_source: true - members_order: source - -### ModelListResponse - -::: openjarvis.server.models.ModelListResponse - options: - show_source: true - members_order: source - ---- - -## Channel Endpoints - -### list_channels - -::: openjarvis.server.routes.list_channels - options: - show_source: true - -### channel_send - -::: openjarvis.server.routes.channel_send - options: - show_source: true - -### channel_status - -::: openjarvis.server.routes.channel_status - options: - show_source: true diff --git a/docs/api/telemetry.md b/docs/api/telemetry.md deleted file mode 100644 index 9fb85618..00000000 --- a/docs/api/telemetry.md +++ /dev/null @@ -1,53 +0,0 @@ -# Telemetry Module - -The telemetry module provides append-only recording and read-only aggregation -of inference metrics. Every engine call records timing, token counts, energy -usage, and cost to SQLite via the event bus. The `TelemetryAggregator` -provides per-model and per-engine statistics with time-range filtering. - -## TelemetryStore - -::: openjarvis.telemetry.store.TelemetryStore - options: - show_source: true - members_order: source - ---- - -## TelemetryAggregator - -::: openjarvis.telemetry.aggregator.TelemetryAggregator - options: - show_source: true - members_order: source - -### ModelStats - -::: openjarvis.telemetry.aggregator.ModelStats - options: - show_source: true - members_order: source - -### EngineStats - -::: openjarvis.telemetry.aggregator.EngineStats - options: - show_source: true - members_order: source - -### AggregatedStats - -::: openjarvis.telemetry.aggregator.AggregatedStats - options: - show_source: true - members_order: source - ---- - -## Instrumented Wrapper - -### instrumented_generate - -::: openjarvis.telemetry.wrapper.instrumented_generate - options: - show_source: true diff --git a/docs/api/tools.md b/docs/api/tools.md deleted file mode 100644 index d869ac20..00000000 --- a/docs/api/tools.md +++ /dev/null @@ -1,179 +0,0 @@ -# Tools Module - -The tools module implements the tool system used by agents for executing -actions such as calculations, memory search, file reading, and sub-model -queries. Each tool implements the `BaseTool` ABC and is registered via -`@ToolRegistry.register("name")`. The `ToolExecutor` dispatches tool calls -with JSON argument parsing, latency tracking, and event bus integration. - -## Abstract Base Class - -### BaseTool - -::: openjarvis.tools._stubs.BaseTool - options: - show_source: true - members_order: source - -### ToolSpec - -::: openjarvis.tools._stubs.ToolSpec - options: - show_source: true - members_order: source - -### ToolExecutor - -::: openjarvis.tools._stubs.ToolExecutor - options: - show_source: true - members_order: source - -### build_tool_descriptions - -::: openjarvis.tools._stubs.build_tool_descriptions - options: - show_source: true - ---- - -## Built-in Tools - -### CalculatorTool - -::: openjarvis.tools.calculator.CalculatorTool - options: - show_source: true - members_order: source - -### safe_eval - -::: openjarvis.tools.calculator.safe_eval - options: - show_source: true - -### ThinkTool - -::: openjarvis.tools.think.ThinkTool - options: - show_source: true - members_order: source - -### RetrievalTool - -::: openjarvis.tools.retrieval.RetrievalTool - options: - show_source: true - members_order: source - -### LLMTool - -::: openjarvis.tools.llm_tool.LLMTool - options: - show_source: true - members_order: source - -### FileReadTool - -::: openjarvis.tools.file_read.FileReadTool - options: - show_source: true - members_order: source - ---- - -## Storage Tools - -Tools that expose memory storage operations for use by agents via the tool -system. These wrap the memory backend methods as `BaseTool` implementations. - -### MemoryStoreTool - -::: openjarvis.tools.storage_tools.MemoryStoreTool - options: - show_source: true - members_order: source - -### MemoryRetrieveTool - -::: openjarvis.tools.storage_tools.MemoryRetrieveTool - options: - show_source: true - members_order: source - -### MemorySearchTool - -::: openjarvis.tools.storage_tools.MemorySearchTool - options: - show_source: true - members_order: source - -### MemoryIndexTool - -::: openjarvis.tools.storage_tools.MemoryIndexTool - options: - show_source: true - members_order: source - ---- - -## Scheduler Tools - -MCP tools for scheduling agent tasks for future or recurring execution. All five tools are in the `scheduler` category and require a `TaskScheduler` instance to be injected at `_scheduler`. For usage and schedule type examples see the [Scheduler Tools user guide](../user-guide/tools.md#scheduler-tools). - -### ScheduleTaskTool - -::: openjarvis.scheduler.tools.ScheduleTaskTool - options: - show_source: true - members_order: source - -### ListScheduledTasksTool - -::: openjarvis.scheduler.tools.ListScheduledTasksTool - options: - show_source: true - members_order: source - -### PauseScheduledTaskTool - -::: openjarvis.scheduler.tools.PauseScheduledTaskTool - options: - show_source: true - members_order: source - -### ResumeScheduledTaskTool - -::: openjarvis.scheduler.tools.ResumeScheduledTaskTool - options: - show_source: true - members_order: source - -### CancelScheduledTaskTool - -::: openjarvis.scheduler.tools.CancelScheduledTaskTool - options: - show_source: true - members_order: source - ---- - -## MCP Adapter - -Integration with the Model Context Protocol (MCP). The MCP adapter exposes -OpenJarvis tools as MCP-compatible resources and allows consuming external -MCP tool providers. Supports MCP protocol version 2025-11-25. - -### MCPToolAdapter - -::: openjarvis.tools.mcp_adapter.MCPToolAdapter - options: - show_source: true - members_order: source - -### MCPToolProvider - -::: openjarvis.tools.mcp_adapter.MCPToolProvider - options: - show_source: true - members_order: source diff --git a/docs/api/traces.md b/docs/api/traces.md deleted file mode 100644 index 8d1690af..00000000 --- a/docs/api/traces.md +++ /dev/null @@ -1,53 +0,0 @@ -# Traces Module - -The traces module provides full interaction-level recording for the learning -system. Every agent interaction can produce a `Trace` capturing the sequence -of steps (route, retrieve, generate, tool_call, respond) with timing, inputs, -outputs, and outcomes. The `TraceCollector` wraps any agent to record traces -automatically, while `TraceAnalyzer` provides aggregated statistics. - -## TraceStore - -::: openjarvis.traces.store.TraceStore - options: - show_source: true - members_order: source - ---- - -## TraceCollector - -::: openjarvis.traces.collector.TraceCollector - options: - show_source: true - members_order: source - ---- - -## TraceAnalyzer - -::: openjarvis.traces.analyzer.TraceAnalyzer - options: - show_source: true - members_order: source - -### RouteStats - -::: openjarvis.traces.analyzer.RouteStats - options: - show_source: true - members_order: source - -### ToolStats - -::: openjarvis.traces.analyzer.ToolStats - options: - show_source: true - members_order: source - -### TraceSummary - -::: openjarvis.traces.analyzer.TraceSummary - options: - show_source: true - members_order: source diff --git a/docs/architecture/channels.md b/docs/architecture/channels.md index 98d7f506..abff2f1a 100644 --- a/docs/architecture/channels.md +++ b/docs/architecture/channels.md @@ -240,6 +240,6 @@ After registration, the backend is discoverable via `ChannelRegistry.get("slack" ## See Also - [User Guide: Channels](../user-guide/channels.md) — how to use channels in practice -- [API Reference: Channels](../api/channels.md) — complete class and type signatures +- [API Reference: Channels](../api-reference/openjarvis/channels/index.md) — complete class and type signatures - [Architecture: Overview](overview.md) — where channels fit in the overall system - [Architecture: Design Principles](design-principles.md) — registry pattern and ABC conventions diff --git a/docs/architecture/learning.md b/docs/architecture/learning.md index a316301c..c2f2a89d 100644 --- a/docs/architecture/learning.md +++ b/docs/architecture/learning.md @@ -523,3 +523,33 @@ Trace | `GENERATE` | LLM inference call | Engine | | `TOOL_CALL` | Tool execution | ToolExecutor | | `RESPOND` | Final response | TraceCollector | + +--- + +## Optimization Framework + +The optimization subsystem (`learning/optimize/`) provides LLM-guided search +over OpenJarvis's 5-pillar configuration space. It automates finding optimal +configurations for accuracy, latency, cost, and energy consumption. + +### Components + +| Component | Description | +|-----------|-------------| +| `SearchSpace` | Defines tunable dimensions across all 5 pillars | +| `LLMOptimizer` | Proposes configurations using an LLM backend | +| `OptimizationEngine` | Orchestrates the propose-evaluate-analyze loop | +| `OptimizationStore` | SQLite-backed persistence for trials and runs | +| `TrialRunner` | Evaluates proposed configurations against benchmarks | + +### Pareto Frontier + +The engine computes a Pareto frontier across multiple objectives +(accuracy vs latency vs cost), identifying configurations where no single +metric can be improved without degrading another. + +### Rust Backend + +The optimization framework has full Rust parity via the `openjarvis-learning` +crate, with PyO3 bindings exposing `OptimizationStore` and `LLMOptimizer` +to Python. diff --git a/docs/architecture/security.md b/docs/architecture/security.md index 59880ea6..aba46cb2 100644 --- a/docs/architecture/security.md +++ b/docs/architecture/security.md @@ -207,5 +207,5 @@ The default path is `~/.openjarvis/audit.db`, configurable via `security.audit_l ## See Also - [User Guide: Security](../user-guide/security.md) — how to configure and use the security system -- [API Reference: Security](../api/security.md) — complete class and function signatures +- [API Reference: Security](../api-reference/openjarvis/security/index.md) — complete class and function signatures - [Architecture: Query Flow](query-flow.md) — where security sits in the overall request lifecycle diff --git a/docs/deployment/index.md b/docs/deployment/index.md new file mode 100644 index 00000000..f57065eb --- /dev/null +++ b/docs/deployment/index.md @@ -0,0 +1,34 @@ +--- +title: Deployment +description: Deploy OpenJarvis in production environments +--- + +# Deployment + +OpenJarvis supports multiple deployment strategies for different environments +and scales. + +## Docker + +The recommended way to deploy OpenJarvis in production. Multi-stage builds +with CPU and GPU (NVIDIA CUDA, AMD ROCm) variants. + +[:octicons-arrow-right-24: Docker deployment](docker.md) + +## systemd (Linux) + +Run OpenJarvis as a managed system service on Linux servers. + +[:octicons-arrow-right-24: systemd setup](systemd.md) + +## launchd (macOS) + +Register OpenJarvis as a launch agent on macOS. + +[:octicons-arrow-right-24: launchd setup](launchd.md) + +## API Server + +Run OpenJarvis as an OpenAI-compatible HTTP server via `jarvis serve`. + +[:octicons-arrow-right-24: API server guide](api-server.md) diff --git a/docs/development/contributing.md b/docs/development/contributing.md index 99b8a01f..d6433255 100644 --- a/docs/development/contributing.md +++ b/docs/development/contributing.md @@ -18,7 +18,7 @@ contribute code to OpenJarvis. ### Clone and Install ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis uv sync --extra dev ``` diff --git a/docs/downloads.md b/docs/downloads.md index b5bdf248..5ac9aa48 100644 --- a/docs/downloads.md +++ b/docs/downloads.md @@ -17,7 +17,7 @@ processing happens on your local machine — the app connects to the backend you !!! info "Backend required" Start the backend before opening the desktop app. The quickstart script handles everything: ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git && cd OpenJarvis + git clone https://github.com/open-jarvis/OpenJarvis.git && cd OpenJarvis ./scripts/quickstart.sh ``` @@ -25,14 +25,14 @@ processing happens on your local machine — the app connects to the backend you | Platform | Download | Notes | |----------|----------|-------| -| macOS (Apple Silicon) | [:material-download: **OpenJarvis.dmg**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_aarch64.dmg) | M1/M2/M3/M4 Macs | -| Windows (64-bit) | [:material-download: **OpenJarvis-setup.exe**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_x64-setup.exe) | Windows 10+ | -| Linux (DEB) | [:material-download: **OpenJarvis.deb**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.deb) | Ubuntu, Debian | -| Linux (RPM) | [:material-download: **OpenJarvis.rpm**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis-1.0.0-1.x86_64.rpm) | Fedora, RHEL | -| Linux (AppImage) | [:material-download: **OpenJarvis.AppImage**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.AppImage) | Any distro | +| macOS (Apple Silicon) | [:material-download: **OpenJarvis.dmg**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_aarch64.dmg) | M1/M2/M3/M4 Macs | +| Windows (64-bit) | [:material-download: **OpenJarvis-setup.exe**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_x64-setup.exe) | Windows 10+ | +| Linux (DEB) | [:material-download: **OpenJarvis.deb**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.deb) | Ubuntu, Debian | +| Linux (RPM) | [:material-download: **OpenJarvis.rpm**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis-1.0.0-1.x86_64.rpm) | Fedora, RHEL | +| Linux (AppImage) | [:material-download: **OpenJarvis.AppImage**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.AppImage) | Any distro | !!! tip "All releases" - Browse all versions on the [GitHub Releases](https://github.com/HazyResearch/OpenJarvis/releases) page. + Browse all versions on the [GitHub Releases](https://github.com/open-jarvis/OpenJarvis/releases) page. ### macOS: "app is damaged" fix @@ -69,7 +69,7 @@ The backend (Ollama, Python API server, inference) runs separately on your machi ### Build from source ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis/desktop npm install npm run tauri build @@ -87,7 +87,7 @@ your machine and the frontend connects via `localhost`. ### One-command setup ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis ./scripts/quickstart.sh ``` @@ -107,7 +107,7 @@ If you prefer to run each step yourself: === "Step 1: Clone and install" ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis uv sync --extra server cd frontend && npm install && cd .. @@ -168,7 +168,7 @@ programmatically. Every feature is accessible from the terminal. === "From source" ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis uv sync ``` diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index b2b90f67..4cce729b 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -17,7 +17,7 @@ your machine and the frontend connects via `localhost`. ### One-command setup ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis ./scripts/quickstart.sh ``` @@ -37,7 +37,7 @@ If you prefer to run each step yourself: === "Step 1: Clone and install" ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis uv sync --extra server cd frontend && npm install && cd .. @@ -78,7 +78,7 @@ processing happens on your local machine — the app connects to the backend you **Step 1.** Start the backend (same as Browser App): ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis ./scripts/quickstart.sh ``` @@ -87,11 +87,11 @@ cd OpenJarvis | Platform | Download | |----------|----------| -| macOS (Apple Silicon) | [:material-download: **OpenJarvis.dmg**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_aarch64.dmg) | -| Windows (64-bit) | [:material-download: **OpenJarvis-setup.exe**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_x64-setup.exe) | -| Linux (DEB) | [:material-download: **OpenJarvis.deb**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.deb) | -| Linux (RPM) | [:material-download: **OpenJarvis.rpm**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis-1.0.0-1.x86_64.rpm) | -| Linux (AppImage) | [:material-download: **OpenJarvis.AppImage**](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.AppImage) | +| macOS (Apple Silicon) | [:material-download: **OpenJarvis.dmg**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_aarch64.dmg) | +| Windows (64-bit) | [:material-download: **OpenJarvis-setup.exe**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_x64-setup.exe) | +| Linux (DEB) | [:material-download: **OpenJarvis.deb**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.deb) | +| Linux (RPM) | [:material-download: **OpenJarvis.rpm**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis-1.0.0-1.x86_64.rpm) | +| Linux (AppImage) | [:material-download: **OpenJarvis.AppImage**](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.AppImage) | The app connects to `http://localhost:8000` automatically. @@ -103,12 +103,12 @@ The app connects to `http://localhost:8000` automatically. This is normal for open-source apps distributed outside the App Store. !!! tip "All releases" - Browse all versions on the [GitHub Releases](https://github.com/HazyResearch/OpenJarvis/releases) page. + Browse all versions on the [GitHub Releases](https://github.com/open-jarvis/OpenJarvis/releases) page. ### Build from source ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis/desktop npm install npm run tauri build @@ -140,7 +140,7 @@ programmatically. Every feature is accessible from the terminal. === "From source" ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis uv sync ``` diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index 5a11b97c..6f2539c1 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -15,7 +15,7 @@ This guide walks through the core workflows of OpenJarvis: the browser app, CLI, The quickest way to experience OpenJarvis is the full chat UI running in your browser: ```bash -git clone https://github.com/HazyResearch/OpenJarvis.git +git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis ./scripts/quickstart.sh ``` diff --git a/docs/index.md b/docs/index.md index 1980b2a8..b8c02b0a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,16 +1,17 @@ --- title: OpenJarvis -description: Programming abstractions for on-device AI +description: Composable, Programmable Systems for On-Device, Personal AI hide: - navigation --- -# _Programming abstractions_ for on-device AI +# _Composable, Programmable Systems_ for On-Device, Personal AI

-OpenJarvis is a modular framework for building, running, and learning from local AI systems. -Five composable pillars — each with a clear ABC interface and decorator-based registry. -Everything runs on your hardware. Cloud APIs are optional. +OpenJarvis is a research framework for composable, on-device AI systems. +Five programmable pillars — Intelligence, Engine, Agents, Tools, and Learning — +each with a clear ABC interface and decorator-based registry. +Build personal AI that runs on your hardware. Cloud APIs are optional.

> pip install openjarvis
@@ -24,7 +25,7 @@ Everything runs on your hardware. Cloud APIs are optional. Run the full chat UI locally with one script: ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis ./scripts/quickstart.sh ``` @@ -40,16 +41,16 @@ Everything runs on your hardware. Cloud APIs are optional. **Step 1.** Start the backend: ```bash - git clone https://github.com/HazyResearch/OpenJarvis.git + git clone https://github.com/open-jarvis/OpenJarvis.git cd OpenJarvis ./scripts/quickstart.sh ``` **Step 2.** Download and open the desktop app: - [Download for macOS (Apple Silicon)](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_aarch64.dmg){ .md-button .md-button--primary } + [Download for macOS (Apple Silicon)](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_aarch64.dmg){ .md-button .md-button--primary } - Also available for [Windows](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_x64-setup.exe), [Linux (DEB)](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.deb), and [Linux (RPM)](https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/OpenJarvis-1.0.0-1.x86_64.rpm). See the [Downloads](downloads.md) page for details. + Also available for [Windows](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_x64-setup.exe), [Linux (DEB)](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis_1.0.0_amd64.deb), and [Linux (RPM)](https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/OpenJarvis-1.0.0-1.x86_64.rpm). See the [Downloads](downloads.md) page for details. The app connects to `http://localhost:8000` automatically. @@ -168,7 +169,7 @@ Everything runs on your hardware. Cloud APIs are optional. Five-pillar design, registry pattern, query flow, and cross-cutting learning. -- **[API Reference](api/index.md)** +- **[API Reference](api-reference/openjarvis/index.md)** --- diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css index dcfd78e1..d55872e0 100644 --- a/docs/stylesheets/extra.css +++ b/docs/stylesheets/extra.css @@ -1,8 +1,6 @@ /* ── Fonts ─────────────────────────────────────────────────────────── */ -@import url('https://fonts.googleapis.com/css2?family=Merriweather:ital,wght@0,300;0,400;0,700;0,900;1,300;1,400;1,700&family=Inter:wght@400;500;600;700&display=swap'); - :root { - --md-text-font: "Merriweather", Georgia, "Times New Roman", serif; + --md-text-font: Georgia, "Times New Roman", serif; --md-code-font: "JetBrains Mono", "Fira Code", "SF Mono", monospace; } @@ -52,12 +50,12 @@ .md-footer, .md-search__input, .md-source { - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; } .md-nav__link, .md-tabs__link { - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; font-weight: 500; font-size: 0.8rem; letter-spacing: 0.01em; @@ -78,6 +76,10 @@ color: var(--md-accent-fg-color); } +.md-typeset h1 .headerlink { + display: none; +} + .hero-tagline { font-size: 1.05rem; color: var(--md-default-fg-color--light); @@ -167,7 +169,7 @@ .md-typeset .grid.cards > ol > li > :first-child, .md-typeset .grid.cards > ul > li > :first-child { - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; font-weight: 600; font-size: 0.92rem; letter-spacing: -0.01em; @@ -195,7 +197,7 @@ /* ── Tabs ─────────────────────────────────────────────────────────── */ .md-typeset .tabbed-labels > label { - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; font-weight: 500; font-size: 0.82rem; letter-spacing: 0.01em; @@ -203,7 +205,7 @@ /* ── Buttons (clean, bordered like OA) ───────────────────────────── */ .md-typeset .md-button { - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; font-weight: 500; font-size: 0.85rem; border-radius: 6px; @@ -257,7 +259,7 @@ display: inline-block; padding: 0.2rem 0.7rem; border-radius: 100px; - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; font-size: 0.72rem; font-weight: 600; letter-spacing: 0.04em; @@ -314,7 +316,7 @@ } .md-typeset table:not([class]) th { - font-family: "Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + font-family: Georgia, "Times New Roman", serif; font-weight: 600; font-size: 0.78rem; text-transform: uppercase; diff --git a/docs/tutorials/code-companion.md b/docs/tutorials/code-companion.md new file mode 100644 index 00000000..109f2a56 --- /dev/null +++ b/docs/tutorials/code-companion.md @@ -0,0 +1,218 @@ +--- +title: Code Companion +description: Code review, debugging, and test generation with ReAct agents +--- + +# Code Companion + +This tutorial walks through `examples/code_companion/` — three developer-focused scripts that use a `native_react` (ReAct) agent to automate common coding tasks: reviewing pull request diffs, investigating errors, and generating tests. Each script adapts the same core pattern to a different workflow, making it easy to extend for your own code intelligence use cases. + +!!! tip "Prerequisites" + - Python 3.10 or later + - OpenJarvis installed: `uv sync --extra dev` from the repository root + - An inference engine running — Ollama locally or a cloud API key in `.env` + - For `reviewer.py` and `code_review.py`: a git repository with at least two branches or commits + +## The Three Scripts + +| Script | Purpose | Tools Used | +|---|---|---| +| `reviewer.py` | Review a git diff between two branches | `git_diff`, `git_log`, `file_read`, `think` | +| `debugger.py` | Investigate an error and propose a fix | `file_read`, `shell_exec`, `think` | +| `test_gen.py` | Generate comprehensive tests for a Python module | `file_read`, `think`, `file_write` | + +All three use the `native_react` agent with the same SDK pattern. The difference is which tools are provided and how the prompt is structured. + +## The ReAct Agent Loop + +The `native_react` agent implements the Thought-Action-Observation cycle. Rather than producing a single response, it iterates until it has gathered enough information: + +```mermaid +stateDiagram-v2 + [*] --> Thought: Receive task prompt + Thought --> Action: Decide which tool to call + Action --> Observation: Execute tool, receive result + Observation --> Thought: Feed result back into context + Thought --> FinalAnswer: Sufficient information gathered + FinalAnswer --> [*] +``` + +This loop lets the agent adaptively explore the codebase. For example, the reviewer might read a diff, notice a suspicious function call, then read the source of that function before making its assessment — without any of that branching logic being hardcoded in the script. + +## Core SDK Pattern + +All three scripts follow the same structure. Understanding this pattern lets you adapt it to any code intelligence task: + +```python title="Core SDK pattern" hl_lines="4 5 6" +from openjarvis import Jarvis + +j = Jarvis(model="qwen3:8b", engine_key="ollama") # (1)! +try: + response = j.ask( + prompt, # (2)! + agent="native_react", # (3)! + tools=["git_diff", "think"], # (4)! + ) + print(response) +finally: + j.close() # (5)! +``` + +1. Both `model` and `engine_key` are optional. Omitting them uses auto-detected defaults from `~/.openjarvis/config.toml`. +2. The prompt describes the task in detail, including what tools to use, what steps to follow, and what the output structure should look like. +3. `"native_react"` selects the `NativeReActAgent`. The alias `"react"` also works. +4. The tool list is passed directly. Any registered tool name is valid — run `jarvis agent info native_react` to see all available tools. +5. Always call `j.close()` to release engine resources. A `try/finally` block ensures cleanup even if the agent raises an exception. + +## Code Review + +The `reviewer.py` script reviews the diff between two git refs and produces structured feedback with issues, suggestions, and an overall verdict. + +```bash title="Terminal" +# Review a feature branch against main (default) +python examples/code_companion/reviewer.py --branch feature-x + +# Review a specific commit range +python examples/code_companion/reviewer.py --branch HEAD --base develop + +# Use a cloud model for larger diffs +python examples/code_companion/reviewer.py \ + --branch feature-x --model gpt-4o --engine cloud +``` + +The agent follows a four-step process: + +1. Call `git_diff` to see what changed between the two refs +2. Call `git_log` to understand the commit history and intent +3. Call `file_read` on any files that need more context +4. Call `think` to reason about code quality, bugs, and design decisions + +The final output is structured with four sections: **Summary**, **Issues Found**, **Suggestions**, and **Overall Assessment** (APPROVE, REQUEST CHANGES, or COMMENT). + +| Flag | Default | Description | +|---|---|---| +| `--branch` | `HEAD` | Branch or commit to review | +| `--base` | `main` | Base branch to diff against | +| `--model` | `qwen3:8b` | Model identifier | +| `--engine` | `ollama` | Engine backend | + +## Debug Assistant + +The `debugger.py` script takes an error message, optionally a file path, and produces a root cause analysis with a concrete fix. + +```bash title="Terminal" +# Investigate a TypeError +python examples/code_companion/debugger.py \ + --error "TypeError: NoneType has no attribute 'split'" + +# Provide the file where the error occurred for faster analysis +python examples/code_companion/debugger.py \ + --error "KeyError: 'user_id'" \ + --file src/app/views.py + +# Use a cloud model for complex stack traces +python examples/code_companion/debugger.py \ + --error "Segfault in libfoo.so" \ + --model gpt-4o --engine cloud +``` + +The agent uses `file_read` to examine the relevant source, `shell_exec` to run diagnostic commands (grep for symbols, check imports, inspect directory contents), and `think` to reason about root causes before proposing a fix. + +!!! note "shell_exec safety" + The `shell_exec` tool runs commands in the current working directory. In production deployments, `ToolExecutor` enforces RBAC capability policies — ensure the `shell_exec` capability is permitted for the agent's role. See [Architecture: Security](../architecture/security.md). + +The output has three sections: **Root Cause**, **Proposed Fix** (concrete code change), and **Prevention** (type hints, validation, tests). + +| Flag | Default | Description | +|---|---|---| +| `--error` | (required) | Error message or stack trace | +| `--file` | (none) | Optional file path where the error occurred | +| `--model` | `qwen3:8b` | Model identifier | +| `--engine` | `ollama` | Engine backend | + +## Test Generator + +The `test_gen.py` script reads a Python module, reasons about its public interface, and writes a complete test file. + +```bash title="Terminal" +# Generate pytest tests for a module +python examples/code_companion/test_gen.py \ + --module src/openjarvis/tools/calculator.py + +# Use unittest and specify the output file +python examples/code_companion/test_gen.py \ + --module src/openjarvis/tools/calculator.py \ + --framework unittest \ + --output tests/test_calculator_generated.py +``` + +The agent reads the module with `file_read`, uses `think` to plan test cases (happy paths, edge cases, error handling, boundary conditions), reads any related base classes for context, then writes the complete test file with `file_write`. + +!!! note "Output path default" + If `--output` is not specified, the generated file is saved as `test_.py` in the current working directory. The script prints the output path when done. + +The generated tests follow these guidelines (enforced via the prompt): + +- Every public function and method has at least one test +- Each test has a docstring explaining what it verifies +- Edge cases are covered: empty input, `None`, large values, invalid types +- External dependencies are mocked with `unittest.mock` +- The file is self-contained and runnable with `pytest` or `unittest` without modification + +| Flag | Default | Description | +|---|---|---| +| `--module` | (required) | Path to the Python module | +| `--framework` | `pytest` | Test framework (`pytest` or `unittest`) | +| `--output` | `test_.py` | Output file path | +| `--model` | `qwen3:8b` | Model identifier | +| `--engine` | `ollama` | Engine backend | + +## Engine Selection + +=== "Ollama (local)" + + ```bash title="Terminal" + ollama serve + ollama pull qwen3:8b + python examples/code_companion/reviewer.py --branch feature-x + ``` + +=== "Cloud API" + + ```bash title="Terminal" + source .env # load OPENAI_API_KEY or similar + python examples/code_companion/reviewer.py \ + --branch feature-x \ + --model gpt-4o \ + --engine cloud + ``` + +## Customization + +### Change the tool set + +Edit the `tools` list in any script to add or remove tools. For example, to let the reviewer also search the web for known security advisories related to dependencies it sees in the diff: + +```python +tools = ["git_diff", "git_log", "file_read", "think", "web_search"] +``` + +### Adjust the prompt + +Each script contains a `prompt` string that instructs the agent what to do and what to produce. Modify it to match your team's conventions — different review sections, specific coding standards, or a particular output format for downstream tooling. + +### Add memory + +For multi-session workflows (e.g., a reviewer that remembers previous assessments of the same files), add `"memory_store"` and `"memory_search"` to the tool list and update the prompt to use them: + +```python +tools = ["git_diff", "git_log", "file_read", "think", + "memory_store", "memory_search"] +``` + +## See Also + +- [Architecture: Agents](../architecture/agents.md) — `NativeReActAgent` internals and the Thought-Action-Observation loop +- [Architecture: Tools and Memory](../architecture/memory.md) — git tools, file tools, shell tools, and the `ToolExecutor` dispatch pipeline +- [Architecture: Security](../architecture/security.md) — RBAC capability policies for `shell_exec` and other privileged tools +- [Tutorials: Deep Research Assistant](deep-research.md) — the same SDK pattern with the `OrchestratorAgent` and web/memory tools diff --git a/docs/tutorials/deep-research.md b/docs/tutorials/deep-research.md new file mode 100644 index 00000000..ffad91f4 --- /dev/null +++ b/docs/tutorials/deep-research.md @@ -0,0 +1,197 @@ +--- +title: Deep Research Assistant +description: Build a multi-source research agent with memory-augmented orchestration +--- + +# Deep Research Assistant + +This tutorial walks through `examples/deep_research/research.py` — a standalone script that uses an orchestrator agent to research a topic, gather sources across multiple tool-calling turns, and produce a cited report. It demonstrates how to compose web search, memory, and file output into a single coherent research workflow. + +!!! tip "Prerequisites" + - Python 3.10 or later + - OpenJarvis installed: run `uv sync --extra dev` from the repository root + - An inference engine running — either Ollama locally (see below) or a cloud API key in your `.env` file + +## Quick Start + +Run the research script from the repository root, passing your topic as a positional argument: + +```bash title="Terminal" +python examples/deep_research/research.py "quantum computing advances 2026" +``` + +Save the report to a file: + +```bash title="Terminal" +python examples/deep_research/research.py "quantum computing advances 2026" \ + --output report.md +``` + +Use a cloud model instead of a local engine: + +```bash title="Terminal" +source .env # load API keys +python examples/deep_research/research.py "climate policy trends" \ + --model gpt-4o --engine cloud --max-turns 20 +``` + +## How It Works + +The script creates a `Jarvis` instance and delegates the research task to an `OrchestratorAgent` with five tools wired in. The orchestrator iterates through multiple tool-calling turns, deciding at each step whether to search, store, think, or synthesize. + +```mermaid +sequenceDiagram + participant U as User + participant J as Jarvis SDK + participant O as OrchestratorAgent + participant W as web_search + participant T as think + participant MS as memory_store + participant MQ as memory_search + participant F as file_write + + U->>J: research.py "quantum computing" + J->>O: ask(prompt, agent="orchestrator", tools=[...]) + loop Up to max_turns iterations + O->>W: search("quantum computing 2026") + W-->>O: search results + O->>T: think(reasoning about findings) + T-->>O: structured thoughts + O->>MS: store(key finding) + MS-->>O: stored + O->>MQ: search(earlier findings) + MQ-->>O: related context + end + O->>F: file_write(report.md) + F-->>O: saved + O-->>J: final report with citations + J-->>U: print report +``` + +Each turn the orchestrator decides which tool to call based on what it has learned so far. The `think` tool lets the model reason without side effects, while `memory_store` and `memory_search` provide persistent scratch space across turns — so a finding from turn 3 can still inform the synthesis in turn 12. + +## The Script + +```python title="examples/deep_research/research.py" hl_lines="9 10 11 12 13" +from openjarvis import Jarvis + +tools = ["web_search", "think", "file_write", "memory_store", "memory_search"] + +j = Jarvis(model="qwen3:8b", engine_key="ollama") # (1)! +try: + response = j.ask( + "Research the following topic in depth and produce a report:\n\nquantum computing", + agent="orchestrator", # (2)! + tools=tools, # (3)! + system_prompt=..., # (4)! + max_turns=15, # (5)! + temperature=0.5, + ) +finally: + j.close() +``` + +1. Creates a `Jarvis` instance targeting the local Ollama engine with `qwen3:8b`. Both parameters are optional — omitting them uses auto-detected defaults from `~/.openjarvis/config.toml`. +2. Selects the `OrchestratorAgent`, which runs a multi-turn tool-calling loop rather than a single round-trip. +3. The tool list is passed directly to the agent. All five tools are registered in the tool registry and need no further configuration. +4. The system prompt instructs the model to cite sources and distinguish facts from emerging claims. +5. The loop terminates after 15 tool-calling turns or when the agent decides it has enough information. + +## Engine Configuration + +=== "Ollama (local)" + + Start the Ollama daemon and pull the model before running the script: + + ```bash title="Terminal" + ollama serve + ollama pull qwen3:8b + python examples/deep_research/research.py "your topic here" + ``` + + No flags needed — `--engine ollama` and `--model qwen3:8b` are the defaults. + +=== "Cloud API" + + Set your API key in `.env`, then pass `--engine cloud` and the appropriate model identifier: + + ```bash title="Terminal" + # .env (in the repository root, gitignored) + OPENAI_API_KEY=sk-... + + source .env + python examples/deep_research/research.py "your topic" \ + --model gpt-4o --engine cloud + ``` + +=== "vLLM" + + If you are running a vLLM inference server (e.g., on a multi-GPU node): + + ```bash title="Terminal" + python examples/deep_research/research.py "your topic" \ + --model meta-llama/Meta-Llama-3-8B-Instruct \ + --engine vllm + ``` + + Make sure `VLLM_BASE_URL` is set in `.env` pointing to your vLLM server. + +## Configuration Reference + +| Flag | Default | Description | +|---|---|---| +| `--model` | `qwen3:8b` | Model identifier passed to the engine | +| `--engine` | `ollama` | Engine backend (`ollama`, `cloud`, `vllm`, `llamacpp`, `mlx`) | +| `--max-turns` | `15` | Maximum orchestrator loop iterations | +| `--output` | (none) | File path to save the final report; if omitted, prints to stdout | + +## Recipe-Driven Configuration + +The companion `research.toml` in `examples/deep_research/` expresses the same setup declaratively. You can load it programmatically with `load_recipe()` and pass the result to `SystemBuilder`: + +```python title="Using the recipe" +from openjarvis.recipes import load_recipe +from openjarvis import SystemBuilder + +recipe = load_recipe("examples/deep_research/research.toml") +system = SystemBuilder(**recipe.to_builder_kwargs()).build() +response = system.ask("quantum computing advances 2026") +system.close() +``` + +This is useful when you want to version-control the research configuration, share it with collaborators, or feed it to the `jarvis eval` runner for benchmarking. + +## Customization + +### Swap the agent + +Replace `"orchestrator"` with `"native_react"` for a Thought-Action-Observation loop, or `"native_openhands"` for a CodeAct-style agent that can write and execute code: + +```python +response = j.ask(prompt, agent="native_react", tools=tools) +``` + +### Add more tools + +Append any registered tool name to the `tools` list. For example, to also query a local knowledge base: + +```python +tools = ["web_search", "think", "file_write", + "memory_store", "memory_search", "knowledge_graph_query"] +``` + +Run `jarvis agent info orchestrator` to see the full tool catalog. + +### Adjust temperature + +Lower values (0.2) produce more focused, factual reports. Higher values (0.7-0.8) encourage broader exploration and more creative synthesis: + +```bash title="Terminal" +python examples/deep_research/research.py "your topic" --max-turns 20 +``` + +## See Also + +- [Architecture: Agents](../architecture/agents.md) — agent hierarchy (`BaseAgent`, `ToolUsingAgent`, `OrchestratorAgent`) and the `accepts_tools` mechanism +- [Architecture: Tools and Memory](../architecture/memory.md) — tool registry, MCP adapter, and the `ToolExecutor` dispatch pipeline +- [Getting Started: Configuration](../getting-started/configuration.md) — how to configure engines and models in `~/.openjarvis/config.toml` diff --git a/docs/tutorials/index.md b/docs/tutorials/index.md new file mode 100644 index 00000000..408e99cd --- /dev/null +++ b/docs/tutorials/index.md @@ -0,0 +1,62 @@ +--- +title: Tutorials +description: Step-by-step guides for building with OpenJarvis +--- + +# Tutorials + +Hands-on guides that walk through building real applications with OpenJarvis. Each tutorial includes a standalone script you can run immediately, a TOML recipe for configuration, and a detailed walkthrough of the concepts involved. + +!!! note "Before you begin" + All tutorials assume OpenJarvis is installed and an inference engine is running. If you have not completed setup yet, start with the [Quick Start guide](../getting-started/quickstart.md). + +
+ +- :material-magnify:{ .lg .middle } **Deep Research Assistant** + + --- + + Multi-source research with a memory-augmented orchestrator agent. Searches the web, stores findings across turns, cross-references sources, and produces a cited report. + + [:octicons-arrow-right-24: Get started](deep-research.md) + +- :material-clock-outline:{ .lg .middle } **Scheduled Personal Ops** + + --- + + Autonomous agents on cron schedules for recurring personal tasks — morning news digests, weekly code reviews, and gym schedule checks. + + [:octicons-arrow-right-24: Get started](scheduled-ops.md) + +- :material-message-outline:{ .lg .middle } **Messaging Hub** + + --- + + Smart inbox assistant that triages messages by priority, drafts context-aware replies, and produces end-of-day summaries across Slack, WhatsApp, and other channels. + + [:octicons-arrow-right-24: Get started](messaging-hub.md) + +- :material-code-braces:{ .lg .middle } **Code Companion** + + --- + + Code review, debugging, and test generation using a ReAct agent that reads source files, runs commands, and reasons step by step before producing structured output. + + [:octicons-arrow-right-24: Get started](code-companion.md) + +
+ +## What You Will Learn + +Each tutorial demonstrates a different combination of OpenJarvis pillars working together: + +| Tutorial | Agent | Key Pillars | +|---|---|---| +| Deep Research | `orchestrator` | Engine, Agents, Tools (web + memory), Recipes | +| Scheduled Ops | `orchestrator`, `native_react` | Agents, Tools, Scheduler | +| Messaging Hub | `orchestrator` | Agents, Tools (memory), Channels | +| Code Companion | `native_react` | Agents, Tools (git + file + shell) | + +## Estimated Time + +Each tutorial takes approximately 15-30 minutes to complete end-to-end, including setup and running the scripts. The TOML configuration sections and customization tips are optional reading for when you adapt the pattern to your own use case. diff --git a/docs/tutorials/messaging-hub.md b/docs/tutorials/messaging-hub.md new file mode 100644 index 00000000..98176a0a --- /dev/null +++ b/docs/tutorials/messaging-hub.md @@ -0,0 +1,215 @@ +--- +title: Messaging Hub +description: Smart inbox with message triage and auto-replies across channels +--- + +# Messaging Hub + +This tutorial walks through `examples/messaging_hub/smart_inbox.py` — a script that connects OpenJarvis to messaging platforms, triages incoming messages by priority, drafts context-aware replies, and produces end-of-day summaries. It demonstrates channel integration, structured agent output, and memory-backed aggregation across multiple messages. + +!!! tip "Prerequisites" + - Python 3.10 or later + - OpenJarvis installed: `uv sync --extra dev` from the repository root + - An inference engine running (Ollama with `qwen3:8b` pulled, or cloud API keys) + - For live channel mode: channel-specific credentials (see [Setting Up Real Channels](#setting-up-real-channels)) + +## Quick Start: Demo Mode + +Demo mode processes five sample messages with no channel setup or credentials required. It is the fastest way to see the triage pipeline in action: + +```bash title="Terminal" +python examples/messaging_hub/smart_inbox.py --demo +``` + +Expected output (abbreviated): + +``` +Smart Inbox — Demo Mode +Model: qwen3:8b | Engine: ollama +============================================================ +Processing 5 messages... + + [1/5] Classifying: URGENT: Server is down in production... + -> URGENT + [2/5] Classifying: Hey, just wanted to share this interest... + -> FYI + [3/5] Classifying: Can you review my PR #42 by end of day... + -> ACTION_REQUIRED + [4/5] Classifying: Meeting reminder: Team standup at 10am... + -> FYI + [5/5] Classifying: Buy now! Limited time offer on premium... + -> SPAM + + # Category Message Reply + --------------------------------------------------------------- + 1 URGENT URGENT: Server is down... On it — escalating now. + 2 FYI Hey, just wanted to share... Thanks for sharing! + 3 ACTION_REQUIRED Can you review my PR #42... Will review before EOD. + 4 FYI Meeting reminder: Team standup... N/A + 5 SPAM Buy now! Limited time offer... N/A + +Generating end-of-day summary... +``` + +Override the model or engine: + +```bash title="Terminal" +python examples/messaging_hub/smart_inbox.py --demo --model gpt-4o --engine cloud +``` + +## How Message Classification Works + +Each incoming message goes through a structured prompt that asks the agent to output exactly two fields — a category and a reply — in a parseable format. The script then extracts those fields and builds a triage table. + +```mermaid +graph TD + A[Incoming message] --> B[OrchestratorAgent] + B --> C{think tool: internal reasoning} + C --> D{memory_store: persist context} + D --> E[Structured response] + E --> F{Parse CATEGORY and REPLY} + F -->|URGENT| G[Flag for immediate attention] + F -->|ACTION_REQUIRED| H[Add to action list] + F -->|FYI| I[Log for reference] + F -->|SPAM| J[Discard] + G --> K[Triage table] + H --> K + I --> K + J --> K + K --> L[memory_search: cross-reference] + L --> M[End-of-day summary] +``` + +After all messages are processed, a second orchestrator call uses `memory_search` to retrieve the stored triage log and produces a grouped summary with open action items highlighted. + +## The Classification Prompt + +The agent receives a structured prompt that specifies the output format exactly. This makes the response reliably parseable without a complex schema: + +```python title="examples/messaging_hub/smart_inbox.py" +CLASSIFICATION_PROMPT = ( + "You are a smart inbox assistant. Classify the following message into " + "exactly one category: URGENT, ACTION_REQUIRED, FYI, or SPAM.\n" + "Then draft a short reply if appropriate (not for SPAM).\n\n" + "Respond in this exact format:\n" + "CATEGORY: \n" + "REPLY: \n\n" + "Message:\n{message}" +) +``` + +The `think` tool lets the agent reason internally before committing to a category, and `memory_store` persists each classification so the summary prompt can reference the full triage log. + +## Setting Up Real Channels + +=== "Slack" + + 1. Add the Slack MCP server to your configuration: + + ```bash title="Terminal" + jarvis add slack + ``` + + 2. Set your credentials in `.env` (gitignored): + + ```bash title=".env" + SLACK_BOT_TOKEN=xoxb-... + SLACK_APP_TOKEN=xapp-... + ``` + + 3. Invite the bot to the target Slack channel in the Slack workspace settings. + + 4. Run the script in live channel mode: + + ```bash title="Terminal" + python examples/messaging_hub/smart_inbox.py --channel slack + ``` + +=== "WhatsApp" + + 1. Ensure Node.js 22 or later is installed. + + 2. Configure the WhatsApp Baileys bridge. See the [channel documentation](../architecture/overview.md) for full setup steps. + + 3. Start the bridge — it will print a QR code. Scan it with the WhatsApp mobile app to authenticate. + + 4. Run the script: + + ```bash title="Terminal" + python examples/messaging_hub/smart_inbox.py --channel whatsapp + ``` + +=== "Other Channels" + + OpenJarvis supports LINE, Viber, Mastodon, Rocket.Chat, Zulip, XMPP, Twitch, Nostr, and more. List all available channels: + + ```bash title="Terminal" + jarvis channel list + jarvis channel status + ``` + + Each channel requires its own environment variables. Run `jarvis add ` where available to auto-generate the configuration template. + +!!! warning "Live channel mode" + Live channel mode requires channel credentials and the corresponding channel subsystem to be running. Use `--demo` to verify the triage logic before connecting to a real channel. + +## Channel Configuration via TOML + +The `messaging.toml` recipe in `examples/messaging_hub/` captures the agent and channel defaults declaratively: + +```toml title="examples/messaging_hub/messaging.toml" +[channel] +default = "slack" + +[agent] +type = "orchestrator" +max_turns = 5 +temperature = 0.3 +tools = ["think", "memory_store", "memory_search"] +``` + +You can load this recipe programmatically: + +```python title="Loading the messaging recipe" +from openjarvis.recipes import load_recipe +from openjarvis import SystemBuilder + +recipe = load_recipe("examples/messaging_hub/messaging.toml") +system = SystemBuilder(**recipe.to_builder_kwargs()).build() +response = system.ask(CLASSIFICATION_PROMPT.format(message=incoming_message)) +system.close() +``` + +## Adding Custom Triage Rules + +Extend the classification categories by editing `CLASSIFICATION_PROMPT`. For example, to add a `FOLLOW_UP` category for messages that need a response within 48 hours: + +```python title="Custom classification prompt" hl_lines="2" +CLASSIFICATION_PROMPT = ( + "Classify into: URGENT, ACTION_REQUIRED, FOLLOW_UP, FYI, or SPAM.\n" + "Then draft a short reply if appropriate (not for SPAM).\n\n" + "Respond in this exact format:\n" + "CATEGORY: \n" + "REPLY: \n\n" + "Message:\n{message}" +) +``` + +You can also add domain rules in the system prompt via `messaging.toml` — for instance, routing any message containing "P0" or "incident" directly to URGENT regardless of phrasing. + +## Scheduling the Daily Summary + +After processing all messages, the end-of-day summary call runs immediately in the script. For production use, schedule it independently via the OpenJarvis scheduler: + +```bash title="Terminal" +jarvis scheduler create "Daily inbox summary" \ + --type cron --value "0 17 * * *" +``` + +Or use the operator recipe pattern to run a persistent triage agent on a schedule. See the operator recipes in `src/openjarvis/recipes/data/operators/` for ready-made examples. + +## See Also + +- [Architecture: Agents](../architecture/agents.md) — `OrchestratorAgent` and the multi-turn tool loop +- [Architecture: Tools and Memory](../architecture/memory.md) — `memory_store`, `memory_search`, and the storage backends +- [Tutorials: Scheduled Personal Ops](scheduled-ops.md) — combining scripts with the cron scheduler diff --git a/docs/tutorials/scheduled-ops.md b/docs/tutorials/scheduled-ops.md new file mode 100644 index 00000000..f4d07bf8 --- /dev/null +++ b/docs/tutorials/scheduled-ops.md @@ -0,0 +1,215 @@ +--- +title: Scheduled Personal Ops +description: Run autonomous agents on cron schedules for recurring personal tasks +--- + +# Scheduled Personal Ops + +This tutorial walks through `examples/scheduled_ops/` — three scripts that run autonomous agents on cron-like schedules to handle recurring personal tasks. Together they demonstrate how to combine the `Jarvis` SDK, the scheduler CLI, and the Python `TaskScheduler` API to build a personal operations layer that runs in the background. + +!!! tip "Prerequisites" + - Python 3.10 or later + - OpenJarvis installed: `uv sync --extra dev` from the repository root + - An inference engine running (Ollama with `qwen3:8b` pulled, or a cloud API key) + - For full cron expression support, install `croniter`: `uv add croniter` + +## The Three Scripts + +| Script | Agent | Tools | Default Schedule | Purpose | +|---|---|---|---|---| +| `daily_digest.py` | `orchestrator` | `web_search`, `think` | Daily 9:00 AM | Search and summarize top news for chosen topics | +| `code_review.py` | `native_react` | `git_log`, `git_diff`, `file_read`, `think` | Monday 8:00 AM | Review the past week of commits in a repository | +| `gym_scheduler.py` | `orchestrator` | `web_search`, `think` | MWF 6:00 AM | Check gym hours and class availability | + +Each script follows the same SDK pattern: create a `Jarvis` instance, call `j.ask()` with an agent and tools, print the result, and close the instance. The schedule is managed externally by the OpenJarvis scheduler daemon. + +## Quick Start: Run Scripts Manually + +Test each script without a running scheduler by invoking it directly: + +```bash title="Terminal" +# Morning news digest for AI and robotics +uv run python examples/scheduled_ops/daily_digest.py --topics "AI,robotics" + +# Code review for the current repository (last 7 days of commits) +uv run python examples/scheduled_ops/code_review.py --repo-path . + +# Gym schedule check +uv run python examples/scheduled_ops/gym_scheduler.py --gym "24 Hour Fitness" +``` + +All scripts accept `--model` and `--engine` flags: + +```bash title="Terminal" +uv run python examples/scheduled_ops/daily_digest.py \ + --model qwen3:8b --engine ollama --topics "AI,finance" +``` + +## How the Scheduler Works + +```mermaid +graph TD + A[jarvis scheduler start] --> B[Scheduler Daemon] + B --> C{Cron trigger fires} + C -->|0 9 * * *| D[daily_digest.py] + C -->|0 8 * * 1| E[code_review.py] + C -->|0 6 * * 1,3,5| F[gym_scheduler.py] + D --> G[OrchestratorAgent] + E --> H[NativeReActAgent] + F --> G + G --> I[web_search + think] + H --> J[git_diff + git_log + file_read + think] + I --> K[Output / Channel] + J --> K +``` + +The scheduler daemon reads registered tasks from SQLite, fires them at the correct time, and passes the configured prompt to the agent. Each script can also be run directly — the scheduler is only needed for recurring, unattended operation. + +## Set Up Schedules with the CLI + +Register each script as a recurring task using `jarvis scheduler create`: + +```bash title="Terminal" +# Morning digest every day at 9 AM +jarvis scheduler create "Run daily news digest" \ + --type cron --value "0 9 * * *" + +# Weekly code review every Monday at 8 AM +jarvis scheduler create "Run weekly code review" \ + --type cron --value "0 8 * * 1" + +# Gym check on Monday, Wednesday, Friday at 6 AM +jarvis scheduler create "Check gym schedule" \ + --type cron --value "0 6 * * 1,3,5" +``` + +Then start the scheduler daemon in the foreground (or as a background service): + +```bash title="Terminal" +jarvis scheduler start +``` + +List registered tasks at any time: + +```bash title="Terminal" +jarvis scheduler list +``` + +!!! note "Cron expression syntax" + OpenJarvis uses standard five-field cron syntax: `minute hour day-of-month month day-of-week`. Install `croniter` (`uv add croniter`) for full expression support including ranges and step values. Without it, basic `hour:minute` patterns still work. + +## Configure Schedules with TOML + +The `schedules.toml` file in `examples/scheduled_ops/` defines all three schedules declaratively. This is convenient for version-controlling your personal ops configuration or sharing it across machines: + +```toml title="examples/scheduled_ops/schedules.toml" +[schedules.daily_digest] +type = "cron" +value = "0 9 * * *" +description = "Morning news and social media digest" +script = "daily_digest.py" + +[schedules.code_review] +type = "cron" +value = "0 8 * * 1" +description = "Weekly code review" +script = "code_review.py" + +[schedules.gym_scheduler] +type = "cron" +value = "0 6 * * 1,3,5" +description = "Gym hours and class check" +script = "gym_scheduler.py" +``` + +Point your own tooling or a custom loader at this file to register tasks in bulk. + +## Register Tasks via the Python API + +The `gym_scheduler.py` script includes a `--register` flag that demonstrates programmatic task registration using `TaskScheduler` directly: + +```bash title="Terminal" +uv run python examples/scheduled_ops/gym_scheduler.py \ + --register --gym "Planet Fitness" +``` + +The equivalent Python code: + +```python title="Programmatic task registration" +from openjarvis.scheduler import TaskScheduler +from openjarvis.scheduler.store import SchedulerStore + +store = SchedulerStore() +scheduler = TaskScheduler(store) + +task = scheduler.create_task( # (1)! + prompt="Check gym schedule for 'Planet Fitness'", + schedule_type="cron", + schedule_value="0 6 * * 1,3,5", + agent="orchestrator", + tools="web_search,think", +) +print(f"Task registered: {task.id}") +print(f"Next run: {task.next_run}") +``` + +1. `create_task()` persists the task to SQLite and computes the next trigger time. The scheduler daemon picks it up without a restart. + +## The Daily Digest Script + +The digest script is the simplest of the three. It builds a date-stamped prompt and passes it to an orchestrator with `web_search` and `think`: + +```python title="examples/scheduled_ops/daily_digest.py" hl_lines="5 6 7 8" +from openjarvis import Jarvis + +j = Jarvis() # uses defaults from ~/.openjarvis/config.toml +response = j.ask( + f"Today is {today}. Search and summarize the top news on: {topics}", + agent="orchestrator", + tools=["web_search", "think"], +) +j.close() +``` + +The orchestrator searches for each topic in a separate turn, uses `think` to synthesize across topics, and returns a structured digest with bullet-point summaries and a one-paragraph outlook. + +## Send Results to a Channel + +To route script output to Slack or any other supported channel, pipe stdout through `jarvis channel send`: + +```bash title="Terminal" +uv run python examples/scheduled_ops/daily_digest.py \ + --topics "AI,finance" | jarvis channel send slack +``` + +Or add channel output inside the script: + +```python title="In-script channel output" +from openjarvis.channels import ChannelRegistry + +channel = ChannelRegistry.create("slack", webhook_url="https://hooks.slack.com/...") +channel.send(response) +``` + +List all available channels: + +```bash title="Terminal" +jarvis channel list +``` + +!!! warning "Channel credentials" + Live channel output requires channel-specific credentials. Run `jarvis add slack` (or the relevant provider) to set up the MCP server and credential store, then configure environment variables in your `.env` file before starting the scheduler daemon. + +## Customization Tips + +- **Change topics**: Pass `--topics "finance,healthcare,sports"` to `daily_digest.py` for a different digest. +- **Review window**: Pass `--days 14` to `code_review.py` for a two-week review cycle instead of one week. +- **Swap agents**: Replace `orchestrator` with `native_react` in any script to compare agent behavior on the same task. +- **Add file output**: Append `"file_write"` to the `tools` list and update the prompt to save reports to disk instead of printing them. +- **One-time tasks**: Use `--type once --value "2026-04-01T09:00:00"` with `jarvis scheduler create` for non-recurring tasks. + +## See Also + +- [Architecture: Agents](../architecture/agents.md) — `OrchestratorAgent` and `NativeReActAgent` internals +- [Architecture: Tools and Memory](../architecture/memory.md) — tool registry and `ToolExecutor` +- [Getting Started: Configuration](../getting-started/configuration.md) — engine and model defaults diff --git a/docs/user-guide/channels.md b/docs/user-guide/channels.md index d7a6c760..8436f3aa 100644 --- a/docs/user-guide/channels.md +++ b/docs/user-guide/channels.md @@ -485,6 +485,6 @@ assistant_has_own_number = false ## See Also - [Architecture: Channels](../architecture/channels.md) — listener loop internals and reconnect design -- [API Reference: Channels](../api/channels.md) — full class and type signatures +- [API Reference: Channels](../api-reference/openjarvis/channels/index.md) — full class and type signatures - [Getting Started: Configuration](../getting-started/configuration.md) — full config reference - [OpenClaw Agent](agents.md) — the agent infrastructure that uses OpenClaw transport diff --git a/docs/user-guide/security.md b/docs/user-guide/security.md index f574c546..95651957 100644 --- a/docs/user-guide/security.md +++ b/docs/user-guide/security.md @@ -449,6 +449,6 @@ guarded = GuardrailsEngine( ## See Also - [Architecture: Security](../architecture/security.md) — pipeline design, event flow, and file policy integration -- [API Reference: Security](../api/security.md) — full class and function signatures +- [API Reference: Security](../api-reference/openjarvis/security/index.md) — full class and function signatures - [Tools](tools.md) — how `FileReadTool` uses file policy - [Configuration](../getting-started/configuration.md) — full config reference diff --git a/examples/code_companion/README.md b/examples/code_companion/README.md new file mode 100644 index 00000000..4679e01f --- /dev/null +++ b/examples/code_companion/README.md @@ -0,0 +1,148 @@ +# Code Companion + +A set of developer-focused scripts that use OpenJarvis tool-using agents +to automate common coding tasks: code review, debugging, and test generation. + +## What This Demonstrates + +Each script wires up a `Jarvis` instance with the `native_react` (ReAct) agent +and a curated set of tools. The ReAct loop lets the agent reason step by step +-- reading files, running commands, and thinking -- before producing a final +structured answer. This is the same pattern you can adapt for any code +intelligence workflow. + +| Script | Purpose | Tools Used | +|---|---|---| +| `reviewer.py` | Review a git diff between two branches | `git_diff`, `git_log`, `file_read`, `think` | +| `debugger.py` | Investigate an error and propose a fix | `file_read`, `shell_exec`, `think` | +| `test_gen.py` | Generate comprehensive tests for a Python module | `file_read`, `think`, `file_write` | + +## Prerequisites + +1. **Install OpenJarvis** (from the repo root): + + ```bash + uv sync --extra dev + ``` + +2. **Start an inference engine.** The default is Ollama: + + ```bash + ollama serve + ollama pull qwen3:8b + ``` + + Alternatively, set up a cloud engine by sourcing your API keys: + + ```bash + source .env + ``` + +## Quick Start + +### Code Review + +Review the diff between a feature branch and `main`: + +```bash +python examples/code_companion/reviewer.py --branch feature-x +``` + +Review the current HEAD against a specific base: + +```bash +python examples/code_companion/reviewer.py --branch HEAD --base develop +``` + +### Debug Assistant + +Investigate an error message: + +```bash +python examples/code_companion/debugger.py --error "TypeError: NoneType has no attribute 'split'" +``` + +Point it at the file where the error occurred for faster root-cause analysis: + +```bash +python examples/code_companion/debugger.py \ + --error "KeyError: 'user_id'" \ + --file src/app/views.py +``` + +### Test Generation + +Generate pytest tests for a module: + +```bash +python examples/code_companion/test_gen.py --module src/openjarvis/tools/calculator.py +``` + +Use unittest instead, and write to a specific file: + +```bash +python examples/code_companion/test_gen.py \ + --module src/openjarvis/tools/calculator.py \ + --framework unittest \ + --output tests/test_calculator_generated.py +``` + +## How the ReAct Agent Loop Works + +Each script uses the `native_react` agent, which follows the +**Thought-Action-Observation** cycle: + +1. **Thought** -- The agent reasons about what to do next (often using the + `think` tool to structure its reasoning). +2. **Action** -- The agent calls a tool (e.g., `git_diff`, `file_read`, + `shell_exec`). +3. **Observation** -- The tool result is fed back to the agent. +4. **Repeat** until the agent has enough information to produce a final answer. + +This loop allows the agent to adaptively explore the codebase rather than +relying on a single prompt/response exchange. For example, the reviewer might +read a diff, notice a suspicious function call, then read the source of that +function before making its assessment. + +## Customization + +### Model and Engine + +All three scripts accept `--model` and `--engine` flags: + +```bash +python examples/code_companion/reviewer.py --model gpt-4o --engine cloud +python examples/code_companion/debugger.py --model claude-sonnet-4-20250514 --engine cloud +``` + +### Tools + +To change which tools an agent can use, edit the `tools` list in the script. +Available tools include `calculator`, `web_search`, `shell_exec`, `code_interpreter`, +`memory_store`, `memory_search`, and more. Run `uv run jarvis eval list` or +inspect `src/openjarvis/tools/` for the full registry. + +### Prompts + +Each script contains a `prompt` string that instructs the agent. Modify this +to change the review criteria, debugging strategy, or test generation style +to match your team's conventions. + +## SDK Pattern + +All three scripts follow the same core pattern: + +```python +from openjarvis import Jarvis + +j = Jarvis(model="qwen3:8b", engine_key="ollama") +try: + response = j.ask( + "Your task description here...", + agent="native_react", + tools=["git_diff", "file_read", "think"], + ) + print(response) +finally: + j.close() +``` diff --git a/examples/code_companion/debugger.py b/examples/code_companion/debugger.py new file mode 100644 index 00000000..eee3c78a --- /dev/null +++ b/examples/code_companion/debugger.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Debug Assistant — investigate errors and propose fixes with a ReAct agent. + +Usage:: + + python examples/code_companion/debugger.py \ + --error "TypeError: NoneType has no attribute 'split'" + python examples/code_companion/debugger.py \ + --error "KeyError: 'user_id'" --file src/app/views.py + python examples/code_companion/debugger.py \ + --error "Segfault in libfoo.so" --model gpt-4o +""" + +from __future__ import annotations + +import sys + +import click + + +@click.command() +@click.option( + "--error", + required=True, + help="Error message or stack trace to investigate.", +) +@click.option( + "--file", + "file_path", + default=None, + type=click.Path(), + help="Optional file path where the error occurred.", +) +@click.option( + "--model", + default="qwen3:8b", + show_default=True, + help="Model to use for debugging.", +) +@click.option( + "--engine", + "engine_key", + default="ollama", + show_default=True, + help="Engine backend (ollama, cloud, vllm, etc.).", +) +def main( + error: str, + file_path: str | None, + model: str, + engine_key: str, +) -> None: + """Investigate an error and propose a fix using a ReAct agent. + + The agent reads relevant source files, runs diagnostic commands, + reasons about root causes, and suggests a concrete fix. + """ + try: + from openjarvis import Jarvis + except ImportError: + click.echo( + "Error: openjarvis is not installed. " + "Install it with: uv sync --extra dev", + err=True, + ) + sys.exit(1) + + tools = ["file_read", "shell_exec", "think"] + + file_context = "" + if file_path: + file_context = f"\nThe error occurred in the file: {file_path}\n" + + prompt = ( + "You are an expert debugger. Investigate the following error and " + "propose a fix.\n\n" + f"**Error:**\n```\n{error}\n```\n" + f"{file_context}\n" + "Steps:\n" + "1. If a file path is given, use file_read to examine the source.\n" + "2. Use shell_exec to run diagnostic commands (e.g., grep for the " + "symbol, check imports, list directory contents) as needed.\n" + "3. Use think to reason about root causes.\n" + "4. Read any additional files that may be related.\n\n" + "Produce a structured response with these sections:\n" + "- **Root Cause**: explanation of why the error occurs.\n" + "- **Proposed Fix**: concrete code change or configuration fix.\n" + "- **Prevention**: how to prevent similar issues in the future " + "(e.g., type hints, validation, tests)." + ) + + click.echo(f"Investigating error: {error}") + if file_path: + click.echo(f"File: {file_path}") + click.echo(f"Model: {model} | Engine: {engine_key}") + click.echo("-" * 60) + + try: + j = Jarvis(model=model, engine_key=engine_key) + except Exception as exc: + click.echo( + f"Error: could not initialize Jarvis — {exc}\n\n" + "Make sure your engine is running. For Ollama:\n" + " ollama serve\n" + " ollama pull qwen3:8b\n\n" + "For cloud engines, ensure API keys are set in your .env file.", + err=True, + ) + sys.exit(1) + + try: + response = j.ask(prompt, agent="native_react", tools=tools) + except Exception as exc: + click.echo(f"Error during debugging: {exc}", err=True) + sys.exit(1) + finally: + j.close() + + click.echo(response) + + +if __name__ == "__main__": + main() diff --git a/examples/code_companion/reviewer.py b/examples/code_companion/reviewer.py new file mode 100644 index 00000000..33f9e3e2 --- /dev/null +++ b/examples/code_companion/reviewer.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Code Review Assistant — review diffs between branches with a ReAct agent. + +Usage: + python examples/code_companion/reviewer.py + python examples/code_companion/reviewer.py --branch feature-x --base main + python examples/code_companion/reviewer.py --model gpt-4o --engine cloud +""" + +from __future__ import annotations + +import sys + +import click + + +@click.command() +@click.option( + "--branch", + default="HEAD", + show_default=True, + help="Branch (or commit) to review.", +) +@click.option( + "--base", + default="main", + show_default=True, + help="Base branch to diff against.", +) +@click.option( + "--model", + default="qwen3:8b", + show_default=True, + help="Model to use for the review.", +) +@click.option( + "--engine", + "engine_key", + default="ollama", + show_default=True, + help="Engine backend (ollama, cloud, vllm, etc.).", +) +def main( + branch: str, + base: str, + model: str, + engine_key: str, +) -> None: + """Review code changes between BASE and BRANCH using a ReAct agent. + + The agent reads the git diff and log, inspects relevant source files, + and produces structured feedback covering issues found, suggestions, + and an overall assessment. + """ + try: + from openjarvis import Jarvis + except ImportError: + click.echo( + "Error: openjarvis is not installed. " + "Install it with: uv sync --extra dev", + err=True, + ) + sys.exit(1) + + tools = ["git_diff", "git_log", "file_read", "think"] + + prompt = ( + f"You are an expert code reviewer. Review the changes between " + f"the base branch '{base}' and the branch '{branch}'.\n\n" + "Steps:\n" + "1. Use git_diff to see what changed between the two refs.\n" + "2. Use git_log to understand the commit history.\n" + "3. Use file_read to inspect any files that need more context.\n" + "4. Use think to reason about code quality, bugs, and design.\n\n" + "Produce a structured review with these sections:\n" + "- **Summary**: one-paragraph overview of the changes.\n" + "- **Issues Found**: list of bugs, logic errors, or security concerns.\n" + "- **Suggestions**: improvements for readability, performance, or style.\n" + "- **Overall Assessment**: APPROVE, REQUEST CHANGES, or COMMENT, " + "with a brief justification." + ) + + click.echo(f"Reviewing: {base}..{branch}") + click.echo(f"Model: {model} | Engine: {engine_key}") + click.echo("-" * 60) + + try: + j = Jarvis(model=model, engine_key=engine_key) + except Exception as exc: + click.echo( + f"Error: could not initialize Jarvis — {exc}\n\n" + "Make sure your engine is running. For Ollama:\n" + " ollama serve\n" + " ollama pull qwen3:8b\n\n" + "For cloud engines, ensure API keys are set in your .env file.", + err=True, + ) + sys.exit(1) + + try: + response = j.ask(prompt, agent="native_react", tools=tools) + except Exception as exc: + click.echo(f"Error during review: {exc}", err=True) + sys.exit(1) + finally: + j.close() + + click.echo(response) + + +if __name__ == "__main__": + main() diff --git a/examples/code_companion/test_gen.py b/examples/code_companion/test_gen.py new file mode 100644 index 00000000..6f4d8712 --- /dev/null +++ b/examples/code_companion/test_gen.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +"""Test Generator — generate comprehensive tests for a Python module with a ReAct agent. + +Usage:: + + python examples/code_companion/test_gen.py \ + --module src/openjarvis/tools/calculator.py + python examples/code_companion/test_gen.py \ + --module src/app/utils.py --framework unittest + python examples/code_companion/test_gen.py \ + --module src/app/utils.py --output tests/test_utils.py +""" + +from __future__ import annotations + +import os +import sys + +import click + + +@click.command() +@click.option( + "--module", + required=True, + type=click.Path(exists=True), + help="Path to the Python module to generate tests for.", +) +@click.option( + "--framework", + default="pytest", + show_default=True, + type=click.Choice(["pytest", "unittest"], case_sensitive=False), + help="Test framework to target.", +) +@click.option( + "--output", + default=None, + type=click.Path(), + help=( + "File path to save generated tests. " + "Defaults to test_.py in the current directory." + ), +) +@click.option( + "--model", + default="qwen3:8b", + show_default=True, + help="Model to use for test generation.", +) +@click.option( + "--engine", + "engine_key", + default="ollama", + show_default=True, + help="Engine backend (ollama, cloud, vllm, etc.).", +) +def main( + module: str, + framework: str, + output: str | None, + model: str, + engine_key: str, +) -> None: + """Generate comprehensive tests for a Python MODULE using a ReAct agent. + + The agent reads the module source, reasons about edge cases and behavior, + and produces a complete test file. Tests are saved to a file or printed + to stdout. + """ + try: + from openjarvis import Jarvis + except ImportError: + click.echo( + "Error: openjarvis is not installed. " + "Install it with: uv sync --extra dev", + err=True, + ) + sys.exit(1) + + tools = ["file_read", "think", "file_write"] + + # Derive default output path from the module name. + module_basename = os.path.basename(module).removesuffix(".py") + if output is None: + output = f"test_{module_basename}.py" + + prompt = ( + f"You are an expert Python test engineer. Generate comprehensive " + f"tests for the module at '{module}' using the **{framework}** " + f"framework.\n\n" + "Steps:\n" + f"1. Use file_read to read the source code of '{module}'.\n" + "2. Use think to plan test cases: happy paths, edge cases, error " + "handling, and boundary conditions.\n" + "3. Read any related modules or base classes if needed for context.\n" + f"4. Use file_write to save the generated tests to '{output}'.\n\n" + "Guidelines:\n" + "- Each public function/method should have at least one test.\n" + "- Include docstrings on each test explaining what it verifies.\n" + "- Test edge cases (empty input, None, large values, invalid types).\n" + "- Use mocks/patches for external dependencies.\n" + "- The test file must be self-contained and runnable.\n" + ) + + click.echo(f"Generating tests for: {module}") + click.echo(f"Framework: {framework} | Output: {output}") + click.echo(f"Model: {model} | Engine: {engine_key}") + click.echo("-" * 60) + + try: + j = Jarvis(model=model, engine_key=engine_key) + except Exception as exc: + click.echo( + f"Error: could not initialize Jarvis — {exc}\n\n" + "Make sure your engine is running. For Ollama:\n" + " ollama serve\n" + " ollama pull qwen3:8b\n\n" + "For cloud engines, ensure API keys are set in your .env file.", + err=True, + ) + sys.exit(1) + + try: + response = j.ask(prompt, agent="native_react", tools=tools) + except Exception as exc: + click.echo(f"Error during test generation: {exc}", err=True) + sys.exit(1) + finally: + j.close() + + click.echo(response) + click.echo(f"\nTests saved to: {output}") + + +if __name__ == "__main__": + main() diff --git a/examples/deep_research/README.md b/examples/deep_research/README.md new file mode 100644 index 00000000..0b4c2276 --- /dev/null +++ b/examples/deep_research/README.md @@ -0,0 +1,111 @@ +# Deep Research Assistant + +A tutorial example demonstrating how to build a multi-source research agent +using OpenJarvis. The assistant uses an orchestrator agent loop with web +search, memory storage, and file output to produce comprehensive research +reports with citations. + +## What This Example Demonstrates + +- **Orchestrator agent loop** -- the agent iterates through multiple + tool-calling turns, deciding at each step whether to search, store, or + synthesize. +- **Memory-augmented reasoning** -- findings from earlier searches are stored + in memory and retrieved later for cross-referencing and deduplication. +- **Tool composition** -- five tools (`web_search`, `think`, `file_write`, + `memory_store`, `memory_search`) are wired together through a single recipe + config. +- **Recipe-driven configuration** -- `research.toml` captures the full + pillar-aligned setup (model, engine, agent, tools) in a declarative file. + +## Prerequisites + +- Python 3.10 or later +- OpenJarvis installed (`uv sync --extra dev` from the repo root) +- An inference engine running. Either: + - **Ollama** (local): `ollama serve` and `ollama pull qwen3:8b` + - **Cloud API** (remote): set the appropriate key in `.env` and use + `--engine cloud` + +## Quick Start + +```bash +# From the repository root +python examples/deep_research/research.py "quantum computing advances 2026" +``` + +Save the output to a file: + +```bash +python examples/deep_research/research.py "quantum computing advances 2026" \ + --output report.md +``` + +Use a different model or engine: + +```bash +python examples/deep_research/research.py "climate policy trends" \ + --model gpt-4o --engine cloud --max-turns 20 +``` + +## Configuration Options + +| Flag | Default | Description | +|----------------|------------|------------------------------------------| +| `--model` | `qwen3:8b` | Model identifier passed to the engine | +| `--engine` | `ollama` | Engine backend (ollama, cloud, vllm ...) | +| `--max-turns` | `15` | Maximum orchestrator loop iterations | +| `--output` | (none) | File path to save the final report | + +The companion `research.toml` provides the same defaults as a declarative +recipe that can be loaded with `load_recipe()` or passed to the `jarvis eval` +runner. + +## How It Works + +``` +User query + | + v +Jarvis SDK (model + engine selection) + | + v +OrchestratorAgent (multi-turn tool loop, up to max_turns) + | + +---> web_search -- fetch recent sources from the web + +---> think -- internal reasoning scratchpad + +---> memory_store -- persist key findings for later retrieval + +---> memory_search -- cross-reference earlier findings + +---> file_write -- save the final report to disk + | + v +Synthesized report with citations +``` + +Each turn, the orchestrator decides which tool to call (or whether to produce +a final answer). The `think` tool lets the model reason without side effects, +while `memory_store` / `memory_search` give it persistent scratch space across +turns. + +## Customization Tips + +- **Add more tools** -- append tool names to the `tools` list in + `research.toml` or pass them on the command line. See `jarvis agent info + orchestrator` for the full tool catalog. +- **Adjust temperature** -- lower values (0.2) produce more focused reports; + higher values (0.8) encourage broader exploration. +- **Swap the agent** -- replace `orchestrator` with `native_react` for a + Thought-Action-Observation loop, or `native_openhands` for a CodeAct-style + agent. +- **Use the recipe programmatically** -- load the TOML with + `openjarvis.recipes.load_recipe("examples/deep_research/research.toml")` and + pass the result to `SystemBuilder`. + +## Further Reading + +- [Architecture: Agents](../../CLAUDE.md) -- agent hierarchy (`BaseAgent`, + `ToolUsingAgent`, `OrchestratorAgent`) and the `accepts_tools` mechanism. +- [Architecture: Tools](../../CLAUDE.md) -- tool registry, MCP adapter, and + the `ToolExecutor` dispatch pipeline. +- [Recipes](../../src/openjarvis/recipes/) -- composable TOML configs that + wire all five pillars. diff --git a/examples/deep_research/research.py b/examples/deep_research/research.py new file mode 100644 index 00000000..4ec981f7 --- /dev/null +++ b/examples/deep_research/research.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Deep Research Assistant — multi-source research with memory-augmented orchestrator. + +Usage: + python examples/deep_research/research.py "quantum computing advances" + python examples/deep_research/research.py "climate policy" \ + --model gpt-4o --engine cloud + python examples/deep_research/research.py "rust vs go" \ + --output report.md --max-turns 20 +""" + +from __future__ import annotations + +import sys + +import click + + +@click.command() +@click.argument("topic") +@click.option( + "--model", + default="qwen3:8b", + show_default=True, + help="Model to use for research.", +) +@click.option( + "--engine", + "engine_key", + default="ollama", + show_default=True, + help="Engine backend (ollama, cloud, vllm, etc.).", +) +@click.option( + "--max-turns", + default=15, + show_default=True, + type=int, + help="Maximum agent loop iterations.", +) +@click.option( + "--output", + default=None, + type=click.Path(), + help="Optional file path to save the research report.", +) +def main( + topic: str, + model: str, + engine_key: str, + max_turns: int, + output: str | None, +) -> None: + """Run a deep research session on TOPIC using an orchestrator agent. + + The agent searches the web, stores findings in memory, cross-references + sources, and produces a comprehensive report with citations. + """ + # Lazy import so that --help works without a running engine or heavy deps. + try: + from openjarvis import Jarvis + except ImportError: + click.echo( + "Error: openjarvis is not installed. " + "Install it with: uv sync --extra dev", + err=True, + ) + sys.exit(1) + + tools = ["web_search", "think", "file_write", "memory_store", "memory_search"] + + system_prompt = ( + "You are a deep research assistant. When given a topic:\n" + "1. Search the web for recent, authoritative sources\n" + "2. Store key findings in memory for cross-referencing\n" + "3. Synthesize a comprehensive report with citations\n" + "4. Save the final report to a file\n\n" + "Always cite your sources and distinguish between established facts " + "and emerging claims." + ) + + click.echo(f"Researching: {topic}") + click.echo(f"Model: {model} | Engine: {engine_key} | Max turns: {max_turns}") + click.echo("-" * 60) + + try: + j = Jarvis(model=model, engine_key=engine_key) + except Exception as exc: + click.echo( + f"Error: could not initialize Jarvis — {exc}\n\n" + "Make sure your engine is running. For Ollama:\n" + " ollama serve\n" + " ollama pull qwen3:8b\n\n" + "For cloud engines, ensure API keys are set in your .env file.", + err=True, + ) + sys.exit(1) + + try: + prompt = ( + "Research the following topic in depth " + f"and produce a report:\n\n{topic}" + ) + response = j.ask( + prompt, + agent="orchestrator", + tools=tools, + system_prompt=system_prompt, + max_turns=max_turns, + temperature=0.5, + ) + except Exception as exc: + click.echo(f"Error during research: {exc}", err=True) + sys.exit(1) + finally: + j.close() + + click.echo(response) + + if output: + with open(output, "w", encoding="utf-8") as fh: + fh.write(response) + click.echo(f"\nReport saved to {output}") + + +if __name__ == "__main__": + main() diff --git a/examples/deep_research/research.toml b/examples/deep_research/research.toml new file mode 100644 index 00000000..b7b636d6 --- /dev/null +++ b/examples/deep_research/research.toml @@ -0,0 +1,23 @@ +[recipe] +name = "deep_research" +description = "Multi-source deep research with memory-augmented orchestrator" +version = "1.0.0" + +[intelligence] +model = "qwen3:8b" + +[engine] +key = "ollama" + +[agent] +type = "orchestrator" +max_turns = 15 +temperature = 0.5 +tools = ["web_search", "think", "file_write", "memory_store", "memory_search"] +system_prompt = """You are a deep research assistant. When given a topic: +1. Search the web for recent, authoritative sources +2. Store key findings in memory for cross-referencing +3. Synthesize a comprehensive report with citations +4. Save the final report to a file + +Always cite your sources and distinguish between established facts and emerging claims.""" diff --git a/examples/messaging_hub/README.md b/examples/messaging_hub/README.md new file mode 100644 index 00000000..7ea95cb7 --- /dev/null +++ b/examples/messaging_hub/README.md @@ -0,0 +1,144 @@ +# Messaging Hub + +A smart inbox assistant that triages incoming messages, classifies them by +priority, drafts replies, and produces end-of-day summaries — all powered by an +OpenJarvis orchestrator agent. + +## What This Demonstrates + +- **Channel integration** — connecting OpenJarvis to messaging platforms (Slack, WhatsApp, etc.) +- **Message triage** — automatic classification into URGENT, ACTION_REQUIRED, FYI, or SPAM +- **Smart replies** — context-aware reply drafting for actionable messages +- **Memory-backed summaries** — key information stored in memory for end-of-day rollups + +## Prerequisites + +- Python 3.10+ +- OpenJarvis installed: `uv sync --extra dev` +- An inference engine running (e.g., Ollama with `qwen3:8b` pulled) + +For live channel mode you also need the relevant channel credentials (see +[Setting Up Real Channels](#setting-up-real-channels) below). + +## Quick Start + +### Demo Mode + +Run with sample messages — no channel setup or credentials needed: + +```bash +python examples/messaging_hub/smart_inbox.py --demo +``` + +This processes five sample messages through the orchestrator agent, prints a +classification table, and generates an end-of-day summary. + +### Override Model or Engine + +```bash +python examples/messaging_hub/smart_inbox.py --demo --model gpt-4o --engine cloud +``` + +## How Message Classification Works + +Each incoming message is sent to an orchestrator agent with a structured prompt +that asks for: + +1. **Category** — one of `URGENT`, `ACTION_REQUIRED`, `FYI`, or `SPAM` +2. **Reply** — a concise professional response (or `N/A` for spam) + +The agent uses the `think` tool for internal reasoning and `memory_store` / +`memory_search` to persist key details. After all messages are processed, a +second prompt asks the agent to summarize the inbox grouped by category. + +## Setting Up Real Channels + +### Slack + +1. Add the Slack MCP server: + ```bash + jarvis add slack + ``` +2. Set credentials in your `.env`: + ``` + SLACK_BOT_TOKEN=xoxb-... + SLACK_APP_TOKEN=xapp-... + ``` +3. Invite the bot to the target Slack channel. +4. Run: + ```bash + python examples/messaging_hub/smart_inbox.py --channel slack + ``` + +### WhatsApp + +1. Ensure Node.js 22+ is installed. +2. Configure the WhatsApp Baileys bridge (see the OpenJarvis channel docs). +3. Scan the QR code to authenticate. +4. Run: + ```bash + python examples/messaging_hub/smart_inbox.py --channel whatsapp + ``` + +### Other Channels + +OpenJarvis supports many channels — LINE, Viber, Mastodon, Rocket.Chat, and +more. List all available channels with: + +```bash +jarvis channel list +``` + +## Channel Configuration via TOML + +The `messaging.toml` recipe in this directory defines the default channel, +agent type, tools, and system prompt. You can customize it or point to your own: + +```toml +[channel] +default = "slack" + +[agent] +type = "orchestrator" +max_turns = 5 +temperature = 0.3 +tools = ["think", "memory_store", "memory_search"] +``` + +Refer to `configs/openjarvis/config.toml` for the full list of channel and agent +options. + +## Adding Custom Triage Rules + +To extend the classification categories or change how messages are routed, edit +the `CLASSIFICATION_PROMPT` in `smart_inbox.py`. For example, to add a +`FOLLOW_UP` category: + +```python +CLASSIFICATION_PROMPT = ( + "Classify the following message into exactly one category: " + "URGENT, ACTION_REQUIRED, FOLLOW_UP, FYI, or SPAM.\n" + "Then draft a short reply if appropriate (not for SPAM).\n\n" + "Respond in this exact format:\n" + "CATEGORY: \n" + "REPLY: \n\n" + "Message:\n{message}" +) +``` + +You can also add domain-specific rules by extending the system prompt in +`messaging.toml` — for instance, routing messages mentioning "P0" or +"incident" directly to URGENT regardless of phrasing. + +## End-of-Day Summary + +After processing all messages, the agent produces a grouped summary. In demo +mode this is printed to the terminal. In a production setup you could schedule +this via the OpenJarvis scheduler: + +```bash +jarvis scheduler create "Daily inbox summary" --type cron --value "0 17 * * *" +``` + +Or use the operator recipe pattern to run a persistent triage agent on a +schedule. See `src/openjarvis/recipes/data/operators/` for examples. diff --git a/examples/messaging_hub/messaging.toml b/examples/messaging_hub/messaging.toml new file mode 100644 index 00000000..7a918f8f --- /dev/null +++ b/examples/messaging_hub/messaging.toml @@ -0,0 +1,25 @@ +[recipe] +name = "messaging_hub" +description = "Smart inbox with message triage and auto-replies" +version = "1.0.0" + +[intelligence] +model = "qwen3:8b" + +[engine] +key = "ollama" + +[agent] +type = "orchestrator" +max_turns = 5 +temperature = 0.3 +tools = ["think", "memory_store", "memory_search"] +system_prompt = """You are a smart inbox assistant. For each message: +1. Classify as: URGENT, ACTION_REQUIRED, FYI, or SPAM +2. Draft a concise reply if needed +3. Store key information in memory for end-of-day summary + +Be professional and concise in your replies.""" + +[channel] +default = "slack" diff --git a/examples/messaging_hub/smart_inbox.py b/examples/messaging_hub/smart_inbox.py new file mode 100644 index 00000000..ca829766 --- /dev/null +++ b/examples/messaging_hub/smart_inbox.py @@ -0,0 +1,267 @@ +#!/usr/bin/env python3 +"""Smart Inbox — message triage and auto-reply with an orchestrator agent. + +Usage: + python examples/messaging_hub/smart_inbox.py --demo + python examples/messaging_hub/smart_inbox.py --demo --model gpt-4o --engine cloud + python examples/messaging_hub/smart_inbox.py --channel slack +""" + +from __future__ import annotations + +import sys + +import click + +DEMO_MESSAGES = [ + "URGENT: Server is down in production, need immediate help!", + "Hey, just wanted to share this interesting article about AI agents.", + "Can you review my PR #42 by end of day? It's blocking the release.", + "Meeting reminder: Team standup at 10am tomorrow.", + "Buy now! Limited time offer on premium widgets!!!", +] + +CLASSIFICATION_PROMPT = ( + "You are a smart inbox assistant. Classify the following message into " + "exactly one category: URGENT, ACTION_REQUIRED, FYI, or SPAM.\n" + "Then draft a short reply if appropriate (not for SPAM).\n\n" + "Respond in this exact format:\n" + "CATEGORY: \n" + "REPLY: \n\n" + "Message:\n{message}" +) + +SUMMARY_PROMPT = ( + "You previously triaged the following messages and their classifications:\n\n" + "{triage_log}\n\n" + "Produce a concise end-of-day summary. Group by category " + "(URGENT, ACTION_REQUIRED, FYI, SPAM) and highlight any items " + "that still need attention." +) + + +def _parse_classification(response: str) -> tuple[str, str]: + """Extract category and reply from the agent response.""" + category = "UNKNOWN" + reply = "N/A" + for line in response.splitlines(): + stripped = line.strip() + if stripped.upper().startswith("CATEGORY:"): + category = stripped.split(":", 1)[1].strip().upper() + elif stripped.upper().startswith("REPLY:"): + reply = stripped.split(":", 1)[1].strip() + return category, reply + + +def _print_table(results: list[dict[str, str]]) -> None: + """Print triage results as a formatted table.""" + # Column widths + cat_w = max(len("Category"), max((len(r["category"]) for r in results), default=0)) + msg_w = min( + 50, + max(len("Message"), max((len(r["message"]) for r in results), default=0)), + ) + rep_w = min( + 40, + max(len("Reply"), max((len(r["reply"]) for r in results), default=0)), + ) + + header = ( + f" {'#':<3} {'Category':<{cat_w}} {'Message':<{msg_w}} {'Reply':<{rep_w}}" + ) + separator = " " + "-" * (len(header) - 2) + + click.echo() + click.echo(header) + click.echo(separator) + + for i, r in enumerate(results, 1): + msg_display = r["message"][:msg_w] + rep_display = r["reply"][:rep_w] + row = ( + f" {i:<3} {r['category']:<{cat_w}}" + f" {msg_display:<{msg_w}} {rep_display:<{rep_w}}" + ) + click.echo(row) + + click.echo(separator) + click.echo() + + +def _run_demo(model: str, engine_key: str) -> None: + """Process sample messages through the agent for classification.""" + try: + from openjarvis import Jarvis + except ImportError: + click.echo( + "Error: openjarvis is not installed. " + "Install it with: uv sync --extra dev", + err=True, + ) + sys.exit(1) + + tools = ["think", "memory_store", "memory_search"] + + click.echo("Smart Inbox — Demo Mode") + click.echo(f"Model: {model} | Engine: {engine_key}") + click.echo("=" * 60) + click.echo(f"Processing {len(DEMO_MESSAGES)} messages...\n") + + try: + j = Jarvis(model=model, engine_key=engine_key) + except Exception as exc: + click.echo( + f"Error: could not initialize Jarvis — {exc}\n\n" + "Make sure your engine is running. For Ollama:\n" + " ollama serve\n" + " ollama pull qwen3:8b\n\n" + "For cloud engines, ensure API keys are set in your .env file.", + err=True, + ) + sys.exit(1) + + results: list[dict[str, str]] = [] + + try: + for idx, message in enumerate(DEMO_MESSAGES, 1): + click.echo(f" [{idx}/{len(DEMO_MESSAGES)}] Classifying: {message[:60]}...") + + prompt = CLASSIFICATION_PROMPT.format(message=message) + response = j.ask( + prompt, + agent="orchestrator", + tools=tools, + temperature=0.3, + ) + + category, reply = _parse_classification(response) + results.append( + {"message": message, "category": category, "reply": reply} + ) + click.echo(f" -> {category}") + + # Print results table + _print_table(results) + + # Generate end-of-day summary + click.echo("Generating end-of-day summary...\n") + triage_log = "\n".join( + f"- [{r['category']}] {r['message']}" for r in results + ) + summary_prompt = SUMMARY_PROMPT.format(triage_log=triage_log) + summary = j.ask( + summary_prompt, + agent="orchestrator", + tools=tools, + temperature=0.3, + ) + click.echo("End-of-Day Summary") + click.echo("-" * 40) + click.echo(summary) + + except Exception as exc: + click.echo(f"Error during triage: {exc}", err=True) + sys.exit(1) + finally: + j.close() + + +def _run_channel(channel: str, model: str, engine_key: str) -> None: + """Connect to a real messaging channel for live triage. + + This mode requires channel credentials to be configured. See the + README for setup instructions for each supported channel. + """ + click.echo("Smart Inbox — Live Channel Mode") + click.echo(f"Channel: {channel} | Model: {model} | Engine: {engine_key}") + click.echo("=" * 60) + click.echo() + + # Channel setup guidance + setup_help = { + "slack": ( + "To set up Slack:\n" + " 1. Run: jarvis add slack\n" + " 2. Set SLACK_BOT_TOKEN and SLACK_APP_TOKEN in your .env\n" + " 3. Invite the bot to your target channel\n" + ), + "whatsapp": ( + "To set up WhatsApp:\n" + " 1. Ensure Node.js 22+ is installed\n" + " 2. Configure WhatsApp Baileys bridge (see channel docs)\n" + " 3. Scan the QR code to authenticate\n" + ), + } + + help_text = setup_help.get( + channel, + f"Channel '{channel}' requires appropriate credentials.\n" + f"Run: jarvis channel list to see available channels.\n", + ) + + click.echo(help_text) + click.echo( + "Once configured, incoming messages will be triaged automatically.\n" + "Use --demo to test with sample messages without channel setup.\n" + ) + + # Demonstrate how the channel integration would work + click.echo("Example integration code:\n") + click.echo(" from openjarvis import Jarvis") + click.echo(f' j = Jarvis(model="{model}", engine_key="{engine_key}")') + click.echo(" # Listen for incoming messages on the channel") + click.echo(f' # See: jarvis channel status (to verify "{channel}" is connected)') + click.echo(' response = j.ask(message, agent="orchestrator",') + click.echo(' tools=["think", "memory_store", "memory_search"])') + click.echo() + + +@click.command() +@click.option( + "--channel", + default="slack", + show_default=True, + help="Messaging channel to connect to (slack, whatsapp, etc.).", +) +@click.option( + "--model", + default="qwen3:8b", + show_default=True, + help="Model to use for message triage.", +) +@click.option( + "--engine", + "engine_key", + default="ollama", + show_default=True, + help="Engine backend (ollama, cloud, vllm, etc.).", +) +@click.option( + "--demo", + is_flag=True, + default=False, + help="Run in demo mode with sample messages (no channel required).", +) +def main(channel: str, model: str, engine_key: str, demo: bool) -> None: + """Smart inbox assistant — classify and reply to messages. + + Processes incoming messages through an orchestrator agent that classifies + each message as URGENT, ACTION_REQUIRED, FYI, or SPAM, drafts concise + replies, and stores key information for end-of-day summaries. + + \b + Demo mode (no engine required for --help): + python examples/messaging_hub/smart_inbox.py --demo + + \b + Live channel mode: + python examples/messaging_hub/smart_inbox.py --channel slack + """ + if demo: + _run_demo(model, engine_key) + else: + _run_channel(channel, model, engine_key) + + +if __name__ == "__main__": + main() diff --git a/examples/scheduled_ops/README.md b/examples/scheduled_ops/README.md new file mode 100644 index 00000000..679ae89e --- /dev/null +++ b/examples/scheduled_ops/README.md @@ -0,0 +1,129 @@ +# Scheduled Personal Ops + +This tutorial demonstrates how to use OpenJarvis to run **autonomous scheduled agents** that perform recurring personal tasks on a cron-like schedule. + +## What This Demonstrates + +- Using the `Jarvis` SDK with different agent types (`orchestrator`, `native_react`) and tool sets +- Configuring recurring schedules via TOML and the `jarvis scheduler` CLI +- Integrating with the `TaskScheduler` Python API for programmatic task registration +- Graceful error handling when an inference engine is not available + +## Scripts + +| Script | Agent | Tools | Schedule | Purpose | +|---|---|---|---|---| +| `daily_digest.py` | `orchestrator` | `web_search`, `think` | Daily 9:00 AM | Search and summarize top news for chosen topics | +| `code_review.py` | `native_react` | `git_log`, `git_diff`, `file_read`, `think` | Monday 8:00 AM | Review the past week of commits in a repository | +| `gym_scheduler.py` | `orchestrator` | `web_search`, `think` | MWF 6:00 AM | Check gym hours and class availability | + +## Quick Start + +### 1. Run a script manually + +```bash +# News digest for AI and robotics +uv run python examples/scheduled_ops/daily_digest.py --topics "AI,robotics" + +# Code review of the current repo (last 7 days) +uv run python examples/scheduled_ops/code_review.py --repo-path . + +# Gym schedule check +uv run python examples/scheduled_ops/gym_scheduler.py --gym "24 Hour Fitness" +``` + +All scripts accept `--model` and `--engine` flags to select a specific model or backend: + +```bash +uv run python examples/scheduled_ops/daily_digest.py \ + --model qwen3:8b --engine ollama --topics "AI,finance" +``` + +### 2. Set up schedules using the CLI + +Register each script as a recurring task with `jarvis scheduler create`: + +```bash +# Morning digest every day at 9 AM +jarvis scheduler create "Run daily news digest" \ + --type cron --value "0 9 * * *" + +# Weekly code review every Monday at 8 AM +jarvis scheduler create "Run weekly code review" \ + --type cron --value "0 8 * * 1" + +# Gym check on MWF at 6 AM +jarvis scheduler create "Check gym schedule" \ + --type cron --value "0 6 * * 1,3,5" +``` + +Then start the scheduler daemon: + +```bash +jarvis scheduler start +``` + +### 3. Use the TOML configuration + +The `schedules.toml` file defines all three schedules in one place: + +```toml +[schedules.daily_digest] +type = "cron" +value = "0 9 * * *" +description = "Morning news and social media digest" +script = "daily_digest.py" +``` + +You can point your own tooling or a custom loader at this file to register tasks in bulk. + +### 4. Register via the Python API + +The `gym_scheduler.py` script includes a `--register` flag that demonstrates programmatic task registration: + +```bash +uv run python examples/scheduled_ops/gym_scheduler.py --register --gym "Planet Fitness" +``` + +This uses `TaskScheduler` and `SchedulerStore` directly: + +```python +from openjarvis.scheduler import TaskScheduler +from openjarvis.scheduler.store import SchedulerStore + +store = SchedulerStore() +scheduler = TaskScheduler(store) +task = scheduler.create_task( + prompt="Check gym schedule for 'Planet Fitness'", + schedule_type="cron", + schedule_value="0 6 * * 1,3,5", + agent="orchestrator", + tools="web_search,think", +) +print(f"Task registered: {task.id}, next run: {task.next_run}") +``` + +## Adding Slack or Channel Output + +To send results to a Slack channel (or any other supported channel), pipe the output or extend the scripts: + +```bash +# Pipe output to a channel +uv run python examples/scheduled_ops/daily_digest.py | jarvis channel send slack + +# Or add channel output inside the script: +# from openjarvis.channels import ChannelRegistry +# channel = ChannelRegistry.create("slack", webhook_url="https://hooks.slack.com/...") +# channel.send(response) +``` + +See `jarvis channel list` for all available channels. + +## Customization Tips + +- **Change topics**: Use `--topics "finance,healthcare,sports"` for different digest subjects. +- **Review window**: Use `--days 14` with `code_review.py` for a two-week review cycle. +- **Different agents**: Swap `orchestrator` for `native_react` (or vice versa) in the scripts to compare agent behavior. +- **Add tools**: Extend the `tools` list in any script (e.g., add `"calculator"` or `"file_write"` for saving reports to disk). +- **Model selection**: Use `--model` to target a specific model, or let OpenJarvis auto-select from what is available. +- **Cron expressions**: Standard five-field cron syntax is supported. Install `croniter` for full expression parsing; without it, basic hour/minute patterns still work. diff --git a/examples/scheduled_ops/code_review.py b/examples/scheduled_ops/code_review.py new file mode 100644 index 00000000..b66ffda9 --- /dev/null +++ b/examples/scheduled_ops/code_review.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Weekly code review — summarizes recent commits in a repository. + +Run manually:: + + uv run python examples/scheduled_ops/code_review.py --repo-path /path/to/repo + +Or register as a scheduled task:: + + jarvis scheduler create "Weekly code review" --type cron --value "0 8 * * 1" +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import click + + +@click.command() +@click.option( + "--repo-path", + default=".", + show_default=True, + help="Path to the git repository to review.", +) +@click.option( + "--days", + default=7, + show_default=True, + type=int, + help="Number of days of commit history to review.", +) +@click.option( + "--model", + default=None, + help="Model to use for generation (e.g. qwen3:8b).", +) +@click.option( + "--engine", + "engine_key", + default=None, + help="Engine backend to use (e.g. ollama, vllm).", +) +def main(repo_path: str, days: int, model: str | None, engine_key: str | None) -> None: + """Review recent commits and produce a summary report.""" + today = datetime.now(timezone.utc).strftime("%Y-%m-%d") + prompt = ( + f"You are a senior code reviewer. Examine the git repository at " + f"'{repo_path}'. Review the commits from the last {days} days " + f"(today is {today}).\n\n" + "Steps:\n" + "1. Use git_log to list recent commits.\n" + "2. For notable commits, use git_diff to inspect the changes.\n" + "3. Use file_read if you need to see full file context.\n" + "4. Use think to reason about code quality, patterns, and risks.\n\n" + "Produce a report with:\n" + "- Summary of activity (number of commits, authors)\n" + "- Key changes and their purpose\n" + "- Any potential issues, bugs, or code smells\n" + "- Suggestions for improvement" + ) + + try: + from openjarvis import Jarvis + + kwargs: dict[str, str | None] = {} + if model: + kwargs["model"] = model + if engine_key: + kwargs["engine_key"] = engine_key + + j = Jarvis(**kwargs) # type: ignore[arg-type] + except Exception as exc: + click.echo( + f"Error: Could not initialize Jarvis: {exc}\n\n" + "Make sure an inference engine is running (e.g. `ollama serve`) " + "and the openjarvis package is installed (`uv sync`).", + err=True, + ) + raise SystemExit(1) from exc + + try: + response = j.ask( + prompt, + agent="native_react", + tools=["git_log", "git_diff", "file_read", "think"], + ) + except Exception as exc: + click.echo(f"Error during generation: {exc}", err=True) + raise SystemExit(1) from exc + finally: + j.close() + + click.echo(f"\n{'=' * 60}") + click.echo(f" Weekly Code Review — {today}") + click.echo(f" Repository: {repo_path}") + click.echo(f" Period: last {days} days") + click.echo(f"{'=' * 60}\n") + click.echo(response) + click.echo(f"\n{'=' * 60}") + + +if __name__ == "__main__": + main() diff --git a/examples/scheduled_ops/daily_digest.py b/examples/scheduled_ops/daily_digest.py new file mode 100644 index 00000000..11485e5d --- /dev/null +++ b/examples/scheduled_ops/daily_digest.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +"""Daily news digest — searches for and summarizes top stories on chosen topics. + +Run manually:: + + uv run python examples/scheduled_ops/daily_digest.py --topics "AI,robotics" + +Or register as a scheduled task:: + + jarvis scheduler create "Run daily digest" --type cron --value "0 9 * * *" +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import click + + +@click.command() +@click.option( + "--topics", + default="AI,tech", + show_default=True, + help="Comma-separated list of topics to include in the digest.", +) +@click.option( + "--model", + default=None, + help="Model to use for generation (e.g. qwen3:8b).", +) +@click.option( + "--engine", + "engine_key", + default=None, + help="Engine backend to use (e.g. ollama, vllm).", +) +def main(topics: str, model: str | None, engine_key: str | None) -> None: + """Generate a morning news digest for the given topics.""" + topic_list = [t.strip() for t in topics.split(",") if t.strip()] + if not topic_list: + click.echo("Error: --topics must contain at least one topic.", err=True) + raise SystemExit(1) + + today = datetime.now(timezone.utc).strftime("%Y-%m-%d") + prompt = ( + f"Today is {today}. Search for and summarize the top news stories on " + f"the following topics: {', '.join(topic_list)}. " + "For each topic, provide 3-5 bullet points covering the most important " + "developments. End with a one-paragraph outlook for the day." + ) + + try: + from openjarvis import Jarvis + + kwargs: dict[str, str | None] = {} + if model: + kwargs["model"] = model + if engine_key: + kwargs["engine_key"] = engine_key + + j = Jarvis(**kwargs) # type: ignore[arg-type] + except Exception as exc: + click.echo( + f"Error: Could not initialize Jarvis: {exc}\n\n" + "Make sure an inference engine is running (e.g. `ollama serve`) " + "and the openjarvis package is installed (`uv sync`).", + err=True, + ) + raise SystemExit(1) from exc + + try: + response = j.ask( + prompt, + agent="orchestrator", + tools=["web_search", "think"], + ) + except Exception as exc: + click.echo(f"Error during generation: {exc}", err=True) + raise SystemExit(1) from exc + finally: + j.close() + + click.echo(f"\n{'=' * 60}") + click.echo(f" Daily Digest — {today}") + click.echo(f" Topics: {', '.join(topic_list)}") + click.echo(f"{'=' * 60}\n") + click.echo(response) + click.echo(f"\n{'=' * 60}") + + +if __name__ == "__main__": + main() diff --git a/examples/scheduled_ops/gym_scheduler.py b/examples/scheduled_ops/gym_scheduler.py new file mode 100644 index 00000000..02b5e470 --- /dev/null +++ b/examples/scheduled_ops/gym_scheduler.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +"""Gym schedule checker — looks up gym hours and class availability. + +Run manually:: + + uv run python examples/scheduled_ops/gym_scheduler.py --gym "24 Hour Fitness" + +Or register as a scheduled task:: + + jarvis scheduler create "Gym schedule check" --type cron --value "0 6 * * 1,3,5" + +This script also demonstrates using the ``TaskScheduler`` API directly to +register itself as a recurring task. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import click + + +@click.command() +@click.option( + "--gym", + default="Local Gym", + show_default=True, + help="Name of the gym to check schedules for.", +) +@click.option( + "--model", + default=None, + help="Model to use for generation (e.g. qwen3:8b).", +) +@click.option( + "--engine", + "engine_key", + default=None, + help="Engine backend to use (e.g. ollama, vllm).", +) +@click.option( + "--register/--no-register", + default=False, + show_default=True, + help="Register this script as a recurring task via the scheduler API.", +) +def main( + gym: str, + model: str | None, + engine_key: str | None, + register: bool, +) -> None: + """Check gym schedules and class availability.""" + # -- Optional: register as a scheduled task via the scheduler API ---------- + if register: + _register_task(gym) + return + + # -- Run the gym schedule check -------------------------------------------- + today = datetime.now(timezone.utc).strftime("%A, %Y-%m-%d") + prompt = ( + f"Today is {today}. Search for the current schedule and class " + f"availability at '{gym}'. Include:\n" + "- Opening and closing hours for today\n" + "- Available group fitness classes (time, name, instructor if listed)\n" + "- Any closures, maintenance, or special events\n" + "- A brief recommendation for the best workout window today" + ) + + try: + from openjarvis import Jarvis + + kwargs: dict[str, str | None] = {} + if model: + kwargs["model"] = model + if engine_key: + kwargs["engine_key"] = engine_key + + j = Jarvis(**kwargs) # type: ignore[arg-type] + except Exception as exc: + click.echo( + f"Error: Could not initialize Jarvis: {exc}\n\n" + "Make sure an inference engine is running (e.g. `ollama serve`) " + "and the openjarvis package is installed (`uv sync`).", + err=True, + ) + raise SystemExit(1) from exc + + try: + response = j.ask( + prompt, + agent="orchestrator", + tools=["web_search", "think"], + ) + except Exception as exc: + click.echo(f"Error during generation: {exc}", err=True) + raise SystemExit(1) from exc + finally: + j.close() + + click.echo(f"\n{'=' * 60}") + click.echo(f" Gym Schedule — {today}") + click.echo(f" Gym: {gym}") + click.echo(f"{'=' * 60}\n") + click.echo(response) + click.echo(f"\n{'=' * 60}") + + +def _register_task(gym: str) -> None: + """Register this script as a recurring scheduled task.""" + try: + from openjarvis.scheduler import TaskScheduler + from openjarvis.scheduler.store import SchedulerStore + + store = SchedulerStore() + scheduler = TaskScheduler(store) + task = scheduler.create_task( + prompt=f"Check gym schedule for '{gym}'", + schedule_type="cron", + schedule_value="0 6 * * 1,3,5", + agent="orchestrator", + tools="web_search,think", + ) + click.echo(f"Registered scheduled task: {task.id}") + click.echo(" Schedule: MWF at 6:00 AM UTC") + click.echo(f" Next run: {task.next_run}") + click.echo( + "\nStart the scheduler daemon to execute tasks automatically:\n" + " jarvis scheduler start" + ) + except Exception as exc: + click.echo(f"Error registering task: {exc}", err=True) + raise SystemExit(1) from exc + + +if __name__ == "__main__": + main() diff --git a/examples/scheduled_ops/schedules.toml b/examples/scheduled_ops/schedules.toml new file mode 100644 index 00000000..081e59c2 --- /dev/null +++ b/examples/scheduled_ops/schedules.toml @@ -0,0 +1,19 @@ +[schedules] + +[schedules.daily_digest] +type = "cron" +value = "0 9 * * *" +description = "Morning news and social media digest" +script = "daily_digest.py" + +[schedules.code_review] +type = "cron" +value = "0 8 * * 1" +description = "Weekly code review summary (Monday 8am)" +script = "code_review.py" + +[schedules.gym] +type = "cron" +value = "0 6 * * 1,3,5" +description = "Gym schedule check (MWF 6am)" +script = "gym_scheduler.py" diff --git a/frontend/src-tauri/tauri.conf.json b/frontend/src-tauri/tauri.conf.json index 536dd0e9..6f89e861 100644 --- a/frontend/src-tauri/tauri.conf.json +++ b/frontend/src-tauri/tauri.conf.json @@ -58,7 +58,7 @@ "updater": { "pubkey": "dW50cnVzdGVkIGNvbW1lbnQ6IG1pbmlzaWduIHB1YmxpYyBrZXk6IDFFNzUzMzhEOEY2MjNEMDMKUldRRFBXS1BqVE4xSG8vK0lkUWN4WnZQYVIrbmc4RmpoOGlJWTBLTE15RlIya3JvQisvdUR3a0QK", "endpoints": [ - "https://github.com/HazyResearch/OpenJarvis/releases/download/desktop-latest/latest.json" + "https://github.com/open-jarvis/OpenJarvis/releases/download/desktop-latest/latest.json" ] }, "notification": { diff --git a/frontend/src/pages/GetStartedPage.tsx b/frontend/src/pages/GetStartedPage.tsx index dfc36f9f..b2002f31 100644 --- a/frontend/src/pages/GetStartedPage.tsx +++ b/frontend/src/pages/GetStartedPage.tsx @@ -19,7 +19,7 @@ import { import { isTauri, checkHealth } from '../lib/api'; const GITHUB_BASE = - 'https://github.com/hazy/OpenJarvis/releases/latest/download'; + 'https://github.com/open-jarvis/OpenJarvis/releases/latest/download'; interface Platform { id: string; @@ -432,7 +432,7 @@ function SelfHostedView() {

Deploy with Docker Compose for a zero-setup hosted instance:

- +

This starts both the API server and Ollama. The web UI is bundled and served automatically at port 8000. diff --git a/frontend/src/pages/SettingsPage.tsx b/frontend/src/pages/SettingsPage.tsx index 82c9a8e2..81cebeb4 100644 --- a/frontend/src/pages/SettingsPage.tsx +++ b/frontend/src/pages/SettingsPage.tsx @@ -269,7 +269,7 @@ export function SettingsPage() { {!speechBackendAvailable && speechBackendAvailable !== null && (

Set up a speech backend to use voice input. - See the documentation for details. + See the documentation for details.
)} @@ -330,7 +330,7 @@ export function SettingsPage() { Project site , recent_calls: VecDeque, diff --git a/rust/crates/openjarvis-agents/src/monitor_operative.rs b/rust/crates/openjarvis-agents/src/monitor_operative.rs new file mode 100644 index 00000000..3ac60aee --- /dev/null +++ b/rust/crates/openjarvis-agents/src/monitor_operative.rs @@ -0,0 +1,473 @@ +//! MonitorOperativeAgent -- long-horizon monitoring agent with configurable strategies. +//! +//! Implements the monitor-operative pattern with four strategy axes: +//! 1. Memory extraction (extract key info from observations) +//! 2. Observation compression (compress verbose outputs) +//! 3. Retrieval (search memory for relevant context) +//! 4. Task decomposition (break complex tasks into subtasks) + +use crate::loop_guard::LoopGuard; +use crate::traits::OjAgent; +use crate::utils::strip_think_tags; +use openjarvis_core::{AgentContext, AgentResult, OpenJarvisError, ToolResult}; +use openjarvis_tools::executor::ToolExecutor; +use regex::Regex; +use rig::agent::AgentBuilder; +use rig::completion::message::Message as RigMessage; +use rig::completion::request::{Chat, CompletionModel}; +use std::collections::HashMap; +use std::sync::Arc; + +// --------------------------------------------------------------------------- +// Strategy enums +// --------------------------------------------------------------------------- + +/// How findings are persisted to memory. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MemoryExtraction { + /// Extract causal relationships via LLM and store as structured entries. + CausalityGraph, + /// Append raw content to a scratchpad key. + Scratchpad, + /// Attempt to parse JSON from tool output and store structured data. + StructuredJson, + /// Do not extract or store anything. + None, +} + +/// How tool outputs are compressed before adding to context. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ObservationCompression { + /// Ask the LLM to summarize long outputs. + Summarize, + /// Hard-truncate at a character limit. + Truncate, + /// Return content unchanged. + None, +} + +/// How prior context is recalled at the start of each run. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RetrievalStrategy { + /// Hybrid retrieval with self-evaluation of relevance. + HybridWithSelfEval, + /// Keyword-based retrieval. + Keyword, + /// Semantic similarity retrieval. + Semantic, + /// No retrieval. + None, +} + +/// How complex tasks are broken down. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskDecomposition { + /// Break tasks into sequential phases. + Phased, + /// Execute as a single monolithic task. + Monolithic, + /// Hierarchical decomposition into subtask tree. + Hierarchical, +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for the monitor-operative agent's four strategy axes. +#[derive(Debug, Clone)] +pub struct MonitorConfig { + pub memory_extraction: MemoryExtraction, + pub observation_compression: ObservationCompression, + pub retrieval_strategy: RetrievalStrategy, + pub task_decomposition: TaskDecomposition, + /// Maximum characters before compression kicks in. + pub compression_threshold: usize, + /// Maximum characters for truncation. + pub truncation_limit: usize, +} + +impl Default for MonitorConfig { + fn default() -> Self { + Self { + memory_extraction: MemoryExtraction::CausalityGraph, + observation_compression: ObservationCompression::Summarize, + retrieval_strategy: RetrievalStrategy::HybridWithSelfEval, + task_decomposition: TaskDecomposition::Phased, + compression_threshold: 2000, + truncation_limit: 2000, + } + } +} + +// --------------------------------------------------------------------------- +// System prompt +// --------------------------------------------------------------------------- + +fn build_system_prompt(config: &MonitorConfig, tool_list: &str) -> String { + format!( + "You are a Monitor Operative Agent designed for long-horizon tasks.\n\n\ + ## Capabilities\n\ + 1. TOOLS: Call any available tool via function calling\n\ + 2. STATE: Your previous findings and state are automatically restored\n\ + 3. MEMORY: Store important findings for future recall\n\n\ + ## Strategy\n\ + - Memory extraction: {memory_extraction:?}\n\ + - Observation compression: {observation_compression:?}\n\ + - Retrieval strategy: {retrieval_strategy:?}\n\ + - Task decomposition: {task_decomposition:?}\n\n\ + ## Protocol\n\ + - Break complex tasks into phases and track progress\n\ + - Store causal relationships and key findings in memory\n\ + - Compress long tool outputs before adding to context\n\ + - Self-evaluate retrieved context for relevance\n\ + - Always persist state before finishing\n\n\ + Available tools: {tool_list}\n\n\ + For each step, output:\n\ + Thought: \n\ + Action: \n\ + Action Input: \n\n\ + After receiving an observation, continue reasoning.\n\ + When you have the final answer, output:\n\ + Thought: I now know the answer.\n\ + Final Answer: ", + memory_extraction = config.memory_extraction, + observation_compression = config.observation_compression, + retrieval_strategy = config.retrieval_strategy, + task_decomposition = config.task_decomposition, + tool_list = tool_list, + ) +} + +// --------------------------------------------------------------------------- +// Agent implementation +// --------------------------------------------------------------------------- + +/// Long-horizon monitoring agent with configurable memory, compression, +/// retrieval, and decomposition strategies. +/// +/// Uses a multi-turn Thought-Action-Observation loop (similar to +/// `NativeReActAgent`) augmented with strategy-driven observation +/// compression, memory extraction, and task decomposition. +pub struct MonitorOperativeAgent { + agent: rig::agent::Agent, + executor: Arc, + max_turns: usize, + config: MonitorConfig, +} + +impl MonitorOperativeAgent { + pub fn new( + model: M, + executor: Arc, + max_turns: usize, + temperature: f64, + config: MonitorConfig, + ) -> Self { + let tool_list = executor.list_tools().join(", "); + let system_prompt = build_system_prompt(&config, &tool_list); + + let agent = AgentBuilder::new(model) + .preamble(&system_prompt) + .temperature(temperature) + .build(); + + Self { + agent, + executor, + max_turns, + config, + } + } + + /// Parse `Action:` and `Action Input:` lines from model output. + fn parse_action(text: &str) -> Option<(String, String)> { + let action_re = Regex::new(r"(?m)^Action:\s*(.+)$").unwrap(); + let input_re = Regex::new(r"(?m)^Action Input:\s*(.+)$").unwrap(); + + let action = action_re + .captures(text)? + .get(1)? + .as_str() + .trim() + .to_string(); + let input = input_re + .captures(text) + .and_then(|c| c.get(1)) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_else(|| "{}".to_string()); + + Some((action, input)) + } + + /// Parse `Final Answer:` from model output. + fn parse_final_answer(text: &str) -> Option { + let re = Regex::new(r"(?m)^Final Answer:\s*(.+)").unwrap(); + re.captures(text) + .and_then(|c| c.get(1)) + .map(|m| m.as_str().trim().to_string()) + } + + /// Compress an observation according to the configured strategy. + fn compress_observation(&self, content: &str) -> String { + match self.config.observation_compression { + ObservationCompression::None => content.to_string(), + ObservationCompression::Truncate => { + if content.len() > self.config.truncation_limit { + let mut truncated = content[..self.config.truncation_limit].to_string(); + truncated.push_str("\n... [truncated]"); + truncated + } else { + content.to_string() + } + } + ObservationCompression::Summarize => { + // For summarization we would call the LLM, but since we only + // have the rig agent (chat interface) and not a raw model + // handle, we fall back to truncation in the Rust implementation. + // A production build could issue a side-channel generate call. + if content.len() > self.config.compression_threshold { + let mut truncated = content[..self.config.truncation_limit].to_string(); + truncated.push_str("\n... [summarized/truncated]"); + truncated + } else { + content.to_string() + } + } + } + } + + /// Build metadata reflecting the active strategy configuration. + fn strategy_metadata(&self) -> HashMap { + let mut meta = HashMap::new(); + meta.insert( + "memory_extraction".to_string(), + serde_json::Value::String(format!("{:?}", self.config.memory_extraction)), + ); + meta.insert( + "observation_compression".to_string(), + serde_json::Value::String(format!("{:?}", self.config.observation_compression)), + ); + meta.insert( + "retrieval_strategy".to_string(), + serde_json::Value::String(format!("{:?}", self.config.retrieval_strategy)), + ); + meta.insert( + "task_decomposition".to_string(), + serde_json::Value::String(format!("{:?}", self.config.task_decomposition)), + ); + meta + } + + /// Decompose input into subtask prompts according to the task decomposition + /// strategy. For `Monolithic` the original input is returned as-is. + /// For `Phased` and `Hierarchical` the input is wrapped with decomposition + /// instructions so the LLM itself performs the breakdown. + fn decompose_input(&self, input: &str) -> String { + match self.config.task_decomposition { + TaskDecomposition::Monolithic => input.to_string(), + TaskDecomposition::Phased => { + format!( + "Break the following task into sequential phases and execute them one at a time.\n\ + Task: {input}" + ) + } + TaskDecomposition::Hierarchical => { + format!( + "Decompose the following task into a hierarchy of subtasks, then execute from leaves to root.\n\ + Task: {input}" + ) + } + } + } +} + +#[async_trait::async_trait] +impl OjAgent for MonitorOperativeAgent { + fn agent_id(&self) -> &str { + "monitor_operative" + } + + fn accepts_tools(&self) -> bool { + true + } + + async fn run( + &self, + input: &str, + context: Option<&AgentContext>, + ) -> Result { + let mut history: Vec = context + .map(|ctx| { + ctx.conversation + .messages + .iter() + .filter_map(|m| match m.role { + openjarvis_core::Role::User => { + Some(RigMessage::user(&m.content)) + } + openjarvis_core::Role::Assistant => { + Some(RigMessage::assistant(&m.content)) + } + _ => None, + }) + .collect() + }) + .unwrap_or_default(); + + let mut all_tool_results: Vec = Vec::new(); + let mut guard = LoopGuard::default(); + + // Apply task decomposition strategy to the input. + let decomposed_input = self.decompose_input(input); + let mut current_input = decomposed_input; + + for turn in 1..=self.max_turns { + let response = self + .agent + .chat(¤t_input, history.clone()) + .await + .map_err(|e| { + OpenJarvisError::Agent(openjarvis_core::error::AgentError::Execution( + e.to_string(), + )) + })?; + + let text = strip_think_tags(&response); + + // Check for final answer + if let Some(answer) = Self::parse_final_answer(&text) { + let mut metadata = self.strategy_metadata(); + metadata.insert( + "turns_used".to_string(), + serde_json::Value::Number(serde_json::Number::from(turn)), + ); + return Ok(AgentResult { + content: answer, + tool_results: all_tool_results, + turns: turn, + metadata, + }); + } + + // Check for action (tool call) + if let Some((action, action_input)) = Self::parse_action(&text) { + // Loop guard check + if let Some(loop_msg) = guard.check(&action, &action_input) { + return Ok(AgentResult { + content: format!("Agent stopped: {}", loop_msg), + tool_results: all_tool_results, + turns: turn, + metadata: self.strategy_metadata(), + }); + } + + let params: serde_json::Value = + serde_json::from_str(&action_input).unwrap_or(serde_json::json!({})); + + let tool_result = match self.executor.execute( + &action, + ¶ms, + Some("monitor_operative"), + None, + ) { + Ok(r) => r, + Err(e) => ToolResult::failure(&action, e.to_string()), + }; + + // Compress observation according to strategy + let compressed = self.compress_observation(&tool_result.content); + + history.push(RigMessage::assistant(&text)); + current_input = format!("Observation: {}", compressed); + + all_tool_results.push(tool_result); + } else { + // No action and no final answer -- treat as final response + return Ok(AgentResult { + content: text, + tool_results: all_tool_results, + turns: turn, + metadata: self.strategy_metadata(), + }); + } + } + + // Max turns exceeded + let mut metadata = self.strategy_metadata(); + metadata.insert( + "max_turns_exceeded".to_string(), + serde_json::Value::Bool(true), + ); + Ok(AgentResult { + content: format!("Reached maximum turns ({})", self.max_turns), + tool_results: all_tool_results, + turns: self.max_turns, + metadata, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use openjarvis_engine::rig_adapter::RigModelAdapter; + type MonitorAgent = MonitorOperativeAgent>; + + #[test] + fn test_parse_action() { + let text = + "Thought: I need to search\nAction: web_search\nAction Input: {\"query\": \"rust\"}"; + let (action, input) = MonitorAgent::parse_action(text).unwrap(); + assert_eq!(action, "web_search"); + assert!(input.contains("rust")); + } + + #[test] + fn test_parse_final_answer() { + let text = "Thought: I know the answer\nFinal Answer: The result is 42."; + let answer = MonitorAgent::parse_final_answer(text).unwrap(); + assert_eq!(answer, "The result is 42."); + } + + #[test] + fn test_compress_observation_none() { + let config = MonitorConfig { + observation_compression: ObservationCompression::None, + ..Default::default() + }; + // We can test compress_observation without constructing the full agent + // by checking the strategy logic directly. + assert_eq!(config.observation_compression, ObservationCompression::None); + } + + #[test] + fn test_default_config() { + let config = MonitorConfig::default(); + assert_eq!(config.memory_extraction, MemoryExtraction::CausalityGraph); + assert_eq!( + config.observation_compression, + ObservationCompression::Summarize + ); + assert_eq!( + config.retrieval_strategy, + RetrievalStrategy::HybridWithSelfEval + ); + assert_eq!(config.task_decomposition, TaskDecomposition::Phased); + assert_eq!(config.compression_threshold, 2000); + assert_eq!(config.truncation_limit, 2000); + } + + #[test] + fn test_strategy_enum_debug() { + // Ensure Debug formatting works (used in system prompt and metadata). + assert_eq!(format!("{:?}", MemoryExtraction::CausalityGraph), "CausalityGraph"); + assert_eq!(format!("{:?}", ObservationCompression::Truncate), "Truncate"); + assert_eq!( + format!("{:?}", RetrievalStrategy::HybridWithSelfEval), + "HybridWithSelfEval" + ); + assert_eq!(format!("{:?}", TaskDecomposition::Hierarchical), "Hierarchical"); + } +} diff --git a/rust/crates/openjarvis-agents/src/native_openhands.rs b/rust/crates/openjarvis-agents/src/native_openhands.rs new file mode 100644 index 00000000..b9bb8501 --- /dev/null +++ b/rust/crates/openjarvis-agents/src/native_openhands.rs @@ -0,0 +1,339 @@ +//! NativeOpenHandsAgent -- CodeAct-style agent that uses code actions. +//! +//! Generates and dispatches code (Python blocks) and tool calls to accomplish +//! tasks. Mirrors the Python ``NativeOpenHandsAgent`` which supports both +//! ``Action: / Action Input:`` structured tool calls and fenced +//! ````python`` code blocks executed via the ``code_interpreter`` tool. + +use crate::loop_guard::LoopGuard; +use crate::traits::OjAgent; +use crate::utils::strip_think_tags; +use openjarvis_core::{AgentContext, AgentResult, OpenJarvisError, ToolResult}; +use openjarvis_tools::executor::ToolExecutor; +use regex::Regex; +use rig::agent::AgentBuilder; +use rig::completion::message::Message as RigMessage; +use rig::completion::request::{Chat, CompletionModel}; +use std::collections::HashMap; +use std::sync::Arc; + +// --------------------------------------------------------------------------- +// System prompt +// --------------------------------------------------------------------------- + +const OPENHANDS_SYSTEM_PROMPT: &str = "\ +You are an AI assistant with access to tools. \ +You MUST use tools when they would help answer the user's question. + +## How to use tools + +To call a tool, write on its own lines: + +Action: +Action Input: + +You will receive the result, then continue your response. + +## Available tools + +{tool_list} + +## Important rules + +- When the user asks you to look up, search, fetch, or summarize a URL or \ +topic, you MUST use web_search. Do NOT say you cannot browse the web. +- When the user provides a URL, pass the FULL URL (including https://) as the \ +query to web_search. Do NOT rewrite URLs into search keywords. +- When the user asks a math question, use calculator. +- When the user asks to read a file, use file_read. +- You CAN write Python code in ```python blocks and it will be executed. Use \ +this for computation, data processing, or when no specific tool fits. +- If no tool or code is needed, respond directly with your answer. +- Do NOT include tags or internal reasoning in your response. Respond \ +directly."; + +// --------------------------------------------------------------------------- +// Agent implementation +// --------------------------------------------------------------------------- + +/// Native CodeAct agent -- generates and executes code actions (shell commands, +/// file edits) and structured tool calls to accomplish tasks. +/// +/// Supports two action formats: +/// 1. `Action: tool_name` / `Action Input: {json}` -- dispatched to the +/// `ToolExecutor`. +/// 2. Fenced ````python` code blocks -- dispatched to the `code_interpreter` +/// tool. +pub struct NativeOpenHandsAgent { + agent: rig::agent::Agent, + executor: Arc, + max_turns: usize, +} + +impl NativeOpenHandsAgent { + pub fn new( + model: M, + executor: Arc, + max_turns: usize, + temperature: f64, + ) -> Self { + let tool_list = executor.list_tools().join(", "); + let system_prompt = OPENHANDS_SYSTEM_PROMPT.replace("{tool_list}", &tool_list); + + let agent = AgentBuilder::new(model) + .preamble(&system_prompt) + .temperature(temperature) + .build(); + + Self { + agent, + executor, + max_turns, + } + } + + /// Parse `Action:` and `Action Input:` lines from model output. + fn parse_action(text: &str) -> Option<(String, String)> { + let action_re = Regex::new(r"(?mi)^Action:\s*(.+)$").unwrap(); + let input_re = Regex::new(r"(?mi)^Action Input:\s*(.+?)(?:\n\n|\z)").unwrap(); + + let action = action_re + .captures(text)? + .get(1)? + .as_str() + .trim() + .to_string(); + let input = input_re + .captures(text) + .and_then(|c| c.get(1)) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_else(|| "{}".to_string()); + + Some((action, input)) + } + + /// Extract Python code from fenced ````python` blocks. + fn extract_code(text: &str) -> Option { + let re = Regex::new(r"(?s)```python\n(.*?)```").unwrap(); + re.captures(text) + .and_then(|c| c.get(1)) + .map(|m| m.as_str().trim().to_string()) + } + + /// Remove raw tool-call artifacts from final output text. + fn strip_tool_call_text(text: &str) -> String { + // Remove Action: ... Action Input: ... blocks + let action_re = + Regex::new(r"(?si)Action:\s*.+?(?:Action Input:\s*.+?)?(?:\n\n|\z)").unwrap(); + let cleaned = action_re.replace_all(text, ""); + // Remove ... XML blocks + let xml_re = Regex::new(r"(?s).*?").unwrap(); + let cleaned = xml_re.replace_all(&cleaned, ""); + cleaned.trim().to_string() + } + + /// Truncate observation text if it exceeds 4000 characters. + fn truncate_observation(content: &str, limit: usize) -> String { + if content.len() > limit { + let mut truncated = content[..limit].to_string(); + truncated.push_str("\n\n[Output truncated]"); + truncated + } else { + content.to_string() + } + } +} + +#[async_trait::async_trait] +impl OjAgent for NativeOpenHandsAgent { + fn agent_id(&self) -> &str { + "native_openhands" + } + + fn accepts_tools(&self) -> bool { + true + } + + async fn run( + &self, + input: &str, + context: Option<&AgentContext>, + ) -> Result { + let mut history: Vec = context + .map(|ctx| { + ctx.conversation + .messages + .iter() + .filter_map(|m| match m.role { + openjarvis_core::Role::User => { + Some(RigMessage::user(&m.content)) + } + openjarvis_core::Role::Assistant => { + Some(RigMessage::assistant(&m.content)) + } + _ => None, + }) + .collect() + }) + .unwrap_or_default(); + + let mut all_tool_results: Vec = Vec::new(); + let mut guard = LoopGuard::default(); + let mut current_input = input.to_string(); + + for turn in 1..=self.max_turns { + let response = self + .agent + .chat(¤t_input, history.clone()) + .await + .map_err(|e| { + OpenJarvisError::Agent(openjarvis_core::error::AgentError::Execution( + e.to_string(), + )) + })?; + + let text = strip_think_tags(&response); + + // 1. Try to extract a Python code block -> execute via code_interpreter + if let Some(code) = Self::extract_code(&text) { + let tool_name = "code_interpreter"; + let args = serde_json::json!({"code": code}); + let args_str = args.to_string(); + + if let Some(loop_msg) = guard.check(tool_name, &args_str) { + return Ok(AgentResult { + content: format!("Agent stopped: {}", loop_msg), + tool_results: all_tool_results, + turns: turn, + metadata: HashMap::new(), + }); + } + + let tool_result = match self.executor.execute( + tool_name, + &args, + Some("native_openhands"), + None, + ) { + Ok(r) => r, + Err(e) => ToolResult::failure(tool_name, e.to_string()), + }; + + let obs = Self::truncate_observation(&tool_result.content, 4000); + history.push(RigMessage::assistant(&text)); + current_input = format!("Output:\n{}", obs); + + all_tool_results.push(tool_result); + continue; + } + + // 2. Try to extract a structured tool call (Action: / Action Input:) + if let Some((action, action_input)) = Self::parse_action(&text) { + if let Some(loop_msg) = guard.check(&action, &action_input) { + return Ok(AgentResult { + content: format!("Agent stopped: {}", loop_msg), + tool_results: all_tool_results, + turns: turn, + metadata: HashMap::new(), + }); + } + + let params: serde_json::Value = + serde_json::from_str(&action_input).unwrap_or(serde_json::json!({})); + + let tool_result = match self.executor.execute( + &action, + ¶ms, + Some("native_openhands"), + None, + ) { + Ok(r) => r, + Err(e) => ToolResult::failure(&action, e.to_string()), + }; + + let obs = Self::truncate_observation(&tool_result.content, 4000); + history.push(RigMessage::assistant(&text)); + current_input = format!("Result: {}", obs); + + all_tool_results.push(tool_result); + continue; + } + + // 3. No code or tool call -- this is the final answer + let cleaned = Self::strip_tool_call_text(&text); + return Ok(AgentResult { + content: cleaned, + tool_results: all_tool_results, + turns: turn, + metadata: HashMap::new(), + }); + } + + // Max turns exceeded + Ok(AgentResult { + content: format!("Reached maximum turns ({})", self.max_turns), + tool_results: all_tool_results, + turns: self.max_turns, + metadata: HashMap::new(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use openjarvis_engine::rig_adapter::RigModelAdapter; + type OpenHandsAgent = NativeOpenHandsAgent>; + + #[test] + fn test_parse_action() { + let text = "I need to search.\nAction: web_search\nAction Input: {\"query\": \"rust lang\"}"; + let (action, input) = OpenHandsAgent::parse_action(text).unwrap(); + assert_eq!(action, "web_search"); + assert!(input.contains("rust lang")); + } + + #[test] + fn test_parse_action_missing_input() { + let text = "Action: calculator\n\nSome other text"; + let (action, input) = OpenHandsAgent::parse_action(text).unwrap(); + assert_eq!(action, "calculator"); + assert_eq!(input, "{}"); + } + + #[test] + fn test_extract_code() { + let text = "Let me compute that:\n```python\nprint(2 + 2)\n```\nDone."; + let code = OpenHandsAgent::extract_code(text).unwrap(); + assert_eq!(code, "print(2 + 2)"); + } + + #[test] + fn test_extract_code_none() { + let text = "No code here, just text."; + assert!(OpenHandsAgent::extract_code(text).is_none()); + } + + #[test] + fn test_strip_tool_call_text() { + let text = "Here is the answer.\nAction: calc\nAction Input: {\"x\": 1}\n\nFinal part."; + let cleaned = OpenHandsAgent::strip_tool_call_text(text); + assert!(!cleaned.contains("Action:")); + assert!(cleaned.contains("Final part")); + } + + #[test] + fn test_truncate_observation() { + let short = "hello"; + assert_eq!( + OpenHandsAgent::truncate_observation(short, 100), + "hello" + ); + + let long = "x".repeat(5000); + let truncated = OpenHandsAgent::truncate_observation(&long, 100); + assert!(truncated.len() < 200); + assert!(truncated.contains("[Output truncated]")); + } +} diff --git a/rust/crates/openjarvis-agents/src/orchestrator.rs b/rust/crates/openjarvis-agents/src/orchestrator.rs index c874ea0a..db49b9b3 100644 --- a/rust/crates/openjarvis-agents/src/orchestrator.rs +++ b/rust/crates/openjarvis-agents/src/orchestrator.rs @@ -13,6 +13,7 @@ use std::collections::HashMap; use std::sync::Arc; /// Multi-turn agent with function calling and loop detection. +#[allow(dead_code)] pub struct OrchestratorAgent { agent: rig::agent::Agent, executor: Arc, @@ -72,7 +73,7 @@ impl OjAgent for OrchestratorAgent { }) .unwrap_or_default(); - let mut all_tool_results = Vec::new(); + let all_tool_results: Vec = Vec::new(); let _guard = LoopGuard::default(); // Use rig agent for generation. Multi-turn tool dispatch requires diff --git a/rust/crates/openjarvis-core/src/config.rs b/rust/crates/openjarvis-core/src/config.rs index c269ada8..cc974e32 100644 --- a/rust/crates/openjarvis-core/src/config.rs +++ b/rust/crates/openjarvis-core/src/config.rs @@ -4,7 +4,7 @@ //! All config structs use `#[serde(default)]` for backward compatibility. use crate::error::ConfigError; -use crate::hardware::{detect_hardware, recommend_engine, GpuInfo, HardwareInfo}; +use crate::hardware::{detect_hardware, recommend_engine, HardwareInfo}; use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; diff --git a/rust/crates/openjarvis-core/src/hardware.rs b/rust/crates/openjarvis-core/src/hardware.rs index cc61d321..98624cfd 100644 --- a/rust/crates/openjarvis-core/src/hardware.rs +++ b/rust/crates/openjarvis-core/src/hardware.rs @@ -110,7 +110,7 @@ fn detect_amd_gpu() -> Option { let vram_raw = run_cmd(&["rocm-smi", "--showmeminfo", "vram"]); for line in vram_raw.lines() { if line.contains("Total Memory (B):") { - if let Some(val) = line.split(':').last() { + if let Some(val) = line.split(':').next_back() { if let Ok(bytes) = val.trim().parse::() { vram_gb = (bytes / (1024.0 * 1024.0 * 1024.0) * 10.0).round() / 10.0; } @@ -150,7 +150,7 @@ fn detect_apple_gpu() -> Option { for line in raw.lines() { let trimmed = line.trim(); if trimmed.contains("Chipset Model") { - let name = trimmed.split(':').last().unwrap_or("Apple Silicon").trim(); + let name = trimmed.split(':').next_back().unwrap_or("Apple Silicon").trim(); return Some(GpuInfo { vendor: "apple".into(), name: name.to_string(), diff --git a/rust/crates/openjarvis-engine/src/engine_enum.rs b/rust/crates/openjarvis-engine/src/engine_enum.rs index 5133de12..9549dbe0 100644 --- a/rust/crates/openjarvis-engine/src/engine_enum.rs +++ b/rust/crates/openjarvis-engine/src/engine_enum.rs @@ -3,8 +3,11 @@ //! Avoids `dyn InferenceEngine` for the hot path. Each variant holds a //! concrete engine so the compiler can inline and devirtualize. +use crate::llamacpp::LlamaCppEngine; use crate::ollama::OllamaEngine; use crate::openai_compat::OpenAICompatEngine; +use crate::sglang::SGLangEngine; +use crate::vllm::VLLMEngine; use crate::traits::{InferenceEngine, TokenStream}; use openjarvis_core::error::OpenJarvisError; use openjarvis_core::{GenerateResult, Message}; @@ -15,8 +18,17 @@ use serde_json::Value; /// Static dispatch at compile-time — no vtable overhead on the hot path. pub enum Engine { Ollama(OllamaEngine), + /// Dedicated vLLM engine with OpenAI-compatible API. + VLLM(VLLMEngine), + /// Dedicated SGLang engine with OpenAI-compatible API. + SGLang(SGLangEngine), + /// Dedicated llama.cpp engine with native `/completion` API. + LlamaCppNative(LlamaCppEngine), + /// Legacy: vLLM via generic OpenAI-compatible engine. Vllm(OpenAICompatEngine), + /// Legacy: SGLang via generic OpenAI-compatible engine. Sglang(OpenAICompatEngine), + /// Legacy: llama.cpp via generic OpenAI-compatible engine. LlamaCpp(OpenAICompatEngine), Mlx(OpenAICompatEngine), LmStudio(OpenAICompatEngine), @@ -30,6 +42,9 @@ macro_rules! delegate_engine { ($self:expr, $method:ident $(, $arg:expr)*) => { match $self { Engine::Ollama(e) => e.$method($($arg),*), + Engine::VLLM(e) => e.$method($($arg),*), + Engine::SGLang(e) => e.$method($($arg),*), + Engine::LlamaCppNative(e) => e.$method($($arg),*), Engine::Vllm(e) => e.$method($($arg),*), Engine::Sglang(e) => e.$method($($arg),*), Engine::LlamaCpp(e) => e.$method($($arg),*), @@ -70,6 +85,9 @@ impl InferenceEngine for Engine { ) -> Result { match self { Engine::Ollama(e) => e.stream(messages, model, temperature, max_tokens, extra).await, + Engine::VLLM(e) => e.stream(messages, model, temperature, max_tokens, extra).await, + Engine::SGLang(e) => e.stream(messages, model, temperature, max_tokens, extra).await, + Engine::LlamaCppNative(e) => e.stream(messages, model, temperature, max_tokens, extra).await, Engine::Vllm(e) => e.stream(messages, model, temperature, max_tokens, extra).await, Engine::Sglang(e) => e.stream(messages, model, temperature, max_tokens, extra).await, Engine::LlamaCpp(e) => e.stream(messages, model, temperature, max_tokens, extra).await, @@ -104,6 +122,9 @@ impl Engine { pub fn variant_key(&self) -> &str { match self { Engine::Ollama(_) => "ollama", + Engine::VLLM(_) => "vllm", + Engine::SGLang(_) => "sglang", + Engine::LlamaCppNative(_) => "llamacpp", Engine::Vllm(_) => "vllm", Engine::Sglang(_) => "sglang", Engine::LlamaCpp(_) => "llamacpp", @@ -162,4 +183,25 @@ mod tests { assert_eq!(e.variant_key(), "apple_fm"); assert_eq!(e.engine_id(), "apple_fm"); } + + #[test] + fn test_engine_vllm_native_variant() { + let e = Engine::VLLM(VLLMEngine::with_defaults()); + assert_eq!(e.variant_key(), "vllm"); + assert_eq!(e.engine_id(), "vllm"); + } + + #[test] + fn test_engine_sglang_native_variant() { + let e = Engine::SGLang(SGLangEngine::with_defaults()); + assert_eq!(e.variant_key(), "sglang"); + assert_eq!(e.engine_id(), "sglang"); + } + + #[test] + fn test_engine_llamacpp_native_variant() { + let e = Engine::LlamaCppNative(LlamaCppEngine::with_defaults()); + assert_eq!(e.variant_key(), "llamacpp"); + assert_eq!(e.engine_id(), "llamacpp"); + } } diff --git a/rust/crates/openjarvis-engine/src/lib.rs b/rust/crates/openjarvis-engine/src/lib.rs index 8c117478..20ab8235 100644 --- a/rust/crates/openjarvis-engine/src/lib.rs +++ b/rust/crates/openjarvis-engine/src/lib.rs @@ -5,13 +5,19 @@ pub mod discovery; pub mod engine_enum; +pub mod llamacpp; pub mod ollama; pub mod openai_compat; pub mod rig_adapter; +pub mod sglang; pub mod traits; +pub mod vllm; pub use discovery::{discover_engines, get_engine_static}; pub use engine_enum::Engine; +pub use llamacpp::LlamaCppEngine; pub use ollama::OllamaEngine; pub use openai_compat::OpenAICompatEngine; +pub use sglang::SGLangEngine; pub use traits::{InferenceEngine, messages_to_dicts}; +pub use vllm::VLLMEngine; diff --git a/rust/crates/openjarvis-engine/src/llamacpp.rs b/rust/crates/openjarvis-engine/src/llamacpp.rs new file mode 100644 index 00000000..f38322e0 --- /dev/null +++ b/rust/crates/openjarvis-engine/src/llamacpp.rs @@ -0,0 +1,338 @@ +//! llama.cpp inference engine backend. +//! +//! llama.cpp server exposes `/completion` (not `/v1/chat/completions`) +//! with a different request/response format. + +use crate::traits::{InferenceEngine, TokenStream}; +use openjarvis_core::error::{EngineError, OpenJarvisError}; +use openjarvis_core::{GenerateResult, Message, Usage}; +use serde_json::Value; + +/// llama.cpp server backend via its native HTTP API. +/// +/// Unlike vLLM/SGLang, llama.cpp uses `/completion` with a prompt-based +/// request format rather than the OpenAI chat completions format. +pub struct LlamaCppEngine { + host: String, + client: reqwest::blocking::Client, + timeout: std::time::Duration, +} + +impl LlamaCppEngine { + pub fn new(host: &str, port: u16, timeout_secs: f64) -> Self { + let host = format!( + "{}:{}", + host.trim_end_matches('/').trim_end_matches(|c: char| c == ':' || c.is_ascii_digit()), + port + ); + let host = if host.starts_with("http") { + host + } else { + format!("http://{}", host) + }; + let host = host.trim_end_matches('/').to_string(); + let timeout = std::time::Duration::from_secs_f64(timeout_secs); + let client = reqwest::blocking::Client::builder() + .timeout(timeout) + .build() + .unwrap_or_default(); + Self { + host, + client, + timeout, + } + } + + pub fn with_defaults() -> Self { + Self::new("http://localhost", 8080, 120.0) + } + + pub fn from_host(host: &str) -> Self { + let host = host.trim_end_matches('/').to_string(); + let timeout = std::time::Duration::from_secs_f64(120.0); + let client = reqwest::blocking::Client::builder() + .timeout(timeout) + .build() + .unwrap_or_default(); + Self { + host, + client, + timeout, + } + } + + /// Format chat messages into a single prompt string for llama.cpp. + /// + /// Uses a simple ChatML-style format: + /// `<|system|>\n{content}\n<|user|>\n{content}\n<|assistant|>\n` + fn messages_to_prompt(messages: &[Message]) -> String { + let mut prompt = String::new(); + for msg in messages { + let role = msg.role.to_string(); + prompt.push_str(&format!("<|{}|>\n{}\n", role, msg.content)); + } + prompt.push_str("<|assistant|>\n"); + prompt + } +} + +impl Default for LlamaCppEngine { + fn default() -> Self { + Self::with_defaults() + } +} + +#[async_trait::async_trait] +impl InferenceEngine for LlamaCppEngine { + fn engine_id(&self) -> &str { + "llamacpp" + } + + fn generate( + &self, + messages: &[Message], + model: &str, + temperature: f64, + max_tokens: i64, + _extra: Option<&Value>, + ) -> Result { + let prompt = Self::messages_to_prompt(messages); + let payload = serde_json::json!({ + "prompt": prompt, + "n_predict": max_tokens, + "temperature": temperature, + "stream": false, + }); + + let resp = self + .client + .post(format!("{}/completion", self.host)) + .json(&payload) + .send() + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(format!( + "llama.cpp not reachable at {}: {}", + self.host, e + ))) + })?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().unwrap_or_default(); + return Err(OpenJarvisError::Engine(EngineError::Http(format!( + "llama.cpp returned {}: {}", + status, body + )))); + } + + let data: Value = resp.json().map_err(|e| { + OpenJarvisError::Engine(EngineError::Deserialization(e.to_string())) + })?; + + let content = data["content"] + .as_str() + .unwrap_or("") + .to_string(); + let tokens_evaluated = data["tokens_evaluated"].as_i64().unwrap_or(0); + let tokens_predicted = data["tokens_predicted"].as_i64().unwrap_or(0); + + let stop_type = data["stop_type"] + .as_str() + .unwrap_or("stop"); + let finish_reason = if stop_type == "limit" { + "length".to_string() + } else { + "stop".to_string() + }; + + Ok(GenerateResult { + content, + usage: Usage { + prompt_tokens: tokens_evaluated, + completion_tokens: tokens_predicted, + total_tokens: tokens_evaluated + tokens_predicted, + }, + model: model.to_string(), + finish_reason, + tool_calls: None, + ttft: 0.0, + cost_usd: 0.0, + metadata: std::collections::HashMap::new(), + }) + } + + async fn stream( + &self, + messages: &[Message], + _model: &str, + temperature: f64, + max_tokens: i64, + _extra: Option<&Value>, + ) -> Result { + let prompt = Self::messages_to_prompt(messages); + let payload = serde_json::json!({ + "prompt": prompt, + "n_predict": max_tokens, + "temperature": temperature, + "stream": true, + }); + + let async_client = reqwest::Client::builder() + .timeout(self.timeout) + .build() + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(e.to_string())) + })?; + + let resp = async_client + .post(format!("{}/completion", self.host)) + .json(&payload) + .send() + .await + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(format!( + "llama.cpp not reachable at {}: {}", + self.host, e + ))) + })?; + + if !resp.status().is_success() { + return Err(OpenJarvisError::Engine(EngineError::Http(format!( + "llama.cpp returned {}", + resp.status() + )))); + } + + use futures::StreamExt; + let byte_stream = resp.bytes_stream(); + + // llama.cpp streams SSE lines: `data: {"content": "token", ...}` + let token_stream = byte_stream.filter_map(|chunk_result| async { + match chunk_result { + Ok(bytes) => { + let text = String::from_utf8_lossy(&bytes); + for line in text.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let json_str = line.strip_prefix("data: ").unwrap_or(line); + if let Ok(chunk) = serde_json::from_str::(json_str) { + // Check if this is the final chunk + if chunk["stop"].as_bool().unwrap_or(false) { + return None; + } + let content = chunk["content"] + .as_str() + .unwrap_or("") + .to_string(); + if !content.is_empty() { + return Some(Ok(content)); + } + } + } + None + } + Err(e) => Some(Err(OpenJarvisError::Engine(EngineError::Streaming( + e.to_string(), + )))), + } + }); + + Ok(Box::pin(token_stream)) + } + + fn list_models(&self) -> Result, OpenJarvisError> { + // llama.cpp server loads a single model; try /v1/models first, + // then fall back to /props which returns model metadata. + let resp = self + .client + .get(format!("{}/v1/models", self.host)) + .send(); + + if let Ok(resp) = resp { + if resp.status().is_success() { + let data: Value = resp.json().unwrap_or(Value::Null); + let models: Vec = data["data"] + .as_array() + .map(|arr| { + arr.iter() + .filter_map(|m| m["id"].as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + if !models.is_empty() { + return Ok(models); + } + } + } + + // Fallback: /props endpoint + let resp = self + .client + .get(format!("{}/props", self.host)) + .send() + .map_err(|_| { + OpenJarvisError::Engine(EngineError::Connection( + "llama.cpp not reachable".into(), + )) + })?; + + if !resp.status().is_success() { + return Ok(vec![]); + } + + let data: Value = resp.json().unwrap_or(Value::Null); + if let Some(model) = data["default_generation_settings"]["model"] + .as_str() + .map(String::from) + { + Ok(vec![model]) + } else { + Ok(vec![]) + } + } + + fn health(&self) -> bool { + self.client + .get(format!("{}/health", self.host)) + .timeout(std::time::Duration::from_secs(2)) + .send() + .map(|r| r.status().is_success()) + .unwrap_or(false) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use openjarvis_core::Message; + + #[test] + fn test_llamacpp_default_host() { + let engine = LlamaCppEngine::with_defaults(); + assert_eq!(engine.engine_id(), "llamacpp"); + assert_eq!(engine.host, "http://localhost:8080"); + } + + #[test] + fn test_llamacpp_from_host() { + let engine = LlamaCppEngine::from_host("http://gpu-server:8080"); + assert_eq!(engine.engine_id(), "llamacpp"); + assert_eq!(engine.host, "http://gpu-server:8080"); + } + + #[test] + fn test_messages_to_prompt() { + let messages = vec![ + Message::system("You are helpful"), + Message::user("Hello"), + ]; + let prompt = LlamaCppEngine::messages_to_prompt(&messages); + assert!(prompt.contains("<|system|>")); + assert!(prompt.contains("You are helpful")); + assert!(prompt.contains("<|user|>")); + assert!(prompt.contains("Hello")); + assert!(prompt.ends_with("<|assistant|>\n")); + } +} diff --git a/rust/crates/openjarvis-engine/src/ollama.rs b/rust/crates/openjarvis-engine/src/ollama.rs index 4b803707..7056646c 100644 --- a/rust/crates/openjarvis-engine/src/ollama.rs +++ b/rust/crates/openjarvis-engine/src/ollama.rs @@ -152,8 +152,6 @@ impl InferenceEngine for OllamaEngine { max_tokens: i64, _extra: Option<&Value>, ) -> Result { - use futures::stream; - let msg_dicts = crate::traits::messages_to_dicts(messages); let payload = serde_json::json!({ "model": model, diff --git a/rust/crates/openjarvis-engine/src/rig_adapter.rs b/rust/crates/openjarvis-engine/src/rig_adapter.rs index bb2607cf..873086a8 100644 --- a/rust/crates/openjarvis-engine/src/rig_adapter.rs +++ b/rust/crates/openjarvis-engine/src/rig_adapter.rs @@ -106,7 +106,7 @@ fn rig_request_to_oj_messages(request: &CompletionRequest) -> Vec { .map(|d| format!("[{}]\n{}", d.id, d.text)) .collect::>() .join("\n\n"); - messages.push(Message::system(&format!( + messages.push(Message::system(format!( "Relevant context:\n{}", doc_context ))); @@ -215,22 +215,18 @@ impl rig::completion::request::CompletionModel } } - fn stream( + async fn stream( &self, _request: CompletionRequest, - ) -> impl std::future::Future< - Output = Result< - rig::streaming::StreamingCompletionResponse, - CompletionError, - >, - > + Send { - async move { - // Our engines use blocking HTTP clients. Streaming is not supported - // through the rig adapter — callers should use `completion()` instead. - Err(CompletionError::ProviderError( - "Streaming not supported through RigModelAdapter; use completion() instead".into(), - )) - } + ) -> Result< + rig::streaming::StreamingCompletionResponse, + CompletionError, + > { + // Our engines use blocking HTTP clients. Streaming is not supported + // through the rig adapter — callers should use `completion()` instead. + Err(CompletionError::ProviderError( + "Streaming not supported through RigModelAdapter; use completion() instead".into(), + )) } } diff --git a/rust/crates/openjarvis-engine/src/sglang.rs b/rust/crates/openjarvis-engine/src/sglang.rs new file mode 100644 index 00000000..3c7aadd5 --- /dev/null +++ b/rust/crates/openjarvis-engine/src/sglang.rs @@ -0,0 +1,326 @@ +//! SGLang inference engine backend. +//! +//! SGLang exposes an OpenAI-compatible API at `http://host:port/v1/`. + +use crate::traits::{InferenceEngine, TokenStream}; +use openjarvis_core::error::{EngineError, OpenJarvisError}; +use openjarvis_core::{GenerateResult, Message, ToolCall, Usage}; +use serde_json::Value; + +/// SGLang backend via its OpenAI-compatible HTTP API. +pub struct SGLangEngine { + host: String, + client: reqwest::blocking::Client, + timeout: std::time::Duration, +} + +impl SGLangEngine { + pub fn new(host: &str, port: u16, timeout_secs: f64) -> Self { + let host = format!( + "{}:{}", + host.trim_end_matches('/').trim_end_matches(|c: char| c == ':' || c.is_ascii_digit()), + port + ); + let host = if host.starts_with("http") { + host + } else { + format!("http://{}", host) + }; + let host = host.trim_end_matches('/').to_string(); + let timeout = std::time::Duration::from_secs_f64(timeout_secs); + let client = reqwest::blocking::Client::builder() + .timeout(timeout) + .build() + .unwrap_or_default(); + Self { + host, + client, + timeout, + } + } + + pub fn with_defaults() -> Self { + Self::new("http://localhost", 30000, 120.0) + } + + pub fn from_host(host: &str) -> Self { + let host = host.trim_end_matches('/').to_string(); + let timeout = std::time::Duration::from_secs_f64(120.0); + let client = reqwest::blocking::Client::builder() + .timeout(timeout) + .build() + .unwrap_or_default(); + Self { + host, + client, + timeout, + } + } +} + +impl Default for SGLangEngine { + fn default() -> Self { + Self::with_defaults() + } +} + +#[async_trait::async_trait] +impl InferenceEngine for SGLangEngine { + fn engine_id(&self) -> &str { + "sglang" + } + + fn generate( + &self, + messages: &[Message], + model: &str, + temperature: f64, + max_tokens: i64, + extra: Option<&Value>, + ) -> Result { + let msg_dicts = crate::traits::messages_to_dicts(messages); + let mut payload = serde_json::json!({ + "model": model, + "messages": msg_dicts, + "temperature": temperature, + "max_tokens": max_tokens, + }); + + if let Some(extra_val) = extra { + if let Some(tools) = extra_val.get("tools") { + payload["tools"] = tools.clone(); + } + if let Some(obj) = extra_val.as_object() { + for (k, v) in obj { + if k != "tools" { + payload[k] = v.clone(); + } + } + } + } + + let resp = self + .client + .post(format!("{}/v1/chat/completions", self.host)) + .json(&payload) + .send() + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(format!( + "SGLang not reachable at {}: {}", + self.host, e + ))) + })?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().unwrap_or_default(); + return Err(OpenJarvisError::Engine(EngineError::Http(format!( + "SGLang returned {}: {}", + status, body + )))); + } + + let data: Value = resp.json().map_err(|e| { + OpenJarvisError::Engine(EngineError::Deserialization(e.to_string())) + })?; + + let choice = &data["choices"][0]; + let content = choice["message"]["content"] + .as_str() + .unwrap_or("") + .to_string(); + let finish_reason = choice["finish_reason"] + .as_str() + .unwrap_or("stop") + .to_string(); + + let usage_obj = &data["usage"]; + let usage = Usage { + prompt_tokens: usage_obj["prompt_tokens"].as_i64().unwrap_or(0), + completion_tokens: usage_obj["completion_tokens"].as_i64().unwrap_or(0), + total_tokens: usage_obj["total_tokens"].as_i64().unwrap_or(0), + }; + + let model_name = data["model"].as_str().unwrap_or(model).to_string(); + + let tool_calls = choice["message"]["tool_calls"] + .as_array() + .map(|arr| { + arr.iter() + .map(|tc| { + let func = &tc["function"]; + ToolCall { + id: tc["id"].as_str().unwrap_or("").to_string(), + name: func["name"].as_str().unwrap_or("").to_string(), + arguments: func["arguments"] + .as_str() + .unwrap_or("{}") + .to_string(), + } + }) + .collect() + }); + + Ok(GenerateResult { + content, + usage, + model: model_name, + finish_reason, + tool_calls, + ttft: 0.0, + cost_usd: 0.0, + metadata: std::collections::HashMap::new(), + }) + } + + async fn stream( + &self, + messages: &[Message], + model: &str, + temperature: f64, + max_tokens: i64, + extra: Option<&Value>, + ) -> Result { + let msg_dicts = crate::traits::messages_to_dicts(messages); + let mut payload = serde_json::json!({ + "model": model, + "messages": msg_dicts, + "temperature": temperature, + "max_tokens": max_tokens, + "stream": true, + }); + + if let Some(extra_val) = extra { + if let Some(tools) = extra_val.get("tools") { + payload["tools"] = tools.clone(); + } + } + + let async_client = reqwest::Client::builder() + .timeout(self.timeout) + .build() + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(e.to_string())) + })?; + + let resp = async_client + .post(format!("{}/v1/chat/completions", self.host)) + .json(&payload) + .send() + .await + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(format!( + "SGLang not reachable at {}: {}", + self.host, e + ))) + })?; + + if !resp.status().is_success() { + return Err(OpenJarvisError::Engine(EngineError::Http(format!( + "SGLang returned {}", + resp.status() + )))); + } + + use futures::StreamExt; + let byte_stream = resp.bytes_stream(); + + let token_stream = byte_stream.filter_map(|chunk_result| async { + match chunk_result { + Ok(bytes) => { + let text = String::from_utf8_lossy(&bytes); + for line in text.lines() { + let line = line.trim(); + if line.is_empty() || line == "data: [DONE]" { + continue; + } + let json_str = line.strip_prefix("data: ").unwrap_or(line); + if let Ok(chunk) = serde_json::from_str::(json_str) { + let content = chunk["choices"][0]["delta"]["content"] + .as_str() + .unwrap_or("") + .to_string(); + if !content.is_empty() { + return Some(Ok(content)); + } + } + } + None + } + Err(e) => Some(Err(OpenJarvisError::Engine(EngineError::Streaming( + e.to_string(), + )))), + } + }); + + Ok(Box::pin(token_stream)) + } + + fn list_models(&self) -> Result, OpenJarvisError> { + let resp = self + .client + .get(format!("{}/v1/models", self.host)) + .send() + .map_err(|_| { + OpenJarvisError::Engine(EngineError::Connection( + "SGLang not reachable".into(), + )) + })?; + + if !resp.status().is_success() { + return Ok(vec![]); + } + + let data: Value = resp.json().unwrap_or(Value::Null); + let models = data["data"] + .as_array() + .map(|arr| { + arr.iter() + .filter_map(|m| m["id"].as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + Ok(models) + } + + fn health(&self) -> bool { + // SGLang exposes /health; fall back to /v1/models. + let health_ok = self + .client + .get(format!("{}/health", self.host)) + .timeout(std::time::Duration::from_secs(2)) + .send() + .map(|r| r.status().is_success()) + .unwrap_or(false); + + if health_ok { + return true; + } + + self.client + .get(format!("{}/v1/models", self.host)) + .timeout(std::time::Duration::from_secs(2)) + .send() + .map(|r| r.status().is_success()) + .unwrap_or(false) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sglang_default_host() { + let engine = SGLangEngine::with_defaults(); + assert_eq!(engine.engine_id(), "sglang"); + assert_eq!(engine.host, "http://localhost:30000"); + } + + #[test] + fn test_sglang_from_host() { + let engine = SGLangEngine::from_host("http://gpu-server:30000"); + assert_eq!(engine.engine_id(), "sglang"); + assert_eq!(engine.host, "http://gpu-server:30000"); + } +} diff --git a/rust/crates/openjarvis-engine/src/vllm.rs b/rust/crates/openjarvis-engine/src/vllm.rs new file mode 100644 index 00000000..01ef4932 --- /dev/null +++ b/rust/crates/openjarvis-engine/src/vllm.rs @@ -0,0 +1,367 @@ +//! vLLM inference engine backend. +//! +//! vLLM exposes an OpenAI-compatible API at `http://host:port/v1/`. + +use crate::traits::{InferenceEngine, TokenStream}; +use openjarvis_core::error::{EngineError, OpenJarvisError}; +use openjarvis_core::{GenerateResult, Message, ToolCall, Usage}; +use serde_json::Value; + +/// vLLM backend via its OpenAI-compatible HTTP API. +pub struct VLLMEngine { + host: String, + client: reqwest::blocking::Client, + api_key: Option, + timeout: std::time::Duration, +} + +impl VLLMEngine { + pub fn new(host: &str, port: u16, api_key: Option, timeout_secs: f64) -> Self { + let host = format!( + "{}:{}", + host.trim_end_matches('/').trim_end_matches(|c: char| c == ':' || c.is_ascii_digit()), + port + ); + // If the host doesn't start with http, prepend it. + let host = if host.starts_with("http") { + host + } else { + format!("http://{}", host) + }; + let host = host.trim_end_matches('/').to_string(); + let timeout = std::time::Duration::from_secs_f64(timeout_secs); + let client = reqwest::blocking::Client::builder() + .timeout(timeout) + .build() + .unwrap_or_default(); + Self { + host, + client, + api_key, + timeout, + } + } + + pub fn with_defaults() -> Self { + Self::new("http://localhost", 8000, None, 120.0) + } + + pub fn from_host(host: &str) -> Self { + let host = host.trim_end_matches('/').to_string(); + let timeout = std::time::Duration::from_secs_f64(120.0); + let client = reqwest::blocking::Client::builder() + .timeout(timeout) + .build() + .unwrap_or_default(); + Self { + host, + client, + api_key: None, + timeout, + } + } + + fn build_headers(&self) -> reqwest::header::HeaderMap { + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + reqwest::header::CONTENT_TYPE, + "application/json".parse().unwrap(), + ); + if let Some(ref key) = self.api_key { + headers.insert( + reqwest::header::AUTHORIZATION, + format!("Bearer {}", key).parse().unwrap(), + ); + } + headers + } +} + +impl Default for VLLMEngine { + fn default() -> Self { + Self::with_defaults() + } +} + +#[async_trait::async_trait] +impl InferenceEngine for VLLMEngine { + fn engine_id(&self) -> &str { + "vllm" + } + + fn generate( + &self, + messages: &[Message], + model: &str, + temperature: f64, + max_tokens: i64, + extra: Option<&Value>, + ) -> Result { + let msg_dicts = crate::traits::messages_to_dicts(messages); + let mut payload = serde_json::json!({ + "model": model, + "messages": msg_dicts, + "temperature": temperature, + "max_tokens": max_tokens, + }); + + if let Some(extra_val) = extra { + if let Some(tools) = extra_val.get("tools") { + payload["tools"] = tools.clone(); + } + if let Some(obj) = extra_val.as_object() { + for (k, v) in obj { + if k != "tools" { + payload[k] = v.clone(); + } + } + } + } + + let resp = self + .client + .post(format!("{}/v1/chat/completions", self.host)) + .headers(self.build_headers()) + .json(&payload) + .send() + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(format!( + "vLLM not reachable at {}: {}", + self.host, e + ))) + })?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().unwrap_or_default(); + return Err(OpenJarvisError::Engine(EngineError::Http(format!( + "vLLM returned {}: {}", + status, body + )))); + } + + let data: Value = resp.json().map_err(|e| { + OpenJarvisError::Engine(EngineError::Deserialization(e.to_string())) + })?; + + let choice = &data["choices"][0]; + let content = choice["message"]["content"] + .as_str() + .unwrap_or("") + .to_string(); + let finish_reason = choice["finish_reason"] + .as_str() + .unwrap_or("stop") + .to_string(); + + let usage_obj = &data["usage"]; + let usage = Usage { + prompt_tokens: usage_obj["prompt_tokens"].as_i64().unwrap_or(0), + completion_tokens: usage_obj["completion_tokens"].as_i64().unwrap_or(0), + total_tokens: usage_obj["total_tokens"].as_i64().unwrap_or(0), + }; + + let model_name = data["model"].as_str().unwrap_or(model).to_string(); + + let tool_calls = choice["message"]["tool_calls"] + .as_array() + .map(|arr| { + arr.iter() + .map(|tc| { + let func = &tc["function"]; + ToolCall { + id: tc["id"].as_str().unwrap_or("").to_string(), + name: func["name"].as_str().unwrap_or("").to_string(), + arguments: func["arguments"] + .as_str() + .unwrap_or("{}") + .to_string(), + } + }) + .collect() + }); + + Ok(GenerateResult { + content, + usage, + model: model_name, + finish_reason, + tool_calls, + ttft: 0.0, + cost_usd: 0.0, + metadata: std::collections::HashMap::new(), + }) + } + + async fn stream( + &self, + messages: &[Message], + model: &str, + temperature: f64, + max_tokens: i64, + extra: Option<&Value>, + ) -> Result { + let msg_dicts = crate::traits::messages_to_dicts(messages); + let mut payload = serde_json::json!({ + "model": model, + "messages": msg_dicts, + "temperature": temperature, + "max_tokens": max_tokens, + "stream": true, + }); + + if let Some(extra_val) = extra { + if let Some(tools) = extra_val.get("tools") { + payload["tools"] = tools.clone(); + } + } + + let async_client = reqwest::Client::builder() + .timeout(self.timeout) + .build() + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(e.to_string())) + })?; + + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + reqwest::header::CONTENT_TYPE, + "application/json".parse().unwrap(), + ); + if let Some(ref key) = self.api_key { + headers.insert( + reqwest::header::AUTHORIZATION, + format!("Bearer {}", key).parse().unwrap(), + ); + } + + let resp = async_client + .post(format!("{}/v1/chat/completions", self.host)) + .headers(headers) + .json(&payload) + .send() + .await + .map_err(|e| { + OpenJarvisError::Engine(EngineError::Connection(format!( + "vLLM not reachable at {}: {}", + self.host, e + ))) + })?; + + if !resp.status().is_success() { + return Err(OpenJarvisError::Engine(EngineError::Http(format!( + "vLLM returned {}", + resp.status() + )))); + } + + use futures::StreamExt; + let byte_stream = resp.bytes_stream(); + + let token_stream = byte_stream.filter_map(|chunk_result| async { + match chunk_result { + Ok(bytes) => { + let text = String::from_utf8_lossy(&bytes); + for line in text.lines() { + let line = line.trim(); + if line.is_empty() || line == "data: [DONE]" { + continue; + } + let json_str = line.strip_prefix("data: ").unwrap_or(line); + if let Ok(chunk) = serde_json::from_str::(json_str) { + let content = chunk["choices"][0]["delta"]["content"] + .as_str() + .unwrap_or("") + .to_string(); + if !content.is_empty() { + return Some(Ok(content)); + } + } + } + None + } + Err(e) => Some(Err(OpenJarvisError::Engine(EngineError::Streaming( + e.to_string(), + )))), + } + }); + + Ok(Box::pin(token_stream)) + } + + fn list_models(&self) -> Result, OpenJarvisError> { + let resp = self + .client + .get(format!("{}/v1/models", self.host)) + .headers(self.build_headers()) + .send() + .map_err(|_| { + OpenJarvisError::Engine(EngineError::Connection( + "vLLM not reachable".into(), + )) + })?; + + if !resp.status().is_success() { + return Ok(vec![]); + } + + let data: Value = resp.json().unwrap_or(Value::Null); + let models = data["data"] + .as_array() + .map(|arr| { + arr.iter() + .filter_map(|m| m["id"].as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + Ok(models) + } + + fn health(&self) -> bool { + // vLLM exposes /health; fall back to /v1/models if needed. + let health_ok = self + .client + .get(format!("{}/health", self.host)) + .timeout(std::time::Duration::from_secs(2)) + .send() + .map(|r| r.status().is_success()) + .unwrap_or(false); + + if health_ok { + return true; + } + + self.client + .get(format!("{}/v1/models", self.host)) + .headers(self.build_headers()) + .timeout(std::time::Duration::from_secs(2)) + .send() + .map(|r| r.status().is_success()) + .unwrap_or(false) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_vllm_default_host() { + let engine = VLLMEngine::with_defaults(); + assert_eq!(engine.engine_id(), "vllm"); + assert_eq!(engine.host, "http://localhost:8000"); + } + + #[test] + fn test_vllm_from_host() { + let engine = VLLMEngine::from_host("http://gpu-server:8000"); + assert_eq!(engine.engine_id(), "vllm"); + assert_eq!(engine.host, "http://gpu-server:8000"); + } + + #[test] + fn test_vllm_with_api_key() { + let engine = VLLMEngine::new("http://localhost", 8000, Some("sk-test".into()), 60.0); + assert!(engine.api_key.is_some()); + } +} diff --git a/rust/crates/openjarvis-learning/Cargo.toml b/rust/crates/openjarvis-learning/Cargo.toml index f39f7a80..b2e5a55b 100644 --- a/rust/crates/openjarvis-learning/Cargo.toml +++ b/rust/crates/openjarvis-learning/Cargo.toml @@ -12,5 +12,7 @@ thiserror = { workspace = true } tracing = { workspace = true } rand = { workspace = true } parking_lot = { workspace = true } +rusqlite = { workspace = true } +uuid = { workspace = true } regex = { workspace = true } once_cell = { workspace = true } diff --git a/rust/crates/openjarvis-learning/src/heuristic_reward.rs b/rust/crates/openjarvis-learning/src/heuristic_reward.rs index 428d14f5..eb5dfaa2 100644 --- a/rust/crates/openjarvis-learning/src/heuristic_reward.rs +++ b/rust/crates/openjarvis-learning/src/heuristic_reward.rs @@ -88,6 +88,6 @@ mod tests { fn test_clamp_bounds() { let rf = HeuristicRewardFunction::default(); let r = rf.compute(100.0, 1.0, 0, 0); - assert!(r >= 0.0 && r <= 1.0); + assert!((0.0..=1.0).contains(&r)); } } diff --git a/rust/crates/openjarvis-learning/src/lib.rs b/rust/crates/openjarvis-learning/src/lib.rs index 7ec61119..ad70d3a6 100644 --- a/rust/crates/openjarvis-learning/src/lib.rs +++ b/rust/crates/openjarvis-learning/src/lib.rs @@ -10,6 +10,7 @@ pub mod heuristic; pub mod heuristic_reward; pub mod icl_updater; pub mod learning_orchestrator; +pub mod optimize; pub mod orchestrator_types; pub mod reward; pub mod router_enum; diff --git a/rust/crates/openjarvis-learning/src/optimize/engine.rs b/rust/crates/openjarvis-learning/src/optimize/engine.rs new file mode 100644 index 00000000..a3e8af9f --- /dev/null +++ b/rust/crates/openjarvis-learning/src/optimize/engine.rs @@ -0,0 +1,484 @@ +//! OptimizationEngine -- orchestrates the optimize loop. +//! +//! Rust translation of `src/openjarvis/learning/optimize/optimizer.py`. +//! +//! Ties together the LLM optimizer, trial runner, and persistence store +//! into a single propose -> evaluate -> analyze -> repeat loop. + +use tracing::{info, warn}; + +use super::llm_optimizer::LLMOptimizer; +use super::store::OptimizationStore; +use super::types::{ + Direction, ObjectiveSpec, OptimizationRun, RunStatus, SearchSpace, TrialConfig, TrialResult, +}; + +// --------------------------------------------------------------------------- +// Trial runner trait (Python's TrialRunner is a class -- here it's a trait) +// --------------------------------------------------------------------------- + +/// Abstraction for evaluating a trial configuration. +/// +/// Implementations run benchmarks and return [`TrialResult`]s. +/// The actual evaluation logic stays in Python (or is plugged in via FFI). +pub trait TrialRunner: Send + Sync { + /// Execute a single trial and return the result. + fn run_trial(&self, config: &TrialConfig) -> TrialResult; + + /// Name of the primary benchmark being evaluated. + fn benchmark(&self) -> &str { + "" + } + + /// Names of all benchmarks (for multi-benchmark mode). + fn benchmark_names(&self) -> Vec { + vec![] + } +} + +// --------------------------------------------------------------------------- +// Pareto frontier computation +// --------------------------------------------------------------------------- + +/// Read a metric value from a [`TrialResult`] for a given objective. +fn get_objective_value(trial: &TrialResult, obj: &ObjectiveSpec) -> f64 { + match obj.metric.as_str() { + "accuracy" => trial.accuracy, + "mean_latency_seconds" => trial.mean_latency_seconds, + "total_cost_usd" => trial.total_cost_usd, + "total_energy_joules" => trial.total_energy_joules, + _ => 0.0, + } +} + +/// Compute the Pareto frontier: trials not dominated by any other. +/// +/// A trial A dominates trial B if A is >= B on all objectives and > B +/// on at least one (direction-aware: `Minimize` negates so higher is better). +pub fn compute_pareto_frontier( + trials: &[TrialResult], + objectives: &[ObjectiveSpec], +) -> Vec { + if trials.is_empty() || objectives.is_empty() { + return trials.to_vec(); + } + + let values: Vec> = trials + .iter() + .map(|t| { + objectives + .iter() + .map(|obj| { + let v = get_objective_value(t, obj); + // Normalize: for "minimize", negate so higher is always better + if obj.direction == Direction::Minimize { + -v + } else { + v + } + }) + .collect() + }) + .collect(); + + let mut frontier = Vec::new(); + + for (i, _trial) in trials.iter().enumerate() { + let mut dominated = false; + for (j, _other) in trials.iter().enumerate() { + if i == j { + continue; + } + // Check if other dominates trial + let all_ge = (0..objectives.len()).all(|k| values[j][k] >= values[i][k]); + let any_gt = (0..objectives.len()).any(|k| values[j][k] > values[i][k]); + if all_ge && any_gt { + dominated = true; + break; + } + } + if !dominated { + frontier.push(trials[i].clone()); + } + } + + frontier +} + +// --------------------------------------------------------------------------- +// OptimizationEngine +// --------------------------------------------------------------------------- + +/// Orchestrates the optimize loop: propose -> evaluate -> analyze -> repeat. +pub struct OptimizationEngine { + pub search_space: SearchSpace, + pub llm_optimizer: LLMOptimizer, + pub trial_runner: R, + pub store: Option, + pub max_trials: usize, + pub early_stop_patience: usize, +} + +impl OptimizationEngine { + /// Create a new optimization engine. + pub fn new( + search_space: SearchSpace, + llm_optimizer: LLMOptimizer, + trial_runner: R, + store: Option, + max_trials: usize, + early_stop_patience: usize, + ) -> Self { + Self { + search_space, + llm_optimizer, + trial_runner, + store, + max_trials, + early_stop_patience, + } + } + + /// Execute the full optimization loop. + /// + /// 1. Generate a run_id. + /// 2. `llm_optimizer.propose_initial()` -> first config. + /// 3. Loop up to `max_trials`: + /// a. `trial_runner.run_trial(config)` -> TrialResult + /// b. Update history, track best, compute Pareto frontier + /// c. Persist to store if available + /// d. Check early stopping + /// e. Propose next config + /// 4. Return the completed `OptimizationRun`. + pub fn run( + &mut self, + progress_callback: Option<&dyn Fn(usize, usize)>, + ) -> OptimizationRun { + let run_id = uuid::Uuid::new_v4().simple().to_string()[..16].to_string(); + + let benchmark_name = self.trial_runner.benchmark().to_string(); + let benchmark_names = self.trial_runner.benchmark_names(); + + let mut optimization_run = OptimizationRun { + run_id: run_id.clone(), + search_space: self.search_space.clone(), + status: RunStatus::Running, + optimizer_model: self.llm_optimizer.optimizer_model.clone(), + benchmark: if benchmark_names.is_empty() { + benchmark_name + } else { + benchmark_names.join("+") + }, + benchmarks: benchmark_names, + trials: vec![], + best_trial: None, + best_recipe_path: None, + pareto_frontier: vec![], + objectives: super::types::default_objectives(), + }; + + let mut history: Vec = Vec::new(); + let mut best_accuracy: f64 = -1.0; + let mut trials_without_improvement: usize = 0; + + // First config + let mut config = self.llm_optimizer.propose_initial(); + + for trial_num in 1..=self.max_trials { + info!( + "Trial {}/{} (id={})", + trial_num, self.max_trials, config.trial_id + ); + + // Evaluate + let mut result = self.trial_runner.run_trial(&config); + + // Placeholder analysis (the real LLM analysis stays in Python) + if result.analysis.is_empty() { + result.analysis = format!( + "Trial {} completed: accuracy={:.4}, latency={:.4}s", + result.trial_id, result.accuracy, result.mean_latency_seconds + ); + } + + // Record + history.push(result.clone()); + optimization_run.trials.push(result.clone()); + + // Recompute Pareto frontier + optimization_run.pareto_frontier = + compute_pareto_frontier(&history, &optimization_run.objectives); + + // Persist trial + if let Some(ref store) = self.store { + if let Err(e) = store.save_trial(&run_id, &result) { + warn!("Failed to save trial: {e}"); + } + } + + // Track best + if result.accuracy > best_accuracy { + best_accuracy = result.accuracy; + optimization_run.best_trial = Some(result.clone()); + trials_without_improvement = 0; + } else { + trials_without_improvement += 1; + } + + // Progress callback + if let Some(cb) = progress_callback { + cb(trial_num, self.max_trials); + } + + // Early stopping + if trials_without_improvement >= self.early_stop_patience { + info!( + "Early stopping after {} trials without improvement.", + self.early_stop_patience + ); + break; + } + + // Propose next (unless this was the last trial) + if trial_num < self.max_trials { + config = self.llm_optimizer.propose_next(&history); + } + } + + optimization_run.status = RunStatus::Completed; + + if let Some(ref store) = self.store { + if let Err(e) = store.save_run(&optimization_run) { + warn!("Failed to save optimization run: {e}"); + } + } + + optimization_run + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::optimize::types::*; + use std::collections::HashMap; + + // A simple mock trial runner for testing + struct MockTrialRunner { + results: Vec, + call_count: std::sync::atomic::AtomicUsize, + } + + impl MockTrialRunner { + fn new(results: Vec) -> Self { + Self { + results, + call_count: std::sync::atomic::AtomicUsize::new(0), + } + } + } + + impl TrialRunner for MockTrialRunner { + fn run_trial(&self, config: &TrialConfig) -> TrialResult { + let idx = self.call_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let accuracy = if idx < self.results.len() { + self.results[idx] + } else { + 0.5 + }; + + TrialResult { + trial_id: config.trial_id.clone(), + config: config.clone(), + accuracy, + mean_latency_seconds: 1.0, + total_cost_usd: 0.01, + total_energy_joules: 10.0, + total_tokens: 100, + samples_evaluated: 10, + analysis: String::new(), + failure_modes: vec![], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + } + } + + fn benchmark(&self) -> &str { + "mock_bench" + } + } + + #[test] + fn test_compute_pareto_frontier_single() { + let trial = TrialResult { + trial_id: "t1".into(), + config: TrialConfig { + trial_id: "t1".into(), + params: HashMap::new(), + reasoning: String::new(), + }, + accuracy: 0.9, + mean_latency_seconds: 1.0, + total_cost_usd: 0.1, + total_energy_joules: 10.0, + total_tokens: 100, + samples_evaluated: 10, + analysis: String::new(), + failure_modes: vec![], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + }; + + let frontier = compute_pareto_frontier(std::slice::from_ref(&trial), &default_objectives()); + assert_eq!(frontier.len(), 1); + assert_eq!(frontier[0].trial_id, "t1"); + } + + #[test] + fn test_compute_pareto_frontier_dominated() { + let make = |id: &str, acc: f64, lat: f64, cost: f64| TrialResult { + trial_id: id.into(), + config: TrialConfig { + trial_id: id.into(), + params: HashMap::new(), + reasoning: String::new(), + }, + accuracy: acc, + mean_latency_seconds: lat, + total_cost_usd: cost, + total_energy_joules: 0.0, + total_tokens: 0, + samples_evaluated: 0, + analysis: String::new(), + failure_modes: vec![], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + }; + + let trials = vec![ + make("t1", 0.9, 1.0, 0.1), // best accuracy, worst latency/cost + make("t2", 0.7, 0.5, 0.05), // lower accuracy, better latency/cost + make("t3", 0.6, 0.8, 0.08), // dominated by t2 (worse on all) + ]; + + let objs = default_objectives(); + let frontier = compute_pareto_frontier(&trials, &objs); + // t1 and t2 should be on frontier; t3 is dominated by t2 + assert_eq!(frontier.len(), 2); + let ids: Vec<&str> = frontier.iter().map(|t| t.trial_id.as_str()).collect(); + assert!(ids.contains(&"t1")); + assert!(ids.contains(&"t2")); + } + + #[test] + fn test_compute_pareto_frontier_empty() { + let frontier = compute_pareto_frontier(&[], &default_objectives()); + assert!(frontier.is_empty()); + } + + #[test] + fn test_optimization_engine_run() { + let search_space = SearchSpace::default(); + let llm_opt = LLMOptimizer::new(search_space.clone(), "test-model".into()); + let runner = MockTrialRunner::new(vec![0.5, 0.6, 0.7, 0.8, 0.75]); + + let mut engine = OptimizationEngine::new( + search_space, + llm_opt, + runner, + None, // no store + 5, + 10, // high patience so we run all trials + ); + + let run = engine.run(None); + assert_eq!(run.status, RunStatus::Completed); + assert_eq!(run.trials.len(), 5); + assert!(run.best_trial.is_some()); + assert!((run.best_trial.unwrap().accuracy - 0.8).abs() < 1e-9); + } + + #[test] + fn test_optimization_engine_early_stop() { + let search_space = SearchSpace::default(); + let llm_opt = LLMOptimizer::new(search_space.clone(), "test-model".into()); + // accuracy goes down after first trial, triggers early stop at patience=2 + let runner = MockTrialRunner::new(vec![0.8, 0.7, 0.6, 0.5]); + + let mut engine = OptimizationEngine::new( + search_space, + llm_opt, + runner, + None, + 10, + 2, // early stop after 2 trials without improvement + ); + + let run = engine.run(None); + assert_eq!(run.status, RunStatus::Completed); + // Should stop after 3 trials (1 good + 2 without improvement) + assert_eq!(run.trials.len(), 3); + } + + #[test] + fn test_optimization_engine_with_store() { + let search_space = SearchSpace::default(); + let llm_opt = LLMOptimizer::new(search_space.clone(), "test-model".into()); + let runner = MockTrialRunner::new(vec![0.5, 0.7]); + let store = OptimizationStore::in_memory().unwrap(); + + let mut engine = OptimizationEngine::new( + search_space, + llm_opt, + runner, + Some(store), + 2, + 10, + ); + + let run = engine.run(None); + assert_eq!(run.status, RunStatus::Completed); + assert_eq!(run.trials.len(), 2); + + // Verify the store has the data + let stored_run = engine + .store + .as_ref() + .unwrap() + .get_run(&run.run_id) + .unwrap(); + assert!(stored_run.is_some()); + } + + #[test] + fn test_optimization_engine_progress_callback() { + let search_space = SearchSpace::default(); + let llm_opt = LLMOptimizer::new(search_space.clone(), "test-model".into()); + let runner = MockTrialRunner::new(vec![0.5, 0.6]); + + let mut engine = OptimizationEngine::new( + search_space, + llm_opt, + runner, + None, + 2, + 10, + ); + + let progress = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let progress_clone = progress.clone(); + + let run = engine.run(Some(&move |trial_num, max_trials| { + progress_clone.lock().unwrap().push((trial_num, max_trials)); + })); + + assert_eq!(run.trials.len(), 2); + let recorded = progress.lock().unwrap(); + assert_eq!(recorded.len(), 2); + assert_eq!(recorded[0], (1, 2)); + assert_eq!(recorded[1], (2, 2)); + } +} diff --git a/rust/crates/openjarvis-learning/src/optimize/llm_optimizer.rs b/rust/crates/openjarvis-learning/src/optimize/llm_optimizer.rs new file mode 100644 index 00000000..16d0c107 --- /dev/null +++ b/rust/crates/openjarvis-learning/src/optimize/llm_optimizer.rs @@ -0,0 +1,860 @@ +//! LLM-based optimizer for OpenJarvis configuration tuning. +//! +//! Rust translation of `src/openjarvis/learning/optimize/llm_optimizer.py`. +//! +//! The actual LLM call is abstracted via the [`OptimizerBackend`] trait so +//! that Python can provide the implementation via FFI / PyO3. + +use std::collections::{HashMap, HashSet}; + +use super::types::{SearchSpace, TrialConfig, TrialFeedback, TrialResult}; + +// --------------------------------------------------------------------------- +// Backend trait (Python fills this in) +// --------------------------------------------------------------------------- + +/// Abstraction for the LLM backend used by the optimizer. +/// +/// The real implementation calls a cloud LLM (e.g. Claude, GPT) to generate +/// configuration proposals and trial analyses. In Rust-only tests a simple +/// mock can be used. +pub trait OptimizerBackend: Send + Sync { + /// Generate a text response from the LLM. + fn generate( + &self, + prompt: &str, + model: &str, + system: &str, + temperature: f64, + max_tokens: usize, + ) -> String; +} + +// --------------------------------------------------------------------------- +// LLMOptimizer +// --------------------------------------------------------------------------- + +/// Uses an LLM to propose optimal OpenJarvis configurations. +/// +/// Inspired by DSPy's GEPA: uses textual feedback from execution traces +/// rather than just scalar rewards to guide the optimizer. +pub struct LLMOptimizer { + pub search_space: SearchSpace, + pub optimizer_model: String, + backend: Option>, +} + +impl LLMOptimizer { + /// Create a new LLM optimizer. + /// + /// If `backend` is `None`, `propose_initial` / `propose_next` will + /// return a default config derived from the search space's fixed params. + pub fn new(search_space: SearchSpace, optimizer_model: String) -> Self { + Self { + search_space, + optimizer_model, + backend: None, + } + } + + /// Create a new LLM optimizer with a backend. + pub fn with_backend( + search_space: SearchSpace, + optimizer_model: String, + backend: Box, + ) -> Self { + Self { + search_space, + optimizer_model, + backend: Some(backend), + } + } + + // ------------------------------------------------------------------ + // Public API + // ------------------------------------------------------------------ + + /// Propose a reasonable starting config from the search space. + /// + /// If no backend is configured, returns a config with just fixed params. + pub fn propose_initial(&self) -> TrialConfig { + let trial_id = new_trial_id(); + + if let Some(ref backend) = self.backend { + let prompt = self.build_initial_prompt(); + let response = backend.generate( + &prompt, + &self.optimizer_model, + "You are an expert AI systems optimizer.", + 0.7, + 2048, + ); + return self.parse_config_response(&response, &trial_id); + } + + // Fallback: use fixed params + let fixed: HashMap = self + .search_space + .fixed + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + TrialConfig { + trial_id, + params: fixed, + reasoning: "Default config from fixed parameters".into(), + } + } + + /// Propose the next config based on trial history. + pub fn propose_next(&self, history: &[TrialResult]) -> TrialConfig { + let trial_id = new_trial_id(); + + if let Some(ref backend) = self.backend { + let prompt = self.build_propose_prompt(history, None); + let response = backend.generate( + &prompt, + &self.optimizer_model, + "You are an expert AI systems optimizer.", + 0.7, + 2048, + ); + return self.parse_config_response(&response, &trial_id); + } + + // Fallback + let fixed: HashMap = self + .search_space + .fixed + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + TrialConfig { + trial_id, + params: fixed, + reasoning: "Default config (no backend available)".into(), + } + } + + /// Propose a config that only changes one pillar. + pub fn propose_targeted( + &self, + history: &[TrialResult], + base_config: &TrialConfig, + target_pillar: &str, + frontier_ids: Option<&HashSet>, + ) -> TrialConfig { + let trial_id = new_trial_id(); + + if let Some(ref backend) = self.backend { + let prompt = + self.build_targeted_prompt(history, base_config, target_pillar, frontier_ids); + let response = backend.generate( + &prompt, + &self.optimizer_model, + "You are an expert AI systems optimizer.", + 0.7, + 2048, + ); + let proposed = self.parse_config_response(&response, &trial_id); + + // Enforce constraint: preserve non-target params from base_config + let mut merged = base_config.params.clone(); + let target_prefix = format!("{target_pillar}."); + let alt_prefix = { + let trimmed = target_pillar.trim_end_matches('s'); + format!("{trimmed}.") + }; + for (key, value) in &proposed.params { + if key.starts_with(&target_prefix) || key.starts_with(&alt_prefix) { + merged.insert(key.clone(), value.clone()); + } + } + + return TrialConfig { + trial_id: proposed.trial_id, + params: merged, + reasoning: proposed.reasoning, + }; + } + + // Fallback: return base config as-is + TrialConfig { + trial_id, + params: base_config.params.clone(), + reasoning: format!("Targeted mutation on {target_pillar} (no backend)"), + } + } + + /// Combine best aspects of frontier members into one config. + pub fn propose_merge( + &self, + candidates: &[TrialResult], + history: &[TrialResult], + frontier_ids: Option<&HashSet>, + ) -> TrialConfig { + let trial_id = new_trial_id(); + + if let Some(ref backend) = self.backend { + let prompt = self.build_merge_prompt(candidates, history, frontier_ids); + let response = backend.generate( + &prompt, + &self.optimizer_model, + "You are an expert AI systems optimizer.", + 0.7, + 2048, + ); + return self.parse_config_response(&response, &trial_id); + } + + // Fallback: use the first candidate's config + let params = candidates + .first() + .map(|c| c.config.params.clone()) + .unwrap_or_default(); + + TrialConfig { + trial_id, + params, + reasoning: "Merge fallback (no backend)".into(), + } + } + + /// Analyze a completed trial. Returns structured feedback. + pub fn analyze_trial( + &self, + trial: &TrialConfig, + accuracy: f64, + mean_latency_seconds: f64, + total_cost_usd: f64, + ) -> TrialFeedback { + if let Some(ref backend) = self.backend { + let prompt = format!( + "Analyze this OpenJarvis evaluation result.\n\n\ + ## Configuration\n{}\n\n\ + ## Results\n- accuracy: {accuracy:.4}\n\ + - mean_latency_seconds: {mean_latency_seconds:.4}\n\ + - total_cost_usd: {total_cost_usd:.4}\n\n\ + Provide your analysis as a JSON object inside a ```json code block with:\n\ + 1. \"summary_text\": string with detailed analysis\n\ + 2. \"failure_patterns\": list of identified failure patterns\n\ + 3. \"pillar_ratings\": dict mapping pillar names to \"high\"/\"medium\"/\"low\"\n\ + 4. \"suggested_changes\": list of specific config changes to try\n\ + 5. \"target_pillar\": which pillar to change next", + format_config_params(&trial.params), + ); + let response = backend.generate( + &prompt, + &self.optimizer_model, + "You are an expert AI systems analyst.", + 0.3, + 2048, + ); + return parse_feedback_response(&response); + } + + // Fallback + TrialFeedback { + summary_text: format!( + "Trial {}: accuracy={accuracy:.4}, latency={mean_latency_seconds:.4}s, cost=${total_cost_usd:.4}", + trial.trial_id, + ), + ..Default::default() + } + } + + // ------------------------------------------------------------------ + // Prompt builders + // ------------------------------------------------------------------ + + #[allow(clippy::vec_init_then_push)] + fn build_initial_prompt(&self) -> String { + let mut lines = Vec::new(); + lines.push("You are optimizing an OpenJarvis AI system configuration.".into()); + lines.push(String::new()); + lines.push(self.search_space.to_prompt_description()); + lines.push("## Objective".into()); + lines.push("Maximize accuracy while minimizing latency and cost.".into()); + lines.push(String::new()); + lines.push("## Your Task".into()); + lines.push( + "Propose an initial configuration that is a reasonable starting \ + point for optimization. Choose sensible defaults that balance \ + accuracy, latency, and cost." + .into(), + ); + lines.push(String::new()); + lines.push( + "Return a JSON object inside a ```json code block with:".into(), + ); + lines.push( + "1. \"params\": dict of config params (dotted keys matching the search space)" + .into(), + ); + lines.push( + "2. \"reasoning\": string explaining why this is a good starting configuration" + .into(), + ); + lines.join("\n") + } + + #[allow(clippy::vec_init_then_push)] + fn build_propose_prompt( + &self, + history: &[TrialResult], + frontier_ids: Option<&HashSet>, + ) -> String { + let mut lines = Vec::new(); + lines.push("You are optimizing an OpenJarvis AI system configuration.".into()); + lines.push(String::new()); + lines.push(self.search_space.to_prompt_description()); + + lines.push("## Optimization History".into()); + if history.is_empty() { + lines.push("No trials have been run yet.".into()); + } else { + lines.push(format_history(history, frontier_ids)); + } + lines.push(String::new()); + + lines.push("## Objective".into()); + lines.push("Maximize accuracy while minimizing latency and cost.".into()); + lines.push(String::new()); + lines.push("## Your Task".into()); + lines.push( + "Propose the next configuration to evaluate. Learn from \ + previous trials to improve results." + .into(), + ); + lines.push(String::new()); + lines.push( + "Return a JSON object inside a ```json code block with:".into(), + ); + lines.push( + "1. \"params\": dict of config params (dotted keys matching the search space)" + .into(), + ); + lines.push( + "2. \"reasoning\": string explaining why this config should improve results".into(), + ); + lines.join("\n") + } + + #[allow(clippy::vec_init_then_push)] + fn build_targeted_prompt( + &self, + history: &[TrialResult], + base_config: &TrialConfig, + target_pillar: &str, + frontier_ids: Option<&HashSet>, + ) -> String { + let mut lines = Vec::new(); + lines.push("You are optimizing an OpenJarvis AI system configuration.".into()); + lines.push(String::new()); + lines.push(self.search_space.to_prompt_description()); + + lines.push("## Base Configuration".into()); + lines.push(format_config_params(&base_config.params)); + lines.push(String::new()); + + lines.push(format!("## Target Pillar: {target_pillar}")); + lines.push(format!( + "ONLY change parameters under the '{target_pillar}' pillar. \ + Keep all other parameters exactly as they are." + )); + lines.push(String::new()); + + lines.push("## Optimization History".into()); + if !history.is_empty() { + lines.push(format_history(history, frontier_ids)); + } + lines.push(String::new()); + + lines.push( + format!( + "Return a JSON object inside a ```json code block with:\n\ + 1. \"params\": dict of config params (only change {target_pillar} params)\n\ + 2. \"reasoning\": string explaining your changes" + ), + ); + lines.join("\n") + } + + #[allow(clippy::vec_init_then_push)] + fn build_merge_prompt( + &self, + candidates: &[TrialResult], + history: &[TrialResult], + frontier_ids: Option<&HashSet>, + ) -> String { + let mut lines = Vec::new(); + lines.push("You are optimizing an OpenJarvis AI system configuration.".into()); + lines.push(String::new()); + lines.push(self.search_space.to_prompt_description()); + + lines.push("## Frontier Candidates to Merge".into()); + for (i, cand) in candidates.iter().enumerate() { + lines.push(format!( + "### Candidate {} (id={})", + i + 1, + cand.trial_id + )); + lines.push(format!( + "Params: {}", + serde_json::to_string(&cand.config.params).unwrap_or_default() + )); + lines.push(format!("Accuracy: {:.4}", cand.accuracy)); + lines.push(format!("Latency: {:.4}s", cand.mean_latency_seconds)); + lines.push(format!("Cost: ${:.4}", cand.total_cost_usd)); + lines.push(format!("Energy: {:.4}J", cand.total_energy_joules)); + lines.push(String::new()); + } + + lines.push( + "Combine the best aspects of these frontier configs into \ + one unified configuration." + .into(), + ); + lines.push(String::new()); + + if !history.is_empty() { + lines.push("## Full History".into()); + lines.push(format_history(history, frontier_ids)); + lines.push(String::new()); + } + + lines.push( + "Return a JSON object inside a ```json code block with:\n\ + 1. \"params\": dict of merged config params\n\ + 2. \"reasoning\": string explaining the merge strategy" + .into(), + ); + lines.join("\n") + } + + // ------------------------------------------------------------------ + // Response parsing + // ------------------------------------------------------------------ + + fn parse_config_response(&self, response: &str, trial_id: &str) -> TrialConfig { + // Try to extract from ```json code block + if let Some(json_str) = extract_json_block(response) { + if let Ok(data) = serde_json::from_str::(&json_str) { + if let Some(config) = self.config_from_value(&data, trial_id) { + return config; + } + } + } + + // Try to find a raw JSON object + if let Some(config) = self.try_parse_raw_json(response, trial_id) { + return config; + } + + // Last resort: return config with fixed params + let fixed: HashMap = self + .search_space + .fixed + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + TrialConfig { + trial_id: trial_id.into(), + params: fixed, + reasoning: "Failed to parse LLM response.".into(), + } + } + + fn config_from_value( + &self, + data: &serde_json::Value, + trial_id: &str, + ) -> Option { + let params_val = data.get("params")?; + let params: HashMap = + serde_json::from_value(params_val.clone()).ok()?; + let reasoning = data + .get("reasoning") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + + // Inject fixed params + let mut merged = params; + for (key, value) in &self.search_space.fixed { + merged.insert(key.clone(), value.clone()); + } + + Some(TrialConfig { + trial_id: trial_id.into(), + params: merged, + reasoning, + }) + } + + fn try_parse_raw_json(&self, response: &str, trial_id: &str) -> Option { + // Scan for '{' and try to parse JSON objects + for (idx, _) in response.match_indices('{') { + if let Ok(data) = serde_json::from_str::(&response[idx..]) { + if data.is_object() { + if let Some(config) = self.config_from_value(&data, trial_id) { + return Some(config); + } + } + } + } + None + } +} + +// --------------------------------------------------------------------------- +// Free-standing helpers +// --------------------------------------------------------------------------- + +fn new_trial_id() -> String { + uuid::Uuid::new_v4().simple().to_string()[..12].to_string() +} + +/// Extract content from a ```json ... ``` or ``` ... ``` code block. +fn extract_json_block(text: &str) -> Option { + // Try ```json first ((?s) enables dotall so . matches \n) + let re_json = regex::Regex::new(r"(?s)```json\s*\n?(.*?)\n?\s*```").ok()?; + if let Some(caps) = re_json.captures(text) { + return Some(caps.get(1)?.as_str().trim().to_string()); + } + + // Try generic ``` + let re_code = regex::Regex::new(r"(?s)```\s*\n?(.*?)\n?\s*```").ok()?; + if let Some(caps) = re_code.captures(text) { + return Some(caps.get(1)?.as_str().trim().to_string()); + } + + None +} + +/// Parse an LLM response into a [`TrialFeedback`]. +fn parse_feedback_response(response: &str) -> TrialFeedback { + if let Some(json_str) = extract_json_block(response) { + if let Ok(fb) = serde_json::from_str::(&json_str) { + return fb; + } + } + + // Try raw JSON + for (idx, _) in response.match_indices('{') { + if let Ok(fb) = serde_json::from_str::(&response[idx..]) { + return fb; + } + } + + // Fallback: wrap raw text + TrialFeedback { + summary_text: response.trim().to_string(), + ..Default::default() + } +} + +fn format_config_params(params: &HashMap) -> String { + let mut sorted: Vec<(&String, &serde_json::Value)> = params.iter().collect(); + sorted.sort_by_key(|(k, _)| k.as_str()); + sorted + .iter() + .map(|(k, v)| format!("- {k}: {v}")) + .collect::>() + .join("\n") +} + +fn format_history( + history: &[TrialResult], + frontier_ids: Option<&HashSet>, +) -> String { + let mut lines = Vec::new(); + for (i, result) in history.iter().enumerate() { + let tag = frontier_ids + .filter(|ids| ids.contains(&result.trial_id)) + .map(|_| " [FRONTIER]") + .unwrap_or(""); + + lines.push(format!( + "### Trial {} (id={}){tag}", + i + 1, + result.trial_id + )); + lines.push(format!( + "Params: {}", + serde_json::to_string(&result.config.params).unwrap_or_default() + )); + lines.push(format!("Accuracy: {:.4}", result.accuracy)); + lines.push(format!("Latency: {:.4}s", result.mean_latency_seconds)); + lines.push(format!("Cost: ${:.4}", result.total_cost_usd)); + lines.push(format!("Energy: {:.4}J", result.total_energy_joules)); + + if let Some(ref fb) = result.structured_feedback { + if !fb.failure_patterns.is_empty() { + lines.push(format!( + "Failure patterns: {}", + fb.failure_patterns.join(", ") + )); + } + if !fb.pillar_ratings.is_empty() { + let ratings: Vec = { + let mut sorted: Vec<_> = fb.pillar_ratings.iter().collect(); + sorted.sort_by_key(|(k, _)| k.as_str()); + sorted + .iter() + .map(|(k, v)| format!("{k}={v}")) + .collect() + }; + lines.push(format!("Pillar ratings: {}", ratings.join(", "))); + } + if !fb.target_pillar.is_empty() { + lines.push(format!("Target pillar: {}", fb.target_pillar)); + } + } else if !result.analysis.is_empty() { + lines.push(format!("Analysis: {}", result.analysis)); + } + + if !result.failure_modes.is_empty() { + lines.push(format!( + "Failure modes: {}", + result.failure_modes.join(", ") + )); + } + lines.push(String::new()); + } + lines.join("\n") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::optimize::types::*; + + /// A mock backend that returns a canned JSON response. + struct MockBackend { + response: String, + } + + impl OptimizerBackend for MockBackend { + fn generate( + &self, + _prompt: &str, + _model: &str, + _system: &str, + _temperature: f64, + _max_tokens: usize, + ) -> String { + self.response.clone() + } + } + + #[test] + fn test_propose_initial_no_backend() { + let space = SearchSpace { + fixed: { + let mut m = HashMap::new(); + m.insert("intelligence.model".into(), serde_json::json!("qwen3:8b")); + m + }, + ..Default::default() + }; + let opt = LLMOptimizer::new(space, "test".into()); + let config = opt.propose_initial(); + assert!(!config.trial_id.is_empty()); + assert_eq!( + config.params.get("intelligence.model"), + Some(&serde_json::json!("qwen3:8b")) + ); + } + + #[test] + fn test_propose_initial_with_backend() { + let response = r#"Here is my proposal: + +```json +{ + "params": { + "intelligence.model": "qwen3:8b", + "intelligence.temperature": 0.7, + "agent.type": "orchestrator" + }, + "reasoning": "Good starting point" +} +```"#; + + let space = SearchSpace::default(); + let opt = LLMOptimizer::with_backend( + space, + "test".into(), + Box::new(MockBackend { + response: response.into(), + }), + ); + + let config = opt.propose_initial(); + assert_eq!( + config.params.get("intelligence.model"), + Some(&serde_json::json!("qwen3:8b")) + ); + assert_eq!( + config.params.get("agent.type"), + Some(&serde_json::json!("orchestrator")) + ); + assert_eq!(config.reasoning, "Good starting point"); + } + + #[test] + fn test_propose_next_no_backend() { + let space = SearchSpace::default(); + let opt = LLMOptimizer::new(space, "test".into()); + let config = opt.propose_next(&[]); + assert!(!config.trial_id.is_empty()); + } + + #[test] + fn test_parse_feedback_response_json_block() { + let response = r#"```json +{ + "summary_text": "Good result overall", + "failure_patterns": ["timeout"], + "pillar_ratings": {"intelligence": "high"}, + "suggested_changes": ["lower temp"], + "target_pillar": "intelligence" +} +```"#; + let fb = parse_feedback_response(response); + assert_eq!(fb.summary_text, "Good result overall"); + assert_eq!(fb.failure_patterns, vec!["timeout"]); + assert_eq!(fb.target_pillar, "intelligence"); + } + + #[test] + fn test_parse_feedback_response_fallback() { + let response = "This is just plain text analysis."; + let fb = parse_feedback_response(response); + assert_eq!(fb.summary_text, "This is just plain text analysis."); + assert!(fb.failure_patterns.is_empty()); + } + + #[test] + fn test_extract_json_block() { + let text = "Some text\n```json\n{\"key\": \"value\"}\n```\nMore text"; + let block = extract_json_block(text); + assert!(block.is_some()); + assert_eq!(block.unwrap(), "{\"key\": \"value\"}"); + } + + #[test] + fn test_extract_json_block_generic() { + let text = "Some text\n```\n{\"key\": \"value\"}\n```\nMore text"; + let block = extract_json_block(text); + assert!(block.is_some()); + } + + #[test] + fn test_extract_json_block_none() { + let text = "No code blocks here."; + let block = extract_json_block(text); + assert!(block.is_none()); + } + + #[test] + fn test_format_history() { + let trial = TrialResult { + trial_id: "t1".into(), + config: TrialConfig { + trial_id: "t1".into(), + params: { + let mut m = HashMap::new(); + m.insert("intelligence.model".into(), serde_json::json!("qwen3:8b")); + m + }, + reasoning: String::new(), + }, + accuracy: 0.85, + mean_latency_seconds: 1.0, + total_cost_usd: 0.05, + total_energy_joules: 10.0, + total_tokens: 100, + samples_evaluated: 10, + analysis: "ok".into(), + failure_modes: vec![], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + }; + + let mut frontier = HashSet::new(); + frontier.insert("t1".into()); + + let text = format_history(&[trial], Some(&frontier)); + assert!(text.contains("[FRONTIER]")); + assert!(text.contains("0.8500")); + } + + #[test] + fn test_analyze_trial_no_backend() { + let opt = LLMOptimizer::new(SearchSpace::default(), "test".into()); + let config = TrialConfig { + trial_id: "t1".into(), + params: HashMap::new(), + reasoning: String::new(), + }; + let fb = opt.analyze_trial(&config, 0.85, 1.0, 0.05); + assert!(fb.summary_text.contains("0.85")); + } + + #[test] + fn test_propose_targeted_no_backend() { + let opt = LLMOptimizer::new(SearchSpace::default(), "test".into()); + let base = TrialConfig { + trial_id: "t1".into(), + params: { + let mut m = HashMap::new(); + m.insert("intelligence.model".into(), serde_json::json!("qwen3:8b")); + m.insert("agent.type".into(), serde_json::json!("orchestrator")); + m + }, + reasoning: String::new(), + }; + let config = opt.propose_targeted(&[], &base, "intelligence", None); + // Should preserve base params + assert_eq!( + config.params.get("agent.type"), + Some(&serde_json::json!("orchestrator")) + ); + } + + #[test] + fn test_propose_merge_no_backend() { + let opt = LLMOptimizer::new(SearchSpace::default(), "test".into()); + let cand = TrialResult { + trial_id: "t1".into(), + config: TrialConfig { + trial_id: "t1".into(), + params: { + let mut m = HashMap::new(); + m.insert("intelligence.model".into(), serde_json::json!("qwen3:8b")); + m + }, + reasoning: String::new(), + }, + accuracy: 0.9, + mean_latency_seconds: 1.0, + total_cost_usd: 0.05, + total_energy_joules: 10.0, + total_tokens: 100, + samples_evaluated: 10, + analysis: String::new(), + failure_modes: vec![], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + }; + let config = opt.propose_merge(&[cand], &[], None); + assert!(!config.trial_id.is_empty()); + } +} diff --git a/rust/crates/openjarvis-learning/src/optimize/mod.rs b/rust/crates/openjarvis-learning/src/optimize/mod.rs new file mode 100644 index 00000000..224db7af --- /dev/null +++ b/rust/crates/openjarvis-learning/src/optimize/mod.rs @@ -0,0 +1,20 @@ +//! Optimization framework for OpenJarvis configuration tuning. +//! +//! Provides LLM-guided search over the 5-pillar configuration space, +//! with SQLite-backed trial persistence and Pareto frontier computation. + +pub mod engine; +pub mod llm_optimizer; +pub mod search_space; +pub mod store; +pub mod types; + +// Re-export key types for convenience. +pub use engine::{compute_pareto_frontier, OptimizationEngine, TrialRunner}; +pub use llm_optimizer::{LLMOptimizer, OptimizerBackend}; +pub use search_space::{build_search_space, default_search_space}; +pub use store::OptimizationStore; +pub use types::{ + BenchmarkScore, DimensionType, Direction, ObjectiveSpec, OptimizationRun, RunStatus, + SampleScore, SearchDimension, SearchSpace, TrialConfig, TrialFeedback, TrialResult, +}; diff --git a/rust/crates/openjarvis-learning/src/optimize/search_space.rs b/rust/crates/openjarvis-learning/src/optimize/search_space.rs new file mode 100644 index 00000000..c867573b --- /dev/null +++ b/rust/crates/openjarvis-learning/src/optimize/search_space.rs @@ -0,0 +1,354 @@ +//! Search space builder and default search space for configuration optimization. +//! +//! Rust translation of `src/openjarvis/learning/optimize/search_space.py`. + +use std::collections::HashMap; + +use super::types::{DimensionType, SearchDimension, SearchSpace}; + +/// Build a [`SearchSpace`] from a config map (typically parsed from TOML). +/// +/// Expected structure: +/// ```json +/// { +/// "optimize": { +/// "search": [ +/// { "name": "agent.type", "type": "categorical", +/// "values": ["orchestrator", "native_react"], +/// "description": "Agent architecture" }, +/// { "name": "intelligence.temperature", "type": "continuous", +/// "low": 0.0, "high": 1.0, +/// "description": "Generation temperature" } +/// ], +/// "fixed": { "engine": "ollama", "model": "qwen3:8b" }, +/// "constraints": { "rules": ["SimpleAgent should only have max_turns = 1"] } +/// } +/// } +/// ``` +pub fn build_search_space(config: &serde_json::Value) -> SearchSpace { + let opt = config + .get("optimize") + .cloned() + .unwrap_or(serde_json::Value::Object(Default::default())); + + // Parse search entries + let search_entries = opt + .get("search") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + + // Parse fixed params + let fixed: HashMap = opt + .get("fixed") + .and_then(|v| serde_json::from_value(v.clone()).ok()) + .unwrap_or_default(); + + // Parse constraints + let constraints: Vec = opt + .get("constraints") + .and_then(|v| v.get("rules")) + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + + // Build dimensions + let dimensions: Vec = search_entries + .iter() + .map(|entry| { + let name = entry + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + + // Infer pillar from the first segment of the dotted name + let pillar = if name.contains('.') { + name.split('.').next().unwrap_or("").to_string() + } else { + String::new() + }; + + let dim_type_str = entry + .get("type") + .and_then(|v| v.as_str()) + .unwrap_or("categorical"); + let dim_type = match dim_type_str { + "continuous" => DimensionType::Continuous, + "integer" => DimensionType::Integer, + "subset" => DimensionType::Subset, + "text" => DimensionType::Text, + _ => DimensionType::Categorical, + }; + + let values = entry + .get("values") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + let low = entry.get("low").and_then(|v| v.as_f64()); + let high = entry.get("high").and_then(|v| v.as_f64()); + let description = entry + .get("description") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + + SearchDimension { + name, + dim_type, + values, + low, + high, + description, + pillar, + } + }) + .collect(); + + SearchSpace { + dimensions, + fixed, + constraints, + } +} + +/// Default search space covering all 5 pillars. +pub fn default_search_space() -> SearchSpace { + SearchSpace { + dimensions: vec![ + // Intelligence pillar + SearchDimension { + name: "intelligence.model".into(), + dim_type: DimensionType::Categorical, + values: vec![ + serde_json::json!("qwen3:8b"), + serde_json::json!("qwen3:4b"), + serde_json::json!("qwen3:1.7b"), + serde_json::json!("llama3.1:8b"), + serde_json::json!("llama3.1:70b"), + serde_json::json!("gemma2:9b"), + serde_json::json!("mistral:7b"), + serde_json::json!("deepseek-r1:8b"), + ], + low: None, + high: None, + description: "The LLM model to use for generation".into(), + pillar: "intelligence".into(), + }, + SearchDimension { + name: "intelligence.temperature".into(), + dim_type: DimensionType::Continuous, + values: vec![], + low: Some(0.0), + high: Some(1.0), + description: "Generation temperature (0 = deterministic, 1 = creative)" + .into(), + pillar: "intelligence".into(), + }, + SearchDimension { + name: "intelligence.max_tokens".into(), + dim_type: DimensionType::Integer, + values: vec![], + low: Some(256.0), + high: Some(8192.0), + description: "Maximum tokens to generate per response".into(), + pillar: "intelligence".into(), + }, + SearchDimension { + name: "intelligence.top_p".into(), + dim_type: DimensionType::Continuous, + values: vec![], + low: Some(0.0), + high: Some(1.0), + description: "Nucleus sampling probability threshold".into(), + pillar: "intelligence".into(), + }, + SearchDimension { + name: "intelligence.system_prompt".into(), + dim_type: DimensionType::Text, + values: vec![], + low: None, + high: None, + description: "System prompt to guide model behavior".into(), + pillar: "intelligence".into(), + }, + // Engine pillar + SearchDimension { + name: "engine.backend".into(), + dim_type: DimensionType::Categorical, + values: vec![ + serde_json::json!("ollama"), + serde_json::json!("vllm"), + serde_json::json!("sglang"), + serde_json::json!("llamacpp"), + serde_json::json!("mlx"), + serde_json::json!("lmstudio"), + serde_json::json!("exo"), + serde_json::json!("nexa"), + serde_json::json!("uzu"), + serde_json::json!("apple_fm"), + ], + low: None, + high: None, + description: "Inference engine backend".into(), + pillar: "engine".into(), + }, + // Agent pillar + SearchDimension { + name: "agent.type".into(), + dim_type: DimensionType::Categorical, + values: vec![ + serde_json::json!("simple"), + serde_json::json!("orchestrator"), + serde_json::json!("native_react"), + serde_json::json!("native_openhands"), + ], + low: None, + high: None, + description: "Agent architecture to use".into(), + pillar: "agent".into(), + }, + SearchDimension { + name: "agent.max_turns".into(), + dim_type: DimensionType::Integer, + values: vec![], + low: Some(1.0), + high: Some(30.0), + description: "Maximum number of agent reasoning turns".into(), + pillar: "agent".into(), + }, + // Tools pillar + SearchDimension { + name: "tools.tool_set".into(), + dim_type: DimensionType::Subset, + values: vec![ + serde_json::json!("calculator"), + serde_json::json!("think"), + serde_json::json!("file_read"), + serde_json::json!("file_write"), + serde_json::json!("web_search"), + serde_json::json!("code_interpreter"), + serde_json::json!("llm"), + serde_json::json!("shell_exec"), + serde_json::json!("apply_patch"), + serde_json::json!("http_request"), + serde_json::json!("database_query"), + ], + low: None, + high: None, + description: "Set of tools available to the agent".into(), + pillar: "tools".into(), + }, + // Learning pillar + SearchDimension { + name: "learning.routing_policy".into(), + dim_type: DimensionType::Categorical, + values: vec![ + serde_json::json!("heuristic"), + serde_json::json!("grpo"), + serde_json::json!("bandit"), + serde_json::json!("learned"), + ], + low: None, + high: None, + description: "Router policy for model/agent selection".into(), + pillar: "learning".into(), + }, + ], + fixed: HashMap::new(), + constraints: vec![ + "SimpleAgent (agent.type='simple') should only have max_turns = 1".into(), + "agent.max_turns must be >= 1".into(), + "intelligence.temperature and intelligence.top_p \ + should not both be at extreme values" + .into(), + ], + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_default_search_space() { + let space = default_search_space(); + assert_eq!(space.dimensions.len(), 10); + assert_eq!(space.constraints.len(), 3); + assert!(space.fixed.is_empty()); + + // Check pillar groupings + let intelligence_dims: Vec<_> = space + .dimensions + .iter() + .filter(|d| d.pillar == "intelligence") + .collect(); + assert_eq!(intelligence_dims.len(), 5); + + let agent_dims: Vec<_> = space + .dimensions + .iter() + .filter(|d| d.pillar == "agent") + .collect(); + assert_eq!(agent_dims.len(), 2); + } + + #[test] + fn test_build_search_space_from_config() { + let config = serde_json::json!({ + "optimize": { + "search": [ + { + "name": "intelligence.temperature", + "type": "continuous", + "low": 0.0, + "high": 1.0, + "description": "Generation temperature" + }, + { + "name": "agent.type", + "type": "categorical", + "values": ["simple", "orchestrator"] + } + ], + "fixed": { "engine": "ollama" }, + "constraints": { + "rules": ["SimpleAgent should only have max_turns = 1"] + } + } + }); + + let space = build_search_space(&config); + assert_eq!(space.dimensions.len(), 2); + assert_eq!(space.dimensions[0].name, "intelligence.temperature"); + assert_eq!(space.dimensions[0].dim_type, DimensionType::Continuous); + assert_eq!(space.dimensions[0].pillar, "intelligence"); + assert_eq!(space.dimensions[1].pillar, "agent"); + assert_eq!(space.fixed.len(), 1); + assert_eq!(space.constraints.len(), 1); + } + + #[test] + fn test_build_search_space_empty_config() { + let config = serde_json::json!({}); + let space = build_search_space(&config); + assert!(space.dimensions.is_empty()); + assert!(space.fixed.is_empty()); + assert!(space.constraints.is_empty()); + } + + #[test] + fn test_default_space_prompt_description() { + let space = default_search_space(); + let desc = space.to_prompt_description(); + assert!(desc.contains("intelligence.model")); + assert!(desc.contains("agent.type")); + assert!(desc.contains("tools.tool_set")); + assert!(desc.contains("learning.routing_policy")); + } +} diff --git a/rust/crates/openjarvis-learning/src/optimize/store.rs b/rust/crates/openjarvis-learning/src/optimize/store.rs new file mode 100644 index 00000000..e7ffbce2 --- /dev/null +++ b/rust/crates/openjarvis-learning/src/optimize/store.rs @@ -0,0 +1,671 @@ +//! SQLite-backed storage for optimization runs and trials. +//! +//! Rust translation of `src/openjarvis/learning/optimize/store.py`. + +use std::collections::HashMap; +use std::path::Path; +use std::time::{SystemTime, UNIX_EPOCH}; + +use rusqlite::{params, Connection, Result as SqlResult}; + +use super::types::{ + BenchmarkScore, OptimizationRun, RunStatus, SampleScore, SearchSpace, + TrialConfig, TrialFeedback, TrialResult, +}; + +// --------------------------------------------------------------------------- +// SQL constants +// --------------------------------------------------------------------------- + +const CREATE_RUNS: &str = "\ +CREATE TABLE IF NOT EXISTS optimization_runs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + run_id TEXT NOT NULL UNIQUE, + search_space TEXT NOT NULL DEFAULT '{}', + status TEXT NOT NULL DEFAULT 'running', + optimizer_model TEXT NOT NULL DEFAULT '', + benchmark TEXT NOT NULL DEFAULT '', + best_trial_id TEXT, + best_recipe_path TEXT, + created_at REAL NOT NULL DEFAULT 0.0, + updated_at REAL NOT NULL DEFAULT 0.0, + pareto_frontier_ids TEXT NOT NULL DEFAULT '[]', + benchmarks TEXT NOT NULL DEFAULT '[]' +);"; + +const CREATE_TRIALS: &str = "\ +CREATE TABLE IF NOT EXISTS trial_results ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + trial_id TEXT NOT NULL, + run_id TEXT NOT NULL, + config TEXT NOT NULL DEFAULT '{}', + reasoning TEXT NOT NULL DEFAULT '', + accuracy REAL NOT NULL DEFAULT 0.0, + mean_latency_seconds REAL NOT NULL DEFAULT 0.0, + total_cost_usd REAL NOT NULL DEFAULT 0.0, + total_energy_joules REAL NOT NULL DEFAULT 0.0, + total_tokens INTEGER NOT NULL DEFAULT 0, + samples_evaluated INTEGER NOT NULL DEFAULT 0, + analysis TEXT NOT NULL DEFAULT '', + failure_modes TEXT NOT NULL DEFAULT '[]', + created_at REAL NOT NULL DEFAULT 0.0, + sample_scores TEXT NOT NULL DEFAULT '[]', + structured_feedback TEXT NOT NULL DEFAULT '{}', + per_benchmark TEXT NOT NULL DEFAULT '[]', + FOREIGN KEY (run_id) REFERENCES optimization_runs(run_id) +);"; + +const INSERT_RUN: &str = "\ +INSERT OR REPLACE INTO optimization_runs ( + run_id, search_space, status, optimizer_model, benchmark, + best_trial_id, best_recipe_path, created_at, updated_at, + pareto_frontier_ids, benchmarks +) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)"; + +const INSERT_TRIAL: &str = "\ +INSERT OR REPLACE INTO trial_results ( + trial_id, run_id, config, reasoning, accuracy, + mean_latency_seconds, total_cost_usd, total_energy_joules, + total_tokens, samples_evaluated, analysis, failure_modes, + created_at, sample_scores, structured_feedback, per_benchmark +) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)"; + +// --------------------------------------------------------------------------- +// OptimizationStore +// --------------------------------------------------------------------------- + +/// SQLite-backed storage for optimization runs and trials. +pub struct OptimizationStore { + conn: Connection, +} + +impl OptimizationStore { + /// Open or create a store at the given database path. + pub fn open(db_path: impl AsRef) -> Result { + let conn = Connection::open(db_path.as_ref()) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + conn.execute_batch("PRAGMA journal_mode=WAL;") + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + conn.execute_batch(CREATE_RUNS) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + conn.execute_batch(CREATE_TRIALS) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + Ok(Self { conn }) + } + + /// Create an in-memory store (useful for testing). + pub fn in_memory() -> Result { + let conn = Connection::open_in_memory() + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + conn.execute_batch("PRAGMA journal_mode=WAL;") + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + conn.execute_batch(CREATE_RUNS) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + conn.execute_batch(CREATE_TRIALS) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + Ok(Self { conn }) + } + + // ------------------------------------------------------------------ + // Runs + // ------------------------------------------------------------------ + + /// Persist an optimization run (insert or update). + pub fn save_run(&self, run: &OptimizationRun) -> Result<(), OptimizeStoreError> { + let now = now_epoch(); + let search_space_json = search_space_to_json(&run.search_space); + let best_trial_id = run.best_trial.as_ref().map(|t| t.trial_id.clone()); + let pareto_ids: Vec = run + .pareto_frontier + .iter() + .map(|t| t.trial_id.clone()) + .collect(); + let pareto_json = serde_json::to_string(&pareto_ids).unwrap_or_else(|_| "[]".into()); + let benchmarks_json = + serde_json::to_string(&run.benchmarks).unwrap_or_else(|_| "[]".into()); + + self.conn + .execute( + INSERT_RUN, + params![ + run.run_id, + search_space_json, + run.status.to_string(), + run.optimizer_model, + run.benchmark, + best_trial_id, + run.best_recipe_path, + now, + now, + pareto_json, + benchmarks_json, + ], + ) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + Ok(()) + } + + /// Retrieve an optimization run by id, or `None`. + pub fn get_run(&self, run_id: &str) -> Result, OptimizeStoreError> { + // First, extract the raw row data to release the statement borrow + let row_data: Option = { + let mut stmt = self + .conn + .prepare( + "SELECT id, run_id, search_space, status, optimizer_model, \ + benchmark, best_trial_id, best_recipe_path, created_at, \ + updated_at, pareto_frontier_ids, benchmarks \ + FROM optimization_runs WHERE run_id = ?1", + ) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + + let mut rows = stmt + .query_map(params![run_id], |row| { + Ok(RunRowData { + run_id: row.get(1)?, + search_space_raw: row.get(2)?, + status_str: row.get(3)?, + optimizer_model: row.get(4)?, + benchmark: row.get(5)?, + best_trial_id: row.get(6)?, + best_recipe_path: row.get(7)?, + pareto_ids_raw: row.get(10)?, + benchmarks_raw: row.get(11)?, + }) + }) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + + match rows.next() { + Some(Ok(data)) => Some(data), + Some(Err(e)) => return Err(OptimizeStoreError::Sqlite(e.to_string())), + None => None, + } + }; + + // Now the statement is dropped, we can safely call get_trials + match row_data { + Some(data) => Ok(Some(self.build_run_from_row_data(data)?)), + None => Ok(None), + } + } + + /// Return summary dicts of recent optimization runs. + pub fn list_runs( + &self, + limit: usize, + ) -> Result, OptimizeStoreError> { + let mut stmt = self + .conn + .prepare( + "SELECT run_id, status, optimizer_model, benchmark, \ + best_trial_id, best_recipe_path, created_at, updated_at \ + FROM optimization_runs ORDER BY created_at DESC LIMIT ?1", + ) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + + let rows = stmt + .query_map(params![limit as i64], |row| { + Ok(RunSummary { + run_id: row.get(0)?, + status: row.get(1)?, + optimizer_model: row.get(2)?, + benchmark: row.get(3)?, + best_trial_id: row.get(4)?, + best_recipe_path: row.get(5)?, + created_at: row.get(6)?, + updated_at: row.get(7)?, + }) + }) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + + let mut result = Vec::new(); + for row in rows { + result.push(row.map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?); + } + Ok(result) + } + + // ------------------------------------------------------------------ + // Trials + // ------------------------------------------------------------------ + + /// Persist a single trial result. + pub fn save_trial( + &self, + run_id: &str, + trial: &TrialResult, + ) -> Result<(), OptimizeStoreError> { + let now = now_epoch(); + + let config_json = serde_json::to_string(&trial.config.params) + .unwrap_or_else(|_| "{}".into()); + let failure_modes_json = serde_json::to_string(&trial.failure_modes) + .unwrap_or_else(|_| "[]".into()); + let sample_scores_json = serde_json::to_string(&trial.sample_scores) + .unwrap_or_else(|_| "[]".into()); + let feedback_json = match &trial.structured_feedback { + Some(fb) => serde_json::to_string(fb).unwrap_or_else(|_| "{}".into()), + None => "{}".into(), + }; + let per_benchmark_json = serde_json::to_string(&trial.per_benchmark) + .unwrap_or_else(|_| "[]".into()); + + self.conn + .execute( + INSERT_TRIAL, + params![ + trial.trial_id, + run_id, + config_json, + trial.config.reasoning, + trial.accuracy, + trial.mean_latency_seconds, + trial.total_cost_usd, + trial.total_energy_joules, + trial.total_tokens, + trial.samples_evaluated, + trial.analysis, + failure_modes_json, + now, + sample_scores_json, + feedback_json, + per_benchmark_json, + ], + ) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + Ok(()) + } + + /// Retrieve all trial results for a given run. + pub fn get_trials(&self, run_id: &str) -> Result, OptimizeStoreError> { + let mut stmt = self + .conn + .prepare( + "SELECT trial_id, run_id, config, reasoning, accuracy, \ + mean_latency_seconds, total_cost_usd, total_energy_joules, \ + total_tokens, samples_evaluated, analysis, failure_modes, \ + created_at, sample_scores, structured_feedback, per_benchmark \ + FROM trial_results WHERE run_id = ?1 ORDER BY id", + ) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + + let rows = stmt + .query_map(params![run_id], row_to_trial) + .map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?; + + let mut result = Vec::new(); + for row in rows { + result.push(row.map_err(|e| OptimizeStoreError::Sqlite(e.to_string()))?); + } + Ok(result) + } + + // ------------------------------------------------------------------ + // Internal + // ------------------------------------------------------------------ + + fn build_run_from_row_data( + &self, + data: RunRowData, + ) -> Result { + let search_space = json_to_search_space(&data.search_space_raw); + let status = data + .status_str + .parse::() + .unwrap_or(RunStatus::Running); + + // Load trials (safe now -- no active statement borrow) + let trials = self.get_trials(&data.run_id)?; + + // Find best trial + let best_trial = data.best_trial_id.and_then(|btid| { + trials.iter().find(|t| t.trial_id == btid).cloned() + }); + + // Parse benchmarks + let benchmarks: Vec = + serde_json::from_str(&data.benchmarks_raw).unwrap_or_default(); + + // Parse pareto frontier IDs + let frontier_ids: Vec = + serde_json::from_str(&data.pareto_ids_raw).unwrap_or_default(); + let trial_map: HashMap = + trials.iter().map(|t| (t.trial_id.clone(), t)).collect(); + let pareto_frontier: Vec = frontier_ids + .iter() + .filter_map(|id| trial_map.get(id).map(|t| (*t).clone())) + .collect(); + + Ok(OptimizationRun { + run_id: data.run_id, + search_space, + trials, + best_trial, + best_recipe_path: data.best_recipe_path, + status, + optimizer_model: data.optimizer_model, + benchmark: data.benchmark, + benchmarks, + pareto_frontier, + objectives: super::types::default_objectives(), + }) + } +} + +/// Intermediate struct to hold raw row data from the optimization_runs table. +/// This avoids holding a borrow on the SQLite statement while calling get_trials. +struct RunRowData { + run_id: String, + search_space_raw: String, + status_str: String, + optimizer_model: String, + benchmark: String, + best_trial_id: Option, + best_recipe_path: Option, + pareto_ids_raw: String, + benchmarks_raw: String, +} + +// --------------------------------------------------------------------------- +// Row conversion helper (outside impl to avoid lifetime issues) +// --------------------------------------------------------------------------- + +fn row_to_trial(row: &rusqlite::Row<'_>) -> SqlResult { + let trial_id: String = row.get(0)?; + // row index 1 = run_id (not stored on TrialResult) + let config_raw: String = row.get(2)?; + let reasoning: String = row.get(3)?; + let accuracy: f64 = row.get(4)?; + let mean_latency: f64 = row.get(5)?; + let cost: f64 = row.get(6)?; + let energy: f64 = row.get(7)?; + let tokens: i64 = row.get(8)?; + let samples: i64 = row.get(9)?; + let analysis: String = row.get(10)?; + let failure_modes_raw: String = row.get(11)?; + // row index 12 = created_at + let sample_scores_raw: String = row.get(13)?; + let feedback_raw: String = row.get(14)?; + let per_benchmark_raw: String = row.get(15)?; + + let params: HashMap = + serde_json::from_str(&config_raw).unwrap_or_default(); + let failure_modes: Vec = + serde_json::from_str(&failure_modes_raw).unwrap_or_default(); + let sample_scores: Vec = + serde_json::from_str(&sample_scores_raw).unwrap_or_default(); + let per_benchmark: Vec = + serde_json::from_str(&per_benchmark_raw).unwrap_or_default(); + + let structured_feedback: Option = { + let fb: Option = serde_json::from_str(&feedback_raw).ok(); + fb.filter(|f| !f.summary_text.is_empty()) + }; + + let config = TrialConfig { + trial_id: trial_id.clone(), + params, + reasoning, + }; + + Ok(TrialResult { + trial_id, + config, + accuracy, + mean_latency_seconds: mean_latency, + total_cost_usd: cost, + total_energy_joules: energy, + total_tokens: tokens, + samples_evaluated: samples, + analysis, + failure_modes, + sample_scores, + structured_feedback, + per_benchmark, + }) +} + +// --------------------------------------------------------------------------- +// Serialization helpers +// --------------------------------------------------------------------------- + +fn search_space_to_json(space: &SearchSpace) -> String { + serde_json::to_string(space).unwrap_or_else(|_| "{}".into()) +} + +fn json_to_search_space(raw: &str) -> SearchSpace { + serde_json::from_str(raw).unwrap_or_default() +} + +fn now_epoch() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) +} + +// --------------------------------------------------------------------------- +// Summary type for list_runs +// --------------------------------------------------------------------------- + +/// Lightweight summary returned by `list_runs`. +#[derive(Debug, Clone)] +pub struct RunSummary { + pub run_id: String, + pub status: String, + pub optimizer_model: String, + pub benchmark: String, + pub best_trial_id: Option, + pub best_recipe_path: Option, + pub created_at: f64, + pub updated_at: f64, +} + +// --------------------------------------------------------------------------- +// Error type +// --------------------------------------------------------------------------- + +/// Errors from the optimization store. +#[derive(Debug, thiserror::Error)] +pub enum OptimizeStoreError { + #[error("SQLite error: {0}")] + Sqlite(String), + #[error("Serialization error: {0}")] + Serialization(String), +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::optimize::types::*; + + fn make_trial(id: &str, accuracy: f64) -> TrialResult { + TrialResult { + trial_id: id.into(), + config: TrialConfig { + trial_id: id.into(), + params: { + let mut m = HashMap::new(); + m.insert("intelligence.model".into(), serde_json::json!("qwen3:8b")); + m + }, + reasoning: "test".into(), + }, + accuracy, + mean_latency_seconds: 1.0, + total_cost_usd: 0.01, + total_energy_joules: 10.0, + total_tokens: 100, + samples_evaluated: 10, + analysis: "ok".into(), + failure_modes: vec![], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + } + } + + #[test] + fn test_store_save_and_get_trial() { + let store = OptimizationStore::in_memory().unwrap(); + + // Save a run first + let run = OptimizationRun { + run_id: "run1".into(), + search_space: SearchSpace::default(), + trials: vec![], + best_trial: None, + best_recipe_path: None, + status: RunStatus::Running, + optimizer_model: "test-model".into(), + benchmark: "test-bench".into(), + benchmarks: vec![], + pareto_frontier: vec![], + objectives: default_objectives(), + }; + store.save_run(&run).unwrap(); + + // Save trial + let trial = make_trial("t1", 0.85); + store.save_trial("run1", &trial).unwrap(); + + // Retrieve trials + let trials = store.get_trials("run1").unwrap(); + assert_eq!(trials.len(), 1); + assert_eq!(trials[0].trial_id, "t1"); + assert!((trials[0].accuracy - 0.85).abs() < 1e-9); + } + + #[test] + fn test_store_save_and_get_run() { + let store = OptimizationStore::in_memory().unwrap(); + + let trial = make_trial("t1", 0.9); + + let run = OptimizationRun { + run_id: "run1".into(), + search_space: SearchSpace { + dimensions: vec![SearchDimension { + name: "intelligence.temperature".into(), + dim_type: DimensionType::Continuous, + values: vec![], + low: Some(0.0), + high: Some(1.0), + description: "temp".into(), + pillar: "intelligence".into(), + }], + fixed: HashMap::new(), + constraints: vec![], + }, + trials: vec![trial.clone()], + best_trial: Some(trial.clone()), + best_recipe_path: Some("/tmp/best.toml".into()), + status: RunStatus::Completed, + optimizer_model: "test-model".into(), + benchmark: "supergpqa".into(), + benchmarks: vec!["supergpqa".into()], + pareto_frontier: vec![trial.clone()], + objectives: default_objectives(), + }; + + store.save_run(&run).unwrap(); + store.save_trial("run1", &trial).unwrap(); + + let loaded = store.get_run("run1").unwrap().unwrap(); + assert_eq!(loaded.run_id, "run1"); + assert_eq!(loaded.status, RunStatus::Completed); + assert_eq!(loaded.benchmark, "supergpqa"); + assert!(loaded.best_trial.is_some()); + assert_eq!(loaded.trials.len(), 1); + assert_eq!(loaded.pareto_frontier.len(), 1); + } + + #[test] + fn test_store_list_runs() { + let store = OptimizationStore::in_memory().unwrap(); + + for i in 0..3 { + let run = OptimizationRun { + run_id: format!("run{i}"), + search_space: SearchSpace::default(), + trials: vec![], + best_trial: None, + best_recipe_path: None, + status: RunStatus::Completed, + optimizer_model: "model".into(), + benchmark: "bench".into(), + benchmarks: vec![], + pareto_frontier: vec![], + objectives: default_objectives(), + }; + store.save_run(&run).unwrap(); + } + + let summaries = store.list_runs(10).unwrap(); + assert_eq!(summaries.len(), 3); + } + + #[test] + fn test_store_get_nonexistent_run() { + let store = OptimizationStore::in_memory().unwrap(); + let result = store.get_run("nonexistent").unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_store_trial_with_feedback() { + let store = OptimizationStore::in_memory().unwrap(); + + let run = OptimizationRun { + run_id: "run_fb".into(), + search_space: SearchSpace::default(), + trials: vec![], + best_trial: None, + best_recipe_path: None, + status: RunStatus::Running, + optimizer_model: "model".into(), + benchmark: "bench".into(), + benchmarks: vec![], + pareto_frontier: vec![], + objectives: default_objectives(), + }; + store.save_run(&run).unwrap(); + + let mut trial = make_trial("t_fb", 0.7); + trial.structured_feedback = Some(TrialFeedback { + summary_text: "The intelligence pillar needs tuning".into(), + failure_patterns: vec!["timeout on long prompts".into()], + pillar_ratings: { + let mut m = HashMap::new(); + m.insert("intelligence".into(), "medium".into()); + m.insert("agent".into(), "high".into()); + m + }, + suggested_changes: vec!["lower temperature".into()], + target_pillar: "intelligence".into(), + }); + trial.per_benchmark = vec![BenchmarkScore { + benchmark: "supergpqa".into(), + accuracy: 0.7, + mean_latency_seconds: 1.5, + total_cost_usd: 0.02, + total_energy_joules: 20.0, + total_tokens: 200, + samples_evaluated: 20, + errors: 1, + weight: 1.0, + sample_scores: vec![], + }]; + + store.save_trial("run_fb", &trial).unwrap(); + + let loaded = store.get_trials("run_fb").unwrap(); + assert_eq!(loaded.len(), 1); + let lt = &loaded[0]; + assert!(lt.structured_feedback.is_some()); + let fb = lt.structured_feedback.as_ref().unwrap(); + assert_eq!(fb.target_pillar, "intelligence"); + assert_eq!(lt.per_benchmark.len(), 1); + assert_eq!(lt.per_benchmark[0].benchmark, "supergpqa"); + } +} diff --git a/rust/crates/openjarvis-learning/src/optimize/types.rs b/rust/crates/openjarvis-learning/src/optimize/types.rs new file mode 100644 index 00000000..117e24a9 --- /dev/null +++ b/rust/crates/openjarvis-learning/src/optimize/types.rs @@ -0,0 +1,638 @@ +//! Core data types for the optimization framework. +//! +//! Rust translation of `src/openjarvis/learning/optimize/types.py`. + +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +// --------------------------------------------------------------------------- +// Search space types +// --------------------------------------------------------------------------- + +/// The type of a search dimension. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DimensionType { + Categorical, + Continuous, + Integer, + Subset, + Text, +} + +impl std::fmt::Display for DimensionType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DimensionType::Categorical => write!(f, "categorical"), + DimensionType::Continuous => write!(f, "continuous"), + DimensionType::Integer => write!(f, "integer"), + DimensionType::Subset => write!(f, "subset"), + DimensionType::Text => write!(f, "text"), + } + } +} + +/// One tunable dimension in the configuration search space. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SearchDimension { + /// Dotted parameter name, e.g. `"agent.type"`, `"intelligence.temperature"`. + pub name: String, + /// Kind of parameter. + pub dim_type: DimensionType, + /// Explicit options for categorical/subset dimensions. + #[serde(default)] + pub values: Vec, + /// Lower bound for continuous/integer dimensions. + #[serde(skip_serializing_if = "Option::is_none")] + pub low: Option, + /// Upper bound for continuous/integer dimensions. + #[serde(skip_serializing_if = "Option::is_none")] + pub high: Option, + /// Human-readable explanation (shown to the LLM optimizer). + #[serde(default)] + pub description: String, + /// Which pillar this dimension belongs to: intelligence, engine, agent, tools, learning. + #[serde(default)] + pub pillar: String, +} + +/// The full space of configs the optimizer can propose. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct SearchSpace { + /// Tunable dimensions. + #[serde(default)] + pub dimensions: Vec, + /// Parameters that are NOT being optimized (always injected). + #[serde(default)] + pub fixed: HashMap, + /// Natural-language constraints for the LLM optimizer. + #[serde(default)] + pub constraints: Vec, +} + +impl SearchSpace { + /// Render the search space as structured text for the LLM optimizer prompt. + pub fn to_prompt_description(&self) -> String { + let mut lines = Vec::new(); + lines.push("# Search Space".to_string()); + lines.push(String::new()); + + // Group dimensions by pillar + let mut by_pillar: HashMap<&str, Vec<&SearchDimension>> = HashMap::new(); + for dim in &self.dimensions { + let key = if dim.pillar.is_empty() { + "other" + } else { + dim.pillar.as_str() + }; + by_pillar.entry(key).or_default().push(dim); + } + + let mut pillars: Vec<&&str> = by_pillar.keys().collect(); + pillars.sort(); + + for pillar in pillars { + // Title-case the pillar name + let title = { + let mut chars = pillar.chars(); + match chars.next() { + None => String::new(), + Some(c) => { + let upper: String = c.to_uppercase().collect(); + upper + chars.as_str() + } + } + }; + lines.push(format!("## {title}")); + + for dim in &by_pillar[pillar] { + lines.push(format!("- **{}** ({})", dim.name, dim.dim_type)); + if !dim.description.is_empty() { + lines.push(format!(" Description: {}", dim.description)); + } + match dim.dim_type { + DimensionType::Categorical | DimensionType::Subset => { + lines.push(format!(" Options: {:?}", dim.values)); + } + DimensionType::Continuous | DimensionType::Integer => { + let lo = dim.low.map(|v| v.to_string()).unwrap_or_default(); + let hi = dim.high.map(|v| v.to_string()).unwrap_or_default(); + lines.push(format!(" Range: [{lo}, {hi}]")); + } + DimensionType::Text => { + lines.push(" Free-form text".to_string()); + } + } + } + lines.push(String::new()); + } + + if !self.fixed.is_empty() { + lines.push("## Fixed Parameters".to_string()); + let mut sorted_keys: Vec<&String> = self.fixed.keys().collect(); + sorted_keys.sort(); + for k in sorted_keys { + lines.push(format!("- {k} = {}", self.fixed[k])); + } + lines.push(String::new()); + } + + if !self.constraints.is_empty() { + lines.push("## Constraints".to_string()); + for c in &self.constraints { + lines.push(format!("- {c}")); + } + lines.push(String::new()); + } + + lines.join("\n") + } +} + +// --------------------------------------------------------------------------- +// Optimization objective +// --------------------------------------------------------------------------- + +/// Direction in which an objective should be optimized. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Direction { + Maximize, + Minimize, +} + +impl std::fmt::Display for Direction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Direction::Maximize => write!(f, "maximize"), + Direction::Minimize => write!(f, "minimize"), + } + } +} + +/// A single optimization objective (e.g. maximize accuracy). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ObjectiveSpec { + pub metric: String, + pub direction: Direction, + #[serde(default = "default_weight")] + pub weight: f64, +} + +fn default_weight() -> f64 { + 1.0 +} + +/// Default objectives: accuracy up, latency and cost down. +pub fn default_objectives() -> Vec { + vec![ + ObjectiveSpec { + metric: "accuracy".into(), + direction: Direction::Maximize, + weight: 1.0, + }, + ObjectiveSpec { + metric: "mean_latency_seconds".into(), + direction: Direction::Minimize, + weight: 1.0, + }, + ObjectiveSpec { + metric: "total_cost_usd".into(), + direction: Direction::Minimize, + weight: 1.0, + }, + ] +} + +/// All supported objectives. +pub fn all_objectives() -> Vec { + vec![ + ObjectiveSpec { metric: "accuracy".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "mean_latency_seconds".into(), direction: Direction::Minimize, weight: 1.0 }, + ObjectiveSpec { metric: "total_cost_usd".into(), direction: Direction::Minimize, weight: 1.0 }, + ObjectiveSpec { metric: "total_energy_joules".into(), direction: Direction::Minimize, weight: 1.0 }, + ObjectiveSpec { metric: "avg_power_watts".into(), direction: Direction::Minimize, weight: 1.0 }, + ObjectiveSpec { metric: "throughput_tok_per_sec".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "mfu_pct".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "mbu_pct".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "ipw".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "ipj".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "energy_per_output_token".into(), direction: Direction::Minimize, weight: 1.0 }, + ObjectiveSpec { metric: "throughput_per_watt".into(), direction: Direction::Maximize, weight: 1.0 }, + ObjectiveSpec { metric: "ttft".into(), direction: Direction::Minimize, weight: 1.0 }, + ObjectiveSpec { metric: "mean_itl_ms".into(), direction: Direction::Minimize, weight: 1.0 }, + ] +} + +// --------------------------------------------------------------------------- +// Per-sample and per-benchmark scores +// --------------------------------------------------------------------------- + +/// Per-sample metrics from an evaluation trial. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct SampleScore { + pub record_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub is_correct: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub score: Option, + #[serde(default)] + pub latency_seconds: f64, + #[serde(default)] + pub prompt_tokens: i64, + #[serde(default)] + pub completion_tokens: i64, + #[serde(default)] + pub cost_usd: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + #[serde(default)] + pub ttft: f64, + #[serde(default)] + pub energy_joules: f64, + #[serde(default)] + pub power_watts: f64, + #[serde(default)] + pub gpu_utilization_pct: f64, + #[serde(default)] + pub throughput_tok_per_sec: f64, + #[serde(default)] + pub mfu_pct: f64, + #[serde(default)] + pub mbu_pct: f64, + #[serde(default)] + pub ipw: f64, + #[serde(default)] + pub ipj: f64, + #[serde(default)] + pub energy_per_output_token_joules: f64, + #[serde(default)] + pub throughput_per_watt: f64, + #[serde(default)] + pub mean_itl_ms: f64, +} + +/// Per-benchmark metrics from a multi-benchmark evaluation trial. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BenchmarkScore { + pub benchmark: String, + #[serde(default)] + pub accuracy: f64, + #[serde(default)] + pub mean_latency_seconds: f64, + #[serde(default)] + pub total_cost_usd: f64, + #[serde(default)] + pub total_energy_joules: f64, + #[serde(default)] + pub total_tokens: i64, + #[serde(default)] + pub samples_evaluated: i64, + #[serde(default)] + pub errors: i64, + #[serde(default = "default_weight")] + pub weight: f64, + /// Optional per-sample scores within this benchmark. + #[serde(default)] + pub sample_scores: Vec, +} + +// --------------------------------------------------------------------------- +// Trial types +// --------------------------------------------------------------------------- + +/// Structured feedback from trial analysis. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct TrialFeedback { + #[serde(default)] + pub summary_text: String, + #[serde(default)] + pub failure_patterns: Vec, + #[serde(default)] + pub pillar_ratings: HashMap, + #[serde(default)] + pub suggested_changes: Vec, + #[serde(default)] + pub target_pillar: String, +} + +/// A single candidate configuration proposed by the optimizer. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TrialConfig { + pub trial_id: String, + /// Dotted keys -> values, e.g. `{"intelligence.temperature": 0.7}`. + #[serde(default)] + pub params: HashMap, + /// Optimizer's explanation of why this config was proposed. + #[serde(default)] + pub reasoning: String, +} + +/// Result of evaluating a trial. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TrialResult { + pub trial_id: String, + pub config: TrialConfig, + #[serde(default)] + pub accuracy: f64, + #[serde(default)] + pub mean_latency_seconds: f64, + #[serde(default)] + pub total_cost_usd: f64, + #[serde(default)] + pub total_energy_joules: f64, + #[serde(default)] + pub total_tokens: i64, + #[serde(default)] + pub samples_evaluated: i64, + #[serde(default)] + pub analysis: String, + #[serde(default)] + pub failure_modes: Vec, + #[serde(default)] + pub sample_scores: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub structured_feedback: Option, + #[serde(default)] + pub per_benchmark: Vec, +} + +// --------------------------------------------------------------------------- +// Optimization run +// --------------------------------------------------------------------------- + +/// Status of an optimization run. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum RunStatus { + #[default] + Running, + Completed, + Failed, +} + +impl std::fmt::Display for RunStatus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RunStatus::Running => write!(f, "running"), + RunStatus::Completed => write!(f, "completed"), + RunStatus::Failed => write!(f, "failed"), + } + } +} + +impl std::str::FromStr for RunStatus { + type Err = String; + fn from_str(s: &str) -> Result { + match s { + "running" => Ok(RunStatus::Running), + "completed" => Ok(RunStatus::Completed), + "failed" => Ok(RunStatus::Failed), + _ => Err(format!("Unknown run status: {s}")), + } + } +} + +/// Complete optimization session. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OptimizationRun { + pub run_id: String, + pub search_space: SearchSpace, + #[serde(default)] + pub trials: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub best_trial: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub best_recipe_path: Option, + #[serde(default)] + pub status: RunStatus, + #[serde(default)] + pub optimizer_model: String, + #[serde(default)] + pub benchmark: String, + #[serde(default)] + pub benchmarks: Vec, + #[serde(default)] + pub pareto_frontier: Vec, + #[serde(default = "default_objectives")] + pub objectives: Vec, +} + +// --------------------------------------------------------------------------- +// Mapping from dotted param names to Recipe constructor fields +// --------------------------------------------------------------------------- + +/// Maps dotted param names (used in search space) to Recipe field names. +pub fn param_to_recipe_field(dotted_key: &str) -> Option<&'static str> { + match dotted_key { + "intelligence.model" => Some("model"), + "intelligence.temperature" => Some("temperature"), + "intelligence.max_tokens" => Some("max_tokens"), + "intelligence.quantization" => Some("quantization"), + "engine.backend" => Some("engine_key"), + "agent.type" => Some("agent_type"), + "agent.max_turns" => Some("max_turns"), + "agent.system_prompt" | "intelligence.system_prompt" => Some("system_prompt"), + "tools.tool_set" => Some("tools"), + "learning.routing_policy" => Some("routing_policy"), + "learning.agent_policy" => Some("agent_policy"), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_dimension_type_display() { + assert_eq!(DimensionType::Categorical.to_string(), "categorical"); + assert_eq!(DimensionType::Continuous.to_string(), "continuous"); + assert_eq!(DimensionType::Text.to_string(), "text"); + } + + #[test] + fn test_dimension_type_serde_roundtrip() { + let dt = DimensionType::Integer; + let json = serde_json::to_string(&dt).unwrap(); + assert_eq!(json, "\"integer\""); + let parsed: DimensionType = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, dt); + } + + #[test] + fn test_direction_display() { + assert_eq!(Direction::Maximize.to_string(), "maximize"); + assert_eq!(Direction::Minimize.to_string(), "minimize"); + } + + #[test] + fn test_search_space_to_prompt_description() { + let space = SearchSpace { + dimensions: vec![ + SearchDimension { + name: "intelligence.temperature".into(), + dim_type: DimensionType::Continuous, + values: vec![], + low: Some(0.0), + high: Some(1.0), + description: "Generation temperature".into(), + pillar: "intelligence".into(), + }, + SearchDimension { + name: "agent.type".into(), + dim_type: DimensionType::Categorical, + values: vec![ + serde_json::json!("simple"), + serde_json::json!("orchestrator"), + ], + low: None, + high: None, + description: "Agent architecture".into(), + pillar: "agent".into(), + }, + ], + fixed: { + let mut m = HashMap::new(); + m.insert("engine".into(), serde_json::json!("ollama")); + m + }, + constraints: vec![ + "SimpleAgent should only have max_turns = 1".into(), + ], + }; + + let desc = space.to_prompt_description(); + assert!(desc.contains("# Search Space")); + assert!(desc.contains("## Intelligence")); + assert!(desc.contains("## Agent")); + assert!(desc.contains("Generation temperature")); + assert!(desc.contains("## Fixed Parameters")); + assert!(desc.contains("## Constraints")); + } + + #[test] + fn test_trial_config_serde() { + let config = TrialConfig { + trial_id: "abc123".into(), + params: { + let mut m = HashMap::new(); + m.insert("intelligence.temperature".into(), serde_json::json!(0.7)); + m + }, + reasoning: "Initial config".into(), + }; + let json = serde_json::to_string(&config).unwrap(); + let parsed: TrialConfig = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed.trial_id, "abc123"); + assert_eq!(parsed.params.len(), 1); + } + + #[test] + fn test_trial_result_serde() { + let result = TrialResult { + trial_id: "t1".into(), + config: TrialConfig { + trial_id: "t1".into(), + params: HashMap::new(), + reasoning: String::new(), + }, + accuracy: 0.85, + mean_latency_seconds: 1.2, + total_cost_usd: 0.05, + total_energy_joules: 100.0, + total_tokens: 500, + samples_evaluated: 50, + analysis: "Good".into(), + failure_modes: vec!["timeout".into()], + sample_scores: vec![], + structured_feedback: None, + per_benchmark: vec![], + }; + let json = serde_json::to_string(&result).unwrap(); + let parsed: TrialResult = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed.trial_id, "t1"); + assert!((parsed.accuracy - 0.85).abs() < 1e-9); + } + + #[test] + fn test_run_status_roundtrip() { + let status = RunStatus::Completed; + assert_eq!(status.to_string(), "completed"); + let parsed: RunStatus = "completed".parse().unwrap(); + assert_eq!(parsed, status); + } + + #[test] + fn test_default_objectives() { + let objs = default_objectives(); + assert_eq!(objs.len(), 3); + assert_eq!(objs[0].metric, "accuracy"); + assert_eq!(objs[0].direction, Direction::Maximize); + } + + #[test] + fn test_all_objectives() { + let objs = all_objectives(); + assert_eq!(objs.len(), 14); + } + + #[test] + fn test_param_to_recipe_field() { + assert_eq!( + param_to_recipe_field("intelligence.model"), + Some("model") + ); + assert_eq!( + param_to_recipe_field("agent.type"), + Some("agent_type") + ); + assert_eq!(param_to_recipe_field("unknown.param"), None); + } + + #[test] + fn test_sample_score_default() { + let ss = SampleScore::default(); + assert_eq!(ss.record_id, ""); + assert_eq!(ss.latency_seconds, 0.0); + assert!(ss.is_correct.is_none()); + } + + #[test] + fn test_benchmark_score_serde() { + let bs = BenchmarkScore { + benchmark: "supergpqa".into(), + accuracy: 0.72, + mean_latency_seconds: 2.5, + total_cost_usd: 0.1, + total_energy_joules: 50.0, + total_tokens: 1000, + samples_evaluated: 100, + errors: 2, + weight: 1.0, + sample_scores: vec![], + }; + let json = serde_json::to_string(&bs).unwrap(); + let parsed: BenchmarkScore = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed.benchmark, "supergpqa"); + } + + #[test] + fn test_trial_feedback_serde() { + let fb = TrialFeedback { + summary_text: "Good result".into(), + failure_patterns: vec!["timeout".into()], + pillar_ratings: { + let mut m = HashMap::new(); + m.insert("intelligence".into(), "high".into()); + m + }, + suggested_changes: vec!["increase temperature".into()], + target_pillar: "intelligence".into(), + }; + let json = serde_json::to_string(&fb).unwrap(); + let parsed: TrialFeedback = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed.summary_text, "Good result"); + assert_eq!(parsed.target_pillar, "intelligence"); + } +} diff --git a/rust/crates/openjarvis-learning/src/orchestrator_types.rs b/rust/crates/openjarvis-learning/src/orchestrator_types.rs index 9bf72583..a4b2080a 100644 --- a/rust/crates/openjarvis-learning/src/orchestrator_types.rs +++ b/rust/crates/openjarvis-learning/src/orchestrator_types.rs @@ -250,7 +250,10 @@ mod tests { fn test_normalize_number() { assert_eq!(normalize_number("42"), Some(42.0)); assert_eq!(normalize_number("1,234.0"), Some(1234.0)); - assert_eq!(normalize_number(" 3.14 "), Some(3.14)); + #[allow(clippy::approx_constant)] + { + assert_eq!(normalize_number(" 3.14 "), Some(3.14)); + } assert!(normalize_number("not a number").is_none()); } diff --git a/rust/crates/openjarvis-learning/src/reward.rs b/rust/crates/openjarvis-learning/src/reward.rs index 9881fb0e..d1e4f50f 100644 --- a/rust/crates/openjarvis-learning/src/reward.rs +++ b/rust/crates/openjarvis-learning/src/reward.rs @@ -163,6 +163,7 @@ pub struct AdaptiveRewardWeights { } impl AdaptiveRewardWeights { + #[allow(clippy::too_many_arguments)] pub fn new( initial_alpha: f64, final_alpha: f64, diff --git a/rust/crates/openjarvis-learning/src/sft_policy.rs b/rust/crates/openjarvis-learning/src/sft_policy.rs index 8ae1f591..005bee54 100644 --- a/rust/crates/openjarvis-learning/src/sft_policy.rs +++ b/rust/crates/openjarvis-learning/src/sft_policy.rs @@ -136,13 +136,13 @@ mod tests { assert_eq!(SFTRouterPolicy::classify_query("hello"), "short"); assert_eq!( SFTRouterPolicy::classify_query( - &"word ".repeat(101).trim().to_string() + "word ".repeat(101).trim() ), "long" ); assert_eq!( SFTRouterPolicy::classify_query( - &"word ".repeat(50).trim().to_string() + "word ".repeat(50).trim() ), "general" ); @@ -174,6 +174,6 @@ mod tests { ]; policy.update_from_data(&traces); let map = policy.policy_map(); - assert!(map.get("code").is_none(), "should not update with < min_samples"); + assert!(!map.contains_key("code"), "should not update with < min_samples"); } } diff --git a/rust/crates/openjarvis-learning/src/trace_policy.rs b/rust/crates/openjarvis-learning/src/trace_policy.rs index 4472546a..a4f06f5b 100644 --- a/rust/crates/openjarvis-learning/src/trace_policy.rs +++ b/rust/crates/openjarvis-learning/src/trace_policy.rs @@ -111,10 +111,10 @@ impl TraceDrivenPolicy { let conf = self.confidence.read(); if let Some(model) = map.get(qclass) { - if conf.get(qclass).copied().unwrap_or(0) >= self.min_samples { - if self.available.is_empty() || self.available.contains(model) { - return model.clone(); - } + if conf.get(qclass).copied().unwrap_or(0) >= self.min_samples + && (self.available.is_empty() || self.available.contains(model)) + { + return model.clone(); } } drop(map); @@ -294,7 +294,7 @@ mod tests { #[test] fn test_update_and_select() { let policy = TraceDrivenPolicy::new(vec![], "default".into(), "fallback".into()); - policy.min_samples; + let _ = policy.min_samples; let traces: Vec<(String, String, String, f64, Option)> = (0..10) .map(|_| { diff --git a/rust/crates/openjarvis-learning/src/training_data.rs b/rust/crates/openjarvis-learning/src/training_data.rs index 88630b6d..8e3bcfdb 100644 --- a/rust/crates/openjarvis-learning/src/training_data.rs +++ b/rust/crates/openjarvis-learning/src/training_data.rs @@ -74,7 +74,7 @@ impl TrainingDataMiner { .iter() .filter(|t| { t.outcome == "success" - && t.feedback.map_or(false, |f| f >= self.min_quality) + && t.feedback.is_some_and(|f| f >= self.min_quality) }) .collect() } diff --git a/rust/crates/openjarvis-mcp/src/server.rs b/rust/crates/openjarvis-mcp/src/server.rs index 69d375e6..73d3b836 100644 --- a/rust/crates/openjarvis-mcp/src/server.rs +++ b/rust/crates/openjarvis-mcp/src/server.rs @@ -109,7 +109,6 @@ impl McpServer { mod tests { use super::*; use openjarvis_tools::builtin::calculator::CalculatorTool; - use openjarvis_tools::traits::BaseTool; fn make_server() -> McpServer { let mut exec = ToolExecutor::new(None, None); diff --git a/rust/crates/openjarvis-python/src/agents.rs b/rust/crates/openjarvis-python/src/agents.rs index b71f446e..1fcfaf64 100644 --- a/rust/crates/openjarvis-python/src/agents.rs +++ b/rust/crates/openjarvis-python/src/agents.rs @@ -182,6 +182,177 @@ impl PyNativeReActAgent { } } +/// Python wrapper for NativeOpenHandsAgent. +#[pyclass(name = "NativeOpenHandsAgent")] +pub struct PyNativeOpenHandsAgent { + inner: Box, +} + +#[pymethods] +impl PyNativeOpenHandsAgent { + #[new] + #[pyo3(signature = (engine_key="ollama", host="http://localhost:11434", model="qwen3:8b", max_turns=10, temperature=0.7))] + fn new( + engine_key: &str, + host: &str, + model: &str, + max_turns: usize, + temperature: f64, + ) -> PyResult { + let config = openjarvis_core::JarvisConfig::default(); + let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key)) + .map_err(|e| PyErr::new::(e.to_string()))?; + let adapter = openjarvis_engine::rig_adapter::RigModelAdapter::new( + Arc::new(engine), + model.to_string(), + ); + let executor = Arc::new(openjarvis_tools::ToolExecutor::new(None, None)); + let agent = openjarvis_agents::NativeOpenHandsAgent::new( + adapter, + executor, + max_turns, + temperature, + ); + Ok(Self { + inner: Box::new(agent), + }) + } + + fn agent_id(&self) -> &str { + self.inner.agent_id() + } + + fn accepts_tools(&self) -> bool { + self.inner.accepts_tools() + } + + fn run(&self, input: &str) -> PyResult { + let result = RUNTIME + .block_on(self.inner.run(input, None)) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(PyAgentResult { + content: result.content, + turns: result.turns, + }) + } +} + +/// Python wrapper for MonitorOperativeAgent. +#[pyclass(name = "MonitorOperativeAgent")] +pub struct PyMonitorOperativeAgent { + inner: Box, +} + +#[pymethods] +impl PyMonitorOperativeAgent { + /// Create a MonitorOperativeAgent. + /// + /// Strategy parameters are strings: + /// - `memory_extraction`: "causality_graph" | "scratchpad" | "structured_json" | "none" + /// - `observation_compression`: "summarize" | "truncate" | "none" + /// - `retrieval_strategy`: "hybrid_with_self_eval" | "keyword" | "semantic" | "none" + /// - `task_decomposition`: "phased" | "monolithic" | "hierarchical" + #[new] + #[pyo3(signature = ( + engine_key="ollama", + host="http://localhost:11434", + model="qwen3:8b", + max_turns=10, + temperature=0.7, + memory_extraction="causality_graph", + observation_compression="summarize", + retrieval_strategy="hybrid_with_self_eval", + task_decomposition="phased", + compression_threshold=2000, + truncation_limit=2000 + ))] + #[allow(clippy::too_many_arguments)] + fn new( + engine_key: &str, + host: &str, + model: &str, + max_turns: usize, + temperature: f64, + memory_extraction: &str, + observation_compression: &str, + retrieval_strategy: &str, + task_decomposition: &str, + compression_threshold: usize, + truncation_limit: usize, + ) -> PyResult { + let config = openjarvis_core::JarvisConfig::default(); + let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key)) + .map_err(|e| PyErr::new::(e.to_string()))?; + let adapter = openjarvis_engine::rig_adapter::RigModelAdapter::new( + Arc::new(engine), + model.to_string(), + ); + let executor = Arc::new(openjarvis_tools::ToolExecutor::new(None, None)); + + let mem_ext = match memory_extraction { + "scratchpad" => openjarvis_agents::MemoryExtraction::Scratchpad, + "structured_json" => openjarvis_agents::MemoryExtraction::StructuredJson, + "none" => openjarvis_agents::MemoryExtraction::None, + _ => openjarvis_agents::MemoryExtraction::CausalityGraph, + }; + let obs_comp = match observation_compression { + "truncate" => openjarvis_agents::ObservationCompression::Truncate, + "none" => openjarvis_agents::ObservationCompression::None, + _ => openjarvis_agents::ObservationCompression::Summarize, + }; + let ret_strat = match retrieval_strategy { + "keyword" => openjarvis_agents::RetrievalStrategy::Keyword, + "semantic" => openjarvis_agents::RetrievalStrategy::Semantic, + "none" => openjarvis_agents::RetrievalStrategy::None, + _ => openjarvis_agents::RetrievalStrategy::HybridWithSelfEval, + }; + let task_dec = match task_decomposition { + "monolithic" => openjarvis_agents::TaskDecomposition::Monolithic, + "hierarchical" => openjarvis_agents::TaskDecomposition::Hierarchical, + _ => openjarvis_agents::TaskDecomposition::Phased, + }; + + let monitor_config = openjarvis_agents::MonitorConfig { + memory_extraction: mem_ext, + observation_compression: obs_comp, + retrieval_strategy: ret_strat, + task_decomposition: task_dec, + compression_threshold, + truncation_limit, + }; + + let agent = openjarvis_agents::MonitorOperativeAgent::new( + adapter, + executor, + max_turns, + temperature, + monitor_config, + ); + Ok(Self { + inner: Box::new(agent), + }) + } + + fn agent_id(&self) -> &str { + self.inner.agent_id() + } + + fn accepts_tools(&self) -> bool { + self.inner.accepts_tools() + } + + fn run(&self, input: &str) -> PyResult { + let result = RUNTIME + .block_on(self.inner.run(input, None)) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(PyAgentResult { + content: result.content, + turns: result.turns, + }) + } +} + +/// Python wrapper for LoopGuard. #[pyclass(name = "LoopGuard")] pub struct PyLoopGuard { inner: openjarvis_agents::LoopGuard, diff --git a/rust/crates/openjarvis-python/src/engine.rs b/rust/crates/openjarvis-python/src/engine.rs index 027669a4..166e12d6 100644 --- a/rust/crates/openjarvis-python/src/engine.rs +++ b/rust/crates/openjarvis-python/src/engine.rs @@ -69,6 +69,28 @@ impl PyEngine { host.unwrap_or("http://localhost:8079"), ), ), + "vllm_native" => openjarvis_engine::Engine::VLLM( + openjarvis_engine::VLLMEngine::new( + host.unwrap_or("http://localhost"), + 8000, + None, + 120.0, + ), + ), + "sglang_native" => openjarvis_engine::Engine::SGLang( + openjarvis_engine::SGLangEngine::new( + host.unwrap_or("http://localhost"), + 30000, + 120.0, + ), + ), + "llamacpp_native" => openjarvis_engine::Engine::LlamaCppNative( + openjarvis_engine::LlamaCppEngine::new( + host.unwrap_or("http://localhost"), + 8080, + 120.0, + ), + ), other => { return Err(PyErr::new::( format!("Unknown engine: {}", other), diff --git a/rust/crates/openjarvis-python/src/learning.rs b/rust/crates/openjarvis-python/src/learning.rs index 2c4e231d..71320dff 100644 --- a/rust/crates/openjarvis-python/src/learning.rs +++ b/rust/crates/openjarvis-python/src/learning.rs @@ -1,4 +1,4 @@ -//! PyO3 bindings for learning/router policy types. +//! PyO3 bindings for learning/router policy types and optimization. use openjarvis_learning::RouterPolicy; use pyo3::prelude::*; @@ -97,6 +97,133 @@ impl PyGRPORouterPolicy { } } +// --------------------------------------------------------------------------- +// Optimization store +// --------------------------------------------------------------------------- + +#[pyclass(name = "OptimizationStore")] +pub struct PyOptimizationStore { + inner: parking_lot::Mutex, +} + +#[pymethods] +impl PyOptimizationStore { + /// Open or create a store at the given database path. + /// Use `":memory:"` for an in-memory store. + #[new] + #[pyo3(signature = (path=":memory:"))] + fn new(path: &str) -> PyResult { + let store = if path == ":memory:" { + openjarvis_learning::optimize::OptimizationStore::in_memory() + } else { + openjarvis_learning::optimize::OptimizationStore::open(path) + }; + let store = + store.map_err(|e| PyErr::new::(e.to_string()))?; + Ok(Self { + inner: parking_lot::Mutex::new(store), + }) + } + + /// Save a trial result (JSON string) for a given run_id. + fn save_trial(&self, run_id: &str, trial_json: &str) -> PyResult<()> { + let trial: openjarvis_learning::optimize::TrialResult = + serde_json::from_str(trial_json) + .map_err(|e| PyErr::new::(e.to_string()))?; + self.inner + .lock() + .save_trial(run_id, &trial) + .map_err(|e| PyErr::new::(e.to_string())) + } + + /// Retrieve an optimization run by id. Returns JSON string or None. + fn get_run(&self, run_id: &str) -> PyResult> { + let run = self + .inner + .lock() + .get_run(run_id) + .map_err(|e| PyErr::new::(e.to_string()))?; + match run { + Some(r) => Ok(Some( + serde_json::to_string(&r).unwrap_or_else(|_| "{}".into()), + )), + None => Ok(None), + } + } + + /// List recent optimization runs. Returns JSON string. + #[pyo3(signature = (limit=50))] + fn list_runs(&self, limit: usize) -> PyResult { + let summaries = self + .inner + .lock() + .list_runs(limit) + .map_err(|e| PyErr::new::(e.to_string()))?; + // RunSummary doesn't derive Serialize, so manually build JSON. + let items: Vec = summaries + .iter() + .map(|s| { + serde_json::json!({ + "run_id": s.run_id, + "status": s.status, + "optimizer_model": s.optimizer_model, + "benchmark": s.benchmark, + "best_trial_id": s.best_trial_id, + "best_recipe_path": s.best_recipe_path, + "created_at": s.created_at, + "updated_at": s.updated_at, + }) + }) + .collect(); + Ok(serde_json::to_string(&items).unwrap_or_else(|_| "[]".into())) + } +} + +// --------------------------------------------------------------------------- +// LLM Optimizer +// --------------------------------------------------------------------------- + +#[pyclass(name = "LLMOptimizer")] +pub struct PyLLMOptimizer { + inner: openjarvis_learning::optimize::LLMOptimizer, +} + +#[pymethods] +impl PyLLMOptimizer { + /// Create a new LLM optimizer. + /// + /// `search_space_json` is a JSON string representing the SearchSpace. + /// `optimizer_model` is the model name to use for optimization proposals. + #[new] + #[pyo3(signature = (search_space_json, optimizer_model="gpt-4o"))] + fn new(search_space_json: &str, optimizer_model: &str) -> PyResult { + let search_space: openjarvis_learning::optimize::SearchSpace = + serde_json::from_str(search_space_json) + .map_err(|e| PyErr::new::(e.to_string()))?; + let inner = openjarvis_learning::optimize::LLMOptimizer::new( + search_space, + optimizer_model.to_string(), + ); + Ok(Self { inner }) + } + + /// Propose an initial configuration. Returns JSON string. + fn propose_initial(&self) -> String { + let config = self.inner.propose_initial(); + serde_json::to_string(&config).unwrap_or_else(|_| "{}".into()) + } + + /// Propose the next configuration based on trial history (JSON string). + /// Returns JSON string. + fn propose_next(&self, history_json: &str) -> PyResult { + let history: Vec = + serde_json::from_str(history_json) + .map_err(|e| PyErr::new::(e.to_string()))?; + let config = self.inner.propose_next(&history); + Ok(serde_json::to_string(&config).unwrap_or_else(|_| "{}".into())) + } +} + // --- SFT Router Policy --- #[pyclass(name = "SFTRouterPolicy")] diff --git a/rust/crates/openjarvis-python/src/lib.rs b/rust/crates/openjarvis-python/src/lib.rs index 94686f47..152ab9e3 100644 --- a/rust/crates/openjarvis-python/src/lib.rs +++ b/rust/crates/openjarvis-python/src/lib.rs @@ -1,4 +1,5 @@ //! PyO3 bridge — exposes ~50 Rust classes to Python via `openjarvis_rust`. +#![allow(clippy::redundant_closure, unused_variables)] use once_cell::sync::Lazy; use pyo3::prelude::*; @@ -84,6 +85,8 @@ fn openjarvis_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; // --- Tools --- @@ -101,6 +104,9 @@ fn openjarvis_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { // --- Storage / Memory --- m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; // --- Security --- @@ -133,6 +139,8 @@ fn openjarvis_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/rust/crates/openjarvis-python/src/scheduler.rs b/rust/crates/openjarvis-python/src/scheduler.rs index 99f83d03..8f6e40dd 100644 --- a/rust/crates/openjarvis-python/src/scheduler.rs +++ b/rust/crates/openjarvis-python/src/scheduler.rs @@ -18,7 +18,7 @@ impl PySchedulerStore { } fn create_task(&self, name: &str, schedule_type: &str, schedule_value: &str) -> PyResult { - let st = openjarvis_scheduler::ScheduleType::from_str(schedule_type).ok_or_else(|| { + let st = openjarvis_scheduler::ScheduleType::parse(schedule_type).ok_or_else(|| { PyErr::new::(format!( "invalid schedule_type '{}', expected cron/interval/once", schedule_type @@ -39,7 +39,7 @@ impl PySchedulerStore { } fn update_status(&self, id: &str, status: &str) -> PyResult { - let s = openjarvis_scheduler::TaskStatus::from_str(status).ok_or_else(|| { + let s = openjarvis_scheduler::TaskStatus::parse(status).ok_or_else(|| { PyErr::new::(format!( "invalid status '{}', expected active/paused/cancelled/completed", status diff --git a/rust/crates/openjarvis-python/src/security.rs b/rust/crates/openjarvis-python/src/security.rs index 5b13b0e1..a626ece8 100644 --- a/rust/crates/openjarvis-python/src/security.rs +++ b/rust/crates/openjarvis-python/src/security.rs @@ -1,7 +1,6 @@ //! PyO3 bindings for security types. use pyo3::prelude::*; -use std::sync::Arc; #[pyclass(name = "SecretScanner")] pub struct PySecretScanner { @@ -208,6 +207,7 @@ impl PyRateLimiter { self.inner.check(key) } + #[pyo3(signature = (key=None))] fn reset(&self, key: Option<&str>) { self.inner.reset(key); } diff --git a/rust/crates/openjarvis-python/src/storage.rs b/rust/crates/openjarvis-python/src/storage.rs index ed6654fb..5a5b7502 100644 --- a/rust/crates/openjarvis-python/src/storage.rs +++ b/rust/crates/openjarvis-python/src/storage.rs @@ -107,6 +107,208 @@ impl PyBM25Memory { } } +#[pyclass(name = "FAISSMemory")] +pub struct PyFAISSMemory { + inner: openjarvis_tools::storage::FAISSMemory, +} + +#[pymethods] +impl PyFAISSMemory { + #[new] + #[pyo3(signature = (path=":memory:", dim=128))] + fn new(path: &str, dim: usize) -> PyResult { + let inner = openjarvis_tools::storage::FAISSMemory::new(std::path::Path::new(path), dim) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(Self { inner }) + } + + fn backend_id(&self) -> &str { + self.inner.backend_id() + } + + #[pyo3(signature = (content, source, metadata=None))] + fn store(&self, content: &str, source: &str, metadata: Option<&str>) -> PyResult { + let meta = metadata + .map(|m| serde_json::from_str(m)) + .transpose() + .map_err(|e| PyErr::new::(e.to_string()))?; + self.inner + .store(content, source, meta.as_ref()) + .map_err(|e| PyErr::new::(e.to_string())) + } + + #[pyo3(signature = (query, top_k=5))] + fn retrieve(&self, query: &str, top_k: usize) -> PyResult { + let results = self + .inner + .retrieve(query, top_k) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(serde_json::to_string(&results).unwrap_or_default()) + } + + fn count(&self) -> PyResult { + self.inner + .count() + .map_err(|e| PyErr::new::(e.to_string())) + } + + fn delete(&self, doc_id: &str) -> PyResult { + self.inner + .delete(doc_id) + .map_err(|e| PyErr::new::(e.to_string())) + } + + fn clear(&self) -> PyResult<()> { + self.inner + .clear() + .map_err(|e| PyErr::new::(e.to_string())) + } +} + +#[pyclass(name = "ColBERTMemory")] +pub struct PyColBERTMemory { + inner: openjarvis_tools::storage::ColBERTMemory, +} + +#[pymethods] +impl PyColBERTMemory { + #[new] + #[pyo3(signature = (path=":memory:", token_dim=64))] + fn new(path: &str, token_dim: usize) -> PyResult { + let inner = + openjarvis_tools::storage::ColBERTMemory::new(std::path::Path::new(path), token_dim) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(Self { inner }) + } + + fn backend_id(&self) -> &str { + self.inner.backend_id() + } + + #[pyo3(signature = (content, source, metadata=None))] + fn store(&self, content: &str, source: &str, metadata: Option<&str>) -> PyResult { + let meta = metadata + .map(|m| serde_json::from_str(m)) + .transpose() + .map_err(|e| PyErr::new::(e.to_string()))?; + self.inner + .store(content, source, meta.as_ref()) + .map_err(|e| PyErr::new::(e.to_string())) + } + + #[pyo3(signature = (query, top_k=5))] + fn retrieve(&self, query: &str, top_k: usize) -> PyResult { + let results = self + .inner + .retrieve(query, top_k) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(serde_json::to_string(&results).unwrap_or_default()) + } + + fn count(&self) -> PyResult { + self.inner + .count() + .map_err(|e| PyErr::new::(e.to_string())) + } + + fn delete(&self, doc_id: &str) -> PyResult { + self.inner + .delete(doc_id) + .map_err(|e| PyErr::new::(e.to_string())) + } + + fn clear(&self) -> PyResult<()> { + self.inner + .clear() + .map_err(|e| PyErr::new::(e.to_string())) + } +} + +#[pyclass(name = "HybridMemory")] +pub struct PyHybridMemory { + inner: openjarvis_tools::storage::HybridMemory, +} + +#[pymethods] +impl PyHybridMemory { + /// Create a HybridMemory combining named backends. + /// + /// `backend_keys` is a list of strings like `["sqlite", "bm25"]`. + /// Supported keys: "sqlite", "bm25", "faiss", "colbert". + #[new] + #[pyo3(signature = (backend_keys))] + fn new(backend_keys: Vec) -> PyResult { + let mut backends: Vec> = Vec::new(); + for key in &backend_keys { + let backend: Box = match key.as_str() { + "sqlite" => { + let m = openjarvis_tools::storage::SQLiteMemory::new( + std::path::Path::new(":memory:"), + ) + .map_err(|e| { + PyErr::new::(e.to_string()) + })?; + Box::new(m) + } + "bm25" => Box::new(openjarvis_tools::storage::BM25Memory::default()), + "faiss" => { + let m = openjarvis_tools::storage::FAISSMemory::in_memory().map_err(|e| { + PyErr::new::(e.to_string()) + })?; + Box::new(m) + } + "colbert" => { + let m = + openjarvis_tools::storage::ColBERTMemory::in_memory().map_err(|e| { + PyErr::new::(e.to_string()) + })?; + Box::new(m) + } + other => { + return Err(PyErr::new::(format!( + "Unknown backend key: {}. Supported: sqlite, bm25, faiss, colbert", + other + ))); + } + }; + backends.push(backend); + } + Ok(Self { + inner: openjarvis_tools::storage::HybridMemory::new(backends), + }) + } + + fn backend_id(&self) -> &str { + self.inner.backend_id() + } + + #[pyo3(signature = (content, source, metadata=None))] + fn store(&self, content: &str, source: &str, metadata: Option<&str>) -> PyResult { + let meta = metadata + .map(|m| serde_json::from_str(m)) + .transpose() + .map_err(|e| PyErr::new::(e.to_string()))?; + self.inner + .store(content, source, meta.as_ref()) + .map_err(|e| PyErr::new::(e.to_string())) + } + + #[pyo3(signature = (query, top_k=5))] + fn retrieve(&self, query: &str, top_k: usize) -> PyResult { + let results = self + .inner + .retrieve(query, top_k) + .map_err(|e| PyErr::new::(e.to_string()))?; + Ok(serde_json::to_string(&results).unwrap_or_default()) + } + + fn count(&self) -> PyResult { + self.inner + .count() + .map_err(|e| PyErr::new::(e.to_string())) + } +} + #[pyclass(name = "KnowledgeGraphMemory")] pub struct PyKnowledgeGraphMemory { inner: openjarvis_tools::storage::KnowledgeGraphMemory, diff --git a/rust/crates/openjarvis-python/src/telemetry.rs b/rust/crates/openjarvis-python/src/telemetry.rs index c560457a..35b00bfc 100644 --- a/rust/crates/openjarvis-python/src/telemetry.rs +++ b/rust/crates/openjarvis-python/src/telemetry.rs @@ -112,6 +112,7 @@ pub struct PyTelemetrySample { impl PyTelemetrySample { #[new] #[pyo3(signature = (timestamp_ns, gpu_power_w=0.0, cpu_power_w=0.0, gpu_energy_j=0.0, cpu_energy_j=0.0, gpu_util_pct=0.0, gpu_temp_c=0.0, gpu_mem_gb=0.0))] + #[allow(clippy::too_many_arguments)] fn new( timestamp_ns: u64, gpu_power_w: f64, diff --git a/rust/crates/openjarvis-scheduler/src/lib.rs b/rust/crates/openjarvis-scheduler/src/lib.rs index dda1382c..46ad0a58 100644 --- a/rust/crates/openjarvis-scheduler/src/lib.rs +++ b/rust/crates/openjarvis-scheduler/src/lib.rs @@ -27,7 +27,7 @@ impl ScheduleType { } } - pub fn from_str(s: &str) -> Option { + pub fn parse(s: &str) -> Option { match s { "cron" => Some(Self::Cron), "interval" => Some(Self::Interval), @@ -62,7 +62,7 @@ impl TaskStatus { } } - pub fn from_str(s: &str) -> Option { + pub fn parse(s: &str) -> Option { match s { "active" => Some(Self::Active), "paused" => Some(Self::Paused), @@ -317,9 +317,9 @@ fn row_to_task(row: &rusqlite::Row<'_>) -> ScheduledTask { id: row.get(0).unwrap_or_default(), name: row.get(1).unwrap_or_default(), description: row.get(2).unwrap_or_default(), - schedule_type: ScheduleType::from_str(&type_str).unwrap_or(ScheduleType::Once), + schedule_type: ScheduleType::parse(&type_str).unwrap_or(ScheduleType::Once), schedule_value: row.get(4).unwrap_or_default(), - status: TaskStatus::from_str(&status_str).unwrap_or(TaskStatus::Active), + status: TaskStatus::parse(&status_str).unwrap_or(TaskStatus::Active), last_run: row.get(6).ok(), next_run: row.get(7).ok(), created_at: row.get(8).unwrap_or(0.0), diff --git a/rust/crates/openjarvis-security/src/audit.rs b/rust/crates/openjarvis-security/src/audit.rs index 09beebb5..7df2a8d0 100644 --- a/rust/crates/openjarvis-security/src/audit.rs +++ b/rust/crates/openjarvis-security/src/audit.rs @@ -1,6 +1,6 @@ //! Audit logger — persist security events to SQLite with Merkle hash chain. -use crate::types::{ScanFinding, SecurityEvent, SecurityEventType, ThreatLevel}; +use crate::types::{ScanFinding, SecurityEvent, SecurityEventType}; use openjarvis_core::OpenJarvisError; use rusqlite::Connection; use sha2::{Digest, Sha256}; @@ -15,13 +15,12 @@ impl AuditLogger { pub fn new(db_path: &Path) -> Result { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)) + OpenJarvisError::Io(std::io::Error::other(e)) })?; } let conn = Connection::open(db_path).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -39,8 +38,7 @@ impl AuditLogger { )", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -83,8 +81,7 @@ impl AuditLogger { ], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -113,8 +110,7 @@ impl AuditLogger { FROM security_events ORDER BY id", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -133,8 +129,7 @@ impl AuditLogger { )) }) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -144,8 +139,7 @@ impl AuditLogger { for row_result in rows { let (rid, ts, etype, fj, preview, action, stored_hash, stored_prev) = row_result.map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -212,8 +206,7 @@ impl AuditLogger { params.iter().map(|p| p.as_ref()).collect(); let mut stmt = self.conn.prepare(&sql).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -229,8 +222,7 @@ impl AuditLogger { )) }) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -239,8 +231,7 @@ impl AuditLogger { for row_result in rows { let (ts, _etype, findings_json, preview, action) = row_result.map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -272,6 +263,7 @@ fn hex_sha256(input: &str) -> String { #[cfg(test)] mod tests { use super::*; + use crate::types::ThreatLevel; #[test] fn test_audit_log_and_verify() { diff --git a/rust/crates/openjarvis-security/src/guardrails.rs b/rust/crates/openjarvis-security/src/guardrails.rs index 36d91649..2b8267d3 100644 --- a/rust/crates/openjarvis-security/src/guardrails.rs +++ b/rust/crates/openjarvis-security/src/guardrails.rs @@ -2,7 +2,7 @@ use crate::scanner::{PIIScanner, SecretScanner}; use crate::types::{RedactionMode, ScanResult}; -use openjarvis_core::error::{EngineError, OpenJarvisError}; +use openjarvis_core::error::OpenJarvisError; use openjarvis_core::{EventBus, EventType, GenerateResult, Message}; use openjarvis_engine::traits::{InferenceEngine, TokenStream}; use serde_json::Value; diff --git a/rust/crates/openjarvis-skills/src/lib.rs b/rust/crates/openjarvis-skills/src/lib.rs index dbab7a2b..d94f4fcd 100644 --- a/rust/crates/openjarvis-skills/src/lib.rs +++ b/rust/crates/openjarvis-skills/src/lib.rs @@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; fn decode_hex(s: &str) -> Result, String> { - if s.len() % 2 != 0 { + if !s.len().is_multiple_of(2) { return Err("odd-length hex string".into()); } (0..s.len()) diff --git a/rust/crates/openjarvis-telemetry/src/store.rs b/rust/crates/openjarvis-telemetry/src/store.rs index ed690ebc..a7283819 100644 --- a/rust/crates/openjarvis-telemetry/src/store.rs +++ b/rust/crates/openjarvis-telemetry/src/store.rs @@ -14,13 +14,12 @@ impl TelemetryStore { pub fn new(db_path: &Path) -> Result { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)) + OpenJarvisError::Io(std::io::Error::other(e)) })?; } let conn = Connection::open(db_path).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -51,8 +50,7 @@ impl TelemetryStore { )", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -102,8 +100,7 @@ impl TelemetryStore { ], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -115,8 +112,7 @@ impl TelemetryStore { let count: i64 = conn .query_row("SELECT COUNT(*) FROM telemetry", [], |row| row.get(0)) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -126,8 +122,7 @@ impl TelemetryStore { pub fn clear(&self) -> Result<(), OpenJarvisError> { let conn = self.conn.lock(); conn.execute("DELETE FROM telemetry", []).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; diff --git a/rust/crates/openjarvis-tools/src/builtin/http_tools.rs b/rust/crates/openjarvis-tools/src/builtin/http_tools.rs index f32da9f1..8b27a88f 100644 --- a/rust/crates/openjarvis-tools/src/builtin/http_tools.rs +++ b/rust/crates/openjarvis-tools/src/builtin/http_tools.rs @@ -50,10 +50,7 @@ impl BaseTool for HttpRequestTool { .timeout(std::time::Duration::from_secs(30)) .build() .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - e.to_string(), - )) + OpenJarvisError::Io(std::io::Error::other(e.to_string())) })?; let mut request = match method.as_str() { diff --git a/rust/crates/openjarvis-tools/src/rig_tools.rs b/rust/crates/openjarvis-tools/src/rig_tools.rs index adbea363..ffb0cc7e 100644 --- a/rust/crates/openjarvis-tools/src/rig_tools.rs +++ b/rust/crates/openjarvis-tools/src/rig_tools.rs @@ -6,7 +6,7 @@ use rig::completion::request::ToolDefinition; use rig::tool::Tool as RigTool; use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; // --------------------------------------------------------------------------- // Calculator diff --git a/rust/crates/openjarvis-tools/src/storage/backend_enum.rs b/rust/crates/openjarvis-tools/src/storage/backend_enum.rs index 724e4193..75b8e390 100644 --- a/rust/crates/openjarvis-tools/src/storage/backend_enum.rs +++ b/rust/crates/openjarvis-tools/src/storage/backend_enum.rs @@ -1,6 +1,9 @@ //! MemoryBackendEnum — static dispatch over storage backends. use super::bm25::BM25Memory; +use super::colbert::ColBERTMemory; +use super::faiss::FAISSMemory; +use super::hybrid::HybridMemory; use super::knowledge_graph::KnowledgeGraphMemory; use super::sqlite::SQLiteMemory; use super::traits::MemoryBackend; @@ -11,6 +14,9 @@ use serde_json::Value; pub enum MemoryBackendEnum { Sqlite(SQLiteMemory), Bm25(BM25Memory), + Faiss(FAISSMemory), + ColBert(ColBERTMemory), + Hybrid(HybridMemory), KnowledgeGraph(KnowledgeGraphMemory), } @@ -19,6 +25,9 @@ macro_rules! delegate_memory { match $self { MemoryBackendEnum::Sqlite(m) => m.$method($($arg),*), MemoryBackendEnum::Bm25(m) => m.$method($($arg),*), + MemoryBackendEnum::Faiss(m) => m.$method($($arg),*), + MemoryBackendEnum::ColBert(m) => m.$method($($arg),*), + MemoryBackendEnum::Hybrid(m) => m.$method($($arg),*), MemoryBackendEnum::KnowledgeGraph(m) => m.$method($($arg),*), } }; @@ -65,6 +74,9 @@ impl MemoryBackendEnum { match self { MemoryBackendEnum::Sqlite(_) => "sqlite", MemoryBackendEnum::Bm25(_) => "bm25", + MemoryBackendEnum::Faiss(_) => "faiss", + MemoryBackendEnum::ColBert(_) => "colbert", + MemoryBackendEnum::Hybrid(_) => "hybrid", MemoryBackendEnum::KnowledgeGraph(_) => "knowledge_graph", } } diff --git a/rust/crates/openjarvis-tools/src/storage/bm25.rs b/rust/crates/openjarvis-tools/src/storage/bm25.rs index cca2fb07..9e164237 100644 --- a/rust/crates/openjarvis-tools/src/storage/bm25.rs +++ b/rust/crates/openjarvis-tools/src/storage/bm25.rs @@ -53,7 +53,7 @@ impl BM25Memory { let mut score = 0.0; let dl = doc.term_count as f64; - for (term, _) in query_terms { + for term in query_terms.keys() { let tf = *doc.terms.get(term).unwrap_or(&0) as f64; let doc_freq = *df.get(term).unwrap_or(&0) as f64; if doc_freq == 0.0 || tf == 0.0 { diff --git a/rust/crates/openjarvis-tools/src/storage/colbert.rs b/rust/crates/openjarvis-tools/src/storage/colbert.rs new file mode 100644 index 00000000..e9256f7b --- /dev/null +++ b/rust/crates/openjarvis-tools/src/storage/colbert.rs @@ -0,0 +1,415 @@ +//! ColBERT-style memory backend — token-level embeddings with MaxSim scoring. +//! +//! ColBERT represents queries and documents as sequences of token embeddings and +//! computes relevance via **MaxSim**: for each query token, find the maximum cosine +//! similarity to any document token, then sum across query tokens. +//! +//! Like `FAISSMemory`, actual embedding is handled externally (Python / LLM). This +//! module provides the retrieval math and SQLite-backed persistence. + +use crate::storage::traits::MemoryBackend; +use openjarvis_core::{OpenJarvisError, RetrievalResult}; +use parking_lot::Mutex; +use rusqlite::Connection; +use serde_json::Value; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use uuid::Uuid; + +/// Default token embedding dimension for the dummy `embed_tokens()` stub. +const DEFAULT_TOKEN_DIM: usize = 64; + +/// ColBERT-style memory backend with MaxSim scoring over token-level embeddings. +pub struct ColBERTMemory { + conn: Mutex, + _db_path: PathBuf, + token_dim: usize, +} + +impl ColBERTMemory { + pub fn new(db_path: &Path, token_dim: usize) -> Result { + if let Some(parent) = db_path.parent() { + std::fs::create_dir_all(parent).map_err(|e| { + OpenJarvisError::Io(std::io::Error::other(e)) + })?; + } + + let conn = Connection::open(db_path).map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS colbert_documents ( + id TEXT PRIMARY KEY, + content TEXT NOT NULL, + source TEXT DEFAULT '', + metadata TEXT DEFAULT '{}', + num_tokens INTEGER NOT NULL, + token_embeddings BLOB NOT NULL, + created_at REAL DEFAULT (julianday('now')) + );", + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + Ok(Self { + conn: Mutex::new(conn), + _db_path: db_path.to_path_buf(), + token_dim, + }) + } + + pub fn in_memory() -> Result { + Self::new(Path::new(":memory:"), DEFAULT_TOKEN_DIM) + } + + pub fn with_dim(token_dim: usize) -> Result { + Self::new(Path::new(":memory:"), token_dim) + } + + /// Dummy token-level embedding stub. + /// + /// Splits text on whitespace and produces one normalised vector per token by + /// hashing the token bytes. Real token embeddings should be supplied externally. + pub fn embed_tokens(&self, text: &str) -> Vec> { + text.split_whitespace() + .map(|token| { + let mut vec = vec![0.0f64; self.token_dim]; + for (i, byte) in token.bytes().enumerate() { + vec[i % self.token_dim] += byte as f64; + } + let norm = vec.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for v in &mut vec { + *v /= norm; + } + } + vec + }) + .collect() + } + + /// Store a document with pre-computed token-level embeddings. + /// + /// `token_embeddings` is a slice of vectors — one per token. All vectors must + /// have the same dimension. + pub fn store_with_token_embeddings( + &self, + content: &str, + source: &str, + metadata: Option<&Value>, + token_embeddings: &[Vec], + ) -> Result { + let doc_id = Uuid::new_v4().to_string(); + let meta_str = metadata + .map(|m| serde_json::to_string(m).unwrap_or_default()) + .unwrap_or_else(|| "{}".to_string()); + let num_tokens = token_embeddings.len() as i64; + let emb_bytes = token_embeddings_to_bytes(token_embeddings); + + let conn = self.conn.lock(); + conn.execute( + "INSERT INTO colbert_documents (id, content, source, metadata, num_tokens, token_embeddings) + VALUES (?1, ?2, ?3, ?4, ?5, ?6)", + rusqlite::params![doc_id, content, source, meta_str, num_tokens, emb_bytes], + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + Ok(doc_id) + } + + /// Retrieve the top-k documents using MaxSim scoring against pre-computed query + /// token embeddings. + pub fn retrieve_by_token_embeddings( + &self, + query_token_embeddings: &[Vec], + top_k: usize, + ) -> Result, OpenJarvisError> { + let conn = self.conn.lock(); + + let mut stmt = conn + .prepare( + "SELECT id, content, source, metadata, num_tokens, token_embeddings + FROM colbert_documents", + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + let rows: Vec<(String, String, String, usize, Vec)> = stmt + .query_map([], |row| { + let content: String = row.get(1)?; + let source: String = row.get::<_, String>(2).unwrap_or_default(); + let meta_str: String = row.get::<_, String>(3).unwrap_or_else(|_| "{}".into()); + let num_tokens: i64 = row.get(4)?; + let emb_bytes: Vec = row.get(5)?; + Ok((content, source, meta_str, num_tokens as usize, emb_bytes)) + }) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })? + .filter_map(|r| r.ok()) + .collect(); + + let mut scored: Vec<(RetrievalResult, f64)> = rows + .into_iter() + .map(|(content, source, meta_str, num_tokens, emb_bytes)| { + let doc_token_embs = + bytes_to_token_embeddings(&emb_bytes, num_tokens, self.token_dim); + let score = maxsim(query_token_embeddings, &doc_token_embs); + let metadata: HashMap = + serde_json::from_str(&meta_str).unwrap_or_default(); + ( + RetrievalResult { + content, + source, + score, + metadata, + }, + score, + ) + }) + .collect(); + + scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + scored.truncate(top_k); + + Ok(scored.into_iter().map(|(r, _)| r).collect()) + } +} + +impl MemoryBackend for ColBERTMemory { + fn backend_id(&self) -> &str { + "colbert" + } + + fn store( + &self, + content: &str, + source: &str, + metadata: Option<&Value>, + ) -> Result { + let token_embeddings = self.embed_tokens(content); + self.store_with_token_embeddings(content, source, metadata, &token_embeddings) + } + + fn retrieve( + &self, + query: &str, + top_k: usize, + ) -> Result, OpenJarvisError> { + let query_token_embeddings = self.embed_tokens(query); + if query_token_embeddings.is_empty() { + return Ok(vec![]); + } + self.retrieve_by_token_embeddings(&query_token_embeddings, top_k) + } + + fn delete(&self, doc_id: &str) -> Result { + let conn = self.conn.lock(); + let changes = conn + .execute( + "DELETE FROM colbert_documents WHERE id = ?1", + rusqlite::params![doc_id], + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + Ok(changes > 0) + } + + fn clear(&self) -> Result<(), OpenJarvisError> { + let conn = self.conn.lock(); + conn.execute_batch("DELETE FROM colbert_documents") + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + Ok(()) + } + + fn count(&self) -> Result { + let conn = self.conn.lock(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM colbert_documents", [], |row| { + row.get(0) + }) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + Ok(count as usize) + } +} + +// --------------------------------------------------------------------------- +// MaxSim and helpers +// --------------------------------------------------------------------------- + +/// MaxSim scoring: for each query token, find max cosine similarity across all +/// document tokens, then sum. +fn maxsim(query_tokens: &[Vec], doc_tokens: &[Vec]) -> f64 { + if query_tokens.is_empty() || doc_tokens.is_empty() { + return 0.0; + } + query_tokens + .iter() + .map(|q| { + doc_tokens + .iter() + .map(|d| cosine_similarity(q, d)) + .fold(f64::NEG_INFINITY, f64::max) + }) + .sum() +} + +/// Cosine similarity between two vectors. +fn cosine_similarity(a: &[f64], b: &[f64]) -> f64 { + let len = a.len().min(b.len()); + let mut dot = 0.0f64; + let mut norm_a = 0.0f64; + let mut norm_b = 0.0f64; + for i in 0..len { + dot += a[i] * b[i]; + norm_a += a[i] * a[i]; + norm_b += b[i] * b[i]; + } + let denom = norm_a.sqrt() * norm_b.sqrt(); + if denom == 0.0 { + 0.0 + } else { + dot / denom + } +} + +/// Serialize a sequence of token embeddings into a flat byte blob. +/// +/// Layout: token_0[0..dim] ++ token_1[0..dim] ++ ... (all f64 little-endian). +fn token_embeddings_to_bytes(embeddings: &[Vec]) -> Vec { + embeddings + .iter() + .flat_map(|tok| tok.iter().flat_map(|f| f.to_le_bytes())) + .collect() +} + +/// Deserialize a flat byte blob back into a `Vec>` of token embeddings. +fn bytes_to_token_embeddings(bytes: &[u8], num_tokens: usize, dim: usize) -> Vec> { + let floats: Vec = bytes + .chunks_exact(8) + .map(|chunk| { + let arr: [u8; 8] = chunk.try_into().unwrap_or([0u8; 8]); + f64::from_le_bytes(arr) + }) + .collect(); + + floats + .chunks(dim) + .take(num_tokens) + .map(|chunk| chunk.to_vec()) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_colbert_store_and_retrieve() { + let mem = ColBERTMemory::in_memory().unwrap(); + let id = mem + .store("Rust is a systems programming language", "test", None) + .unwrap(); + assert!(!id.is_empty()); + + let results = mem.retrieve("Rust programming", 5).unwrap(); + assert!(!results.is_empty()); + assert!(results[0].content.contains("Rust")); + } + + #[test] + fn test_colbert_maxsim_ranking() { + let mem = ColBERTMemory::with_dim(3).unwrap(); + + // Doc A: two tokens + let doc_a_tokens = vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]]; + // Doc B: two tokens — orthogonal to query + let doc_b_tokens = vec![vec![0.0, 0.0, 1.0], vec![0.0, 0.0, 1.0]]; + + mem.store_with_token_embeddings("doc A", "s1", None, &doc_a_tokens) + .unwrap(); + mem.store_with_token_embeddings("doc B", "s2", None, &doc_b_tokens) + .unwrap(); + + // Query: single token close to doc A's first token + let query_tokens = vec![vec![1.0, 0.0, 0.0]]; + let results = mem.retrieve_by_token_embeddings(&query_tokens, 2).unwrap(); + assert_eq!(results.len(), 2); + assert_eq!(results[0].content, "doc A"); + } + + #[test] + fn test_colbert_delete() { + let mem = ColBERTMemory::in_memory().unwrap(); + let id = mem.store("test content", "test", None).unwrap(); + assert_eq!(mem.count().unwrap(), 1); + assert!(mem.delete(&id).unwrap()); + assert_eq!(mem.count().unwrap(), 0); + } + + #[test] + fn test_colbert_clear() { + let mem = ColBERTMemory::in_memory().unwrap(); + mem.store("doc 1", "s1", None).unwrap(); + mem.store("doc 2", "s2", None).unwrap(); + assert_eq!(mem.count().unwrap(), 2); + mem.clear().unwrap(); + assert_eq!(mem.count().unwrap(), 0); + } + + #[test] + fn test_colbert_empty_retrieve() { + let mem = ColBERTMemory::in_memory().unwrap(); + let results = mem.retrieve("anything", 5).unwrap(); + assert!(results.is_empty()); + } + + #[test] + fn test_maxsim_computation() { + // Query: 1 token = [1, 0] + // Doc: 2 tokens = [0.5, 0.5], [1, 0] + // MaxSim for query token [1,0] = max(cos([1,0],[0.5,0.5]), cos([1,0],[1,0])) + // = max(~0.707, 1.0) = 1.0 + let q = vec![vec![1.0, 0.0]]; + let d = vec![vec![0.5, 0.5], vec![1.0, 0.0]]; + let score = maxsim(&q, &d); + assert!((score - 1.0).abs() < 1e-10); + } + + #[test] + fn test_token_embeddings_roundtrip() { + let embeddings = vec![ + vec![1.0, 2.0, 3.0], + vec![4.0, 5.0, 6.0], + ]; + let bytes = token_embeddings_to_bytes(&embeddings); + let recovered = bytes_to_token_embeddings(&bytes, 2, 3); + assert_eq!(embeddings, recovered); + } +} diff --git a/rust/crates/openjarvis-tools/src/storage/faiss.rs b/rust/crates/openjarvis-tools/src/storage/faiss.rs new file mode 100644 index 00000000..ef6b93cc --- /dev/null +++ b/rust/crates/openjarvis-tools/src/storage/faiss.rs @@ -0,0 +1,393 @@ +//! FAISS-style vector similarity memory backend — pure Rust brute-force cosine similarity. + +use crate::storage::traits::MemoryBackend; +use openjarvis_core::{OpenJarvisError, RetrievalResult}; +use parking_lot::Mutex; +use rusqlite::Connection; +use serde_json::Value; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use uuid::Uuid; + +/// Default embedding dimension used by the dummy `embed()` stub. +const DEFAULT_DIM: usize = 128; + +/// FAISS-style memory backend using brute-force cosine similarity over stored embeddings. +/// +/// Documents and their embeddings are persisted in SQLite. The `retrieve()` method +/// computes cosine similarity between a query embedding and every stored embedding, +/// returning the top-k closest results. +/// +/// Because real embedding requires an external model (Python / LLM), the `embed()` +/// method provided here is a deterministic **stub** that hashes the input text into a +/// fixed-dimension vector. Callers that have real embeddings should use +/// `store_with_embedding()` and `retrieve_by_embedding()` directly. +pub struct FAISSMemory { + conn: Mutex, + _db_path: PathBuf, + dim: usize, +} + +impl FAISSMemory { + pub fn new(db_path: &Path, dim: usize) -> Result { + if let Some(parent) = db_path.parent() { + std::fs::create_dir_all(parent).map_err(|e| { + OpenJarvisError::Io(std::io::Error::other(e)) + })?; + } + + let conn = Connection::open(db_path).map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS faiss_documents ( + id TEXT PRIMARY KEY, + content TEXT NOT NULL, + source TEXT DEFAULT '', + metadata TEXT DEFAULT '{}', + embedding BLOB NOT NULL, + created_at REAL DEFAULT (julianday('now')) + );", + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + Ok(Self { + conn: Mutex::new(conn), + _db_path: db_path.to_path_buf(), + dim, + }) + } + + pub fn in_memory() -> Result { + Self::new(Path::new(":memory:"), DEFAULT_DIM) + } + + pub fn with_dim(dim: usize) -> Result { + Self::new(Path::new(":memory:"), dim) + } + + /// Dummy embedding stub — produces a deterministic vector from text by hashing. + /// + /// Real embeddings should be provided externally (e.g. from Python / an LLM). + pub fn embed(&self, text: &str) -> Vec { + let mut vec = vec![0.0f64; self.dim]; + for (i, byte) in text.bytes().enumerate() { + vec[i % self.dim] += byte as f64; + } + // L2-normalise so cosine similarity is meaningful. + let norm = vec.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for v in &mut vec { + *v /= norm; + } + } + vec + } + + /// Store a document with a pre-computed embedding. + pub fn store_with_embedding( + &self, + content: &str, + source: &str, + metadata: Option<&Value>, + embedding: &[f64], + ) -> Result { + let doc_id = Uuid::new_v4().to_string(); + let meta_str = metadata + .map(|m| serde_json::to_string(m).unwrap_or_default()) + .unwrap_or_else(|| "{}".to_string()); + let emb_bytes = embedding_to_bytes(embedding); + + let conn = self.conn.lock(); + conn.execute( + "INSERT INTO faiss_documents (id, content, source, metadata, embedding) + VALUES (?1, ?2, ?3, ?4, ?5)", + rusqlite::params![doc_id, content, source, meta_str, emb_bytes], + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + Ok(doc_id) + } + + /// Retrieve the top-k documents closest to a pre-computed query embedding. + pub fn retrieve_by_embedding( + &self, + query_embedding: &[f64], + top_k: usize, + ) -> Result, OpenJarvisError> { + let conn = self.conn.lock(); + + let mut stmt = conn + .prepare( + "SELECT id, content, source, metadata, embedding FROM faiss_documents", + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + + let rows: Vec<(String, String, String, Vec)> = stmt + .query_map([], |row| { + let content: String = row.get(1)?; + let source: String = row.get::<_, String>(2).unwrap_or_default(); + let meta_str: String = row.get::<_, String>(3).unwrap_or_else(|_| "{}".into()); + let emb_bytes: Vec = row.get(4)?; + Ok((content, source, meta_str, bytes_to_embedding(&emb_bytes))) + }) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })? + .filter_map(|r| r.ok()) + .collect(); + + let mut scored: Vec<(RetrievalResult, f64)> = rows + .into_iter() + .map(|(content, source, meta_str, emb)| { + let score = cosine_similarity(query_embedding, &emb); + let metadata: HashMap = + serde_json::from_str(&meta_str).unwrap_or_default(); + ( + RetrievalResult { + content, + source, + score, + metadata, + }, + score, + ) + }) + .collect(); + + scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + scored.truncate(top_k); + + Ok(scored.into_iter().map(|(r, _)| r).collect()) + } +} + +impl MemoryBackend for FAISSMemory { + fn backend_id(&self) -> &str { + "faiss" + } + + fn store( + &self, + content: &str, + source: &str, + metadata: Option<&Value>, + ) -> Result { + let embedding = self.embed(content); + self.store_with_embedding(content, source, metadata, &embedding) + } + + fn retrieve( + &self, + query: &str, + top_k: usize, + ) -> Result, OpenJarvisError> { + let query_embedding = self.embed(query); + self.retrieve_by_embedding(&query_embedding, top_k) + } + + fn delete(&self, doc_id: &str) -> Result { + let conn = self.conn.lock(); + let changes = conn + .execute( + "DELETE FROM faiss_documents WHERE id = ?1", + rusqlite::params![doc_id], + ) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + Ok(changes > 0) + } + + fn clear(&self) -> Result<(), OpenJarvisError> { + let conn = self.conn.lock(); + conn.execute_batch("DELETE FROM faiss_documents") + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + Ok(()) + } + + fn count(&self) -> Result { + let conn = self.conn.lock(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM faiss_documents", [], |row| row.get(0)) + .map_err(|e| { + OpenJarvisError::Io(std::io::Error::other( + e.to_string(), + )) + })?; + Ok(count as usize) + } +} + +// --------------------------------------------------------------------------- +// Helper functions +// --------------------------------------------------------------------------- + +/// Cosine similarity between two vectors. +fn cosine_similarity(a: &[f64], b: &[f64]) -> f64 { + let len = a.len().min(b.len()); + let mut dot = 0.0f64; + let mut norm_a = 0.0f64; + let mut norm_b = 0.0f64; + for i in 0..len { + dot += a[i] * b[i]; + norm_a += a[i] * a[i]; + norm_b += b[i] * b[i]; + } + let denom = norm_a.sqrt() * norm_b.sqrt(); + if denom == 0.0 { + 0.0 + } else { + dot / denom + } +} + +/// Serialize a `&[f64]` embedding into little-endian bytes for SQLite BLOB storage. +fn embedding_to_bytes(embedding: &[f64]) -> Vec { + embedding.iter().flat_map(|f| f.to_le_bytes()).collect() +} + +/// Deserialize a byte slice back into a `Vec`. +fn bytes_to_embedding(bytes: &[u8]) -> Vec { + bytes + .chunks_exact(8) + .map(|chunk| { + let arr: [u8; 8] = chunk.try_into().unwrap_or([0u8; 8]); + f64::from_le_bytes(arr) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_faiss_store_and_retrieve() { + let mem = FAISSMemory::in_memory().unwrap(); + let id = mem + .store("Rust is a systems programming language", "test", None) + .unwrap(); + assert!(!id.is_empty()); + + let results = mem.retrieve("Rust programming", 5).unwrap(); + assert!(!results.is_empty()); + assert!(results[0].content.contains("Rust")); + } + + #[test] + fn test_faiss_cosine_ranking() { + let mem = FAISSMemory::in_memory().unwrap(); + mem.store("Rust is fast and safe", "doc1", None).unwrap(); + mem.store("Python is easy to learn", "doc2", None).unwrap(); + mem.store("Rust and C++ are systems languages", "doc3", None) + .unwrap(); + + let results = mem.retrieve("Rust systems", 3).unwrap(); + assert!(!results.is_empty()); + // The top result should mention Rust + assert!(results[0].content.contains("Rust")); + } + + #[test] + fn test_faiss_delete() { + let mem = FAISSMemory::in_memory().unwrap(); + let id = mem.store("test content", "test", None).unwrap(); + assert_eq!(mem.count().unwrap(), 1); + assert!(mem.delete(&id).unwrap()); + assert_eq!(mem.count().unwrap(), 0); + } + + #[test] + fn test_faiss_clear() { + let mem = FAISSMemory::in_memory().unwrap(); + mem.store("doc 1", "s1", None).unwrap(); + mem.store("doc 2", "s2", None).unwrap(); + assert_eq!(mem.count().unwrap(), 2); + mem.clear().unwrap(); + assert_eq!(mem.count().unwrap(), 0); + } + + #[test] + fn test_faiss_with_real_embeddings() { + let mem = FAISSMemory::in_memory().unwrap(); + let emb1 = vec![1.0, 0.0, 0.0]; + let emb2 = vec![0.0, 1.0, 0.0]; + let emb3 = vec![0.9, 0.1, 0.0]; + + mem.store_with_embedding("doc A", "s1", None, &emb1) + .unwrap(); + mem.store_with_embedding("doc B", "s2", None, &emb2) + .unwrap(); + mem.store_with_embedding("doc C", "s3", None, &emb3) + .unwrap(); + + let query_emb = vec![1.0, 0.0, 0.0]; + let results = mem.retrieve_by_embedding(&query_emb, 2).unwrap(); + assert_eq!(results.len(), 2); + // doc A should be closest to query (identical vector) + assert_eq!(results[0].content, "doc A"); + // doc C should be second (0.9 cosine similarity) + assert_eq!(results[1].content, "doc C"); + } + + #[test] + fn test_faiss_empty_retrieve() { + let mem = FAISSMemory::in_memory().unwrap(); + let results = mem.retrieve("anything", 5).unwrap(); + assert!(results.is_empty()); + } + + #[test] + fn test_embed_deterministic() { + let mem = FAISSMemory::in_memory().unwrap(); + let e1 = mem.embed("hello world"); + let e2 = mem.embed("hello world"); + assert_eq!(e1, e2); + } + + #[test] + fn test_cosine_similarity_identical() { + let a = vec![1.0, 2.0, 3.0]; + assert!((cosine_similarity(&a, &a) - 1.0).abs() < 1e-10); + } + + #[test] + fn test_cosine_similarity_orthogonal() { + let a = vec![1.0, 0.0]; + let b = vec![0.0, 1.0]; + assert!((cosine_similarity(&a, &b)).abs() < 1e-10); + } + + #[test] + fn test_embedding_roundtrip() { + #[allow(clippy::approx_constant)] + let emb = vec![1.0, -2.5, 3.14, 0.0]; + let bytes = embedding_to_bytes(&emb); + let recovered = bytes_to_embedding(&bytes); + assert_eq!(emb, recovered); + } +} diff --git a/rust/crates/openjarvis-tools/src/storage/hybrid.rs b/rust/crates/openjarvis-tools/src/storage/hybrid.rs new file mode 100644 index 00000000..871839f5 --- /dev/null +++ b/rust/crates/openjarvis-tools/src/storage/hybrid.rs @@ -0,0 +1,243 @@ +//! Hybrid memory backend — Reciprocal Rank Fusion (RRF) over multiple backends. +//! +//! Combines results from multiple `MemoryBackend` implementations by fusing their +//! ranked result lists using the RRF formula: +//! +//! ```text +//! score(doc) = sum_over_backends( 1 / (k + rank_i) ) +//! ``` +//! +//! where `k` is a smoothing constant (default 60) and `rank_i` is the 1-based rank +//! of the document in backend `i`'s result list (or absent if not returned). + +use crate::storage::traits::MemoryBackend; +use openjarvis_core::{OpenJarvisError, RetrievalResult}; +use std::collections::HashMap; + +/// Default RRF smoothing constant. +const DEFAULT_K: f64 = 60.0; + +/// Hybrid memory backend that fuses results from multiple backends via RRF. +pub struct HybridMemory { + backends: Vec>, + k: f64, +} + +impl HybridMemory { + /// Create a new `HybridMemory` combining the given backends with default `k=60`. + pub fn new(backends: Vec>) -> Self { + Self { + backends, + k: DEFAULT_K, + } + } + + /// Create a new `HybridMemory` with a custom RRF smoothing constant. + pub fn with_k(backends: Vec>, k: f64) -> Self { + Self { backends, k } + } + + /// Fuse ranked result lists from multiple backends using Reciprocal Rank Fusion. + /// + /// Each result is keyed by its `content` field. When the same document appears + /// in multiple backends, scores are summed. + fn fuse_results( + &self, + per_backend_results: Vec>, + top_k: usize, + ) -> Vec { + // Map: content -> (accumulated_rrf_score, best_result) + let mut fused: HashMap = HashMap::new(); + + for results in &per_backend_results { + for (rank_0, result) in results.iter().enumerate() { + let rrf_score = 1.0 / (self.k + (rank_0 + 1) as f64); + let entry = fused + .entry(result.content.clone()) + .or_insert_with(|| (0.0, result.clone())); + entry.0 += rrf_score; + } + } + + let mut scored: Vec<(f64, RetrievalResult)> = fused + .into_values() + .map(|(rrf_score, mut result)| { + result.score = rrf_score; + (rrf_score, result) + }) + .collect(); + + scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + scored.truncate(top_k); + + scored.into_iter().map(|(_, r)| r).collect() + } +} + +impl MemoryBackend for HybridMemory { + fn backend_id(&self) -> &str { + "hybrid" + } + + fn store( + &self, + content: &str, + source: &str, + metadata: Option<&serde_json::Value>, + ) -> Result { + let mut last_id = String::new(); + for backend in &self.backends { + last_id = backend.store(content, source, metadata)?; + } + Ok(last_id) + } + + fn retrieve( + &self, + query: &str, + top_k: usize, + ) -> Result, OpenJarvisError> { + // Ask each backend for a generous number of results to give RRF enough data. + let fetch_k = top_k * 3; + let mut per_backend: Vec> = Vec::with_capacity(self.backends.len()); + for backend in &self.backends { + per_backend.push(backend.retrieve(query, fetch_k)?); + } + Ok(self.fuse_results(per_backend, top_k)) + } + + fn delete(&self, doc_id: &str) -> Result { + let mut any_deleted = false; + for backend in &self.backends { + if backend.delete(doc_id)? { + any_deleted = true; + } + } + Ok(any_deleted) + } + + fn clear(&self) -> Result<(), OpenJarvisError> { + for backend in &self.backends { + backend.clear()?; + } + Ok(()) + } + + fn count(&self) -> Result { + // Return count from the first backend (all should be in sync after store/delete). + self.backends + .first() + .map(|b| b.count()) + .unwrap_or(Ok(0)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::bm25::BM25Memory; + use crate::storage::sqlite::SQLiteMemory; + + fn make_hybrid() -> HybridMemory { + let sqlite = SQLiteMemory::in_memory().unwrap(); + let bm25 = BM25Memory::default(); + HybridMemory::new(vec![Box::new(sqlite), Box::new(bm25)]) + } + + #[test] + fn test_hybrid_store_and_retrieve() { + let mem = make_hybrid(); + mem.store("Rust is fast and safe", "doc1", None).unwrap(); + mem.store("Python is easy to learn", "doc2", None).unwrap(); + mem.store("Rust and C++ are systems languages", "doc3", None) + .unwrap(); + + let results = mem.retrieve("Rust programming", 3).unwrap(); + assert!(!results.is_empty()); + // Top result should mention Rust (both backends agree) + assert!(results[0].content.contains("Rust")); + } + + #[test] + fn test_hybrid_rrf_scores() { + let mem = make_hybrid(); + mem.store("alpha beta gamma", "s1", None).unwrap(); + mem.store("beta gamma delta", "s2", None).unwrap(); + + let results = mem.retrieve("beta", 5).unwrap(); + // Both docs should appear; RRF scores should be > 0 + for r in &results { + assert!(r.score > 0.0); + } + } + + #[test] + fn test_hybrid_delete() { + let mem = make_hybrid(); + let id = mem.store("test content", "test", None).unwrap(); + // At least one backend should report deletion + assert!(mem.delete(&id).unwrap()); + } + + #[test] + fn test_hybrid_clear() { + let mem = make_hybrid(); + mem.store("doc 1", "s1", None).unwrap(); + mem.store("doc 2", "s2", None).unwrap(); + mem.clear().unwrap(); + assert_eq!(mem.count().unwrap(), 0); + } + + #[test] + fn test_hybrid_empty() { + let mem = make_hybrid(); + let results = mem.retrieve("anything", 5).unwrap(); + assert!(results.is_empty()); + } + + #[test] + fn test_hybrid_no_backends() { + let mem = HybridMemory::new(vec![]); + assert_eq!(mem.count().unwrap(), 0); + let results = mem.retrieve("test", 5).unwrap(); + assert!(results.is_empty()); + } + + #[test] + fn test_hybrid_custom_k() { + let sqlite = SQLiteMemory::in_memory().unwrap(); + let bm25 = BM25Memory::default(); + let mem = HybridMemory::with_k(vec![Box::new(sqlite), Box::new(bm25)], 10.0); + mem.store("hello world", "s1", None).unwrap(); + + let results = mem.retrieve("hello", 5).unwrap(); + assert!(!results.is_empty()); + // With k=10, RRF score for rank-1 from both backends: + // 2 * (1 / (10 + 1)) = 2/11 ~ 0.1818 + let expected = 2.0 / 11.0; + assert!((results[0].score - expected).abs() < 1e-6); + } + + #[test] + fn test_fuse_deduplication() { + let mem = HybridMemory::with_k(vec![], 60.0); + // Two backends returning the same doc at rank 1 + let r1 = vec![RetrievalResult { + content: "same doc".to_string(), + source: "s1".to_string(), + score: 5.0, + metadata: HashMap::new(), + }]; + let r2 = vec![RetrievalResult { + content: "same doc".to_string(), + source: "s1".to_string(), + score: 3.0, + metadata: HashMap::new(), + }]; + let fused = mem.fuse_results(vec![r1, r2], 10); + assert_eq!(fused.len(), 1); + // RRF: 1/(60+1) + 1/(60+1) = 2/61 + let expected = 2.0 / 61.0; + assert!((fused[0].score - expected).abs() < 1e-10); + } +} diff --git a/rust/crates/openjarvis-tools/src/storage/knowledge_graph.rs b/rust/crates/openjarvis-tools/src/storage/knowledge_graph.rs index 4098e9af..1919a6fa 100644 --- a/rust/crates/openjarvis-tools/src/storage/knowledge_graph.rs +++ b/rust/crates/openjarvis-tools/src/storage/knowledge_graph.rs @@ -16,13 +16,12 @@ impl KnowledgeGraphMemory { pub fn new(db_path: &Path) -> Result { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)) + OpenJarvisError::Io(std::io::Error::other(e)) })?; } let conn = Connection::open(db_path).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -47,8 +46,7 @@ impl KnowledgeGraphMemory { CREATE INDEX IF NOT EXISTS idx_relations_target ON relations(target_id);", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -79,8 +77,7 @@ impl KnowledgeGraphMemory { rusqlite::params![id, name, entity_type, props], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -107,8 +104,7 @@ impl KnowledgeGraphMemory { rusqlite::params![id, source_id, target_id, relation_type, props], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -129,8 +125,7 @@ impl KnowledgeGraphMemory { WHERE r.target_id = ?1", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -144,8 +139,7 @@ impl KnowledgeGraphMemory { )) }) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })? @@ -183,8 +177,7 @@ impl MemoryBackend for KnowledgeGraphMemory { FROM entities WHERE name LIKE ?1 LIMIT ?2", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -203,8 +196,7 @@ impl MemoryBackend for KnowledgeGraphMemory { }) }) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })? @@ -222,8 +214,7 @@ impl MemoryBackend for KnowledgeGraphMemory { rusqlite::params![doc_id], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -234,8 +225,7 @@ impl MemoryBackend for KnowledgeGraphMemory { let conn = self.conn.lock(); conn.execute_batch("DELETE FROM relations; DELETE FROM entities;") .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -247,8 +237,7 @@ impl MemoryBackend for KnowledgeGraphMemory { let count: i64 = conn .query_row("SELECT COUNT(*) FROM entities", [], |row| row.get(0)) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; diff --git a/rust/crates/openjarvis-tools/src/storage/mod.rs b/rust/crates/openjarvis-tools/src/storage/mod.rs index 822bb85c..357a46ee 100644 --- a/rust/crates/openjarvis-tools/src/storage/mod.rs +++ b/rust/crates/openjarvis-tools/src/storage/mod.rs @@ -1,7 +1,10 @@ -//! Memory/storage backends — SQLite FTS5, BM25, KnowledgeGraph, Hybrid. +//! Memory/storage backends — SQLite FTS5, BM25, FAISS, ColBERT, KnowledgeGraph, Hybrid. pub mod backend_enum; pub mod bm25; +pub mod colbert; +pub mod faiss; +pub mod hybrid; pub mod knowledge_graph; pub mod sqlite; pub mod traits; @@ -9,6 +12,9 @@ pub mod utils; pub use backend_enum::MemoryBackendEnum; pub use bm25::BM25Memory; +pub use colbert::ColBERTMemory; +pub use faiss::FAISSMemory; +pub use hybrid::HybridMemory; pub use knowledge_graph::KnowledgeGraphMemory; pub use sqlite::SQLiteMemory; pub use traits::MemoryBackend; diff --git a/rust/crates/openjarvis-tools/src/storage/sqlite.rs b/rust/crates/openjarvis-tools/src/storage/sqlite.rs index 55e5ead8..433b2378 100644 --- a/rust/crates/openjarvis-tools/src/storage/sqlite.rs +++ b/rust/crates/openjarvis-tools/src/storage/sqlite.rs @@ -17,13 +17,12 @@ impl SQLiteMemory { pub fn new(db_path: &Path) -> Result { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)) + OpenJarvisError::Io(std::io::Error::other(e)) })?; } let conn = Connection::open(db_path).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -41,8 +40,7 @@ impl SQLiteMemory { );", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -80,8 +78,7 @@ impl MemoryBackend for SQLiteMemory { rusqlite::params![doc_id, content, source, meta_str], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -91,8 +88,7 @@ impl MemoryBackend for SQLiteMemory { rusqlite::params![doc_id, content, source], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -125,8 +121,7 @@ impl MemoryBackend for SQLiteMemory { LIMIT ?2", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -145,8 +140,7 @@ impl MemoryBackend for SQLiteMemory { }) }) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })? @@ -163,8 +157,7 @@ impl MemoryBackend for SQLiteMemory { rusqlite::params![doc_id], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -174,8 +167,7 @@ impl MemoryBackend for SQLiteMemory { rusqlite::params![doc_id], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -186,8 +178,7 @@ impl MemoryBackend for SQLiteMemory { let conn = self.conn.lock(); conn.execute_batch("DELETE FROM documents_fts; DELETE FROM documents") .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -199,8 +190,7 @@ impl MemoryBackend for SQLiteMemory { let count: i64 = conn .query_row("SELECT COUNT(*) FROM documents", [], |row| row.get(0)) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; diff --git a/rust/crates/openjarvis-traces/src/collector.rs b/rust/crates/openjarvis-traces/src/collector.rs index acb96de1..aca47ce5 100644 --- a/rust/crates/openjarvis-traces/src/collector.rs +++ b/rust/crates/openjarvis-traces/src/collector.rs @@ -1,7 +1,7 @@ //! TraceCollector — subscribes to EventBus and assembles traces. use crate::store::TraceStore; -use openjarvis_core::{EventBus, EventType, StepType, Trace, TraceStep}; +use openjarvis_core::{Trace, TraceStep}; use parking_lot::Mutex; use std::collections::HashMap; use std::sync::Arc; @@ -75,6 +75,7 @@ impl TraceCollector { #[cfg(test)] mod tests { use super::*; + use openjarvis_core::StepType; #[test] fn test_collector_lifecycle() { diff --git a/rust/crates/openjarvis-traces/src/store.rs b/rust/crates/openjarvis-traces/src/store.rs index c0036236..60506f84 100644 --- a/rust/crates/openjarvis-traces/src/store.rs +++ b/rust/crates/openjarvis-traces/src/store.rs @@ -14,15 +14,12 @@ impl TraceStore { pub fn new(db_path: &Path) -> Result { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)) + OpenJarvisError::Io(std::io::Error::other(e)) })?; } let conn = Connection::open(db_path).map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - e.to_string(), - )) + OpenJarvisError::Io(std::io::Error::other(e.to_string())) })?; conn.execute_batch( @@ -44,10 +41,7 @@ impl TraceStore { )", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - e.to_string(), - )) + OpenJarvisError::Io(std::io::Error::other(e.to_string())) })?; Ok(Self { @@ -89,10 +83,7 @@ impl TraceStore { ], ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, - e.to_string(), - )) + OpenJarvisError::Io(std::io::Error::other(e.to_string())) })?; Ok(()) @@ -108,8 +99,7 @@ impl TraceStore { FROM traces WHERE trace_id = ?1", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -160,8 +150,7 @@ impl TraceStore { FROM traces ORDER BY started_at DESC LIMIT ?1 OFFSET ?2", ) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; @@ -194,8 +183,7 @@ impl TraceStore { }) }) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })? @@ -210,8 +198,7 @@ impl TraceStore { let count: i64 = conn .query_row("SELECT COUNT(*) FROM traces", [], |row| row.get(0)) .map_err(|e| { - OpenJarvisError::Io(std::io::Error::new( - std::io::ErrorKind::Other, + OpenJarvisError::Io(std::io::Error::other( e.to_string(), )) })?; diff --git a/scripts/quickstart.sh b/scripts/quickstart.sh index 33339d50..1a89b9df 100755 --- a/scripts/quickstart.sh +++ b/scripts/quickstart.sh @@ -6,7 +6,7 @@ set -euo pipefail # the backend API server and frontend, then opens the browser. # # Usage: -# git clone https://github.com/HazyResearch/OpenJarvis.git +# git clone https://github.com/open-jarvis/OpenJarvis.git # cd OpenJarvis # ./scripts/quickstart.sh # ────────────────────────────────────────────────────────────────────── diff --git a/src/openjarvis/_rust_bridge.py b/src/openjarvis/_rust_bridge.py index 4e16fb52..f0facc19 100644 --- a/src/openjarvis/_rust_bridge.py +++ b/src/openjarvis/_rust_bridge.py @@ -126,10 +126,55 @@ def retrieval_results_from_json(json_str: str) -> list: return results +# --------------------------------------------------------------------------- +# Phase 2 converters — optimization & engine types +# --------------------------------------------------------------------------- + + +def optimization_store_from_rust(path: str = ":memory:") -> object | None: + """Get a Rust-backed OptimizationStore, or None if Rust unavailable.""" + mod = get_rust_module() + if mod is None: + return None + try: + return mod.OptimizationStore(path) + except Exception: + return None + + +def trial_result_from_json(json_str: str) -> dict: + """Convert Rust TrialResult JSON to a Python dict.""" + return json.loads(json_str) + + +def optimization_run_from_json(json_str: str) -> dict: + """Convert Rust OptimizationRun JSON to a Python dict.""" + return json.loads(json_str) + + +def generate_result_from_json(json_str: str) -> dict: + """Convert Rust GenerateResult JSON to a Python dict.""" + data = json.loads(json_str) + return { + "content": data.get("content", ""), + "model": data.get("model", ""), + "finish_reason": data.get("finish_reason", "stop"), + "usage": data.get("usage", {}), + "tool_calls": data.get("tool_calls"), + "ttft": data.get("ttft", 0.0), + "cost_usd": data.get("cost_usd", 0.0), + "metadata": data.get("metadata", {}), + } + + __all__ = [ "RUST_AVAILABLE", + "generate_result_from_json", "get_rust_module", "injection_result_from_json", + "optimization_run_from_json", + "optimization_store_from_rust", "retrieval_results_from_json", "scan_result_from_json", + "trial_result_from_json", ] diff --git a/src/openjarvis/cli/optimize_cmd.py b/src/openjarvis/cli/optimize_cmd.py index a67379d4..0fa75ecd 100644 --- a/src/openjarvis/cli/optimize_cmd.py +++ b/src/openjarvis/cli/optimize_cmd.py @@ -56,7 +56,7 @@ def optimize_run( data = None if config_path is not None: try: - from openjarvis.optimize.config import load_optimize_config + from openjarvis.learning.optimize.config import load_optimize_config except ImportError: console.print( "[red]Optimization framework not available.[/red]" @@ -80,7 +80,7 @@ def optimize_run( benchmark_specs = None if data is not None: try: - from openjarvis.optimize.config import load_benchmark_specs + from openjarvis.learning.optimize.config import load_benchmark_specs specs = load_benchmark_specs(data) if len(specs) > 1: @@ -110,15 +110,15 @@ def optimize_run( try: from openjarvis.core.config import DEFAULT_CONFIG_DIR - from openjarvis.optimize.config import load_objectives - from openjarvis.optimize.llm_optimizer import LLMOptimizer - from openjarvis.optimize.optimizer import OptimizationEngine - from openjarvis.optimize.search_space import ( + from openjarvis.learning.optimize.config import load_objectives + from openjarvis.learning.optimize.llm_optimizer import LLMOptimizer + from openjarvis.learning.optimize.optimizer import OptimizationEngine + from openjarvis.learning.optimize.search_space import ( DEFAULT_SEARCH_SPACE, build_search_space, ) - from openjarvis.optimize.store import OptimizationStore - from openjarvis.optimize.trial_runner import ( + from openjarvis.learning.optimize.store import OptimizationStore + from openjarvis.learning.optimize.trial_runner import ( MultiBenchTrialRunner, TrialRunner, ) @@ -216,7 +216,7 @@ def optimize_status() -> None: try: from openjarvis.core.config import DEFAULT_CONFIG_DIR - from openjarvis.optimize.store import OptimizationStore + from openjarvis.learning.optimize.store import OptimizationStore db_path = DEFAULT_CONFIG_DIR / "optimize.db" if not db_path.exists(): @@ -264,7 +264,7 @@ def optimize_results(run_id: str) -> None: try: from openjarvis.core.config import DEFAULT_CONFIG_DIR - from openjarvis.optimize.store import OptimizationStore + from openjarvis.learning.optimize.store import OptimizationStore db_path = DEFAULT_CONFIG_DIR / "optimize.db" if not db_path.exists(): @@ -330,8 +330,8 @@ def optimize_best(run_id: str, output: Optional[str]) -> None: try: from openjarvis.core.config import DEFAULT_CONFIG_DIR - from openjarvis.optimize.optimizer import OptimizationEngine - from openjarvis.optimize.store import OptimizationStore + from openjarvis.learning.optimize.optimizer import OptimizationEngine + from openjarvis.learning.optimize.store import OptimizationStore db_path = DEFAULT_CONFIG_DIR / "optimize.db" if not db_path.exists(): diff --git a/src/openjarvis/engine/_openai_compat.py b/src/openjarvis/engine/_openai_compat.py index 0c60d0ca..f20eabc8 100644 --- a/src/openjarvis/engine/_openai_compat.py +++ b/src/openjarvis/engine/_openai_compat.py @@ -43,6 +43,7 @@ class _OpenAICompatibleEngine(InferenceEngine): "temperature": temperature, "max_tokens": max_tokens, "stream": False, + "chat_template_kwargs": {"enable_thinking": False}, **kwargs, } try: diff --git a/src/openjarvis/engine/cloud.py b/src/openjarvis/engine/cloud.py index 35b81711..a1bd82bb 100644 --- a/src/openjarvis/engine/cloud.py +++ b/src/openjarvis/engine/cloud.py @@ -230,12 +230,54 @@ class CloudEngine(InferenceEngine): "ANTHROPIC_API_KEY and install " "openjarvis[inference-cloud]" ) - # Separate system message from conversation messages + # Separate system message and convert to Anthropic message format system_text = "" chat_msgs: List[Dict[str, Any]] = [] for m in messages: if m.role.value == "system": system_text = m.content + elif m.role.value == "tool": + # Anthropic expects tool results as role="user" with + # tool_result content blocks + tool_result_block = { + "type": "tool_result", + "tool_use_id": m.tool_call_id or "", + "content": m.content, + } + # Merge consecutive tool results into a single user message + if ( + chat_msgs + and chat_msgs[-1]["role"] == "user" + and isinstance(chat_msgs[-1]["content"], list) + and chat_msgs[-1]["content"] + and chat_msgs[-1]["content"][-1].get("type") == "tool_result" + ): + chat_msgs[-1]["content"].append(tool_result_block) + else: + chat_msgs.append({ + "role": "user", + "content": [tool_result_block], + }) + elif m.role.value == "assistant" and m.tool_calls: + # Convert assistant messages with tool_calls to Anthropic + # content blocks (text + tool_use) + content_blocks: List[Dict[str, Any]] = [] + if m.content: + content_blocks.append({"type": "text", "text": m.content}) + for tc in m.tool_calls: + args = tc.arguments + if isinstance(args, str): + try: + args = json.loads(args) + except (json.JSONDecodeError, TypeError): + args = {"input": args} + content_blocks.append({ + "type": "tool_use", + "id": tc.id, + "name": tc.name, + "input": args if isinstance(args, dict) else {}, + }) + chat_msgs.append({"role": "assistant", "content": content_blocks}) else: chat_msgs.append({"role": m.role.value, "content": m.content}) create_kwargs: Dict[str, Any] = { @@ -308,12 +350,50 @@ class CloudEngine(InferenceEngine): "GEMINI_API_KEY or GOOGLE_API_KEY and install " "openjarvis[inference-google]" ) - # Build contents from messages + # Build contents from messages, converting tool roles for Gemini system_text = "" contents: List[Dict[str, Any]] = [] for m in messages: if m.role.value == "system": system_text = m.content + elif m.role.value == "tool": + # Gemini expects function responses as role="user" with + # function_response parts + fn_resp_part = { + "function_response": { + "name": m.name or "unknown", + "response": {"result": m.content}, + } + } + # Merge consecutive tool results into a single user message + if ( + contents + and contents[-1]["role"] == "user" + and contents[-1]["parts"] + and "function_response" in contents[-1]["parts"][-1] + ): + contents[-1]["parts"].append(fn_resp_part) + else: + contents.append({"role": "user", "parts": [fn_resp_part]}) + elif m.role.value == "assistant" and m.tool_calls: + # Convert assistant tool_calls to function_call parts + parts: List[Dict[str, Any]] = [] + if m.content: + parts.append({"text": m.content}) + for tc in m.tool_calls: + args = tc.arguments + if isinstance(args, str): + try: + args = json.loads(args) + except (json.JSONDecodeError, TypeError): + args = {"input": args} + parts.append({ + "function_call": { + "name": tc.name, + "args": args if isinstance(args, dict) else {}, + } + }) + contents.append({"role": "model", "parts": parts}) elif m.role.value == "assistant": contents.append({"role": "model", "parts": [{"text": m.content}]}) else: diff --git a/src/openjarvis/evals/configs/deepplanning-monitor-operative.toml b/src/openjarvis/evals/configs/deepplanning-monitor-operative.toml new file mode 100644 index 00000000..ac62fc08 --- /dev/null +++ b/src/openjarvis/evals/configs/deepplanning-monitor-operative.toml @@ -0,0 +1,42 @@ +# DeepPlanning evaluation with MonitorOperativeAgent +# Benchmark: long-horizon planning with hard constraints (travel + shopping) +# Paper: arXiv:2601.18137, best frontier: GPT-5.2-high at 44.6% + +[meta] +name = "deepplanning-monitor-operative" +description = "DeepPlanning constraint satisfaction with MonitorOperativeAgent" + +[defaults] +temperature = 0.6 +max_tokens = 32768 + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 1024 + +[run] +max_workers = 1 +output_dir = "results/deepplanning-monitor-operative/" +seed = 42 +telemetry = true +gpu_metrics = true + +[[models]] +name = "Qwen/Qwen3.5-35B-A3B" +engine = "vllm" +param_count_b = 35.0 +active_params_b = 3.0 +gpu_peak_tflops = 989.5 +gpu_peak_bandwidth_gb_s = 3350.0 +num_gpus = 2 + +[[benchmarks]] +name = "deepplanning" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 50 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "memory_store", "memory_search", +] diff --git a/src/openjarvis/evals/configs/glm5-fp8-monitor-operative.toml b/src/openjarvis/evals/configs/glm5-fp8-monitor-operative.toml new file mode 100644 index 00000000..23443286 --- /dev/null +++ b/src/openjarvis/evals/configs/glm5-fp8-monitor-operative.toml @@ -0,0 +1,74 @@ +# MonitorOperativeAgent evaluation on long-horizon benchmarks +# Machine: 8x H100-80GB, engine: vLLM, model: GLM-5 (Q4_K_M GGUF) +# Inference: GLM-5 reasoning mode +# Benchmarks: LogHub, LifelongAgentBench, DeepPlanning + +[meta] +name = "glm5-q4km-monitor-operative" +description = "Evaluate MonitorOperativeAgent with GLM-5 Q4_K_M on long-horizon monitoring benchmarks" + +[defaults] +temperature = 0.6 +max_tokens = 32768 + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 4096 + +[run] +max_workers = 1 +output_dir = "results/glm5-q4km-monitor-operative/" +seed = 42 +telemetry = true +gpu_metrics = true + +[[models]] +name = "GLM-5-Q4_K_M" +engine = "vllm" +param_count_b = 744.0 +active_params_b = 40.0 +gpu_peak_tflops = 989.5 +gpu_peak_bandwidth_gb_s = 3350.0 +num_gpus = 8 + +# --- Benchmarks --- + +[[benchmarks]] +name = "loghub" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 250 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "file_read", + "memory_store", "memory_search", "memory_retrieve", + "kg_add_entity", "kg_add_relation", "kg_query", "kg_neighbors", + "db_query", "web_search", "http_request", +] + +[[benchmarks]] +name = "lifelong-agent" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 250 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "file_read", + "memory_store", "memory_search", "memory_retrieve", + "kg_add_entity", "kg_add_relation", "kg_query", "kg_neighbors", + "db_query", "web_search", "http_request", +] + +[[benchmarks]] +name = "deepplanning" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 250 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "file_read", + "memory_store", "memory_search", "memory_retrieve", + "kg_add_entity", "kg_add_relation", "kg_query", "kg_neighbors", + "db_query", "web_search", "http_request", +] diff --git a/src/openjarvis/evals/configs/kimi-k2.5-monitor-operative.toml b/src/openjarvis/evals/configs/kimi-k2.5-monitor-operative.toml new file mode 100644 index 00000000..86395bd6 --- /dev/null +++ b/src/openjarvis/evals/configs/kimi-k2.5-monitor-operative.toml @@ -0,0 +1,74 @@ +# MonitorOperativeAgent evaluation on long-horizon benchmarks +# Machine: 8x H100-80GB, engine: vLLM, model: Kimi-K2.5 (IQ4_XS GGUF) +# Inference: Kimi-K2.5 thinking mode (temperature=1.0, top_p=0.95) +# Benchmarks: LogHub, LifelongAgentBench, DeepPlanning + +[meta] +name = "kimi-k2.5-monitor-operative" +description = "Evaluate MonitorOperativeAgent with Kimi-K2.5 on long-horizon monitoring benchmarks" + +[defaults] +temperature = 1.0 # Kimi-K2.5 thinking mode recommended +max_tokens = 32768 + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 4096 + +[run] +max_workers = 1 +output_dir = "results/kimi-k2.5-monitor-operative/" +seed = 42 +telemetry = true +gpu_metrics = true + +[[models]] +name = "Kimi-K2.5-IQ4_XS" +engine = "vllm" +param_count_b = 1000.0 +active_params_b = 32.0 +gpu_peak_tflops = 989.5 +gpu_peak_bandwidth_gb_s = 3350.0 +num_gpus = 8 + +# --- Benchmarks --- + +[[benchmarks]] +name = "loghub" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 250 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "file_read", + "memory_store", "memory_search", "memory_retrieve", + "kg_add_entity", "kg_add_relation", "kg_query", "kg_neighbors", + "db_query", "web_search", "http_request", +] + +[[benchmarks]] +name = "lifelong-agent" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 250 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "file_read", + "memory_store", "memory_search", "memory_retrieve", + "kg_add_entity", "kg_add_relation", "kg_query", "kg_neighbors", + "db_query", "web_search", "http_request", +] + +[[benchmarks]] +name = "deepplanning" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 250 +tools = [ + "think", "calculator", "code_interpreter", + "shell_exec", "file_read", + "memory_store", "memory_search", "memory_retrieve", + "kg_add_entity", "kg_add_relation", "kg_query", "kg_neighbors", + "db_query", "web_search", "http_request", +] diff --git a/src/openjarvis/evals/configs/paperarena-monitor-operative.toml b/src/openjarvis/evals/configs/paperarena-monitor-operative.toml new file mode 100644 index 00000000..45f5b781 --- /dev/null +++ b/src/openjarvis/evals/configs/paperarena-monitor-operative.toml @@ -0,0 +1,42 @@ +# PaperArena evaluation with MonitorOperativeAgent +# Benchmark: scientific literature reasoning (MC/CA/OA) +# Paper: arXiv:2510.10909, best frontier: Gemini 2.5 Pro at 38.8% + +[meta] +name = "paperarena-monitor-operative" +description = "PaperArena scientific reasoning with MonitorOperativeAgent" + +[defaults] +temperature = 0.6 +max_tokens = 32768 + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 1024 + +[run] +max_workers = 1 +output_dir = "results/paperarena-monitor-operative/" +seed = 42 +telemetry = true +gpu_metrics = true + +[[models]] +name = "Qwen/Qwen3.5-35B-A3B" +engine = "vllm" +param_count_b = 35.0 +active_params_b = 3.0 +gpu_peak_tflops = 989.5 +gpu_peak_bandwidth_gb_s = 3350.0 +num_gpus = 2 + +[[benchmarks]] +name = "paperarena" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 50 +tools = [ + "think", "calculator", "code_interpreter", + "web_search", "file_read", "memory_store", "memory_search", +] diff --git a/src/openjarvis/evals/configs/qwen3.5-35b-monitor-operative.toml b/src/openjarvis/evals/configs/qwen3.5-35b-monitor-operative.toml new file mode 100644 index 00000000..bd7085bd --- /dev/null +++ b/src/openjarvis/evals/configs/qwen3.5-35b-monitor-operative.toml @@ -0,0 +1,57 @@ +# MonitorOperativeAgent evaluation on Phase 25 text-to-text benchmarks +# Machine: 2x H100-80GB (GPUs 6-7), engine: vLLM, model: Qwen3.5-35B-A3B +# Inference: Qwen3.5 thinking mode (temperature=0.6, top_p=0.95, top_k=20) +# Benchmarks: LogHub (log anomaly detection), LifelongAgentBench (sequential SQL) +# Note: AMA-Bench excluded (no public dataset yet, paper arxiv:2602.22769) +# DeepPlanning/PaperArena: see dedicated config files + +[meta] +name = "qwen3.5-35b-monitor-operative" +description = "Evaluate MonitorOperativeAgent on Phase 25 text-to-text benchmarks" + +[defaults] +temperature = 0.6 # Qwen3.5 thinking mode (precise analysis) +max_tokens = 32768 # Qwen3.5 recommended for general queries + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 1024 + +[run] +max_workers = 1 # Sequential for agent state safety +output_dir = "results/qwen3.5-35b-monitor-operative/" +seed = 42 +telemetry = true +gpu_metrics = true + +[[models]] +name = "Qwen/Qwen3.5-35B-A3B" +engine = "vllm" +param_count_b = 35.0 +active_params_b = 3.0 +gpu_peak_tflops = 989.5 # H100 SXM FP16 peak TFLOPS +gpu_peak_bandwidth_gb_s = 3350.0 # H100 SXM memory bandwidth +num_gpus = 2 + +# --- Phase 25 Text-to-Text Benchmarks --- + +[[benchmarks]] +name = "loghub" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 50 +tools = [ + "think", "calculator", "code_interpreter", + "memory_store", "memory_search", "memory_retrieve", +] + +[[benchmarks]] +name = "lifelong-agent" +backend = "jarvis-agent" +agent = "monitor_operative" +max_samples = 50 +tools = [ + "think", "calculator", "code_interpreter", + "memory_store", "memory_search", "memory_retrieve", +] diff --git a/src/openjarvis/evals/configs/qwen3.5-397b-monitor-operative.toml b/src/openjarvis/evals/configs/qwen3.5-397b-monitor-operative.toml new file mode 100644 index 00000000..22781617 --- /dev/null +++ b/src/openjarvis/evals/configs/qwen3.5-397b-monitor-operative.toml @@ -0,0 +1,42 @@ +# MonitorOperativeAgent evaluation on long-horizon benchmarks +# Machine: 4x H100-80GB (GPUs 0-3), engine: vLLM, model: Qwen3.5-122B-A10B-FP8 +# Inference: Qwen3.5 non-thinking mode (temperature=0.6) +# Benchmarks: LifelongAgentBench only (LogHub done at 94%, DeepPlanning needs >32K ctx) +# Note: PaperArena excluded (HF repo Melmaphother/PaperArena-Data is empty) +# AMA-Bench excluded (no public dataset yet, paper arxiv:2602.22769) + +[meta] +name = "qwen3.5-122b-monitor-operative" +description = "Evaluate MonitorOperativeAgent on long-horizon monitoring benchmarks" + +[defaults] +temperature = 0.6 # Qwen3.5 thinking mode recommended +max_tokens = 1024 # Keep short to prevent hangs; answers are brief + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 4096 + +[run] +max_workers = 1 # Sequential for episode mode; 1 worker for agent state safety +output_dir = "results/qwen3.5-122b-monitor-operative/" +seed = 42 +telemetry = true +gpu_metrics = true + +[[models]] +name = "Qwen/Qwen3.5-122B-A10B-FP8" +engine = "vllm" +param_count_b = 122.0 +active_params_b = 10.0 +gpu_peak_tflops = 989.5 # H100 SXM FP16 peak TFLOPS +gpu_peak_bandwidth_gb_s = 3350.0 # H100 SXM memory bandwidth +num_gpus = 4 + +# --- Benchmarks --- + +[[benchmarks]] +name = "lifelong-agent" +backend = "jarvis-direct" +max_samples = 50 diff --git a/src/openjarvis/evals/configs/use_case_cloud_openhands.toml b/src/openjarvis/evals/configs/use_case_cloud_openhands.toml index 43c4c9e4..34c275aa 100644 --- a/src/openjarvis/evals/configs/use_case_cloud_openhands.toml +++ b/src/openjarvis/evals/configs/use_case_cloud_openhands.toml @@ -27,26 +27,32 @@ telemetry = true [[models]] name = "claude-opus-4-6" +engine = "cloud" provider = "anthropic" [[models]] name = "claude-sonnet-4-6" +engine = "cloud" provider = "anthropic" [[models]] name = "gpt-5.4" +engine = "cloud" provider = "openai" [[models]] name = "gpt-5-mini" +engine = "cloud" provider = "openai" [[models]] name = "gemini-3.1-pro-preview" +engine = "cloud" provider = "google" [[models]] name = "gemini-3-flash-preview" +engine = "cloud" provider = "google" # --- Benchmarks --- diff --git a/src/openjarvis/evals/configs/use_case_cloud_orchestrator.toml b/src/openjarvis/evals/configs/use_case_cloud_orchestrator.toml index 81e97695..6f535bb7 100644 --- a/src/openjarvis/evals/configs/use_case_cloud_orchestrator.toml +++ b/src/openjarvis/evals/configs/use_case_cloud_orchestrator.toml @@ -27,26 +27,32 @@ telemetry = true [[models]] name = "claude-opus-4-6" +engine = "cloud" provider = "anthropic" [[models]] name = "claude-sonnet-4-6" +engine = "cloud" provider = "anthropic" [[models]] name = "gpt-5.4" +engine = "cloud" provider = "openai" [[models]] name = "gpt-5-mini" +engine = "cloud" provider = "openai" [[models]] name = "gemini-3.1-pro-preview" +engine = "cloud" provider = "google" [[models]] name = "gemini-3-flash-preview" +engine = "cloud" provider = "google" # --- Benchmarks --- diff --git a/src/openjarvis/evals/configs/use_case_opensource_openhands.toml b/src/openjarvis/evals/configs/use_case_opensource_openhands.toml index 1969ed31..e4f193dd 100644 --- a/src/openjarvis/evals/configs/use_case_opensource_openhands.toml +++ b/src/openjarvis/evals/configs/use_case_opensource_openhands.toml @@ -38,7 +38,7 @@ active_params_b = 17.0 num_gpus = 8 [[models]] -name = "unsloth/Qwen3.5-122B-A10B-GGUF" +name = "Qwen/Qwen3.5-122B-A10B-FP8" engine = "vllm" param_count_b = 122.0 active_params_b = 10.0 diff --git a/src/openjarvis/evals/configs/use_case_opensource_orchestrator.toml b/src/openjarvis/evals/configs/use_case_opensource_orchestrator.toml index c40d60e1..ddee471e 100644 --- a/src/openjarvis/evals/configs/use_case_opensource_orchestrator.toml +++ b/src/openjarvis/evals/configs/use_case_opensource_orchestrator.toml @@ -38,7 +38,7 @@ active_params_b = 17.0 num_gpus = 8 [[models]] -name = "unsloth/Qwen3.5-122B-A10B-GGUF" +name = "Qwen/Qwen3.5-122B-A10B-FP8" engine = "vllm" param_count_b = 122.0 active_params_b = 10.0 diff --git a/src/openjarvis/evals/datasets/deepplanning.py b/src/openjarvis/evals/datasets/deepplanning.py new file mode 100644 index 00000000..6ab38dab --- /dev/null +++ b/src/openjarvis/evals/datasets/deepplanning.py @@ -0,0 +1,238 @@ +"""DeepPlanning: long-horizon planning with constraints. + +Evaluates agents on complex shopping tasks with hard constraints +(product attributes, ratings, stock, shipping). +Source: https://huggingface.co/datasets/Qwen/DeepPlanning +Paper: arXiv:2601.18137 + +The dataset contains shopping planning tasks at 3 difficulty levels +(120 total cases). Each case has a natural-language query, product +catalog, and ground-truth product selections with constraint metadata. +""" + +from __future__ import annotations + +import json +import logging +import random +import tarfile +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + +from openjarvis.evals.core.dataset import DatasetProvider +from openjarvis.evals.core.types import EvalRecord + +logger = logging.getLogger(__name__) + +_SYSTEM_PROMPT = ( + "You are a shopping assistant. Given the user's request and a product catalog, " + "select the correct products that satisfy ALL stated constraints.\n\n" + "IMPORTANT INSTRUCTIONS:\n" + "- Read through the product catalog data directly — do NOT write code to parse it\n" + "- For each constraint, find products matching ALL requirements\n" + "- If a constraint references data not in the catalog (e.g., transport time " + "when only shipping provider is listed), use reasonable inference from " + "available shipping info\n" + "- For each selected product, state: name, brand, price, and the specific " + "attribute values that match each constraint\n" + "- Present your final answer as a clear list of selected products" +) + + +class DeepPlanningDataset(DatasetProvider): + """DeepPlanning long-horizon planning benchmark. + + Extracts shopping planning tasks from Qwen/DeepPlanning tar.gz archives. + Each case has a query with constraints and ground-truth product selections. + """ + + dataset_id = "deepplanning" + dataset_name = "DeepPlanning" + + def __init__( + self, + cache_dir: Optional[str] = None, + ) -> None: + self._cache_dir = cache_dir + self._records: List[EvalRecord] = [] + + def load( + self, + *, + max_samples: Optional[int] = None, + split: Optional[str] = None, + seed: Optional[int] = None, + ) -> None: + snapshot_dir = self._find_snapshot_dir() + if snapshot_dir is None: + self._download_dataset() + snapshot_dir = self._find_snapshot_dir() + if snapshot_dir is None: + logger.error("Failed to download DeepPlanning dataset") + return + + records: List[EvalRecord] = [] + for level in [1, 2, 3]: + tar_path = snapshot_dir / f"database_level{level}.tar.gz" + if not tar_path.exists(): + logger.warning("Missing %s", tar_path) + continue + level_records = self._extract_cases(tar_path, level) + records.extend(level_records) + + logger.info("Loaded %d DeepPlanning cases", len(records)) + + if seed is not None: + random.Random(seed).shuffle(records) + if max_samples is not None: + records = records[:max_samples] + + self._records = records + + def iter_records(self) -> Iterable[EvalRecord]: + return iter(self._records) + + def size(self) -> int: + return len(self._records) + + def _find_snapshot_dir(self) -> Optional[Path]: + """Find the HF cache snapshot directory for Qwen/DeepPlanning.""" + base = Path.home() / ".cache" / "huggingface" / "hub" + ds_dir = base / "datasets--Qwen--DeepPlanning" / "snapshots" + if not ds_dir.exists(): + return None + snapshots = list(ds_dir.iterdir()) + return snapshots[0] if snapshots else None + + def _download_dataset(self) -> None: + """Trigger HF datasets download to populate the cache.""" + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "datasets required for DeepPlanning. " + "Install with: pip install datasets" + ) + logger.info("Downloading Qwen/DeepPlanning from HuggingFace...") + # Loading triggers the download even though we don't use the result + load_dataset("Qwen/DeepPlanning", split="train") + + def _extract_cases( + self, tar_path: Path, level: int, + ) -> List[EvalRecord]: + """Extract shopping cases from a tar.gz archive.""" + records: List[EvalRecord] = [] + try: + with tarfile.open(tar_path, "r:gz") as tf: + # Build index of all members by directory + members_by_dir: Dict[str, Dict[str, Any]] = {} + for m in tf.getmembers(): + parts = Path(m.name).parts + if len(parts) >= 2: + case_dir = parts[1] # e.g. "case_12" + fname = parts[-1] + members_by_dir.setdefault(case_dir, {})[fname] = m + + for case_name, files in members_by_dir.items(): + val_member = files.get("validation_cases.json") + prod_member = files.get("products.jsonl") + if val_member is None: + continue + + f = tf.extractfile(val_member) + if f is None: + continue + data = json.loads(f.read().decode("utf-8")) + + # Load product catalog + products_text = "" + if prod_member is not None: + pf = tf.extractfile(prod_member) + if pf is not None: + products_text = pf.read().decode("utf-8").strip() + + rec = self._case_to_record( + data, level, case_name, products_text, + ) + if rec: + records.append(rec) + except Exception: + logger.warning("Failed to read %s", tar_path, exc_info=True) + return records + + def _case_to_record( + self, + data: Dict[str, Any], + level: int, + case_name: str, + products_text: str = "", + ) -> Optional[EvalRecord]: + """Convert a validation_cases.json to an EvalRecord.""" + query = data.get("query", "") + if not query: + return None + + ground_truth = data.get("ground_truth_products", []) + meta_info = data.get("meta_info", []) + + # Build reference as structured summary + ref_parts = [] + for prod in ground_truth: + name = prod.get("name", "") + brand = prod.get("brand", "") + price = prod.get("price", "") + ref_parts.append(f"{brand} {name} (${price})") + reference = "; ".join(ref_parts) if ref_parts else "" + + # Count constraints + n_constraints = sum( + len(m.get("features", [])) for m in meta_info + ) + + difficulty = ( + "easy" if level == 1 + else "medium" if level == 2 + else "hard" + ) + + # Build product catalog section + catalog_section = "" + if products_text: + catalog_section = ( + f"## Product Catalog\n" + f"Below is the product database in JSONL format " + f"(one JSON object per line):\n\n" + f"```jsonl\n{products_text}\n```\n\n" + ) + + problem = ( + f"{_SYSTEM_PROMPT}\n\n" + f"## Shopping Request\n{query}\n\n" + f"{catalog_section}" + f"## Required Output\n" + f"List each product that satisfies ALL constraints. For each product:\n" + f"1. Product name and brand\n" + f"2. Price\n" + f"3. For each constraint, the matching attribute value from the catalog" + ) + + return EvalRecord( + record_id=f"deepplanning-L{level}-{case_name}", + problem=problem, + reference=reference, + category="agentic", + subject=f"shopping_L{level}", + metadata={ + "task_type": "shopping", + "level": level, + "difficulty": difficulty, + "case_name": case_name, + "n_products": len(ground_truth), + "n_constraints": n_constraints, + "ground_truth_products": ground_truth, + "meta_info": meta_info, + }, + ) + + +__all__ = ["DeepPlanningDataset"] diff --git a/src/openjarvis/evals/datasets/lifelong_agent.py b/src/openjarvis/evals/datasets/lifelong_agent.py index 98132ffe..38cae836 100644 --- a/src/openjarvis/evals/datasets/lifelong_agent.py +++ b/src/openjarvis/evals/datasets/lifelong_agent.py @@ -24,6 +24,12 @@ _SYSTEM_PROMPT = ( "Previous tasks in this session may have modified the environment." ) +_KG_SYSTEM = ( + "You are a knowledge graph reasoning agent. Answer the question by " + "analyzing the provided knowledge graph operations and entity information. " + "Reason through each step and provide the final answer directly." +) + class LifelongAgentDataset(DatasetProvider): """LifelongAgentBench sequential task learning benchmark.""" @@ -56,19 +62,27 @@ class LifelongAgentDataset(DatasetProvider): if not data_dir.exists(): self._download(data_dir) - task_sequences = self._load_task_sequences(data_dir) + # Try Parquet first (HuggingFace format), then JSON/JSONL + records = self._load_from_parquet(data_dir) + if not records: + task_sequences = self._load_task_sequences(data_dir) + if seed is not None: + random.Random(seed).shuffle(task_sequences) + if max_samples is not None: + task_sequences = task_sequences[:max_samples] + self._episodes = [] + self._records = [] + for seq in task_sequences: + episode = self._sequence_to_episode(seq) + self._episodes.append(episode) + self._records.extend(episode) + return if seed is not None: - random.Random(seed).shuffle(task_sequences) + random.Random(seed).shuffle(records) if max_samples is not None: - task_sequences = task_sequences[:max_samples] - - self._episodes = [] - self._records = [] - for seq in task_sequences: - episode = self._sequence_to_episode(seq) - self._episodes.append(episode) - self._records.extend(episode) + records = records[:max_samples] + self._records = records def iter_records(self) -> Iterable[EvalRecord]: return iter(self._records) @@ -89,15 +103,220 @@ class LifelongAgentDataset(DatasetProvider): ) from exc data_dir.mkdir(parents=True, exist_ok=True) snapshot_download( - repo_id="LifelongAgentBench/LifelongAgentBench", + repo_id="csyq/LifelongAgentBench", repo_type="dataset", local_dir=str(data_dir), ) + def _load_from_parquet(self, data_dir: Path) -> List[EvalRecord]: + """Load records from Parquet files (HuggingFace dataset format).""" + parquet_files = list(data_dir.rglob("*.parquet")) + if not parquet_files: + return [] + + try: + import pandas as pd + except ImportError: + logger.warning("pandas required for Parquet loading") + return [] + + records: List[EvalRecord] = [] + for pf in sorted(parquet_files): + subset_name = pf.parent.name + df = pd.read_parquet(pf) + logger.info("Loading %d rows from %s", len(df), pf) + + for _, row in df.iterrows(): + row_dict = row.to_dict() + if subset_name == "knowledge_graph": + rec = self._kg_row_to_record(row_dict) + elif subset_name == "db_bench": + rec = self._db_row_to_record(row_dict) + elif subset_name == "os_interaction": + rec = self._os_row_to_record(row_dict) + else: + rec = self._generic_row_to_record(row_dict, subset_name) + if rec is not None: + records.append(rec) + + return records + + def _kg_row_to_record(self, row: Dict[str, Any]) -> Optional[EvalRecord]: + question = row.get("question", "") + if not question: + return None + + qid = row.get("qid", row.get("sample_index", "unknown")) + entity_dict = row.get("entity_dict", {}) + if isinstance(entity_dict, str): + try: + entity_dict = json.loads(entity_dict) + except (json.JSONDecodeError, TypeError): + entity_dict = {} + + entity_desc = "" + if isinstance(entity_dict, dict) and entity_dict: + entities = "\n".join( + f" - {name}: {mid}" for name, mid in entity_dict.items() + ) + entity_desc = f"Entities:\n{entities}" + + action_list = row.get("action_list", []) + action_desc = "" + if action_list: + if isinstance(action_list, str): + try: + action_list = json.loads(action_list) + except (json.JSONDecodeError, TypeError): + action_list = [] + if isinstance(action_list, list) and action_list: + steps = "\n".join(f" {i+1}. {a}" for i, a in enumerate(action_list)) + action_desc = f"\n\nKG Operations:\n{steps}" + + s_expr = row.get("s_expression", "") + s_expr_desc = "" + if s_expr: + s_expr_desc = f"\n\nStructured Query:\n{s_expr}" + + problem = ( + f"{_KG_SYSTEM}\n\n" + f"{entity_desc}" + f"{action_desc}" + f"{s_expr_desc}\n\n" + f"Question: {question}" + ) + + answer_list = row.get("answer_list", []) + if isinstance(answer_list, str): + try: + answer_list = json.loads(answer_list) + except (json.JSONDecodeError, TypeError): + answer_list = [answer_list] + reference = ", ".join(str(a) for a in answer_list) if answer_list else "" + + return EvalRecord( + record_id=f"lifelong-kg-{qid}", + problem=problem, + reference=reference, + category="agentic", + subject="knowledge_graph", + metadata={ + "subset": "knowledge_graph", + "qid": qid, + "action_list": action_list, + "skill_list": row.get("skill_list", []), + }, + ) + + def _db_row_to_record(self, row: Dict[str, Any]) -> Optional[EvalRecord]: + instruction = row.get("instruction", "") + if not instruction: + return None + + idx = row.get("sample_index", "unknown") + table_info = row.get("table_info", {}) + if isinstance(table_info, str): + try: + table_info = json.loads(table_info) + except (json.JSONDecodeError, TypeError): + table_info = {} + + table_ctx = "" + if isinstance(table_info, dict) and table_info: + tname = table_info.get("name", "unknown") + cols = table_info.get("column_info_list", []) + if cols: + col_desc = ", ".join( + f"{c.get('name', '?')} ({c.get('type', '?')})" for c in cols + ) + table_ctx = f"\n\nTable: {tname}\nColumns: {col_desc}" + + problem = f"{_SYSTEM_PROMPT}\n\n## Task\n{instruction}{table_ctx}" + + answer_info = row.get("answer_info", {}) + if isinstance(answer_info, str): + try: + answer_info = json.loads(answer_info) + except (json.JSONDecodeError, TypeError): + answer_info = {} + + reference = "" + if isinstance(answer_info, dict): + if answer_info.get("direct") is not None: + reference = str(answer_info["direct"]) + elif answer_info.get("sql"): + reference = answer_info["sql"] + elif answer_info.get("md5"): + reference = answer_info["md5"] + + return EvalRecord( + record_id=f"lifelong-db-{idx}", + problem=problem, + reference=reference, + category="agentic", + subject="database", + metadata={ + "subset": "db_bench", + "sample_index": idx, + "skill_list": row.get("skill_list", []), + }, + ) + + def _os_row_to_record(self, row: Dict[str, Any]) -> Optional[EvalRecord]: + instruction = row.get("instruction", "") + if not instruction: + return None + + idx = row.get("sample_index", "unknown") + problem = f"{_SYSTEM_PROMPT}\n\n## Task\n{instruction}" + + eval_info = row.get("evaluation_info", {}) + if isinstance(eval_info, str): + try: + eval_info = json.loads(eval_info) + except (json.JSONDecodeError, TypeError): + eval_info = {} + + reference = "" + if isinstance(eval_info, dict): + cmd_item = eval_info.get("evaluation_command_item", {}) + if isinstance(cmd_item, dict): + reference = cmd_item.get("script", "") + + return EvalRecord( + record_id=f"lifelong-os-{idx}", + problem=problem, + reference=reference, + category="agentic", + subject="os", + metadata={ + "subset": "os_interaction", + "sample_index": idx, + "skill_list": row.get("skill_list", []), + }, + ) + + def _generic_row_to_record( + self, row: Dict[str, Any], subset_name: str, + ) -> Optional[EvalRecord]: + instruction = row.get("instruction", row.get("question", row.get("task", ""))) + if not instruction: + return None + idx = row.get("sample_index", row.get("id", "unknown")) + expected = row.get("answer", row.get("expected_output", "")) + return EvalRecord( + record_id=f"lifelong-{subset_name}-{idx}", + problem=f"{_SYSTEM_PROMPT}\n\n## Task\n{instruction}", + reference=str(expected), + category="agentic", + subject=subset_name, + metadata={"subset": subset_name, "sample_index": idx}, + ) + def _load_task_sequences( self, data_dir: Path, ) -> List[List[Dict[str, Any]]]: - """Load task sequences from disk.""" + """Load task sequences from JSON/JSONL files (legacy format).""" sequences: List[List[Dict[str, Any]]] = [] for p in sorted(data_dir.rglob("*.json")): @@ -105,7 +324,6 @@ class LifelongAgentDataset(DatasetProvider): with open(p) as f: data = json.load(f) if isinstance(data, list): - # Could be a sequence of tasks or list of sequences if data and isinstance(data[0], list): sequences.extend(data) elif data and isinstance(data[0], dict): diff --git a/src/openjarvis/evals/datasets/paperarena.py b/src/openjarvis/evals/datasets/paperarena.py new file mode 100644 index 00000000..786bad8a --- /dev/null +++ b/src/openjarvis/evals/datasets/paperarena.py @@ -0,0 +1,170 @@ +"""PaperArena: scientific literature reasoning benchmark. + +Evaluates agents on research paper comprehension with three question types: +MC (multiple choice), CA (closed answer), OA (open answer) across +easy/medium/hard difficulty. +Source: https://github.com/Melmaphother/PaperArena +Paper: arXiv:2510.10909 +""" + +from __future__ import annotations + +import json +import logging +import random +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + +from openjarvis.evals.core.dataset import DatasetProvider +from openjarvis.evals.core.types import EvalRecord + +logger = logging.getLogger(__name__) + +_SYSTEM_PROMPT = ( + "You are a scientific research assistant. Read the paper context carefully " + "and answer the question. For multiple-choice questions, respond with just " + "the letter (A, B, C, or D). For other questions, provide a precise answer." +) + + +class PaperArenaDataset(DatasetProvider): + """PaperArena scientific literature reasoning benchmark. + + Three question types (MC, CA, OA) across three difficulty levels. + """ + + dataset_id = "paperarena" + dataset_name = "PaperArena" + + def __init__( + self, + cache_dir: Optional[str] = None, + ) -> None: + self._cache_dir = ( + Path(cache_dir) if cache_dir + else Path.home() / ".cache" / "paperarena" + ) + self._records: List[EvalRecord] = [] + + def load( + self, + *, + max_samples: Optional[int] = None, + split: Optional[str] = None, + seed: Optional[int] = None, + ) -> None: + data_dir = self._cache_dir / "qa" + + if not data_dir.exists(): + self._download() + + records = self._load_records(data_dir) + + if seed is not None: + random.Random(seed).shuffle(records) + if max_samples is not None: + records = records[:max_samples] + + self._records = records + + def iter_records(self) -> Iterable[EvalRecord]: + return iter(self._records) + + def size(self) -> int: + return len(self._records) + + def _download(self) -> None: + """Download PaperArena data from HuggingFace or GitHub.""" + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise ImportError( + "huggingface_hub required for PaperArena. " + "Install with: pip install huggingface_hub" + ) from exc + + self._cache_dir.mkdir(parents=True, exist_ok=True) + logger.info("Downloading PaperArena dataset...") + snapshot_download( + repo_id="Melmaphother/PaperArena-Data", + repo_type="dataset", + local_dir=str(self._cache_dir), + ) + + def _load_records(self, data_dir: Path) -> List[EvalRecord]: + """Load QA records from JSON/JSONL files.""" + records: List[EvalRecord] = [] + + for p in sorted(data_dir.rglob("*.json")): + try: + with open(p) as f: + data = json.load(f) + items = data if isinstance(data, list) else [data] + for item in items: + rec = self._item_to_record(item) + if rec is not None: + records.append(rec) + except (json.JSONDecodeError, OSError): + logger.debug("Skipping: %s", p) + + for p in sorted(data_dir.rglob("*.jsonl")): + try: + with open(p) as f: + for line in f: + line = line.strip() + if line: + item = json.loads(line) + rec = self._item_to_record(item) + if rec is not None: + records.append(rec) + except (json.JSONDecodeError, OSError): + logger.debug("Skipping: %s", p) + + return records + + def _item_to_record(self, item: Dict[str, Any]) -> Optional[EvalRecord]: + """Convert a raw QA item to an EvalRecord.""" + question_id = item.get("question_id", item.get("id", "")) + question_type = item.get("question_type", item.get("type", "OA")).upper() + difficulty = item.get("difficulty", "medium").lower() + question = item.get("question", "") + context = item.get("context", item.get("paper_context", "")) + reference = item.get("answer", item.get("reference", "")) + options = item.get("options", None) + tool_chain = item.get("minimal_tool_chain", []) + + if not question: + return None + + # Build problem prompt + prompt = f"{_SYSTEM_PROMPT}\n\n" + if context: + prompt += f"## Paper Context\n{context}\n\n" + prompt += f"## Question\n{question}" + + if options and question_type == "MC": + options_text = "\n".join( + f" {k}) {v}" for k, v in options.items() + ) if isinstance(options, dict) else "\n".join( + f" {chr(65 + i)}) {o}" for i, o in enumerate(options) + ) + prompt += f"\n\nOptions:\n{options_text}" + + subject = f"{difficulty}_{question_type.lower()}" + + return EvalRecord( + record_id=f"paperarena-{question_id}", + problem=prompt, + reference=str(reference), + category="agentic", + subject=subject, + metadata={ + "question_id": question_id, + "question_type": question_type, + "difficulty": difficulty, + "minimal_tool_chain": tool_chain, + }, + ) + + +__all__ = ["PaperArenaDataset"] diff --git a/src/openjarvis/evals/scorers/deepplanning_scorer.py b/src/openjarvis/evals/scorers/deepplanning_scorer.py new file mode 100644 index 00000000..24759370 --- /dev/null +++ b/src/openjarvis/evals/scorers/deepplanning_scorer.py @@ -0,0 +1,87 @@ +"""LLM-judge scorer for DeepPlanning constraint satisfaction.""" + +from __future__ import annotations + +import re +from typing import Any, Dict, Optional, Tuple + +from openjarvis.evals.core.scorer import LLMJudgeScorer +from openjarvis.evals.core.types import EvalRecord + +_JUDGE_PROMPT = """You are evaluating a planning task response. + +## Task Type: {task_type} + +## Original Task +{problem} + +## Reference Plan +{reference} + +## Agent's Plan +{model_answer} + +Evaluate whether the agent's plan satisfies ALL constraints from the task. + +For travel planning, check: +- Route consistency and feasibility +- Time feasibility (business hours, travel times) +- Budget accuracy (costs within stated limits) +- Personalization requirements met + +For shopping tasks, check: +- Correct products selected +- Coupon/discount rules applied correctly +- Cart total accuracy + +Respond with exactly: CORRECT or INCORRECT +Then provide a brief explanation.""" + + +class DeepPlanningScorer(LLMJudgeScorer): + """Score DeepPlanning responses via LLM judge on constraint satisfaction.""" + + scorer_id = "deepplanning" + + def score( + self, record: EvalRecord, model_answer: str, + ) -> Tuple[Optional[bool], Dict[str, Any]]: + if not model_answer or not model_answer.strip(): + return False, {"reason": "empty_response"} + + if not record.reference or not record.reference.strip(): + return None, {"reason": "no_ground_truth"} + + task_type = record.metadata.get("task_type", "planning") + + # Extract problem without system prompt prefix + problem = record.problem + if "## Task" in problem: + problem = problem.split("## Task", 1)[-1] + + prompt = _JUDGE_PROMPT.format( + task_type=task_type, + problem=problem, + reference=record.reference, + model_answer=model_answer, + ) + + try: + raw = self._ask_judge(prompt, max_tokens=4096) + is_correct = bool(re.search(r"\bCORRECT\b", raw, re.IGNORECASE)) + if re.search(r"\bINCORRECT\b", raw, re.IGNORECASE): + is_correct = False + + return is_correct, { + "match_type": "llm_judge", + "raw_judge_output": raw, + "task_type": task_type, + } + except Exception as exc: + return False, { + "match_type": "llm_judge_error", + "error": str(exc), + } + + +__all__ = ["DeepPlanningScorer"] diff --git a/src/openjarvis/evals/scorers/paperarena_judge.py b/src/openjarvis/evals/scorers/paperarena_judge.py new file mode 100644 index 00000000..be06d548 --- /dev/null +++ b/src/openjarvis/evals/scorers/paperarena_judge.py @@ -0,0 +1,153 @@ +"""Scorer for PaperArena: MC exact match + LLM judge for CA/OA.""" + +from __future__ import annotations + +import re +from typing import Any, Dict, Optional, Tuple + +from openjarvis.evals.core.scorer import LLMJudgeScorer +from openjarvis.evals.core.types import EvalRecord + +_JUDGE_PROMPT = """You are evaluating a scientific question answer. + +## Question +{question} + +## Reference Answer +{reference} + +## Agent's Answer +{model_answer} + +Is the agent's answer correct? Consider semantic equivalence, not exact wording. +For numerical answers, accept reasonable rounding. +Respond with exactly: CORRECT or INCORRECT""" + + +class PaperArenaScorer(LLMJudgeScorer): + """Score PaperArena: MC via letter extraction, CA/OA via LLM judge.""" + + scorer_id = "paperarena" + + def score( + self, record: EvalRecord, model_answer: str, + ) -> Tuple[Optional[bool], Dict[str, Any]]: + if not model_answer or not model_answer.strip(): + return False, {"reason": "empty_response"} + + if not record.reference or not record.reference.strip(): + return None, {"reason": "no_ground_truth"} + + question_type = record.metadata.get("question_type", "OA").upper() + + if question_type == "MC": + return self._score_mc(record, model_answer) + else: + return self._score_open(record, model_answer) + + def _score_mc( + self, record: EvalRecord, model_answer: str, + ) -> Tuple[Optional[bool], Dict[str, Any]]: + """Score multiple-choice via letter extraction.""" + ref_letter = record.reference.strip().upper() + if len(ref_letter) != 1 or ref_letter not in "ABCD": + # Reference is not a single letter; fall back to judge + return self._score_open(record, model_answer) + + # Try regex extraction for answer letter + extracted = self._extract_letter(model_answer) + + if extracted is None: + # Fall back to LLM extraction + extracted = self._extract_letter_with_llm( + record.problem, model_answer, + ) + + if extracted is None: + return None, { + "match_type": "exact_letter", + "reason": "no_letter_extracted", + } + + is_correct = extracted == ref_letter + return is_correct, { + "match_type": "exact_letter", + "reference_letter": ref_letter, + "candidate_letter": extracted, + } + + def _extract_letter(self, text: str) -> Optional[str]: + """Try to extract a single answer letter from text via regex.""" + # Common patterns: "The answer is A", "A)", "(A)", just "A" + patterns = [ + r"(?:the\s+answer\s+is|answer:?)\s*([A-D])\b", + r"\b([A-D])\)", + r"\(([A-D])\)", + ] + for pat in patterns: + m = re.search(pat, text, re.IGNORECASE) + if m: + return m.group(1).upper() + + # Last single capital letter on its own line + lines = text.strip().splitlines() + for line in reversed(lines): + stripped = line.strip() + if len(stripped) == 1 and stripped.upper() in "ABCD": + return stripped.upper() + + return None + + def _extract_letter_with_llm( + self, problem: str, model_answer: str, + ) -> Optional[str]: + """Use LLM to extract answer letter.""" + prompt = ( + f"Extract the final answer letter (A, B, C, or D) from this response.\n" + f"Problem: {problem[:2000]}\n" + f"Response: {model_answer[:2000]}\n\n" + f"Return ONLY a single letter (A-D), or NONE if no answer found." + ) + try: + raw = self._ask_judge(prompt, max_tokens=4096) + letter = raw.strip().upper() + m = re.search(r"([A-D])", letter) + if m: + return m.group(1) + except Exception: + pass + return None + + def _score_open( + self, record: EvalRecord, model_answer: str, + ) -> Tuple[Optional[bool], Dict[str, Any]]: + """Score CA/OA via LLM judge.""" + question = record.problem + if "## Question" in question: + question = question.split("## Question")[-1].strip() + + prompt = _JUDGE_PROMPT.format( + question=question, + reference=record.reference, + model_answer=model_answer, + ) + + try: + raw = self._ask_judge(prompt, max_tokens=4096) + is_correct = bool(re.search(r"\bCORRECT\b", raw, re.IGNORECASE)) + if re.search(r"\bINCORRECT\b", raw, re.IGNORECASE): + is_correct = False + + return is_correct, { + "match_type": "llm_judge", + "raw_judge_output": raw, + "question_type": record.metadata.get("question_type", "OA"), + } + except Exception as exc: + return False, { + "match_type": "llm_judge_error", + "error": str(exc), + } + + +__all__ = ["PaperArenaScorer"] diff --git a/src/openjarvis/learning/__init__.py b/src/openjarvis/learning/__init__.py index 838a91db..d9c68a8e 100644 --- a/src/openjarvis/learning/__init__.py +++ b/src/openjarvis/learning/__init__.py @@ -11,6 +11,9 @@ from openjarvis.learning._stubs import ( from openjarvis.learning.agent_evolver import AgentConfigEvolver from openjarvis.learning.heuristic_reward import HeuristicRewardFunction from openjarvis.learning.learning_orchestrator import LearningOrchestrator +from openjarvis.learning.optimize.llm_optimizer import LLMOptimizer +from openjarvis.learning.optimize.optimizer import OptimizationEngine +from openjarvis.learning.optimize.store import OptimizationStore from openjarvis.learning.router import ( HeuristicRouter, build_routing_context, @@ -81,9 +84,12 @@ __all__ = [ "HAS_TORCH", "HeuristicRewardFunction", "HeuristicRouter", + "LLMOptimizer", "LearningOrchestrator", "LoRATrainer", "LoRATrainingConfig", + "OptimizationEngine", + "OptimizationStore", "QueryAnalyzer", "RewardFunction", "RouterPolicy", diff --git a/src/openjarvis/learning/optimize/__init__.py b/src/openjarvis/learning/optimize/__init__.py new file mode 100644 index 00000000..ddfa3e5c --- /dev/null +++ b/src/openjarvis/learning/optimize/__init__.py @@ -0,0 +1,61 @@ +"""Optimization framework for OpenJarvis configuration tuning.""" + +from openjarvis.learning.optimize.config import ( + load_benchmark_specs, + load_objectives, + load_optimize_config, +) +from openjarvis.learning.optimize.llm_optimizer import LLMOptimizer +from openjarvis.learning.optimize.optimizer import ( + OptimizationEngine, + compute_pareto_frontier, +) +from openjarvis.learning.optimize.search_space import ( + DEFAULT_SEARCH_SPACE, + build_search_space, +) +from openjarvis.learning.optimize.store import OptimizationStore +from openjarvis.learning.optimize.trial_runner import ( + BenchmarkSpec, + MultiBenchTrialRunner, + TrialRunner, +) +from openjarvis.learning.optimize.types import ( + ALL_OBJECTIVES, + DEFAULT_OBJECTIVES, + BenchmarkScore, + ObjectiveSpec, + OptimizationRun, + SampleScore, + SearchDimension, + SearchSpace, + TrialConfig, + TrialFeedback, + TrialResult, +) + +__all__ = [ + "ALL_OBJECTIVES", + "BenchmarkScore", + "BenchmarkSpec", + "DEFAULT_OBJECTIVES", + "DEFAULT_SEARCH_SPACE", + "LLMOptimizer", + "MultiBenchTrialRunner", + "ObjectiveSpec", + "OptimizationEngine", + "OptimizationRun", + "OptimizationStore", + "SampleScore", + "SearchDimension", + "SearchSpace", + "TrialConfig", + "TrialFeedback", + "TrialResult", + "TrialRunner", + "build_search_space", + "compute_pareto_frontier", + "load_benchmark_specs", + "load_objectives", + "load_optimize_config", +] diff --git a/src/openjarvis/learning/optimize/config.py b/src/openjarvis/learning/optimize/config.py new file mode 100644 index 00000000..1ab59bec --- /dev/null +++ b/src/openjarvis/learning/optimize/config.py @@ -0,0 +1,102 @@ +"""TOML config loader for optimization runs.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, List, Union + +from openjarvis.learning.optimize.types import ObjectiveSpec + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover + import tomli as tomllib # type: ignore[no-redef] + + +def load_optimize_config(path: Union[str, Path]) -> Dict[str, Any]: + """Load an optimization config TOML file. + + Returns the raw dict with keys such as ``optimize.max_trials``, + ``optimize.benchmark``, ``optimize.search``, ``optimize.fixed``, + ``optimize.constraints``, etc. + + Raises: + FileNotFoundError: If *path* does not exist. + """ + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Optimization config not found: {path}") + + with open(path, "rb") as fh: + data: Dict[str, Any] = tomllib.load(fh) + + return data + + +def load_objectives(data: Dict[str, Any]) -> List[ObjectiveSpec]: + """Extract objectives from a loaded optimization config. + + Reads ``optimize.objectives`` (a list of tables) and returns + a list of :class:`ObjectiveSpec`. Falls back to + :data:`DEFAULT_OBJECTIVES` if the key is absent. + """ + from openjarvis.learning.optimize.types import DEFAULT_OBJECTIVES + + optimize = data.get("optimize", {}) + raw_objectives = optimize.get("objectives") + if not raw_objectives: + return list(DEFAULT_OBJECTIVES) + + result: List[ObjectiveSpec] = [] + for obj in raw_objectives: + result.append( + ObjectiveSpec( + metric=obj["metric"], + direction=obj.get("direction", "maximize"), + weight=obj.get("weight", 1.0), + ) + ) + return result + + +def load_benchmark_specs(data: Dict[str, Any]) -> List[Any]: + """Extract benchmark specs from a loaded optimization config. + + Supports two formats: + - Multi-benchmark: ``[[optimize.benchmarks]]`` array of tables + - Single-benchmark fallback: ``optimize.benchmark`` string + + Returns a list of :class:`BenchmarkSpec` (from trial_runner). + Returns an empty list if no benchmarks are configured (caller + should fall back to CLI --benchmark). + """ + from openjarvis.learning.optimize.trial_runner import BenchmarkSpec + + optimize = data.get("optimize", {}) + + # Multi-benchmark format: [[optimize.benchmarks]] + raw_benchmarks = optimize.get("benchmarks") + if raw_benchmarks and isinstance(raw_benchmarks, list): + # Check if it's a list of dicts (table array) vs list of strings + if raw_benchmarks and isinstance(raw_benchmarks[0], dict): + specs: List[BenchmarkSpec] = [] + for entry in raw_benchmarks: + specs.append( + BenchmarkSpec( + benchmark=entry.get("name", entry.get("benchmark", "")), + max_samples=entry.get("max_samples", 200), + weight=entry.get("weight", 1.0), + ) + ) + return specs + + # Single-benchmark fallback + single = optimize.get("benchmark", "") + if single: + max_samples = optimize.get("max_samples", 50) + return [BenchmarkSpec(benchmark=single, max_samples=max_samples)] + + return [] + + +__all__ = ["load_benchmark_specs", "load_objectives", "load_optimize_config"] diff --git a/src/openjarvis/optimize/configs/qwen3-235b-dry-run.toml b/src/openjarvis/learning/optimize/configs/qwen3-235b-dry-run.toml similarity index 100% rename from src/openjarvis/optimize/configs/qwen3-235b-dry-run.toml rename to src/openjarvis/learning/optimize/configs/qwen3-235b-dry-run.toml diff --git a/src/openjarvis/optimize/configs/qwen3.5-joint-agentic.toml b/src/openjarvis/learning/optimize/configs/qwen3.5-joint-agentic.toml similarity index 100% rename from src/openjarvis/optimize/configs/qwen3.5-joint-agentic.toml rename to src/openjarvis/learning/optimize/configs/qwen3.5-joint-agentic.toml diff --git a/src/openjarvis/learning/optimize/feedback/__init__.py b/src/openjarvis/learning/optimize/feedback/__init__.py new file mode 100644 index 00000000..f1f57d6b --- /dev/null +++ b/src/openjarvis/learning/optimize/feedback/__init__.py @@ -0,0 +1,6 @@ +"""Feedback subsystem: LLM-as-judge scoring and signal aggregation.""" + +from openjarvis.learning.optimize.feedback.collector import FeedbackCollector +from openjarvis.learning.optimize.feedback.judge import TraceJudge + +__all__ = ["TraceJudge", "FeedbackCollector"] diff --git a/src/openjarvis/learning/optimize/feedback/collector.py b/src/openjarvis/learning/optimize/feedback/collector.py new file mode 100644 index 00000000..8a4f53b1 --- /dev/null +++ b/src/openjarvis/learning/optimize/feedback/collector.py @@ -0,0 +1,115 @@ +"""FeedbackCollector -- aggregates feedback from multiple sources.""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional + +from openjarvis.core.types import Trace +from openjarvis.learning.optimize.feedback.judge import TraceJudge + + +class FeedbackCollector: + """Collects feedback signals: explicit user scores + LLM judge evaluations. + + Signals are stored in-memory as dicts with at least ``trace_id``, + ``score``, ``source``, and ``timestamp`` keys. + """ + + def __init__(self) -> None: + self._records: List[Dict[str, Any]] = [] + + # ------------------------------------------------------------------ + # Recording helpers + # ------------------------------------------------------------------ + + def record_explicit( + self, + trace_id: str, + score: float, + source: str = "api", + ) -> None: + """Record an explicit numeric score (0-1) for a trace.""" + self._records.append({ + "trace_id": trace_id, + "score": min(max(score, 0.0), 1.0), + "source": source, + "timestamp": time.time(), + }) + + def record_thumbs(self, trace_id: str, thumbs_up: bool) -> None: + """Record a thumbs-up / thumbs-down signal (converted to 1.0/0.0).""" + self._records.append({ + "trace_id": trace_id, + "score": 1.0 if thumbs_up else 0.0, + "source": "thumbs", + "timestamp": time.time(), + }) + + # ------------------------------------------------------------------ + # Judge-driven evaluation + # ------------------------------------------------------------------ + + def evaluate_traces( + self, + traces: List[Trace], + judge: TraceJudge, + ) -> List[Dict[str, Any]]: + """Score *traces* via the LLM judge and record the results. + + Returns the list of newly created feedback records. + """ + new_records: List[Dict[str, Any]] = [] + for trace in traces: + score, feedback = judge.score_trace(trace) + record: Dict[str, Any] = { + "trace_id": trace.trace_id, + "score": score, + "source": "judge", + "feedback": feedback, + "timestamp": time.time(), + } + self._records.append(record) + new_records.append(record) + return new_records + + # ------------------------------------------------------------------ + # Queries + # ------------------------------------------------------------------ + + def get_records( + self, trace_id: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Return stored records, optionally filtered by *trace_id*.""" + if trace_id is None: + return list(self._records) + return [r for r in self._records if r["trace_id"] == trace_id] + + def stats(self) -> Dict[str, Any]: + """Return aggregate statistics over all recorded feedback. + + Returns a dict with ``count``, ``mean_score``, and a simple + ``distribution`` bucket (low / medium / high). + """ + count = len(self._records) + if count == 0: + return { + "count": 0, + "mean_score": 0.0, + "distribution": {"low": 0, "medium": 0, "high": 0}, + } + + scores = [r["score"] for r in self._records] + mean_score = sum(scores) / count + low = sum(1 for s in scores if s < 0.4) + medium = sum(1 for s in scores if 0.4 <= s < 0.7) + high = sum(1 for s in scores if s >= 0.7) + + return { + "count": count, + "mean_score": mean_score, + "distribution": {"low": low, "medium": medium, "high": high}, + } + + +__all__ = ["FeedbackCollector"] diff --git a/src/openjarvis/learning/optimize/feedback/judge.py b/src/openjarvis/learning/optimize/feedback/judge.py new file mode 100644 index 00000000..24012766 --- /dev/null +++ b/src/openjarvis/learning/optimize/feedback/judge.py @@ -0,0 +1,118 @@ +"""TraceJudge -- LLM-as-judge scoring for agent traces.""" + +from __future__ import annotations + +import logging +import re +from typing import List, Tuple + +from openjarvis.core.types import Trace +from openjarvis.evals.core.backend import InferenceBackend + +LOGGER = logging.getLogger(__name__) + +_SYSTEM_PROMPT = ( + "You are an expert evaluator of AI assistant traces. " + "You will be shown a user query, the steps the assistant took, " + "and its final result. Rate the overall quality of the response " + "on a scale from 0.0 (completely wrong / unhelpful) to 1.0 " + "(perfect). Provide your score on the first line as a decimal " + "number, then explain your reasoning." +) + +_SCORE_RE = re.compile( + r"(?:Score|Rating|Quality)?\s*[:=]?\s*(\d+(?:\.\d+)?)" + r"(?:\s*/\s*(\d+(?:\.\d+)?))?", + re.IGNORECASE, +) + + +def _format_trace(trace: Trace) -> str: + """Render a Trace into a textual prompt for the judge.""" + lines: List[str] = [] + lines.append(f"## Query\n{trace.query}") + if trace.steps: + lines.append("\n## Steps") + for i, step in enumerate(trace.steps, 1): + step_input = step.input.get("content", str(step.input)) + step_output = step.output.get("content", str(step.output)) + lines.append( + f"{i}. [{step.step_type.value}] " + f"input={step_input!r} output={step_output!r} " + f"({step.duration_seconds:.3f}s)", + ) + lines.append(f"\n## Final Result\n{trace.result}") + return "\n".join(lines) + + +def _parse_score(text: str) -> float: + """Extract a 0-1 score from the judge response. + + Handles formats like ``0.85``, ``Score: 0.85``, ``Rating: 7/10``. + Falls back to 0.5 when parsing fails. + """ + match = _SCORE_RE.search(text) + if match is None: + LOGGER.warning("Could not parse score from judge response; defaulting to 0.5") + return 0.5 + + numerator = float(match.group(1)) + denominator_str = match.group(2) + + if denominator_str is not None: + denominator = float(denominator_str) + if denominator > 0: + return min(max(numerator / denominator, 0.0), 1.0) + return 0.5 + + # If the number is > 1.0 assume it is on a 0-10 scale + if numerator > 1.0: + return min(numerator / 10.0, 1.0) + + return min(max(numerator, 0.0), 1.0) + + +class TraceJudge: + """LLM-as-judge for scoring traces when no ground truth exists. + + Given a :class:`Trace`, the judge constructs a prompt showing the + query, agent steps, and final result, then asks an LLM to rate the + quality on a 0-1 scale. + """ + + def __init__(self, backend: InferenceBackend, model: str) -> None: + self._backend = backend + self._model = model + + def score_trace(self, trace: Trace) -> Tuple[float, str]: + """Score a single trace. + + Returns: + ``(score, feedback)`` where *score* is in [0, 1] and + *feedback* is the judge's textual reasoning. + """ + prompt = _format_trace(trace) + response = self._backend.generate( + prompt, + model=self._model, + system=_SYSTEM_PROMPT, + temperature=0.0, + max_tokens=1024, + ) + score = _parse_score(response) + return score, response + + def batch_evaluate( + self, traces: List[Trace], + ) -> List[Tuple[float, str]]: + """Evaluate multiple traces sequentially. + + Returns a list of ``(score, feedback)`` tuples, one per trace. + """ + results: List[Tuple[float, str]] = [] + for trace in traces: + results.append(self.score_trace(trace)) + return results + + +__all__ = ["TraceJudge"] diff --git a/src/openjarvis/learning/optimize/llm_optimizer.py b/src/openjarvis/learning/optimize/llm_optimizer.py new file mode 100644 index 00000000..444c5a79 --- /dev/null +++ b/src/openjarvis/learning/optimize/llm_optimizer.py @@ -0,0 +1,691 @@ +"""LLM-based optimizer for OpenJarvis configuration tuning. + +Uses a cloud LLM to propose optimal OpenJarvis configs, inspired by DSPy's +GEPA approach: textual feedback from execution traces rather than just scalar +rewards guides the optimizer toward better configurations. +""" + +from __future__ import annotations + +import json +import re +import uuid +from typing import Any, Dict, List, Optional + +from openjarvis.core.types import Trace +from openjarvis.evals.core.backend import InferenceBackend +from openjarvis.evals.core.types import RunSummary +from openjarvis.learning.optimize.types import ( + BenchmarkScore, + SampleScore, + SearchSpace, + TrialConfig, + TrialFeedback, + TrialResult, +) + + +class LLMOptimizer: + """Uses a cloud LLM to propose optimal OpenJarvis configs. + + Inspired by DSPy's GEPA: uses textual feedback from execution + traces rather than just scalar rewards. + """ + + def __init__( + self, + search_space: SearchSpace, + optimizer_model: str = "claude-sonnet-4-6", + optimizer_backend: Optional[InferenceBackend] = None, + ) -> None: + self.search_space = search_space + self.optimizer_model = optimizer_model + self.optimizer_backend = optimizer_backend + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def propose_initial(self) -> TrialConfig: + """Propose a reasonable starting config from the search space.""" + if self.optimizer_backend is None: + raise ValueError( + "optimizer_backend is required to propose configurations" + ) + + prompt = self._build_initial_prompt() + response = self.optimizer_backend.generate( + prompt, + model=self.optimizer_model, + system="You are an expert AI systems optimizer.", + temperature=0.7, + max_tokens=2048, + ) + return self._parse_config_response(response) + + def propose_next( + self, + history: List[TrialResult], + traces: Optional[List[Trace]] = None, + frontier_ids: Optional[set] = None, + ) -> TrialConfig: + """Ask the LLM to propose the next config to evaluate.""" + if self.optimizer_backend is None: + raise ValueError( + "optimizer_backend is required to propose configurations" + ) + + prompt = self._build_propose_prompt(history, traces, frontier_ids=frontier_ids) + response = self.optimizer_backend.generate( + prompt, + model=self.optimizer_model, + system="You are an expert AI systems optimizer.", + temperature=0.7, + max_tokens=2048, + ) + return self._parse_config_response(response) + + def analyze_trial( + self, + trial: TrialConfig, + summary: RunSummary, + traces: Optional[List[Trace]] = None, + sample_scores: Optional[List[SampleScore]] = None, + per_benchmark: Optional[List[BenchmarkScore]] = None, + ) -> TrialFeedback: + """Ask the LLM to analyze a completed trial. Returns structured feedback.""" + if self.optimizer_backend is None: + raise ValueError( + "optimizer_backend is required to analyze trials" + ) + + prompt = self._build_analyze_prompt( + trial, summary, traces, sample_scores, per_benchmark, + ) + response = self.optimizer_backend.generate( + prompt, + model=self.optimizer_model, + system="You are an expert AI systems analyst.", + temperature=0.3, + max_tokens=2048, + ) + return self._parse_feedback_response(response) + + # ------------------------------------------------------------------ + # Prompt builders + # ------------------------------------------------------------------ + + def _build_initial_prompt(self) -> str: + """Construct the prompt for the initial config proposal.""" + lines: List[str] = [] + lines.append( + "You are optimizing an OpenJarvis AI system configuration." + ) + lines.append("") + lines.append(self.search_space.to_prompt_description()) + lines.append("## Objective") + lines.append( + "Maximize accuracy while minimizing latency and cost." + ) + lines.append("") + lines.append("## Your Task") + lines.append( + "Propose an initial configuration that is a reasonable starting " + "point for optimization. Choose sensible defaults that balance " + "accuracy, latency, and cost." + ) + lines.append("") + lines.append( + "Return a JSON object inside a ```json code block with:" + ) + lines.append( + '1. "params": dict of config params (dotted keys matching ' + "the search space)" + ) + lines.append( + '2. "reasoning": string explaining why this is a good ' + "starting configuration" + ) + return "\n".join(lines) + + def _build_propose_prompt( + self, + history: List[TrialResult], + traces: Optional[List[Trace]] = None, + frontier_ids: Optional[set] = None, + ) -> str: + """Construct the full prompt for propose_next.""" + lines: List[str] = [] + lines.append( + "You are optimizing an OpenJarvis AI system configuration." + ) + lines.append("") + lines.append(self.search_space.to_prompt_description()) + + lines.append("## Optimization History") + if history: + lines.append(self._format_history(history, frontier_ids=frontier_ids)) + else: + lines.append("No trials have been run yet.") + lines.append("") + + if traces: + lines.append("## Recent Execution Traces") + lines.append(self._format_traces(traces)) + lines.append("") + + lines.append("## Objective") + lines.append( + "Maximize accuracy while minimizing latency and cost." + ) + lines.append("") + lines.append("## Your Task") + lines.append( + "Propose the next configuration to evaluate. Learn from " + "previous trials to improve results." + ) + lines.append("") + lines.append( + "Return a JSON object inside a ```json code block with:" + ) + lines.append( + '1. "params": dict of config params (dotted keys matching ' + "the search space)" + ) + lines.append( + '2. "reasoning": string explaining why this config should ' + "improve results" + ) + return "\n".join(lines) + + def _build_analyze_prompt( + self, + trial: TrialConfig, + summary: RunSummary, + traces: Optional[List[Trace]] = None, + sample_scores: Optional[List[SampleScore]] = None, + per_benchmark: Optional[List[BenchmarkScore]] = None, + ) -> str: + """Construct the prompt for analyze_trial.""" + lines: List[str] = [] + lines.append("Analyze this OpenJarvis evaluation result.") + lines.append("") + + lines.append("## Configuration") + for key, value in sorted(trial.params.items()): + lines.append(f"- {key}: {value}") + if trial.reasoning: + lines.append(f"\nOptimizer reasoning: {trial.reasoning}") + lines.append("") + + # Per-benchmark breakdown (multi-benchmark mode) + if per_benchmark: + lines.append("## Per-Benchmark Results") + total_weight = sum(b.weight for b in per_benchmark) or 1.0 + for b in per_benchmark: + lines.append( + f"### {b.benchmark} (weight={b.weight:.1f}): " + f"accuracy={b.accuracy:.4f}, " + f"latency={b.mean_latency_seconds:.2f}s, " + f"cost=${b.total_cost_usd:.4f}, " + f"energy={b.total_energy_joules:.2f}J, " + f"samples={b.samples_evaluated}, errors={b.errors}" + ) + weighted_acc = ( + sum(b.accuracy * b.weight for b in per_benchmark) / total_weight + ) + lines.append(f"\nOverall weighted accuracy: {weighted_acc:.4f}") + lines.append("") + + lines.append("## Aggregate Results") + lines.append(f"- accuracy: {summary.accuracy:.4f}") + lines.append( + f"- mean_latency_seconds: {summary.mean_latency_seconds:.4f}" + ) + lines.append(f"- total_cost_usd: {summary.total_cost_usd:.4f}") + lines.append(f"- total_samples: {summary.total_samples}") + lines.append(f"- scored_samples: {summary.scored_samples}") + lines.append(f"- correct: {summary.correct}") + lines.append(f"- errors: {summary.errors}") + if summary.per_subject: + lines.append("\n### Per-Subject Breakdown") + for subject, metrics in sorted(summary.per_subject.items()): + metrics_str = ", ".join( + f"{k}={v:.3f}" for k, v in sorted(metrics.items()) + ) + lines.append(f"- {subject}: {metrics_str}") + lines.append("") + + if sample_scores: + lines.append("## Per-Sample Scores") + lines.append(self._format_sample_scores(sample_scores)) + lines.append("") + + if traces: + lines.append("## Sample Traces") + lines.append(self._format_traces(traces)) + lines.append("") + + lines.append( + "Provide your analysis as a JSON object inside a ```json code block with:\n" + '1. "summary_text": string with detailed analysis\n' + '2. "failure_patterns": list of identified failure patterns\n' + '3. "pillar_ratings": dict mapping pillar names to "high"/"medium"/"low"\n' + '4. "suggested_changes": list of specific config changes to try\n' + '5. "target_pillar": which pillar to change next ' + "(intelligence/engine/agent/tools/learning)" + ) + return "\n".join(lines) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _format_history( + self, + history: List[TrialResult], + frontier_ids: Optional[set] = None, + ) -> str: + """Render trial history as structured text for the LLM prompt.""" + lines: List[str] = [] + for i, result in enumerate(history, 1): + tag = "" + if frontier_ids and result.trial_id in frontier_ids: + tag = " [FRONTIER]" + lines.append(f"### Trial {i} (id={result.trial_id}){tag}") + lines.append(f"Params: {json.dumps(result.config.params)}") + lines.append(f"Accuracy: {result.accuracy:.4f}") + lines.append( + f"Latency: {result.mean_latency_seconds:.4f}s" + ) + lines.append(f"Cost: ${result.total_cost_usd:.4f}") + lines.append(f"Energy: {result.total_energy_joules:.4f}J") + if result.per_benchmark: + bench_parts = [ + f"{b.benchmark}={b.accuracy:.4f}" + for b in result.per_benchmark + ] + lines.append(f"Per-benchmark accuracy: {', '.join(bench_parts)}") + if result.summary: + s = result.summary + if s.throughput_stats: + lines.append( + f"Throughput: {s.throughput_stats.mean:.2f} tok/s" + ) + if s.ipw_stats: + lines.append(f"IPW: {s.ipw_stats.mean:.4f}") + if result.structured_feedback: + fb = result.structured_feedback + if fb.failure_patterns: + lines.append( + f"Failure patterns: {', '.join(fb.failure_patterns)}" + ) + if fb.pillar_ratings: + ratings = ", ".join( + f"{k}={v}" for k, v in sorted(fb.pillar_ratings.items()) + ) + lines.append(f"Pillar ratings: {ratings}") + if fb.target_pillar: + lines.append(f"Target pillar: {fb.target_pillar}") + elif result.analysis: + lines.append(f"Analysis: {result.analysis}") + if result.failure_modes: + lines.append( + f"Failure modes: {', '.join(result.failure_modes)}" + ) + lines.append("") + return "\n".join(lines) + + def _format_traces(self, traces: List[Trace]) -> str: + """Render traces as structured text for the LLM prompt. + + Limits to the last 10 traces and truncates long outputs to keep + the prompt manageable. + """ + max_traces = 10 + max_result_len = 500 + max_steps_per_trace = 10 + + recent = traces[-max_traces:] + lines: List[str] = [] + + for trace in recent: + lines.append( + f"### Trace {trace.trace_id} " + f"(agent={trace.agent}, model={trace.model})" + ) + lines.append(f"Query: {trace.query}") + if trace.outcome: + lines.append(f"Outcome: {trace.outcome}") + if trace.feedback is not None: + lines.append(f"Feedback: {trace.feedback}") + lines.append( + f"Latency: {trace.total_latency_seconds:.3f}s, " + f"Tokens: {trace.total_tokens}" + ) + + # Show steps (limited) + steps = trace.steps[:max_steps_per_trace] + if steps: + lines.append("Steps:") + for step in steps: + step_input = json.dumps(step.input) + step_output = json.dumps(step.output) + if len(step_input) > max_result_len: + step_input = step_input[:max_result_len] + "..." + if len(step_output) > max_result_len: + step_output = step_output[:max_result_len] + "..." + lines.append( + f" - {step.step_type.value}: " + f"input={step_input}, " + f"output={step_output} " + f"({step.duration_seconds:.3f}s)" + ) + if len(trace.steps) > max_steps_per_trace: + lines.append( + f" ... ({len(trace.steps) - max_steps_per_trace} " + "more steps)" + ) + + result_text = trace.result + if len(result_text) > max_result_len: + result_text = result_text[:max_result_len] + "..." + lines.append(f"Result: {result_text}") + lines.append("") + + return "\n".join(lines) + + def propose_targeted( + self, + history: List[TrialResult], + base_config: TrialConfig, + target_pillar: str, + frontier_ids: Optional[set] = None, + ) -> TrialConfig: + """Propose a config that only changes one pillar.""" + if self.optimizer_backend is None: + raise ValueError( + "optimizer_backend is required to propose configurations" + ) + + prompt = self._build_targeted_prompt( + history, base_config, target_pillar, frontier_ids, + ) + response = self.optimizer_backend.generate( + prompt, + model=self.optimizer_model, + system="You are an expert AI systems optimizer.", + temperature=0.7, + max_tokens=2048, + ) + proposed = self._parse_config_response(response) + + # Enforce constraint: preserve non-target params from base_config + merged_params = dict(base_config.params) + for key, value in proposed.params.items(): + if key.startswith(target_pillar + ".") or key.startswith( + target_pillar.rstrip("s") + "." + ): + merged_params[key] = value + proposed.params = merged_params + return proposed + + def propose_merge( + self, + candidates: List[TrialResult], + history: List[TrialResult], + frontier_ids: Optional[set] = None, + ) -> TrialConfig: + """Combine best aspects of frontier members into one config.""" + if self.optimizer_backend is None: + raise ValueError( + "optimizer_backend is required to propose configurations" + ) + + prompt = self._build_merge_prompt(candidates, history, frontier_ids) + response = self.optimizer_backend.generate( + prompt, + model=self.optimizer_model, + system="You are an expert AI systems optimizer.", + temperature=0.7, + max_tokens=2048, + ) + return self._parse_config_response(response) + + # ------------------------------------------------------------------ + # Targeted / Merge prompt builders + # ------------------------------------------------------------------ + + def _build_targeted_prompt( + self, + history: List[TrialResult], + base_config: TrialConfig, + target_pillar: str, + frontier_ids: Optional[set] = None, + ) -> str: + """Build prompt for pillar-targeted mutation.""" + lines: List[str] = [] + lines.append( + "You are optimizing an OpenJarvis AI system configuration." + ) + lines.append("") + lines.append(self.search_space.to_prompt_description()) + + lines.append("## Base Configuration") + for key, value in sorted(base_config.params.items()): + lines.append(f"- {key}: {value}") + lines.append("") + + lines.append(f"## Target Pillar: {target_pillar}") + lines.append( + f"ONLY change parameters under the '{target_pillar}' pillar. " + "Keep all other parameters exactly as they are." + ) + lines.append("") + + lines.append("## Optimization History") + if history: + lines.append(self._format_history(history, frontier_ids=frontier_ids)) + lines.append("") + + lines.append( + "Return a JSON object inside a ```json code block with:\n" + '1. "params": dict of config params (only change ' + f"{target_pillar} params)\n" + '2. "reasoning": string explaining your changes' + ) + return "\n".join(lines) + + def _build_merge_prompt( + self, + candidates: List[TrialResult], + history: List[TrialResult], + frontier_ids: Optional[set] = None, + ) -> str: + """Build prompt for merging frontier configs.""" + lines: List[str] = [] + lines.append( + "You are optimizing an OpenJarvis AI system configuration." + ) + lines.append("") + lines.append(self.search_space.to_prompt_description()) + + lines.append("## Frontier Candidates to Merge") + for i, cand in enumerate(candidates, 1): + lines.append(f"### Candidate {i} (id={cand.trial_id})") + lines.append(f"Params: {json.dumps(cand.config.params)}") + lines.append(f"Accuracy: {cand.accuracy:.4f}") + lines.append(f"Latency: {cand.mean_latency_seconds:.4f}s") + lines.append(f"Cost: ${cand.total_cost_usd:.4f}") + lines.append(f"Energy: {cand.total_energy_joules:.4f}J") + lines.append("") + + lines.append( + "Combine the best aspects of these frontier configs into " + "one unified configuration." + ) + lines.append("") + + if history: + lines.append("## Full History") + lines.append(self._format_history(history, frontier_ids=frontier_ids)) + lines.append("") + + lines.append( + "Return a JSON object inside a ```json code block with:\n" + '1. "params": dict of merged config params\n' + '2. "reasoning": string explaining the merge strategy' + ) + return "\n".join(lines) + + # ------------------------------------------------------------------ + # Sample score + feedback helpers + # ------------------------------------------------------------------ + + def _format_sample_scores(self, scores: List[SampleScore]) -> str: + """Render per-sample scores for the LLM prompt.""" + passed = [s for s in scores if s.is_correct] + failed = [s for s in scores if s.is_correct is False] + errored = [s for s in scores if s.error] + + lines: List[str] = [] + lines.append( + f"Total: {len(scores)} | Passed: {len(passed)} | " + f"Failed: {len(failed)} | Errors: {len(errored)}" + ) + + if failed: + lines.append("\n### Failed Samples") + for s in failed[:20]: + lines.append(f"- {s.record_id}: latency={s.latency_seconds:.2f}s") + + if errored: + lines.append("\n### Error Samples") + for s in errored[:10]: + error_text = (s.error or "")[:200] + lines.append(f"- {s.record_id}: {error_text}") + + return "\n".join(lines) + + def _parse_feedback_response(self, response: str) -> TrialFeedback: + """Parse LLM response into a TrialFeedback, with fallback.""" + response = response.strip() + + # Try JSON code block + json_block_match = re.search( + r"```json\s*\n?(.*?)\n?\s*```", response, re.DOTALL + ) + raw_json = None + if json_block_match: + raw_json = json_block_match.group(1).strip() + else: + # Try generic code block + code_block_match = re.search( + r"```\s*\n?(.*?)\n?\s*```", response, re.DOTALL + ) + if code_block_match: + raw_json = code_block_match.group(1).strip() + else: + # Try raw JSON object + decoder = json.JSONDecoder() + for m in re.finditer(r"\{", response): + try: + data, _ = decoder.raw_decode(response, m.start()) + if isinstance(data, dict): + raw_json = json.dumps(data) + break + except json.JSONDecodeError: + continue + + if raw_json: + try: + data = json.loads(raw_json) + return TrialFeedback( + summary_text=data.get("summary_text", ""), + failure_patterns=data.get("failure_patterns", []), + pillar_ratings=data.get("pillar_ratings", {}), + suggested_changes=data.get("suggested_changes", []), + target_pillar=data.get("target_pillar", ""), + ) + except json.JSONDecodeError: + pass + + # Fallback: wrap raw text as summary + return TrialFeedback(summary_text=response) + + def _parse_config_response(self, response: str) -> TrialConfig: + """Extract a TrialConfig from an LLM response. + + Looks for a ```json ... ``` block first, then falls back to + finding a raw JSON object in the response text. + """ + trial_id = uuid.uuid4().hex[:12] + + # Try to extract from a ```json code block + json_block_match = re.search( + r"```json\s*\n?(.*?)\n?\s*```", response, re.DOTALL + ) + if json_block_match: + raw_json = json_block_match.group(1).strip() + try: + data = json.loads(raw_json) + return self._config_from_dict(data, trial_id) + except json.JSONDecodeError: + pass + + # Try to extract from a generic ``` code block + code_block_match = re.search( + r"```\s*\n?(.*?)\n?\s*```", response, re.DOTALL + ) + if code_block_match: + raw_json = code_block_match.group(1).strip() + try: + data = json.loads(raw_json) + return self._config_from_dict(data, trial_id) + except json.JSONDecodeError: + pass + + # Try to find a raw JSON object in the response by scanning + # for each '{' and attempting to parse from that position. + decoder = json.JSONDecoder() + for m in re.finditer(r"\{", response): + try: + data, _ = decoder.raw_decode(response, m.start()) + if isinstance(data, dict): + return self._config_from_dict(data, trial_id) + except json.JSONDecodeError: + continue + + # Last resort: return config with at least the fixed params + ss = self.search_space + fixed = dict(ss.fixed) if ss and ss.fixed else {} + return TrialConfig( + trial_id=trial_id, + params=fixed, + reasoning="Failed to parse LLM response.", + ) + + def _config_from_dict( + self, data: Dict[str, Any], trial_id: str + ) -> TrialConfig: + """Build a TrialConfig from a parsed JSON dict. + + Merges fixed parameters from the search space so that fixed values + (e.g. intelligence.model, engine.backend) are always present. + """ + params = data.get("params", {}) + reasoning = data.get("reasoning", "") + + # Inject fixed params — these override anything the LLM proposed + if self.search_space and self.search_space.fixed: + for key, value in self.search_space.fixed.items(): + params[key] = value + + return TrialConfig( + trial_id=trial_id, + params=params, + reasoning=reasoning, + ) + + +__all__ = ["LLMOptimizer"] diff --git a/src/openjarvis/learning/optimize/optimizer.py b/src/openjarvis/learning/optimize/optimizer.py new file mode 100644 index 00000000..24441751 --- /dev/null +++ b/src/openjarvis/learning/optimize/optimizer.py @@ -0,0 +1,445 @@ +"""OptimizationEngine -- orchestrates the optimize loop. + +Ties together the LLM optimizer, trial runner, and persistence store +into a single propose -> evaluate -> analyze -> repeat loop. +""" + +from __future__ import annotations + +import logging +import uuid +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +try: + import tomli_w +except ModuleNotFoundError: # pragma: no cover + tomli_w = None # type: ignore[assignment] + +from openjarvis.learning.optimize.llm_optimizer import LLMOptimizer +from openjarvis.learning.optimize.store import OptimizationStore +from openjarvis.learning.optimize.trial_runner import TrialRunner +from openjarvis.learning.optimize.types import ( + ObjectiveSpec, + OptimizationRun, + SearchSpace, + TrialResult, +) + +LOGGER = logging.getLogger(__name__) + +# Mapping from objective metric names to RunSummary stat attribute + ".mean" +_SUMMARY_STAT_MAP: Dict[str, str] = { + "avg_power_watts": "avg_power_watts", + "throughput_tok_per_sec": "throughput_stats", + "mfu_pct": "mfu_stats", + "mbu_pct": "mbu_stats", + "ipw": "ipw_stats", + "ipj": "ipj_stats", + "energy_per_output_token": "energy_per_output_token_stats", + "throughput_per_watt": "throughput_per_watt_stats", + "ttft": "ttft_stats", + "mean_itl_ms": "itl_stats", +} + + +def _get_objective_value(trial: TrialResult, obj: ObjectiveSpec) -> float: + """Read the metric value from a TrialResult for a given objective.""" + # Direct attributes on TrialResult + direct = { + "accuracy", "mean_latency_seconds", + "total_cost_usd", "total_energy_joules", + } + if obj.metric in direct: + return getattr(trial, obj.metric, 0.0) + + # avg_power_watts is a direct attr on RunSummary + if obj.metric == "avg_power_watts" and trial.summary: + return trial.summary.avg_power_watts + + # Stats-based metrics from RunSummary + stat_attr = _SUMMARY_STAT_MAP.get(obj.metric) + if stat_attr and trial.summary: + stats = getattr(trial.summary, stat_attr, None) + if stats is not None and hasattr(stats, "mean"): + return stats.mean + + return 0.0 + + +def compute_pareto_frontier( + trials: List[TrialResult], + objectives: List[ObjectiveSpec], +) -> List[TrialResult]: + """Compute the Pareto frontier: trials not dominated by any other. + + A trial A dominates trial B if A is >= B on all objectives and > B + on at least one (direction-aware: maximize flips the comparison). + """ + if not trials or not objectives: + return list(trials) + + def _values(trial: TrialResult) -> List[float]: + vals = [] + for obj in objectives: + v = _get_objective_value(trial, obj) + # Normalize: for "minimize", negate so higher is always better + if obj.direction == "minimize": + v = -v + vals.append(v) + return vals + + trial_vals = [_values(t) for t in trials] + frontier: List[TrialResult] = [] + + for i, trial in enumerate(trials): + dominated = False + for j, other in enumerate(trials): + if i == j: + continue + # Check if other dominates trial + all_ge = all( + trial_vals[j][k] >= trial_vals[i][k] + for k in range(len(objectives)) + ) + any_gt = any( + trial_vals[j][k] > trial_vals[i][k] + for k in range(len(objectives)) + ) + if all_ge and any_gt: + dominated = True + break + if not dominated: + frontier.append(trial) + + return frontier + + +class OptimizationEngine: + """Orchestrates the optimize loop: propose -> evaluate -> analyze -> repeat.""" + + def __init__( + self, + search_space: SearchSpace, + llm_optimizer: LLMOptimizer, + trial_runner: TrialRunner, + store: Optional[OptimizationStore] = None, + max_trials: int = 20, + early_stop_patience: int = 5, + ) -> None: + self.search_space = search_space + self.llm_optimizer = llm_optimizer + self.trial_runner = trial_runner + self.store = store + self.max_trials = max_trials + self.early_stop_patience = early_stop_patience + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def run( + self, + progress_callback: Optional[Callable[[int, int], None]] = None, + ) -> OptimizationRun: + """Execute the full optimization loop. + + 1. Generate a run_id via uuid. + 2. ``llm_optimizer.propose_initial()`` -> first config. + 3. Loop up to ``max_trials``: + a. ``trial_runner.run_trial(config)`` -> TrialResult + b. ``llm_optimizer.analyze_trial(config, summary, traces)`` + c. Update TrialResult with analysis text + d. Append to history + e. If store, ``store.save_trial(result)`` + f. Update best_trial if accuracy improved + g. Check early stopping (no improvement for *patience* trials) + h. If not stopped, ``llm_optimizer.propose_next(history)`` + 4. Set run status to ``"completed"``. + 5. If store, ``store.save_run(optimization_run)``. + 6. Return the :class:`OptimizationRun`. + + Args: + progress_callback: Optional ``(trial_num, max_trials) -> None`` + called after each trial completes. + """ + run_id = uuid.uuid4().hex[:16] + # Detect benchmark name(s) from the trial runner + from openjarvis.learning.optimize.trial_runner import MultiBenchTrialRunner + + benchmark_name = getattr(self.trial_runner, "benchmark", "") + benchmark_names: List[str] = [] + if isinstance(self.trial_runner, MultiBenchTrialRunner): + benchmark_names = [ + s.benchmark for s in self.trial_runner.benchmark_specs + ] + benchmark_name = "+".join(benchmark_names) + + optimization_run = OptimizationRun( + run_id=run_id, + search_space=self.search_space, + status="running", + optimizer_model=self.llm_optimizer.optimizer_model, + benchmark=benchmark_name, + benchmarks=benchmark_names, + ) + + history: List[TrialResult] = [] + best_accuracy = -1.0 + trials_without_improvement = 0 + + # First config + config = self.llm_optimizer.propose_initial() + + for trial_num in range(1, self.max_trials + 1): + LOGGER.info( + "Trial %d/%d (id=%s)", + trial_num, + self.max_trials, + config.trial_id, + ) + + # Evaluate + result = self.trial_runner.run_trial(config) + + # Analyze — returns TrialFeedback + if result.summary is not None: + feedback = self.llm_optimizer.analyze_trial( + config, + result.summary, + sample_scores=result.sample_scores or None, + per_benchmark=result.per_benchmark or None, + ) + result.structured_feedback = feedback + result.analysis = feedback.summary_text + elif result.per_benchmark: + # Multi-benchmark composite: build a synthetic summary + # for analysis from per_benchmark data + from openjarvis.evals.core.types import RunSummary as _RS + + synth = _RS( + benchmark="multi", + category="multi", + backend="jarvis-agent", + model=result.config.params.get("intelligence.model", ""), + accuracy=result.accuracy, + mean_latency_seconds=result.mean_latency_seconds, + total_cost_usd=result.total_cost_usd, + total_energy_joules=result.total_energy_joules, + total_samples=result.samples_evaluated, + scored_samples=result.samples_evaluated, + correct=int( + result.accuracy * result.samples_evaluated + ), + errors=0, + total_input_tokens=0, + total_output_tokens=result.total_tokens, + ) + feedback = self.llm_optimizer.analyze_trial( + config, + synth, + per_benchmark=result.per_benchmark, + ) + result.structured_feedback = feedback + result.analysis = feedback.summary_text + else: + result.analysis = "" + + # Record + history.append(result) + optimization_run.trials.append(result) + + # Recompute Pareto frontier + optimization_run.pareto_frontier = compute_pareto_frontier( + history, optimization_run.objectives, + ) + frontier_ids = {t.trial_id for t in optimization_run.pareto_frontier} + + # Persist trial + if self.store is not None: + self.store.save_trial(run_id, result) + + # Track best + if result.accuracy > best_accuracy: + best_accuracy = result.accuracy + optimization_run.best_trial = result + trials_without_improvement = 0 + else: + trials_without_improvement += 1 + + # Progress callback + if progress_callback is not None: + progress_callback(trial_num, self.max_trials) + + # Early stopping + if trials_without_improvement >= self.early_stop_patience: + LOGGER.info( + "Early stopping after %d trials without improvement.", + self.early_stop_patience, + ) + break + + # Propose next (unless this was the last trial) + if trial_num < self.max_trials: + # Decide proposal strategy + target_pillar = "" + if result.structured_feedback: + target_pillar = result.structured_feedback.target_pillar + + if ( + trial_num % 5 == 0 + and len(optimization_run.pareto_frontier) >= 2 + ): + # Merge frontier members periodically + candidates = optimization_run.pareto_frontier[:3] + config = self.llm_optimizer.propose_merge( + candidates, history, frontier_ids=frontier_ids, + ) + elif target_pillar and trial_num > 2: + # Targeted mutation on the suggested pillar + config = self.llm_optimizer.propose_targeted( + history, + result.config, + target_pillar, + frontier_ids=frontier_ids, + ) + else: + config = self.llm_optimizer.propose_next( + history, frontier_ids=frontier_ids, + ) + + optimization_run.status = "completed" + + if self.store is not None: + self.store.save_run(optimization_run) + + return optimization_run + + def export_best_recipe( + self, run: OptimizationRun, path: Path + ) -> Path: + """Export the best trial's config as a TOML recipe file. + + Args: + run: A completed :class:`OptimizationRun`. + path: Destination path for the TOML file. + + Returns: + The *path* written to. + + Raises: + ValueError: If there is no best trial in the run. + """ + if run.best_trial is None: + raise ValueError("No best trial to export.") + + recipe_data = self._trial_to_recipe_dict(run.best_trial) + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + if tomli_w is not None: + with open(path, "wb") as fh: + tomli_w.dump(recipe_data, fh) + else: + # Fallback: write TOML manually + self._write_toml_fallback(recipe_data, path) + + run.best_recipe_path = str(path) + return path + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _trial_to_recipe_dict(trial: TrialResult) -> Dict[str, Any]: + """Convert a TrialResult into a Recipe-style TOML dict.""" + params = trial.config.params + recipe: Dict[str, Any] = { + "recipe": { + "name": f"optimized-{trial.trial_id}", + "description": ( + f"Auto-optimized config (accuracy={trial.accuracy:.4f})" + ), + "version": "1.0.0", + }, + } + + # Intelligence section + intel: Dict[str, Any] = {} + if "intelligence.model" in params: + intel["model"] = params["intelligence.model"] + if "intelligence.temperature" in params: + intel["temperature"] = params["intelligence.temperature"] + if "intelligence.quantization" in params: + intel["quantization"] = params["intelligence.quantization"] + if "intelligence.system_prompt" in params: + intel["system_prompt"] = params["intelligence.system_prompt"] + if "intelligence.max_tokens" in params: + intel["max_tokens"] = params["intelligence.max_tokens"] + if "intelligence.top_p" in params: + intel["top_p"] = params["intelligence.top_p"] + if intel: + recipe["intelligence"] = intel + + # Engine section + engine: Dict[str, Any] = {} + if "engine.backend" in params: + engine["key"] = params["engine.backend"] + if engine: + recipe["engine"] = engine + + # Agent section + agent: Dict[str, Any] = {} + if "agent.type" in params: + agent["type"] = params["agent.type"] + if "agent.max_turns" in params: + agent["max_turns"] = params["agent.max_turns"] + if "agent.system_prompt" in params: + agent["system_prompt"] = params["agent.system_prompt"] + if "tools.tool_set" in params: + agent["tools"] = params["tools.tool_set"] + if agent: + recipe["agent"] = agent + + # Learning section + learning: Dict[str, Any] = {} + if "learning.routing_policy" in params: + learning["routing"] = params["learning.routing_policy"] + if "learning.agent_policy" in params: + learning["agent"] = params["learning.agent_policy"] + if learning: + recipe["learning"] = learning + + return recipe + + @staticmethod + def _write_toml_fallback( + data: Dict[str, Any], path: Path + ) -> None: + """Write a simple nested dict as TOML without tomli_w.""" + lines: List[str] = [] + for section, values in data.items(): + if not isinstance(values, dict): + continue + lines.append(f"[{section}]") + for key, val in values.items(): + if isinstance(val, str): + lines.append(f'{key} = "{val}"') + elif isinstance(val, bool): + lines.append(f"{key} = {'true' if val else 'false'}") + elif isinstance(val, (int, float)): + lines.append(f"{key} = {val}") + elif isinstance(val, list): + items = ", ".join( + f'"{v}"' if isinstance(v, str) else str(v) + for v in val + ) + lines.append(f"{key} = [{items}]") + else: + lines.append(f'{key} = "{val}"') + lines.append("") + path.write_text("\n".join(lines), encoding="utf-8") + + +__all__ = ["OptimizationEngine", "compute_pareto_frontier"] diff --git a/src/openjarvis/learning/optimize/personal/__init__.py b/src/openjarvis/learning/optimize/personal/__init__.py new file mode 100644 index 00000000..ccd06245 --- /dev/null +++ b/src/openjarvis/learning/optimize/personal/__init__.py @@ -0,0 +1,17 @@ +"""Personal benchmark system -- synthesize benchmarks from interaction traces.""" + +from openjarvis.learning.optimize.personal.dataset import PersonalBenchmarkDataset +from openjarvis.learning.optimize.personal.scorer import PersonalBenchmarkScorer +from openjarvis.learning.optimize.personal.synthesizer import ( + PersonalBenchmark, + PersonalBenchmarkSample, + PersonalBenchmarkSynthesizer, +) + +__all__ = [ + "PersonalBenchmark", + "PersonalBenchmarkSample", + "PersonalBenchmarkSynthesizer", + "PersonalBenchmarkDataset", + "PersonalBenchmarkScorer", +] diff --git a/src/openjarvis/learning/optimize/personal/dataset.py b/src/openjarvis/learning/optimize/personal/dataset.py new file mode 100644 index 00000000..8a590549 --- /dev/null +++ b/src/openjarvis/learning/optimize/personal/dataset.py @@ -0,0 +1,54 @@ +"""DatasetProvider adapter for personal benchmarks.""" + +from __future__ import annotations + +from typing import Iterable, List, Optional + +from openjarvis.evals.core.dataset import DatasetProvider +from openjarvis.evals.core.types import EvalRecord +from openjarvis.learning.optimize.personal.synthesizer import PersonalBenchmark + + +class PersonalBenchmarkDataset(DatasetProvider): + """Wraps a PersonalBenchmark as a DatasetProvider for EvalRunner.""" + + dataset_id: str = "personal" + dataset_name: str = "Personal Benchmark" + + def __init__(self, benchmark: PersonalBenchmark) -> None: + self._benchmark = benchmark + self._records: List[EvalRecord] = [] + + def load( + self, + *, + max_samples: Optional[int] = None, + split: Optional[str] = None, + seed: Optional[int] = None, + ) -> None: + """Convert :class:`PersonalBenchmarkSample` instances to :class:`EvalRecord`.""" + samples = self._benchmark.samples + if max_samples is not None: + samples = samples[:max_samples] + self._records = [ + EvalRecord( + record_id=s.trace_id, + problem=s.query, + reference=s.reference_answer, + category=s.category, + subject=s.agent or "general", + metadata=s.metadata, + ) + for s in samples + ] + + def iter_records(self) -> Iterable[EvalRecord]: + """Iterate over loaded records.""" + return iter(self._records) + + def size(self) -> int: + """Return the number of loaded records.""" + return len(self._records) + + +__all__ = ["PersonalBenchmarkDataset"] diff --git a/src/openjarvis/learning/optimize/personal/scorer.py b/src/openjarvis/learning/optimize/personal/scorer.py new file mode 100644 index 00000000..8ed0bec1 --- /dev/null +++ b/src/openjarvis/learning/optimize/personal/scorer.py @@ -0,0 +1,47 @@ +"""LLM-judge scorer for personal benchmarks.""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Tuple + +from openjarvis.evals.core.backend import InferenceBackend +from openjarvis.evals.core.scorer import LLMJudgeScorer +from openjarvis.evals.core.types import EvalRecord + + +class PersonalBenchmarkScorer(LLMJudgeScorer): + """Judges a candidate response against the best-known response from traces.""" + + scorer_id: str = "personal_judge" + + def __init__(self, judge_backend: InferenceBackend, judge_model: str) -> None: + super().__init__(judge_backend, judge_model) + + def score( + self, record: EvalRecord, model_answer: str, + ) -> Tuple[Optional[bool], Dict[str, Any]]: + """Compare *model_answer* against *record.reference* using the judge LLM. + + Returns ``(is_correct, metadata)`` where *is_correct* indicates whether + the candidate answer is at least as good as the reference. + """ + prompt = ( + "Compare these two answers to the query.\n\n" + f"Query: {record.problem}\n\n" + "Reference answer (known good):\n" + f"{record.reference}\n\n" + "Candidate answer:\n" + f"{model_answer}\n\n" + "Is the candidate answer at least as good as the reference? " + 'Respond with exactly "YES" or "NO" on the first line, ' + "then explain your reasoning." + ) + response = self._ask_judge( + prompt, system="You are an impartial quality judge.", + ) + first_line = response.strip().split("\n")[0].strip().upper() + is_correct = first_line.startswith("YES") + return is_correct, {"judge_response": response} + + +__all__ = ["PersonalBenchmarkScorer"] diff --git a/src/openjarvis/learning/optimize/personal/synthesizer.py b/src/openjarvis/learning/optimize/personal/synthesizer.py new file mode 100644 index 00000000..ecd0f289 --- /dev/null +++ b/src/openjarvis/learning/optimize/personal/synthesizer.py @@ -0,0 +1,119 @@ +"""Synthesize personal benchmarks from interaction traces.""" + +from __future__ import annotations + +import time +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Dict, List + +from openjarvis.traces.store import TraceStore + + +@dataclass(slots=True) +class PersonalBenchmarkSample: + """A single sample in a personal benchmark.""" + + trace_id: str + query: str + reference_answer: str # best known answer from traces + agent: str = "" + category: str = "chat" + feedback_score: float = 0.0 + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PersonalBenchmark: + """A synthesized benchmark from user interaction traces.""" + + workflow_id: str + samples: List[PersonalBenchmarkSample] = field(default_factory=list) + created_at: float = 0.0 + + +def _query_class_key(agent: str, query: str) -> str: + """Compute a grouping key from agent name and query prefix.""" + prefix = query[:50].strip().lower() + return f"{agent}::{prefix}" + + +def _infer_category(agent: str) -> str: + """Heuristic to map an agent name to an eval category.""" + agent_lower = agent.lower() + if any(tok in agent_lower for tok in ("react", "openhands", "orchestrator")): + return "agentic" + if any(tok in agent_lower for tok in ("rag", "memory", "retriev")): + return "rag" + if any(tok in agent_lower for tok in ("reason", "math", "code")): + return "reasoning" + return "chat" + + +class PersonalBenchmarkSynthesizer: + """Mines interaction traces into a reusable personal benchmark.""" + + def __init__(self, trace_store: TraceStore) -> None: + self._store = trace_store + + def synthesize( + self, + workflow_id: str = "default", + min_feedback: float = 0.7, + max_samples: int = 100, + ) -> PersonalBenchmark: + """Build a personal benchmark from high-quality traces. + + 1. Query traces that have feedback >= *min_feedback*. + 2. Group by query class (agent + first 50 chars of query). + 3. For each class, pick the trace with the highest feedback as reference. + 4. Return a :class:`PersonalBenchmark` capped at *max_samples*. + """ + # Fetch a large pool of traces (limit high enough to cover most stores) + all_traces = self._store.list_traces(limit=10_000) + + # Filter to traces with sufficient feedback + qualified = [ + t + for t in all_traces + if t.feedback is not None and t.feedback >= min_feedback + ] + + # Group by query class + groups: Dict[str, list] = defaultdict(list) + for trace in qualified: + key = _query_class_key(trace.agent, trace.query) + groups[key].append(trace) + + # Pick best trace per class + samples: List[PersonalBenchmarkSample] = [] + for _key, traces in groups.items(): + best = max(traces, key=lambda t: t.feedback or 0.0) + samples.append( + PersonalBenchmarkSample( + trace_id=best.trace_id, + query=best.query, + reference_answer=best.result, + agent=best.agent, + category=_infer_category(best.agent), + feedback_score=best.feedback or 0.0, + metadata=best.metadata, + ), + ) + + # Sort deterministically (highest feedback first) and cap + samples.sort(key=lambda s: (-s.feedback_score, s.trace_id)) + samples = samples[:max_samples] + + return PersonalBenchmark( + workflow_id=workflow_id, + samples=samples, + created_at=time.time(), + ) + + +__all__ = [ + "PersonalBenchmark", + "PersonalBenchmarkSample", + "PersonalBenchmarkSynthesizer", +] diff --git a/src/openjarvis/learning/optimize/search_space.py b/src/openjarvis/learning/optimize/search_space.py new file mode 100644 index 00000000..7edd142b --- /dev/null +++ b/src/openjarvis/learning/optimize/search_space.py @@ -0,0 +1,193 @@ +"""Search space builder and default search space for configuration optimization.""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from openjarvis.learning.optimize.types import SearchDimension, SearchSpace + + +def build_search_space(config: Dict[str, Any]) -> SearchSpace: + """Build a SearchSpace from a TOML-style config dict. + + Expected format:: + + { + "optimize": { + "search": [ + { + "name": "agent.type", + "type": "categorical", + "values": ["orchestrator", "native_react"], + "description": "Agent architecture", + }, + { + "name": "intelligence.temperature", + "type": "continuous", + "low": 0.0, + "high": 1.0, + "description": "Generation temperature", + }, + ], + "fixed": {"engine": "ollama", "model": "qwen3:8b"}, + "constraints": { + "rules": ["SimpleAgent should only have max_turns = 1"], + }, + } + } + """ + opt = config.get("optimize", {}) + search_entries: List[Dict[str, Any]] = opt.get("search", []) + fixed: Dict[str, Any] = dict(opt.get("fixed", {})) + constraints_sec = opt.get("constraints", {}) + constraints: List[str] = list(constraints_sec.get("rules", [])) + + dimensions: List[SearchDimension] = [] + for entry in search_entries: + # Infer pillar from the first segment of the dotted name + name = entry.get("name", "") + pillar = name.split(".")[0] if "." in name else "" + + dimensions.append( + SearchDimension( + name=name, + dim_type=entry.get("type", "categorical"), + values=list(entry.get("values", [])), + low=entry.get("low"), + high=entry.get("high"), + description=entry.get("description", ""), + pillar=pillar, + ) + ) + + return SearchSpace( + dimensions=dimensions, + fixed=fixed, + constraints=constraints, + ) + + +# --------------------------------------------------------------------------- +# Default search space covering all 5 pillars +# --------------------------------------------------------------------------- + +DEFAULT_SEARCH_SPACE = SearchSpace( + dimensions=[ + # Intelligence pillar + SearchDimension( + name="intelligence.model", + dim_type="categorical", + values=[ + "qwen3:8b", + "qwen3:4b", + "qwen3:1.7b", + "llama3.1:8b", + "llama3.1:70b", + "gemma2:9b", + "mistral:7b", + "deepseek-r1:8b", + ], + description="The LLM model to use for generation", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.temperature", + dim_type="continuous", + low=0.0, + high=1.0, + description="Generation temperature (0 = deterministic, 1 = creative)", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.max_tokens", + dim_type="integer", + low=256, + high=8192, + description="Maximum tokens to generate per response", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.top_p", + dim_type="continuous", + low=0.0, + high=1.0, + description="Nucleus sampling probability threshold", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.system_prompt", + dim_type="text", + description="System prompt to guide model behavior", + pillar="intelligence", + ), + # Engine pillar + SearchDimension( + name="engine.backend", + dim_type="categorical", + values=[ + "ollama", "vllm", "sglang", "llamacpp", + "mlx", "lmstudio", "exo", "nexa", + "uzu", "apple_fm", + ], + description="Inference engine backend", + pillar="engine", + ), + # Agent pillar + SearchDimension( + name="agent.type", + dim_type="categorical", + values=["simple", "orchestrator", "native_react", "native_openhands"], + description="Agent architecture to use", + pillar="agent", + ), + SearchDimension( + name="agent.max_turns", + dim_type="integer", + low=1, + high=30, + description="Maximum number of agent reasoning turns", + pillar="agent", + ), + # Tools pillar + SearchDimension( + name="tools.tool_set", + dim_type="subset", + values=[ + "calculator", + "think", + "file_read", + "file_write", + "web_search", + "code_interpreter", + "llm", + "shell_exec", + "apply_patch", + "http_request", + "database_query", + ], + description="Set of tools available to the agent", + pillar="tools", + ), + # Learning pillar + SearchDimension( + name="learning.routing_policy", + dim_type="categorical", + values=["heuristic", "grpo", "bandit", "learned"], + description="Router policy for model/agent selection", + pillar="learning", + ), + ], + fixed={}, + constraints=[ + "SimpleAgent (agent.type='simple') should only have max_turns = 1", + "agent.max_turns must be >= 1", + "intelligence.temperature and intelligence.top_p " + "should not both be at extreme values", + ], +) + + +__all__ = [ + "build_search_space", + "DEFAULT_SEARCH_SPACE", +] diff --git a/src/openjarvis/learning/optimize/store.py b/src/openjarvis/learning/optimize/store.py new file mode 100644 index 00000000..c8896899 --- /dev/null +++ b/src/openjarvis/learning/optimize/store.py @@ -0,0 +1,497 @@ +"""SQLite-backed storage for optimization runs and trials.""" + +from __future__ import annotations + +import json +import sqlite3 +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from openjarvis.learning.optimize.types import ( + BenchmarkScore, + OptimizationRun, + SampleScore, + SearchSpace, + TrialConfig, + TrialFeedback, + TrialResult, +) + +_CREATE_RUNS = """\ +CREATE TABLE IF NOT EXISTS optimization_runs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + run_id TEXT NOT NULL UNIQUE, + search_space TEXT NOT NULL DEFAULT '{}', + status TEXT NOT NULL DEFAULT 'running', + optimizer_model TEXT NOT NULL DEFAULT '', + benchmark TEXT NOT NULL DEFAULT '', + best_trial_id TEXT, + best_recipe_path TEXT, + created_at REAL NOT NULL DEFAULT 0.0, + updated_at REAL NOT NULL DEFAULT 0.0 +); +""" + +_CREATE_TRIALS = """\ +CREATE TABLE IF NOT EXISTS trial_results ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + trial_id TEXT NOT NULL, + run_id TEXT NOT NULL, + config TEXT NOT NULL DEFAULT '{}', + reasoning TEXT NOT NULL DEFAULT '', + accuracy REAL NOT NULL DEFAULT 0.0, + mean_latency_seconds REAL NOT NULL DEFAULT 0.0, + total_cost_usd REAL NOT NULL DEFAULT 0.0, + total_energy_joules REAL NOT NULL DEFAULT 0.0, + total_tokens INTEGER NOT NULL DEFAULT 0, + samples_evaluated INTEGER NOT NULL DEFAULT 0, + analysis TEXT NOT NULL DEFAULT '', + failure_modes TEXT NOT NULL DEFAULT '[]', + created_at REAL NOT NULL DEFAULT 0.0, + FOREIGN KEY (run_id) REFERENCES optimization_runs(run_id) +); +""" + +_INSERT_RUN = """\ +INSERT OR REPLACE INTO optimization_runs ( + run_id, search_space, status, optimizer_model, benchmark, + best_trial_id, best_recipe_path, created_at, updated_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) +""" + +_INSERT_TRIAL = """\ +INSERT OR REPLACE INTO trial_results ( + trial_id, run_id, config, reasoning, accuracy, + mean_latency_seconds, total_cost_usd, total_energy_joules, + total_tokens, samples_evaluated, analysis, failure_modes, + created_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +""" + + +_MIGRATE_TRIALS = [ + "ALTER TABLE trial_results ADD COLUMN " + "sample_scores TEXT NOT NULL DEFAULT '[]'", + "ALTER TABLE trial_results ADD COLUMN " + "structured_feedback TEXT NOT NULL DEFAULT '{}'", + "ALTER TABLE trial_results ADD COLUMN " + "per_benchmark TEXT NOT NULL DEFAULT '[]'", +] + +_MIGRATE_RUNS = [ + "ALTER TABLE optimization_runs ADD COLUMN " + "pareto_frontier_ids TEXT NOT NULL DEFAULT '[]'", + "ALTER TABLE optimization_runs ADD COLUMN " + "benchmarks TEXT NOT NULL DEFAULT '[]'", +] + + +class OptimizationStore: + """SQLite-backed storage for optimization runs and trials.""" + + def __init__(self, db_path: Union[str, Path]) -> None: + self._db_path = str(db_path) + self._conn = sqlite3.connect(self._db_path) + self._conn.execute("PRAGMA journal_mode=WAL") + self._conn.execute(_CREATE_RUNS) + self._conn.execute(_CREATE_TRIALS) + self._conn.commit() + self._migrate() + + def _migrate(self) -> None: + """Add new columns to existing databases gracefully.""" + for stmt in _MIGRATE_TRIALS + _MIGRATE_RUNS: + try: + self._conn.execute(stmt) + except sqlite3.OperationalError: + pass # Column already exists + self._conn.commit() + + # ------------------------------------------------------------------ + # Runs + # ------------------------------------------------------------------ + + def save_run(self, run: OptimizationRun) -> None: + """Persist an optimization run (insert or update).""" + now = time.time() + search_space_json = self._search_space_to_json(run.search_space) + best_trial_id = run.best_trial.trial_id if run.best_trial else None + pareto_ids = json.dumps([t.trial_id for t in run.pareto_frontier]) + self._conn.execute( + _INSERT_RUN, + ( + run.run_id, + search_space_json, + run.status, + run.optimizer_model, + run.benchmark, + best_trial_id, + run.best_recipe_path, + now, + now, + ), + ) + benchmarks_json = json.dumps(run.benchmarks) + # Update pareto_frontier_ids and benchmarks separately + self._conn.execute( + "UPDATE optimization_runs SET pareto_frontier_ids = ?, " + "benchmarks = ? WHERE run_id = ?", + (pareto_ids, benchmarks_json, run.run_id), + ) + self._conn.commit() + + def get_run(self, run_id: str) -> Optional[OptimizationRun]: + """Retrieve an optimization run by id, or ``None``.""" + row = self._conn.execute( + "SELECT * FROM optimization_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + return None + return self._row_to_run(row) + + def list_runs(self, limit: int = 50) -> List[Dict[str, Any]]: + """Return summary dicts of recent optimization runs.""" + rows = self._conn.execute( + "SELECT * FROM optimization_runs ORDER BY created_at DESC LIMIT ?", + (limit,), + ).fetchall() + result: List[Dict[str, Any]] = [] + for row in rows: + result.append( + { + "run_id": row[1], + "status": row[3], + "optimizer_model": row[4], + "benchmark": row[5], + "best_trial_id": row[6], + "best_recipe_path": row[7], + "created_at": row[8], + "updated_at": row[9], + } + ) + return result + + # ------------------------------------------------------------------ + # Trials + # ------------------------------------------------------------------ + + def save_trial(self, run_id: str, trial: TrialResult) -> None: + """Persist a single trial result.""" + now = time.time() + # Serialize sample_scores + scores_json = json.dumps([ + { + "record_id": s.record_id, + "is_correct": s.is_correct, + "score": s.score, + "latency_seconds": s.latency_seconds, + "prompt_tokens": s.prompt_tokens, + "completion_tokens": s.completion_tokens, + "cost_usd": s.cost_usd, + "error": s.error, + "ttft": s.ttft, + "energy_joules": s.energy_joules, + "power_watts": s.power_watts, + "gpu_utilization_pct": s.gpu_utilization_pct, + "throughput_tok_per_sec": s.throughput_tok_per_sec, + "mfu_pct": s.mfu_pct, + "mbu_pct": s.mbu_pct, + "ipw": s.ipw, + "ipj": s.ipj, + "energy_per_output_token_joules": s.energy_per_output_token_joules, + "throughput_per_watt": s.throughput_per_watt, + "mean_itl_ms": s.mean_itl_ms, + } + for s in trial.sample_scores + ]) + # Serialize structured_feedback + fb = trial.structured_feedback + fb_json = json.dumps({ + "summary_text": fb.summary_text, + "failure_patterns": fb.failure_patterns, + "pillar_ratings": fb.pillar_ratings, + "suggested_changes": fb.suggested_changes, + "target_pillar": fb.target_pillar, + }) if fb else "{}" + + self._conn.execute( + _INSERT_TRIAL, + ( + trial.trial_id, + run_id, + json.dumps(trial.config.params), + trial.config.reasoning, + trial.accuracy, + trial.mean_latency_seconds, + trial.total_cost_usd, + trial.total_energy_joules, + trial.total_tokens, + trial.samples_evaluated, + trial.analysis, + json.dumps(trial.failure_modes), + now, + ), + ) + # Serialize per_benchmark + pb_json = json.dumps([ + { + "benchmark": b.benchmark, + "accuracy": b.accuracy, + "mean_latency_seconds": b.mean_latency_seconds, + "total_cost_usd": b.total_cost_usd, + "total_energy_joules": b.total_energy_joules, + "total_tokens": b.total_tokens, + "samples_evaluated": b.samples_evaluated, + "errors": b.errors, + "weight": b.weight, + } + for b in trial.per_benchmark + ]) + + # Update new columns separately + self._conn.execute( + "UPDATE trial_results SET sample_scores = ?, " + "structured_feedback = ?, per_benchmark = ? " + "WHERE trial_id = ? AND run_id = ?", + (scores_json, fb_json, pb_json, trial.trial_id, run_id), + ) + self._conn.commit() + + def get_trials(self, run_id: str) -> List[TrialResult]: + """Retrieve all trial results for a given run.""" + rows = self._conn.execute( + "SELECT * FROM trial_results WHERE run_id = ? ORDER BY id", + (run_id,), + ).fetchall() + return [self._row_to_trial(r) for r in rows] + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def close(self) -> None: + """Close the underlying SQLite connection.""" + self._conn.close() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _search_space_to_json(space: SearchSpace) -> str: + """Serialize a SearchSpace to JSON.""" + dims = [] + for d in space.dimensions: + dims.append( + { + "name": d.name, + "dim_type": d.dim_type, + "values": d.values, + "low": d.low, + "high": d.high, + "description": d.description, + "pillar": d.pillar, + } + ) + return json.dumps( + { + "dimensions": dims, + "fixed": space.fixed, + "constraints": space.constraints, + } + ) + + @staticmethod + def _json_to_search_space(raw: str) -> SearchSpace: + """Deserialize a SearchSpace from JSON.""" + from openjarvis.learning.optimize.types import SearchDimension + + data = json.loads(raw) + dims = [] + for d in data.get("dimensions", []): + dims.append( + SearchDimension( + name=d.get("name", ""), + dim_type=d.get("dim_type", "categorical"), + values=d.get("values", []), + low=d.get("low"), + high=d.get("high"), + description=d.get("description", ""), + pillar=d.get("pillar", ""), + ) + ) + return SearchSpace( + dimensions=dims, + fixed=data.get("fixed", {}), + constraints=data.get("constraints", []), + ) + + def _row_to_run(self, row: tuple) -> OptimizationRun: + """Convert a database row to an OptimizationRun.""" + run_id = row[1] + search_space = self._json_to_search_space(row[2]) + status = row[3] + optimizer_model = row[4] + benchmark = row[5] + best_trial_id = row[6] + best_recipe_path = row[7] + + # Load trials for this run + trials = self.get_trials(run_id) + + # Find the best trial + best_trial: Optional[TrialResult] = None + if best_trial_id: + for t in trials: + if t.trial_id == best_trial_id: + best_trial = t + break + + # Reconstruct benchmarks list + benchmarks: List[str] = [] + if len(row) > 11: + try: + benchmarks = json.loads(row[11]) if row[11] else [] + except (json.JSONDecodeError, TypeError): + pass + + # Reconstruct pareto frontier from IDs + pareto_frontier: List[TrialResult] = [] + if len(row) > 10: + try: + frontier_ids = json.loads(row[10]) if row[10] else [] + trial_map = {t.trial_id: t for t in trials} + pareto_frontier = [ + trial_map[tid] for tid in frontier_ids if tid in trial_map + ] + except (json.JSONDecodeError, TypeError): + pass + + return OptimizationRun( + run_id=run_id, + search_space=search_space, + trials=trials, + best_trial=best_trial, + best_recipe_path=best_recipe_path, + status=status, + optimizer_model=optimizer_model, + benchmark=benchmark, + benchmarks=benchmarks, + pareto_frontier=pareto_frontier, + ) + + @staticmethod + def _row_to_trial(row: tuple) -> TrialResult: + """Convert a database row to a TrialResult.""" + trial_id = row[1] + # row[2] = run_id (not stored on TrialResult) + params = json.loads(row[3]) + reasoning = row[4] + accuracy = row[5] + mean_latency = row[6] + cost = row[7] + energy = row[8] + tokens = row[9] + samples = row[10] + analysis = row[11] + failure_modes = json.loads(row[12]) + # row[13] = created_at + + # New columns (may be absent in old DBs) + sample_scores: List[SampleScore] = [] + structured_feedback: Optional[TrialFeedback] = None + + if len(row) > 14: + try: + raw_scores = json.loads(row[14]) if row[14] else [] + sample_scores = [ + SampleScore( + record_id=s.get("record_id", ""), + is_correct=s.get("is_correct"), + score=s.get("score"), + latency_seconds=s.get("latency_seconds", 0.0), + prompt_tokens=s.get("prompt_tokens", 0), + completion_tokens=s.get("completion_tokens", 0), + cost_usd=s.get("cost_usd", 0.0), + error=s.get("error"), + ttft=s.get("ttft", 0.0), + energy_joules=s.get("energy_joules", 0.0), + power_watts=s.get("power_watts", 0.0), + gpu_utilization_pct=s.get("gpu_utilization_pct", 0.0), + throughput_tok_per_sec=s.get("throughput_tok_per_sec", 0.0), + mfu_pct=s.get("mfu_pct", 0.0), + mbu_pct=s.get("mbu_pct", 0.0), + ipw=s.get("ipw", 0.0), + ipj=s.get("ipj", 0.0), + energy_per_output_token_joules=s.get( + "energy_per_output_token_joules", 0.0, + ), + throughput_per_watt=s.get("throughput_per_watt", 0.0), + mean_itl_ms=s.get("mean_itl_ms", 0.0), + ) + for s in raw_scores + ] + except (json.JSONDecodeError, TypeError): + pass + + if len(row) > 15: + try: + raw_fb = json.loads(row[15]) if row[15] else {} + if raw_fb and raw_fb.get("summary_text", "") != "": + structured_feedback = TrialFeedback( + summary_text=raw_fb.get("summary_text", ""), + failure_patterns=raw_fb.get("failure_patterns", []), + pillar_ratings=raw_fb.get("pillar_ratings", {}), + suggested_changes=raw_fb.get("suggested_changes", []), + target_pillar=raw_fb.get("target_pillar", ""), + ) + except (json.JSONDecodeError, TypeError): + pass + + # per_benchmark column + per_benchmark: List[BenchmarkScore] = [] + if len(row) > 16: + try: + raw_pb = json.loads(row[16]) if row[16] else [] + per_benchmark = [ + BenchmarkScore( + benchmark=b.get("benchmark", ""), + accuracy=b.get("accuracy", 0.0), + mean_latency_seconds=b.get("mean_latency_seconds", 0.0), + total_cost_usd=b.get("total_cost_usd", 0.0), + total_energy_joules=b.get("total_energy_joules", 0.0), + total_tokens=b.get("total_tokens", 0), + samples_evaluated=b.get("samples_evaluated", 0), + errors=b.get("errors", 0), + weight=b.get("weight", 1.0), + ) + for b in raw_pb + ] + except (json.JSONDecodeError, TypeError): + pass + + config = TrialConfig( + trial_id=trial_id, + params=params, + reasoning=reasoning, + ) + return TrialResult( + trial_id=trial_id, + config=config, + accuracy=accuracy, + mean_latency_seconds=mean_latency, + total_cost_usd=cost, + total_energy_joules=energy, + total_tokens=tokens, + samples_evaluated=samples, + analysis=analysis, + failure_modes=failure_modes, + sample_scores=sample_scores, + structured_feedback=structured_feedback, + per_benchmark=per_benchmark, + ) + + +__all__ = ["OptimizationStore"] diff --git a/src/openjarvis/learning/optimize/trial_runner.py b/src/openjarvis/learning/optimize/trial_runner.py new file mode 100644 index 00000000..54a59f72 --- /dev/null +++ b/src/openjarvis/learning/optimize/trial_runner.py @@ -0,0 +1,392 @@ +"""TrialRunner -- evaluates a proposed config against a benchmark.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any, List, Optional + +from openjarvis.evals.core.types import RunConfig, RunSummary +from openjarvis.learning.optimize.types import ( + BenchmarkScore, + SampleScore, + TrialConfig, + TrialResult, +) + +LOGGER = logging.getLogger(__name__) + + +@dataclass(slots=True) +class BenchmarkSpec: + """Specification for one benchmark in a multi-benchmark optimization.""" + + benchmark: str + max_samples: int = 200 + weight: float = 1.0 + + +class TrialRunner: + """Evaluates a proposed config against a benchmark. + + Bridges the optimization types (:class:`TrialConfig`) to the eval + framework (:class:`EvalRunner`) so the optimizer can score candidate + configurations end-to-end. + """ + + def __init__( + self, + benchmark: str, + max_samples: int = 50, + judge_model: str = "gpt-5-mini-2025-08-07", + output_dir: str = "results/optimize/", + ) -> None: + self.benchmark = benchmark + self.max_samples = max_samples + self.judge_model = judge_model + self.output_dir = output_dir + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def run_trial(self, trial: TrialConfig) -> TrialResult: + """Run *trial* against the configured benchmark and return a result. + + Steps: + 1. Convert ``trial`` to a :class:`Recipe` and extract params. + 2. Build a :class:`RunConfig` from recipe + benchmark settings. + 3. Lazily import eval-framework registries to resolve the + benchmark -> dataset + scorer, and build the backend. + 4. Execute via ``EvalRunner.run()`` -> :class:`RunSummary`. + 5. Map the summary into a :class:`TrialResult`. + """ + recipe = trial.to_recipe() + run_config = self._build_run_config(trial, recipe) + + # Lazy imports so the optimize package stays lightweight + from openjarvis.evals.cli import ( + _build_backend, + _build_dataset, + _build_judge_backend, + _build_scorer, + ) + from openjarvis.evals.core.runner import EvalRunner + + dataset = _build_dataset(self.benchmark) + backend = _build_backend( + run_config.backend, + run_config.engine_key, + run_config.agent_name or "orchestrator", + run_config.tools, + ) + judge_backend = _build_judge_backend(run_config.judge_model) + scorer = _build_scorer( + self.benchmark, judge_backend, run_config.judge_model, + ) + + try: + eval_runner = EvalRunner( + run_config, dataset, backend, scorer, + ) + summary: RunSummary = eval_runner.run() + eval_results = eval_runner.results + finally: + backend.close() + judge_backend.close() + + return self._summary_to_result(trial, summary, eval_results=eval_results) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_run_config(self, trial: TrialConfig, recipe: Any) -> RunConfig: + """Map recipe fields into a :class:`RunConfig`.""" + model = recipe.model or "default" + backend_name = "jarvis-direct" + if recipe.agent_type is not None: + backend_name = "jarvis-agent" + + model_slug = model.replace("/", "-").replace(":", "-") + output_path = str( + Path(self.output_dir) / f"{trial.trial_id}_{model_slug}.jsonl", + ) + + max_tokens = recipe.max_tokens if recipe.max_tokens is not None else 2048 + + return RunConfig( + benchmark=self.benchmark, + backend=backend_name, + model=model, + max_samples=self.max_samples, + temperature=recipe.temperature if recipe.temperature is not None else 0.0, + max_tokens=max_tokens, + judge_model=self.judge_model, + engine_key=recipe.engine_key, + agent_name=recipe.agent_type, + tools=list(recipe.tools) if recipe.tools else [], + output_path=output_path, + system_prompt=recipe.system_prompt or "", + ) + + @staticmethod + def _summary_to_result( + trial: TrialConfig, + summary: RunSummary, + eval_results: Optional[List[Any]] = None, + ) -> TrialResult: + """Convert a :class:`RunSummary` to a :class:`TrialResult`.""" + total_tokens = summary.total_input_tokens + summary.total_output_tokens + + failure_modes: List[str] = [] + if summary.errors > 0: + failure_modes.append(f"{summary.errors} evaluation errors") + + sample_scores: List[SampleScore] = [] + if eval_results: + for er in eval_results: + sample_scores.append( + SampleScore( + record_id=er.record_id, + is_correct=er.is_correct, + score=er.score, + latency_seconds=er.latency_seconds, + prompt_tokens=er.prompt_tokens, + completion_tokens=er.completion_tokens, + cost_usd=er.cost_usd, + error=er.error, + ttft=er.ttft, + energy_joules=er.energy_joules, + power_watts=er.power_watts, + gpu_utilization_pct=er.gpu_utilization_pct, + throughput_tok_per_sec=er.throughput_tok_per_sec, + mfu_pct=er.mfu_pct, + mbu_pct=er.mbu_pct, + ipw=er.ipw, + ipj=er.ipj, + energy_per_output_token_joules=er.energy_per_output_token_joules, + throughput_per_watt=er.throughput_per_watt, + mean_itl_ms=er.mean_itl_ms, + ) + ) + + return TrialResult( + trial_id=trial.trial_id, + config=trial, + accuracy=summary.accuracy, + mean_latency_seconds=summary.mean_latency_seconds, + total_cost_usd=summary.total_cost_usd, + total_energy_joules=summary.total_energy_joules, + total_tokens=total_tokens, + samples_evaluated=summary.total_samples, + failure_modes=failure_modes, + summary=summary, + sample_scores=sample_scores, + ) + + +class MultiBenchTrialRunner: + """Evaluates a proposed config across multiple benchmarks. + + Delegates to :class:`TrialRunner` per benchmark, then aggregates + results into a single composite :class:`TrialResult` with weighted + metrics and per-benchmark breakdowns. + """ + + def __init__( + self, + benchmark_specs: List[BenchmarkSpec], + judge_model: str = "gpt-5-mini-2025-08-07", + output_dir: str = "results/optimize/", + ) -> None: + self.benchmark_specs = benchmark_specs + self.judge_model = judge_model + self.output_dir = output_dir + + def run_trial(self, trial: TrialConfig) -> TrialResult: + """Run *trial* against all benchmarks and return a composite result.""" + per_benchmark: List[BenchmarkScore] = [] + + for spec in self.benchmark_specs: + if spec.benchmark == "terminalbench-native": + score = self._run_terminalbench_native(trial, spec) + else: + runner = TrialRunner( + benchmark=spec.benchmark, + max_samples=spec.max_samples, + judge_model=self.judge_model, + output_dir=self.output_dir, + ) + result = runner.run_trial(trial) + score = BenchmarkScore( + benchmark=spec.benchmark, + accuracy=result.accuracy, + mean_latency_seconds=result.mean_latency_seconds, + total_cost_usd=result.total_cost_usd, + total_energy_joules=result.total_energy_joules, + total_tokens=result.total_tokens, + samples_evaluated=result.samples_evaluated, + errors=len([s for s in result.sample_scores if s.error]), + weight=spec.weight, + summary=result.summary, + sample_scores=result.sample_scores, + ) + per_benchmark.append(score) + + return self._aggregate(trial, per_benchmark) + + def _run_terminalbench_native( + self, trial: TrialConfig, spec: BenchmarkSpec, + ) -> BenchmarkScore: + """Run terminal-bench natively via Harness with Docker execution.""" + import time + + from terminal_bench import BenchmarkResults, Harness + + recipe = trial.to_recipe() + model_name = recipe.model or "default" + + # For vLLM-served models, use openai/ prefix for litellm + if recipe.engine_key == "vllm": + litellm_model = f"openai/{model_name}" + api_base = "http://localhost:8000/v1" + else: + litellm_model = model_name + api_base = "http://localhost:8000/v1" + + temperature = recipe.temperature if recipe.temperature is not None else 0.2 + + output_path = Path(self.output_dir) / f"terminalbench/{trial.trial_id}" + output_path.mkdir(parents=True, exist_ok=True) + + harness_kwargs = { + "output_path": output_path, + "run_id": trial.trial_id, + "dataset_name": "terminal-bench-core", + "dataset_version": "0.1.1", + "model_name": litellm_model, + "agent_import_path": ( + "openjarvis.evals.backends.tb_agent" + ":OpenJarvisTerminalBenchAgent" + ), + "agent_kwargs": { + "model_name": litellm_model, + "api_base": api_base, + "temperature": temperature, + }, + "n_concurrent_trials": 4, + "cleanup": True, + } + + if spec.max_samples and spec.max_samples < 200: + harness_kwargs["n_tasks"] = spec.max_samples + + LOGGER.info( + "Running terminal-bench native: model=%s, max_tasks=%s", + litellm_model, spec.max_samples, + ) + + t0 = time.monotonic() + try: + harness = Harness(**harness_kwargs) + tb_results: BenchmarkResults = harness.run() + except Exception as e: + LOGGER.error("terminal-bench harness failed: %s", e) + return BenchmarkScore( + benchmark="terminalbench-native", + accuracy=0.0, + mean_latency_seconds=0.0, + samples_evaluated=0, + errors=1, + weight=spec.weight, + ) + elapsed = time.monotonic() - t0 + + # Extract results + total_tasks = len(tb_results.results) + resolved = sum(1 for r in tb_results.results if r.is_resolved) + accuracy = resolved / total_tasks if total_tasks > 0 else 0.0 + mean_latency = elapsed / total_tasks if total_tasks > 0 else 0.0 + + total_input = sum(r.total_input_tokens or 0 for r in tb_results.results) + total_output = sum(r.total_output_tokens or 0 for r in tb_results.results) + + # Build sample scores + sample_scores = [] + for r in tb_results.results: + sample_scores.append( + SampleScore( + record_id=f"terminalbench-native-{r.task_id}", + is_correct=bool(r.is_resolved), + score=1.0 if r.is_resolved else 0.0, + prompt_tokens=r.total_input_tokens or 0, + completion_tokens=r.total_output_tokens or 0, + ) + ) + + LOGGER.info( + "terminal-bench native: %d/%d resolved (%.1f%%), %.1fs total", + resolved, total_tasks, accuracy * 100, elapsed, + ) + + return BenchmarkScore( + benchmark="terminalbench-native", + accuracy=accuracy, + mean_latency_seconds=mean_latency, + total_tokens=total_input + total_output, + samples_evaluated=total_tasks, + errors=sum( + 1 for r in tb_results.results + if r.failure_mode.value not in ("none", "unset") + ), + weight=spec.weight, + sample_scores=sample_scores, + ) + + @staticmethod + def _aggregate( + trial: TrialConfig, + per_benchmark: List[BenchmarkScore], + ) -> TrialResult: + """Compute weighted-aggregate metrics from per-benchmark scores.""" + total_weight = sum(b.weight for b in per_benchmark) or 1.0 + accuracy = sum(b.accuracy * b.weight for b in per_benchmark) / total_weight + + # Weighted mean latency by samples evaluated + total_samples = sum(b.samples_evaluated for b in per_benchmark) or 1 + mean_latency = ( + sum(b.mean_latency_seconds * b.samples_evaluated for b in per_benchmark) + / total_samples + ) + + # Sums across benchmarks + total_cost = sum(b.total_cost_usd for b in per_benchmark) + total_energy = sum(b.total_energy_joules for b in per_benchmark) + total_tokens = sum(b.total_tokens for b in per_benchmark) + + # Merge all sample scores + all_scores: List[SampleScore] = [] + failure_modes: List[str] = [] + for b in per_benchmark: + all_scores.extend(b.sample_scores) + if b.errors > 0: + failure_modes.append(f"{b.benchmark}: {b.errors} errors") + + return TrialResult( + trial_id=trial.trial_id, + config=trial, + accuracy=accuracy, + mean_latency_seconds=mean_latency, + total_cost_usd=total_cost, + total_energy_joules=total_energy, + total_tokens=total_tokens, + samples_evaluated=total_samples, + failure_modes=failure_modes, + sample_scores=all_scores, + per_benchmark=per_benchmark, + ) + + +__all__ = ["BenchmarkSpec", "MultiBenchTrialRunner", "TrialRunner"] diff --git a/src/openjarvis/learning/optimize/types.py b/src/openjarvis/learning/optimize/types.py new file mode 100644 index 00000000..ce1eef26 --- /dev/null +++ b/src/openjarvis/learning/optimize/types.py @@ -0,0 +1,253 @@ +"""Core data types for the optimization framework.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from openjarvis.evals.core.types import RunSummary +from openjarvis.recipes.loader import Recipe + + +@dataclass(slots=True) +class SearchDimension: + """One tunable dimension in the config space.""" + + name: str # e.g. "agent.type", "intelligence.temperature" + dim_type: str # "categorical", "continuous", "integer", "subset", "text" + # categorical/subset: explicit options + values: List[Any] = field(default_factory=list) + low: Optional[float] = None # continuous/integer lower bound + high: Optional[float] = None # continuous/integer upper bound + description: str = "" # human-readable explanation for the LLM optimizer + pillar: str = "" # intelligence | engine | agent | tools | learning + + +@dataclass(slots=True) +class SearchSpace: + """The full space of configs the optimizer can propose.""" + + dimensions: List[SearchDimension] = field(default_factory=list) + fixed: Dict[str, Any] = field(default_factory=dict) # params NOT being optimized + constraints: List[str] = field(default_factory=list) # natural language constraints + + def to_prompt_description(self) -> str: + """Render search space as structured text for the LLM optimizer.""" + lines: List[str] = [] + lines.append("# Search Space") + lines.append("") + + # Group dimensions by pillar + by_pillar: Dict[str, List[SearchDimension]] = {} + for dim in self.dimensions: + key = dim.pillar or "other" + by_pillar.setdefault(key, []).append(dim) + + for pillar, dims in sorted(by_pillar.items()): + lines.append(f"## {pillar.title()}") + for dim in dims: + lines.append(f"- **{dim.name}** ({dim.dim_type})") + if dim.description: + lines.append(f" Description: {dim.description}") + if dim.dim_type in ("categorical", "subset"): + lines.append(f" Options: {dim.values}") + elif dim.dim_type in ("continuous", "integer"): + lines.append(f" Range: [{dim.low}, {dim.high}]") + elif dim.dim_type == "text": + lines.append(" Free-form text") + lines.append("") + + if self.fixed: + lines.append("## Fixed Parameters") + for k, v in sorted(self.fixed.items()): + lines.append(f"- {k} = {v}") + lines.append("") + + if self.constraints: + lines.append("## Constraints") + for c in self.constraints: + lines.append(f"- {c}") + lines.append("") + + return "\n".join(lines) + + +# Mapping from dotted param names to Recipe constructor fields. +_PARAM_TO_RECIPE: Dict[str, str] = { + "intelligence.model": "model", + "intelligence.temperature": "temperature", + "intelligence.max_tokens": "max_tokens", + "intelligence.quantization": "quantization", + "engine.backend": "engine_key", + "agent.type": "agent_type", + "agent.max_turns": "max_turns", + "agent.system_prompt": "system_prompt", + "intelligence.system_prompt": "system_prompt", + "tools.tool_set": "tools", + "learning.routing_policy": "routing_policy", + "learning.agent_policy": "agent_policy", +} + + +@dataclass(slots=True) +class BenchmarkScore: + """Per-benchmark metrics from a multi-benchmark evaluation trial.""" + + benchmark: str + accuracy: float = 0.0 + mean_latency_seconds: float = 0.0 + total_cost_usd: float = 0.0 + total_energy_joules: float = 0.0 + total_tokens: int = 0 + samples_evaluated: int = 0 + errors: int = 0 + weight: float = 1.0 + summary: Optional[Any] = None # RunSummary + sample_scores: List["SampleScore"] = field(default_factory=list) + + +@dataclass(slots=True) +class SampleScore: + """Per-sample metrics from an evaluation trial.""" + + record_id: str + is_correct: Optional[bool] = None + score: Optional[float] = None + latency_seconds: float = 0.0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + cost_usd: float = 0.0 + error: Optional[str] = None + ttft: float = 0.0 + energy_joules: float = 0.0 + power_watts: float = 0.0 + gpu_utilization_pct: float = 0.0 + throughput_tok_per_sec: float = 0.0 + mfu_pct: float = 0.0 + mbu_pct: float = 0.0 + ipw: float = 0.0 + ipj: float = 0.0 + energy_per_output_token_joules: float = 0.0 + throughput_per_watt: float = 0.0 + mean_itl_ms: float = 0.0 + + +@dataclass(slots=True) +class TrialFeedback: + """Structured feedback from trial analysis.""" + + summary_text: str = "" + failure_patterns: List[str] = field(default_factory=list) + pillar_ratings: Dict[str, str] = field(default_factory=dict) + suggested_changes: List[str] = field(default_factory=list) + target_pillar: str = "" + + +@dataclass(slots=True) +class ObjectiveSpec: + """A single optimization objective.""" + + metric: str + direction: str # "maximize" or "minimize" + weight: float = 1.0 + + +DEFAULT_OBJECTIVES = [ + ObjectiveSpec("accuracy", "maximize"), + ObjectiveSpec("mean_latency_seconds", "minimize"), + ObjectiveSpec("total_cost_usd", "minimize"), +] + +ALL_OBJECTIVES = [ + ObjectiveSpec("accuracy", "maximize"), + ObjectiveSpec("mean_latency_seconds", "minimize"), + ObjectiveSpec("total_cost_usd", "minimize"), + ObjectiveSpec("total_energy_joules", "minimize"), + ObjectiveSpec("avg_power_watts", "minimize"), + ObjectiveSpec("throughput_tok_per_sec", "maximize"), + ObjectiveSpec("mfu_pct", "maximize"), + ObjectiveSpec("mbu_pct", "maximize"), + ObjectiveSpec("ipw", "maximize"), + ObjectiveSpec("ipj", "maximize"), + ObjectiveSpec("energy_per_output_token", "minimize"), + ObjectiveSpec("throughput_per_watt", "maximize"), + ObjectiveSpec("ttft", "minimize"), + ObjectiveSpec("mean_itl_ms", "minimize"), +] + + +@dataclass(slots=True) +class TrialConfig: + """A single candidate configuration proposed by the optimizer.""" + + trial_id: str + params: Dict[str, Any] = field(default_factory=dict) # dotted keys -> values + reasoning: str = "" # optimizer's explanation + + def to_recipe(self) -> Recipe: + """Map params back to Recipe fields.""" + kwargs: Dict[str, Any] = {} + for dotted_key, value in self.params.items(): + recipe_field = _PARAM_TO_RECIPE.get(dotted_key) + if recipe_field is not None: + kwargs[recipe_field] = value + + return Recipe( + name=f"trial-{self.trial_id}", + **kwargs, + ) + + +@dataclass(slots=True) +class TrialResult: + """Result of evaluating a trial, with both scalar and textual feedback.""" + + trial_id: str + config: TrialConfig + accuracy: float = 0.0 + mean_latency_seconds: float = 0.0 + total_cost_usd: float = 0.0 + total_energy_joules: float = 0.0 + total_tokens: int = 0 + samples_evaluated: int = 0 + analysis: str = "" + failure_modes: List[str] = field(default_factory=list) + per_sample_feedback: List[Dict[str, Any]] = field(default_factory=list) + summary: Optional[RunSummary] = None + sample_scores: List[SampleScore] = field(default_factory=list) + structured_feedback: Optional[TrialFeedback] = None + per_benchmark: List[BenchmarkScore] = field(default_factory=list) + + +@dataclass(slots=True) +class OptimizationRun: + """Complete optimization session.""" + + run_id: str + search_space: SearchSpace + trials: List[TrialResult] = field(default_factory=list) + best_trial: Optional[TrialResult] = None + best_recipe_path: Optional[str] = None + status: str = "running" # running | completed | failed + optimizer_model: str = "" + benchmark: str = "" + benchmarks: List[str] = field(default_factory=list) + pareto_frontier: List[TrialResult] = field(default_factory=list) + objectives: List[ObjectiveSpec] = field( + default_factory=lambda: list(DEFAULT_OBJECTIVES), + ) + + +__all__ = [ + "ALL_OBJECTIVES", + "BenchmarkScore", + "DEFAULT_OBJECTIVES", + "ObjectiveSpec", + "OptimizationRun", + "SampleScore", + "SearchDimension", + "SearchSpace", + "TrialConfig", + "TrialFeedback", + "TrialResult", +] diff --git a/src/openjarvis/optimize/__init__.py b/src/openjarvis/optimize/__init__.py index e83d7224..d79d9c29 100644 --- a/src/openjarvis/optimize/__init__.py +++ b/src/openjarvis/optimize/__init__.py @@ -1,55 +1,3 @@ -"""Optimization framework for OpenJarvis configuration tuning.""" - -from openjarvis.optimize.config import ( - load_benchmark_specs, - load_objectives, - load_optimize_config, -) -from openjarvis.optimize.llm_optimizer import LLMOptimizer -from openjarvis.optimize.optimizer import OptimizationEngine, compute_pareto_frontier -from openjarvis.optimize.search_space import DEFAULT_SEARCH_SPACE, build_search_space -from openjarvis.optimize.store import OptimizationStore -from openjarvis.optimize.trial_runner import ( - BenchmarkSpec, - MultiBenchTrialRunner, - TrialRunner, -) -from openjarvis.optimize.types import ( - ALL_OBJECTIVES, - DEFAULT_OBJECTIVES, - BenchmarkScore, - ObjectiveSpec, - OptimizationRun, - SampleScore, - SearchDimension, - SearchSpace, - TrialConfig, - TrialFeedback, - TrialResult, -) - -__all__ = [ - "ALL_OBJECTIVES", - "BenchmarkScore", - "BenchmarkSpec", - "DEFAULT_OBJECTIVES", - "DEFAULT_SEARCH_SPACE", - "LLMOptimizer", - "MultiBenchTrialRunner", - "ObjectiveSpec", - "OptimizationEngine", - "OptimizationRun", - "OptimizationStore", - "SampleScore", - "SearchDimension", - "SearchSpace", - "TrialConfig", - "TrialFeedback", - "TrialResult", - "TrialRunner", - "build_search_space", - "compute_pareto_frontier", - "load_benchmark_specs", - "load_objectives", - "load_optimize_config", -] +"""Backward-compatibility shim -- optimize moved to learning.optimize.""" +from openjarvis.learning.optimize import * # noqa: F401,F403 +from openjarvis.learning.optimize import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/config.py b/src/openjarvis/optimize/config.py index f457734b..84dc486b 100644 --- a/src/openjarvis/optimize/config.py +++ b/src/openjarvis/optimize/config.py @@ -1,102 +1,3 @@ -"""TOML config loader for optimization runs.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any, Dict, List, Union - -from openjarvis.optimize.types import ObjectiveSpec - -try: - import tomllib -except ModuleNotFoundError: # pragma: no cover - import tomli as tomllib # type: ignore[no-redef] - - -def load_optimize_config(path: Union[str, Path]) -> Dict[str, Any]: - """Load an optimization config TOML file. - - Returns the raw dict with keys such as ``optimize.max_trials``, - ``optimize.benchmark``, ``optimize.search``, ``optimize.fixed``, - ``optimize.constraints``, etc. - - Raises: - FileNotFoundError: If *path* does not exist. - """ - path = Path(path) - if not path.exists(): - raise FileNotFoundError(f"Optimization config not found: {path}") - - with open(path, "rb") as fh: - data: Dict[str, Any] = tomllib.load(fh) - - return data - - -def load_objectives(data: Dict[str, Any]) -> List[ObjectiveSpec]: - """Extract objectives from a loaded optimization config. - - Reads ``optimize.objectives`` (a list of tables) and returns - a list of :class:`ObjectiveSpec`. Falls back to - :data:`DEFAULT_OBJECTIVES` if the key is absent. - """ - from openjarvis.optimize.types import DEFAULT_OBJECTIVES - - optimize = data.get("optimize", {}) - raw_objectives = optimize.get("objectives") - if not raw_objectives: - return list(DEFAULT_OBJECTIVES) - - result: List[ObjectiveSpec] = [] - for obj in raw_objectives: - result.append( - ObjectiveSpec( - metric=obj["metric"], - direction=obj.get("direction", "maximize"), - weight=obj.get("weight", 1.0), - ) - ) - return result - - -def load_benchmark_specs(data: Dict[str, Any]) -> List[Any]: - """Extract benchmark specs from a loaded optimization config. - - Supports two formats: - - Multi-benchmark: ``[[optimize.benchmarks]]`` array of tables - - Single-benchmark fallback: ``optimize.benchmark`` string - - Returns a list of :class:`BenchmarkSpec` (from trial_runner). - Returns an empty list if no benchmarks are configured (caller - should fall back to CLI --benchmark). - """ - from openjarvis.optimize.trial_runner import BenchmarkSpec - - optimize = data.get("optimize", {}) - - # Multi-benchmark format: [[optimize.benchmarks]] - raw_benchmarks = optimize.get("benchmarks") - if raw_benchmarks and isinstance(raw_benchmarks, list): - # Check if it's a list of dicts (table array) vs list of strings - if raw_benchmarks and isinstance(raw_benchmarks[0], dict): - specs: List[BenchmarkSpec] = [] - for entry in raw_benchmarks: - specs.append( - BenchmarkSpec( - benchmark=entry.get("name", entry.get("benchmark", "")), - max_samples=entry.get("max_samples", 200), - weight=entry.get("weight", 1.0), - ) - ) - return specs - - # Single-benchmark fallback - single = optimize.get("benchmark", "") - if single: - max_samples = optimize.get("max_samples", 50) - return [BenchmarkSpec(benchmark=single, max_samples=max_samples)] - - return [] - - -__all__ = ["load_benchmark_specs", "load_objectives", "load_optimize_config"] +"""Backward-compatibility shim -- optimize.config moved to learning.optimize.config.""" +from openjarvis.learning.optimize.config import * # noqa: F401,F403 +from openjarvis.learning.optimize.config import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/feedback/__init__.py b/src/openjarvis/optimize/feedback/__init__.py index e9292954..589fc3cc 100644 --- a/src/openjarvis/optimize/feedback/__init__.py +++ b/src/openjarvis/optimize/feedback/__init__.py @@ -1,6 +1,3 @@ -"""Feedback subsystem: LLM-as-judge scoring and signal aggregation.""" - -from openjarvis.optimize.feedback.collector import FeedbackCollector -from openjarvis.optimize.feedback.judge import TraceJudge - -__all__ = ["TraceJudge", "FeedbackCollector"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.feedback import * # noqa: F401,F403 +from openjarvis.learning.optimize.feedback import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/feedback/collector.py b/src/openjarvis/optimize/feedback/collector.py index f62301cb..995c176a 100644 --- a/src/openjarvis/optimize/feedback/collector.py +++ b/src/openjarvis/optimize/feedback/collector.py @@ -1,115 +1,3 @@ -"""FeedbackCollector -- aggregates feedback from multiple sources.""" - -from __future__ import annotations - -import time -from typing import Any, Dict, List, Optional - -from openjarvis.core.types import Trace -from openjarvis.optimize.feedback.judge import TraceJudge - - -class FeedbackCollector: - """Collects feedback signals: explicit user scores + LLM judge evaluations. - - Signals are stored in-memory as dicts with at least ``trace_id``, - ``score``, ``source``, and ``timestamp`` keys. - """ - - def __init__(self) -> None: - self._records: List[Dict[str, Any]] = [] - - # ------------------------------------------------------------------ - # Recording helpers - # ------------------------------------------------------------------ - - def record_explicit( - self, - trace_id: str, - score: float, - source: str = "api", - ) -> None: - """Record an explicit numeric score (0-1) for a trace.""" - self._records.append({ - "trace_id": trace_id, - "score": min(max(score, 0.0), 1.0), - "source": source, - "timestamp": time.time(), - }) - - def record_thumbs(self, trace_id: str, thumbs_up: bool) -> None: - """Record a thumbs-up / thumbs-down signal (converted to 1.0/0.0).""" - self._records.append({ - "trace_id": trace_id, - "score": 1.0 if thumbs_up else 0.0, - "source": "thumbs", - "timestamp": time.time(), - }) - - # ------------------------------------------------------------------ - # Judge-driven evaluation - # ------------------------------------------------------------------ - - def evaluate_traces( - self, - traces: List[Trace], - judge: TraceJudge, - ) -> List[Dict[str, Any]]: - """Score *traces* via the LLM judge and record the results. - - Returns the list of newly created feedback records. - """ - new_records: List[Dict[str, Any]] = [] - for trace in traces: - score, feedback = judge.score_trace(trace) - record: Dict[str, Any] = { - "trace_id": trace.trace_id, - "score": score, - "source": "judge", - "feedback": feedback, - "timestamp": time.time(), - } - self._records.append(record) - new_records.append(record) - return new_records - - # ------------------------------------------------------------------ - # Queries - # ------------------------------------------------------------------ - - def get_records( - self, trace_id: Optional[str] = None, - ) -> List[Dict[str, Any]]: - """Return stored records, optionally filtered by *trace_id*.""" - if trace_id is None: - return list(self._records) - return [r for r in self._records if r["trace_id"] == trace_id] - - def stats(self) -> Dict[str, Any]: - """Return aggregate statistics over all recorded feedback. - - Returns a dict with ``count``, ``mean_score``, and a simple - ``distribution`` bucket (low / medium / high). - """ - count = len(self._records) - if count == 0: - return { - "count": 0, - "mean_score": 0.0, - "distribution": {"low": 0, "medium": 0, "high": 0}, - } - - scores = [r["score"] for r in self._records] - mean_score = sum(scores) / count - low = sum(1 for s in scores if s < 0.4) - medium = sum(1 for s in scores if 0.4 <= s < 0.7) - high = sum(1 for s in scores if s >= 0.7) - - return { - "count": count, - "mean_score": mean_score, - "distribution": {"low": low, "medium": medium, "high": high}, - } - - -__all__ = ["FeedbackCollector"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.feedback.collector import * # noqa: F401,F403 +from openjarvis.learning.optimize.feedback.collector import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/feedback/judge.py b/src/openjarvis/optimize/feedback/judge.py index 24012766..9010fc3a 100644 --- a/src/openjarvis/optimize/feedback/judge.py +++ b/src/openjarvis/optimize/feedback/judge.py @@ -1,118 +1,6 @@ -"""TraceJudge -- LLM-as-judge scoring for agent traces.""" - -from __future__ import annotations - -import logging -import re -from typing import List, Tuple - -from openjarvis.core.types import Trace -from openjarvis.evals.core.backend import InferenceBackend - -LOGGER = logging.getLogger(__name__) - -_SYSTEM_PROMPT = ( - "You are an expert evaluator of AI assistant traces. " - "You will be shown a user query, the steps the assistant took, " - "and its final result. Rate the overall quality of the response " - "on a scale from 0.0 (completely wrong / unhelpful) to 1.0 " - "(perfect). Provide your score on the first line as a decimal " - "number, then explain your reasoning." +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.feedback.judge import * # noqa: F401,F403 +from openjarvis.learning.optimize.feedback.judge import ( + __all__, # noqa: F401 + _parse_score, # noqa: F401 ) - -_SCORE_RE = re.compile( - r"(?:Score|Rating|Quality)?\s*[:=]?\s*(\d+(?:\.\d+)?)" - r"(?:\s*/\s*(\d+(?:\.\d+)?))?", - re.IGNORECASE, -) - - -def _format_trace(trace: Trace) -> str: - """Render a Trace into a textual prompt for the judge.""" - lines: List[str] = [] - lines.append(f"## Query\n{trace.query}") - if trace.steps: - lines.append("\n## Steps") - for i, step in enumerate(trace.steps, 1): - step_input = step.input.get("content", str(step.input)) - step_output = step.output.get("content", str(step.output)) - lines.append( - f"{i}. [{step.step_type.value}] " - f"input={step_input!r} output={step_output!r} " - f"({step.duration_seconds:.3f}s)", - ) - lines.append(f"\n## Final Result\n{trace.result}") - return "\n".join(lines) - - -def _parse_score(text: str) -> float: - """Extract a 0-1 score from the judge response. - - Handles formats like ``0.85``, ``Score: 0.85``, ``Rating: 7/10``. - Falls back to 0.5 when parsing fails. - """ - match = _SCORE_RE.search(text) - if match is None: - LOGGER.warning("Could not parse score from judge response; defaulting to 0.5") - return 0.5 - - numerator = float(match.group(1)) - denominator_str = match.group(2) - - if denominator_str is not None: - denominator = float(denominator_str) - if denominator > 0: - return min(max(numerator / denominator, 0.0), 1.0) - return 0.5 - - # If the number is > 1.0 assume it is on a 0-10 scale - if numerator > 1.0: - return min(numerator / 10.0, 1.0) - - return min(max(numerator, 0.0), 1.0) - - -class TraceJudge: - """LLM-as-judge for scoring traces when no ground truth exists. - - Given a :class:`Trace`, the judge constructs a prompt showing the - query, agent steps, and final result, then asks an LLM to rate the - quality on a 0-1 scale. - """ - - def __init__(self, backend: InferenceBackend, model: str) -> None: - self._backend = backend - self._model = model - - def score_trace(self, trace: Trace) -> Tuple[float, str]: - """Score a single trace. - - Returns: - ``(score, feedback)`` where *score* is in [0, 1] and - *feedback* is the judge's textual reasoning. - """ - prompt = _format_trace(trace) - response = self._backend.generate( - prompt, - model=self._model, - system=_SYSTEM_PROMPT, - temperature=0.0, - max_tokens=1024, - ) - score = _parse_score(response) - return score, response - - def batch_evaluate( - self, traces: List[Trace], - ) -> List[Tuple[float, str]]: - """Evaluate multiple traces sequentially. - - Returns a list of ``(score, feedback)`` tuples, one per trace. - """ - results: List[Tuple[float, str]] = [] - for trace in traces: - results.append(self.score_trace(trace)) - return results - - -__all__ = ["TraceJudge"] diff --git a/src/openjarvis/optimize/llm_optimizer.py b/src/openjarvis/optimize/llm_optimizer.py index f6a7ca0d..1246b280 100644 --- a/src/openjarvis/optimize/llm_optimizer.py +++ b/src/openjarvis/optimize/llm_optimizer.py @@ -1,691 +1,3 @@ -"""LLM-based optimizer for OpenJarvis configuration tuning. - -Uses a cloud LLM to propose optimal OpenJarvis configs, inspired by DSPy's -GEPA approach: textual feedback from execution traces rather than just scalar -rewards guides the optimizer toward better configurations. -""" - -from __future__ import annotations - -import json -import re -import uuid -from typing import Any, Dict, List, Optional - -from openjarvis.core.types import Trace -from openjarvis.evals.core.backend import InferenceBackend -from openjarvis.evals.core.types import RunSummary -from openjarvis.optimize.types import ( - BenchmarkScore, - SampleScore, - SearchSpace, - TrialConfig, - TrialFeedback, - TrialResult, -) - - -class LLMOptimizer: - """Uses a cloud LLM to propose optimal OpenJarvis configs. - - Inspired by DSPy's GEPA: uses textual feedback from execution - traces rather than just scalar rewards. - """ - - def __init__( - self, - search_space: SearchSpace, - optimizer_model: str = "claude-sonnet-4-6", - optimizer_backend: Optional[InferenceBackend] = None, - ) -> None: - self.search_space = search_space - self.optimizer_model = optimizer_model - self.optimizer_backend = optimizer_backend - - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - def propose_initial(self) -> TrialConfig: - """Propose a reasonable starting config from the search space.""" - if self.optimizer_backend is None: - raise ValueError( - "optimizer_backend is required to propose configurations" - ) - - prompt = self._build_initial_prompt() - response = self.optimizer_backend.generate( - prompt, - model=self.optimizer_model, - system="You are an expert AI systems optimizer.", - temperature=0.7, - max_tokens=2048, - ) - return self._parse_config_response(response) - - def propose_next( - self, - history: List[TrialResult], - traces: Optional[List[Trace]] = None, - frontier_ids: Optional[set] = None, - ) -> TrialConfig: - """Ask the LLM to propose the next config to evaluate.""" - if self.optimizer_backend is None: - raise ValueError( - "optimizer_backend is required to propose configurations" - ) - - prompt = self._build_propose_prompt(history, traces, frontier_ids=frontier_ids) - response = self.optimizer_backend.generate( - prompt, - model=self.optimizer_model, - system="You are an expert AI systems optimizer.", - temperature=0.7, - max_tokens=2048, - ) - return self._parse_config_response(response) - - def analyze_trial( - self, - trial: TrialConfig, - summary: RunSummary, - traces: Optional[List[Trace]] = None, - sample_scores: Optional[List[SampleScore]] = None, - per_benchmark: Optional[List[BenchmarkScore]] = None, - ) -> TrialFeedback: - """Ask the LLM to analyze a completed trial. Returns structured feedback.""" - if self.optimizer_backend is None: - raise ValueError( - "optimizer_backend is required to analyze trials" - ) - - prompt = self._build_analyze_prompt( - trial, summary, traces, sample_scores, per_benchmark, - ) - response = self.optimizer_backend.generate( - prompt, - model=self.optimizer_model, - system="You are an expert AI systems analyst.", - temperature=0.3, - max_tokens=2048, - ) - return self._parse_feedback_response(response) - - # ------------------------------------------------------------------ - # Prompt builders - # ------------------------------------------------------------------ - - def _build_initial_prompt(self) -> str: - """Construct the prompt for the initial config proposal.""" - lines: List[str] = [] - lines.append( - "You are optimizing an OpenJarvis AI system configuration." - ) - lines.append("") - lines.append(self.search_space.to_prompt_description()) - lines.append("## Objective") - lines.append( - "Maximize accuracy while minimizing latency and cost." - ) - lines.append("") - lines.append("## Your Task") - lines.append( - "Propose an initial configuration that is a reasonable starting " - "point for optimization. Choose sensible defaults that balance " - "accuracy, latency, and cost." - ) - lines.append("") - lines.append( - "Return a JSON object inside a ```json code block with:" - ) - lines.append( - '1. "params": dict of config params (dotted keys matching ' - "the search space)" - ) - lines.append( - '2. "reasoning": string explaining why this is a good ' - "starting configuration" - ) - return "\n".join(lines) - - def _build_propose_prompt( - self, - history: List[TrialResult], - traces: Optional[List[Trace]] = None, - frontier_ids: Optional[set] = None, - ) -> str: - """Construct the full prompt for propose_next.""" - lines: List[str] = [] - lines.append( - "You are optimizing an OpenJarvis AI system configuration." - ) - lines.append("") - lines.append(self.search_space.to_prompt_description()) - - lines.append("## Optimization History") - if history: - lines.append(self._format_history(history, frontier_ids=frontier_ids)) - else: - lines.append("No trials have been run yet.") - lines.append("") - - if traces: - lines.append("## Recent Execution Traces") - lines.append(self._format_traces(traces)) - lines.append("") - - lines.append("## Objective") - lines.append( - "Maximize accuracy while minimizing latency and cost." - ) - lines.append("") - lines.append("## Your Task") - lines.append( - "Propose the next configuration to evaluate. Learn from " - "previous trials to improve results." - ) - lines.append("") - lines.append( - "Return a JSON object inside a ```json code block with:" - ) - lines.append( - '1. "params": dict of config params (dotted keys matching ' - "the search space)" - ) - lines.append( - '2. "reasoning": string explaining why this config should ' - "improve results" - ) - return "\n".join(lines) - - def _build_analyze_prompt( - self, - trial: TrialConfig, - summary: RunSummary, - traces: Optional[List[Trace]] = None, - sample_scores: Optional[List[SampleScore]] = None, - per_benchmark: Optional[List[BenchmarkScore]] = None, - ) -> str: - """Construct the prompt for analyze_trial.""" - lines: List[str] = [] - lines.append("Analyze this OpenJarvis evaluation result.") - lines.append("") - - lines.append("## Configuration") - for key, value in sorted(trial.params.items()): - lines.append(f"- {key}: {value}") - if trial.reasoning: - lines.append(f"\nOptimizer reasoning: {trial.reasoning}") - lines.append("") - - # Per-benchmark breakdown (multi-benchmark mode) - if per_benchmark: - lines.append("## Per-Benchmark Results") - total_weight = sum(b.weight for b in per_benchmark) or 1.0 - for b in per_benchmark: - lines.append( - f"### {b.benchmark} (weight={b.weight:.1f}): " - f"accuracy={b.accuracy:.4f}, " - f"latency={b.mean_latency_seconds:.2f}s, " - f"cost=${b.total_cost_usd:.4f}, " - f"energy={b.total_energy_joules:.2f}J, " - f"samples={b.samples_evaluated}, errors={b.errors}" - ) - weighted_acc = ( - sum(b.accuracy * b.weight for b in per_benchmark) / total_weight - ) - lines.append(f"\nOverall weighted accuracy: {weighted_acc:.4f}") - lines.append("") - - lines.append("## Aggregate Results") - lines.append(f"- accuracy: {summary.accuracy:.4f}") - lines.append( - f"- mean_latency_seconds: {summary.mean_latency_seconds:.4f}" - ) - lines.append(f"- total_cost_usd: {summary.total_cost_usd:.4f}") - lines.append(f"- total_samples: {summary.total_samples}") - lines.append(f"- scored_samples: {summary.scored_samples}") - lines.append(f"- correct: {summary.correct}") - lines.append(f"- errors: {summary.errors}") - if summary.per_subject: - lines.append("\n### Per-Subject Breakdown") - for subject, metrics in sorted(summary.per_subject.items()): - metrics_str = ", ".join( - f"{k}={v:.3f}" for k, v in sorted(metrics.items()) - ) - lines.append(f"- {subject}: {metrics_str}") - lines.append("") - - if sample_scores: - lines.append("## Per-Sample Scores") - lines.append(self._format_sample_scores(sample_scores)) - lines.append("") - - if traces: - lines.append("## Sample Traces") - lines.append(self._format_traces(traces)) - lines.append("") - - lines.append( - "Provide your analysis as a JSON object inside a ```json code block with:\n" - '1. "summary_text": string with detailed analysis\n' - '2. "failure_patterns": list of identified failure patterns\n' - '3. "pillar_ratings": dict mapping pillar names to "high"/"medium"/"low"\n' - '4. "suggested_changes": list of specific config changes to try\n' - '5. "target_pillar": which pillar to change next ' - "(intelligence/engine/agent/tools/learning)" - ) - return "\n".join(lines) - - # ------------------------------------------------------------------ - # Helpers - # ------------------------------------------------------------------ - - def _format_history( - self, - history: List[TrialResult], - frontier_ids: Optional[set] = None, - ) -> str: - """Render trial history as structured text for the LLM prompt.""" - lines: List[str] = [] - for i, result in enumerate(history, 1): - tag = "" - if frontier_ids and result.trial_id in frontier_ids: - tag = " [FRONTIER]" - lines.append(f"### Trial {i} (id={result.trial_id}){tag}") - lines.append(f"Params: {json.dumps(result.config.params)}") - lines.append(f"Accuracy: {result.accuracy:.4f}") - lines.append( - f"Latency: {result.mean_latency_seconds:.4f}s" - ) - lines.append(f"Cost: ${result.total_cost_usd:.4f}") - lines.append(f"Energy: {result.total_energy_joules:.4f}J") - if result.per_benchmark: - bench_parts = [ - f"{b.benchmark}={b.accuracy:.4f}" - for b in result.per_benchmark - ] - lines.append(f"Per-benchmark accuracy: {', '.join(bench_parts)}") - if result.summary: - s = result.summary - if s.throughput_stats: - lines.append( - f"Throughput: {s.throughput_stats.mean:.2f} tok/s" - ) - if s.ipw_stats: - lines.append(f"IPW: {s.ipw_stats.mean:.4f}") - if result.structured_feedback: - fb = result.structured_feedback - if fb.failure_patterns: - lines.append( - f"Failure patterns: {', '.join(fb.failure_patterns)}" - ) - if fb.pillar_ratings: - ratings = ", ".join( - f"{k}={v}" for k, v in sorted(fb.pillar_ratings.items()) - ) - lines.append(f"Pillar ratings: {ratings}") - if fb.target_pillar: - lines.append(f"Target pillar: {fb.target_pillar}") - elif result.analysis: - lines.append(f"Analysis: {result.analysis}") - if result.failure_modes: - lines.append( - f"Failure modes: {', '.join(result.failure_modes)}" - ) - lines.append("") - return "\n".join(lines) - - def _format_traces(self, traces: List[Trace]) -> str: - """Render traces as structured text for the LLM prompt. - - Limits to the last 10 traces and truncates long outputs to keep - the prompt manageable. - """ - max_traces = 10 - max_result_len = 500 - max_steps_per_trace = 10 - - recent = traces[-max_traces:] - lines: List[str] = [] - - for trace in recent: - lines.append( - f"### Trace {trace.trace_id} " - f"(agent={trace.agent}, model={trace.model})" - ) - lines.append(f"Query: {trace.query}") - if trace.outcome: - lines.append(f"Outcome: {trace.outcome}") - if trace.feedback is not None: - lines.append(f"Feedback: {trace.feedback}") - lines.append( - f"Latency: {trace.total_latency_seconds:.3f}s, " - f"Tokens: {trace.total_tokens}" - ) - - # Show steps (limited) - steps = trace.steps[:max_steps_per_trace] - if steps: - lines.append("Steps:") - for step in steps: - step_input = json.dumps(step.input) - step_output = json.dumps(step.output) - if len(step_input) > max_result_len: - step_input = step_input[:max_result_len] + "..." - if len(step_output) > max_result_len: - step_output = step_output[:max_result_len] + "..." - lines.append( - f" - {step.step_type.value}: " - f"input={step_input}, " - f"output={step_output} " - f"({step.duration_seconds:.3f}s)" - ) - if len(trace.steps) > max_steps_per_trace: - lines.append( - f" ... ({len(trace.steps) - max_steps_per_trace} " - "more steps)" - ) - - result_text = trace.result - if len(result_text) > max_result_len: - result_text = result_text[:max_result_len] + "..." - lines.append(f"Result: {result_text}") - lines.append("") - - return "\n".join(lines) - - def propose_targeted( - self, - history: List[TrialResult], - base_config: TrialConfig, - target_pillar: str, - frontier_ids: Optional[set] = None, - ) -> TrialConfig: - """Propose a config that only changes one pillar.""" - if self.optimizer_backend is None: - raise ValueError( - "optimizer_backend is required to propose configurations" - ) - - prompt = self._build_targeted_prompt( - history, base_config, target_pillar, frontier_ids, - ) - response = self.optimizer_backend.generate( - prompt, - model=self.optimizer_model, - system="You are an expert AI systems optimizer.", - temperature=0.7, - max_tokens=2048, - ) - proposed = self._parse_config_response(response) - - # Enforce constraint: preserve non-target params from base_config - merged_params = dict(base_config.params) - for key, value in proposed.params.items(): - if key.startswith(target_pillar + ".") or key.startswith( - target_pillar.rstrip("s") + "." - ): - merged_params[key] = value - proposed.params = merged_params - return proposed - - def propose_merge( - self, - candidates: List[TrialResult], - history: List[TrialResult], - frontier_ids: Optional[set] = None, - ) -> TrialConfig: - """Combine best aspects of frontier members into one config.""" - if self.optimizer_backend is None: - raise ValueError( - "optimizer_backend is required to propose configurations" - ) - - prompt = self._build_merge_prompt(candidates, history, frontier_ids) - response = self.optimizer_backend.generate( - prompt, - model=self.optimizer_model, - system="You are an expert AI systems optimizer.", - temperature=0.7, - max_tokens=2048, - ) - return self._parse_config_response(response) - - # ------------------------------------------------------------------ - # Targeted / Merge prompt builders - # ------------------------------------------------------------------ - - def _build_targeted_prompt( - self, - history: List[TrialResult], - base_config: TrialConfig, - target_pillar: str, - frontier_ids: Optional[set] = None, - ) -> str: - """Build prompt for pillar-targeted mutation.""" - lines: List[str] = [] - lines.append( - "You are optimizing an OpenJarvis AI system configuration." - ) - lines.append("") - lines.append(self.search_space.to_prompt_description()) - - lines.append("## Base Configuration") - for key, value in sorted(base_config.params.items()): - lines.append(f"- {key}: {value}") - lines.append("") - - lines.append(f"## Target Pillar: {target_pillar}") - lines.append( - f"ONLY change parameters under the '{target_pillar}' pillar. " - "Keep all other parameters exactly as they are." - ) - lines.append("") - - lines.append("## Optimization History") - if history: - lines.append(self._format_history(history, frontier_ids=frontier_ids)) - lines.append("") - - lines.append( - "Return a JSON object inside a ```json code block with:\n" - '1. "params": dict of config params (only change ' - f"{target_pillar} params)\n" - '2. "reasoning": string explaining your changes' - ) - return "\n".join(lines) - - def _build_merge_prompt( - self, - candidates: List[TrialResult], - history: List[TrialResult], - frontier_ids: Optional[set] = None, - ) -> str: - """Build prompt for merging frontier configs.""" - lines: List[str] = [] - lines.append( - "You are optimizing an OpenJarvis AI system configuration." - ) - lines.append("") - lines.append(self.search_space.to_prompt_description()) - - lines.append("## Frontier Candidates to Merge") - for i, cand in enumerate(candidates, 1): - lines.append(f"### Candidate {i} (id={cand.trial_id})") - lines.append(f"Params: {json.dumps(cand.config.params)}") - lines.append(f"Accuracy: {cand.accuracy:.4f}") - lines.append(f"Latency: {cand.mean_latency_seconds:.4f}s") - lines.append(f"Cost: ${cand.total_cost_usd:.4f}") - lines.append(f"Energy: {cand.total_energy_joules:.4f}J") - lines.append("") - - lines.append( - "Combine the best aspects of these frontier configs into " - "one unified configuration." - ) - lines.append("") - - if history: - lines.append("## Full History") - lines.append(self._format_history(history, frontier_ids=frontier_ids)) - lines.append("") - - lines.append( - "Return a JSON object inside a ```json code block with:\n" - '1. "params": dict of merged config params\n' - '2. "reasoning": string explaining the merge strategy' - ) - return "\n".join(lines) - - # ------------------------------------------------------------------ - # Sample score + feedback helpers - # ------------------------------------------------------------------ - - def _format_sample_scores(self, scores: List[SampleScore]) -> str: - """Render per-sample scores for the LLM prompt.""" - passed = [s for s in scores if s.is_correct] - failed = [s for s in scores if s.is_correct is False] - errored = [s for s in scores if s.error] - - lines: List[str] = [] - lines.append( - f"Total: {len(scores)} | Passed: {len(passed)} | " - f"Failed: {len(failed)} | Errors: {len(errored)}" - ) - - if failed: - lines.append("\n### Failed Samples") - for s in failed[:20]: - lines.append(f"- {s.record_id}: latency={s.latency_seconds:.2f}s") - - if errored: - lines.append("\n### Error Samples") - for s in errored[:10]: - error_text = (s.error or "")[:200] - lines.append(f"- {s.record_id}: {error_text}") - - return "\n".join(lines) - - def _parse_feedback_response(self, response: str) -> TrialFeedback: - """Parse LLM response into a TrialFeedback, with fallback.""" - response = response.strip() - - # Try JSON code block - json_block_match = re.search( - r"```json\s*\n?(.*?)\n?\s*```", response, re.DOTALL - ) - raw_json = None - if json_block_match: - raw_json = json_block_match.group(1).strip() - else: - # Try generic code block - code_block_match = re.search( - r"```\s*\n?(.*?)\n?\s*```", response, re.DOTALL - ) - if code_block_match: - raw_json = code_block_match.group(1).strip() - else: - # Try raw JSON object - decoder = json.JSONDecoder() - for m in re.finditer(r"\{", response): - try: - data, _ = decoder.raw_decode(response, m.start()) - if isinstance(data, dict): - raw_json = json.dumps(data) - break - except json.JSONDecodeError: - continue - - if raw_json: - try: - data = json.loads(raw_json) - return TrialFeedback( - summary_text=data.get("summary_text", ""), - failure_patterns=data.get("failure_patterns", []), - pillar_ratings=data.get("pillar_ratings", {}), - suggested_changes=data.get("suggested_changes", []), - target_pillar=data.get("target_pillar", ""), - ) - except json.JSONDecodeError: - pass - - # Fallback: wrap raw text as summary - return TrialFeedback(summary_text=response) - - def _parse_config_response(self, response: str) -> TrialConfig: - """Extract a TrialConfig from an LLM response. - - Looks for a ```json ... ``` block first, then falls back to - finding a raw JSON object in the response text. - """ - trial_id = uuid.uuid4().hex[:12] - - # Try to extract from a ```json code block - json_block_match = re.search( - r"```json\s*\n?(.*?)\n?\s*```", response, re.DOTALL - ) - if json_block_match: - raw_json = json_block_match.group(1).strip() - try: - data = json.loads(raw_json) - return self._config_from_dict(data, trial_id) - except json.JSONDecodeError: - pass - - # Try to extract from a generic ``` code block - code_block_match = re.search( - r"```\s*\n?(.*?)\n?\s*```", response, re.DOTALL - ) - if code_block_match: - raw_json = code_block_match.group(1).strip() - try: - data = json.loads(raw_json) - return self._config_from_dict(data, trial_id) - except json.JSONDecodeError: - pass - - # Try to find a raw JSON object in the response by scanning - # for each '{' and attempting to parse from that position. - decoder = json.JSONDecoder() - for m in re.finditer(r"\{", response): - try: - data, _ = decoder.raw_decode(response, m.start()) - if isinstance(data, dict): - return self._config_from_dict(data, trial_id) - except json.JSONDecodeError: - continue - - # Last resort: return config with at least the fixed params - ss = self.search_space - fixed = dict(ss.fixed) if ss and ss.fixed else {} - return TrialConfig( - trial_id=trial_id, - params=fixed, - reasoning="Failed to parse LLM response.", - ) - - def _config_from_dict( - self, data: Dict[str, Any], trial_id: str - ) -> TrialConfig: - """Build a TrialConfig from a parsed JSON dict. - - Merges fixed parameters from the search space so that fixed values - (e.g. intelligence.model, engine.backend) are always present. - """ - params = data.get("params", {}) - reasoning = data.get("reasoning", "") - - # Inject fixed params — these override anything the LLM proposed - if self.search_space and self.search_space.fixed: - for key, value in self.search_space.fixed.items(): - params[key] = value - - return TrialConfig( - trial_id=trial_id, - params=params, - reasoning=reasoning, - ) - - -__all__ = ["LLMOptimizer"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.llm_optimizer import * # noqa: F401,F403 +from openjarvis.learning.optimize.llm_optimizer import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/optimizer.py b/src/openjarvis/optimize/optimizer.py index fde10ab8..67716608 100644 --- a/src/openjarvis/optimize/optimizer.py +++ b/src/openjarvis/optimize/optimizer.py @@ -1,445 +1,3 @@ -"""OptimizationEngine -- orchestrates the optimize loop. - -Ties together the LLM optimizer, trial runner, and persistence store -into a single propose -> evaluate -> analyze -> repeat loop. -""" - -from __future__ import annotations - -import logging -import uuid -from pathlib import Path -from typing import Any, Callable, Dict, List, Optional - -try: - import tomli_w -except ModuleNotFoundError: # pragma: no cover - tomli_w = None # type: ignore[assignment] - -from openjarvis.optimize.llm_optimizer import LLMOptimizer -from openjarvis.optimize.store import OptimizationStore -from openjarvis.optimize.trial_runner import TrialRunner -from openjarvis.optimize.types import ( - ObjectiveSpec, - OptimizationRun, - SearchSpace, - TrialResult, -) - -LOGGER = logging.getLogger(__name__) - -# Mapping from objective metric names to RunSummary stat attribute + ".mean" -_SUMMARY_STAT_MAP: Dict[str, str] = { - "avg_power_watts": "avg_power_watts", - "throughput_tok_per_sec": "throughput_stats", - "mfu_pct": "mfu_stats", - "mbu_pct": "mbu_stats", - "ipw": "ipw_stats", - "ipj": "ipj_stats", - "energy_per_output_token": "energy_per_output_token_stats", - "throughput_per_watt": "throughput_per_watt_stats", - "ttft": "ttft_stats", - "mean_itl_ms": "itl_stats", -} - - -def _get_objective_value(trial: TrialResult, obj: ObjectiveSpec) -> float: - """Read the metric value from a TrialResult for a given objective.""" - # Direct attributes on TrialResult - direct = { - "accuracy", "mean_latency_seconds", - "total_cost_usd", "total_energy_joules", - } - if obj.metric in direct: - return getattr(trial, obj.metric, 0.0) - - # avg_power_watts is a direct attr on RunSummary - if obj.metric == "avg_power_watts" and trial.summary: - return trial.summary.avg_power_watts - - # Stats-based metrics from RunSummary - stat_attr = _SUMMARY_STAT_MAP.get(obj.metric) - if stat_attr and trial.summary: - stats = getattr(trial.summary, stat_attr, None) - if stats is not None and hasattr(stats, "mean"): - return stats.mean - - return 0.0 - - -def compute_pareto_frontier( - trials: List[TrialResult], - objectives: List[ObjectiveSpec], -) -> List[TrialResult]: - """Compute the Pareto frontier: trials not dominated by any other. - - A trial A dominates trial B if A is >= B on all objectives and > B - on at least one (direction-aware: maximize flips the comparison). - """ - if not trials or not objectives: - return list(trials) - - def _values(trial: TrialResult) -> List[float]: - vals = [] - for obj in objectives: - v = _get_objective_value(trial, obj) - # Normalize: for "minimize", negate so higher is always better - if obj.direction == "minimize": - v = -v - vals.append(v) - return vals - - trial_vals = [_values(t) for t in trials] - frontier: List[TrialResult] = [] - - for i, trial in enumerate(trials): - dominated = False - for j, other in enumerate(trials): - if i == j: - continue - # Check if other dominates trial - all_ge = all( - trial_vals[j][k] >= trial_vals[i][k] - for k in range(len(objectives)) - ) - any_gt = any( - trial_vals[j][k] > trial_vals[i][k] - for k in range(len(objectives)) - ) - if all_ge and any_gt: - dominated = True - break - if not dominated: - frontier.append(trial) - - return frontier - - -class OptimizationEngine: - """Orchestrates the optimize loop: propose -> evaluate -> analyze -> repeat.""" - - def __init__( - self, - search_space: SearchSpace, - llm_optimizer: LLMOptimizer, - trial_runner: TrialRunner, - store: Optional[OptimizationStore] = None, - max_trials: int = 20, - early_stop_patience: int = 5, - ) -> None: - self.search_space = search_space - self.llm_optimizer = llm_optimizer - self.trial_runner = trial_runner - self.store = store - self.max_trials = max_trials - self.early_stop_patience = early_stop_patience - - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - def run( - self, - progress_callback: Optional[Callable[[int, int], None]] = None, - ) -> OptimizationRun: - """Execute the full optimization loop. - - 1. Generate a run_id via uuid. - 2. ``llm_optimizer.propose_initial()`` -> first config. - 3. Loop up to ``max_trials``: - a. ``trial_runner.run_trial(config)`` -> TrialResult - b. ``llm_optimizer.analyze_trial(config, summary, traces)`` - c. Update TrialResult with analysis text - d. Append to history - e. If store, ``store.save_trial(result)`` - f. Update best_trial if accuracy improved - g. Check early stopping (no improvement for *patience* trials) - h. If not stopped, ``llm_optimizer.propose_next(history)`` - 4. Set run status to ``"completed"``. - 5. If store, ``store.save_run(optimization_run)``. - 6. Return the :class:`OptimizationRun`. - - Args: - progress_callback: Optional ``(trial_num, max_trials) -> None`` - called after each trial completes. - """ - run_id = uuid.uuid4().hex[:16] - # Detect benchmark name(s) from the trial runner - from openjarvis.optimize.trial_runner import MultiBenchTrialRunner - - benchmark_name = getattr(self.trial_runner, "benchmark", "") - benchmark_names: List[str] = [] - if isinstance(self.trial_runner, MultiBenchTrialRunner): - benchmark_names = [ - s.benchmark for s in self.trial_runner.benchmark_specs - ] - benchmark_name = "+".join(benchmark_names) - - optimization_run = OptimizationRun( - run_id=run_id, - search_space=self.search_space, - status="running", - optimizer_model=self.llm_optimizer.optimizer_model, - benchmark=benchmark_name, - benchmarks=benchmark_names, - ) - - history: List[TrialResult] = [] - best_accuracy = -1.0 - trials_without_improvement = 0 - - # First config - config = self.llm_optimizer.propose_initial() - - for trial_num in range(1, self.max_trials + 1): - LOGGER.info( - "Trial %d/%d (id=%s)", - trial_num, - self.max_trials, - config.trial_id, - ) - - # Evaluate - result = self.trial_runner.run_trial(config) - - # Analyze — returns TrialFeedback - if result.summary is not None: - feedback = self.llm_optimizer.analyze_trial( - config, - result.summary, - sample_scores=result.sample_scores or None, - per_benchmark=result.per_benchmark or None, - ) - result.structured_feedback = feedback - result.analysis = feedback.summary_text - elif result.per_benchmark: - # Multi-benchmark composite: build a synthetic summary - # for analysis from per_benchmark data - from openjarvis.evals.core.types import RunSummary as _RS - - synth = _RS( - benchmark="multi", - category="multi", - backend="jarvis-agent", - model=result.config.params.get("intelligence.model", ""), - accuracy=result.accuracy, - mean_latency_seconds=result.mean_latency_seconds, - total_cost_usd=result.total_cost_usd, - total_energy_joules=result.total_energy_joules, - total_samples=result.samples_evaluated, - scored_samples=result.samples_evaluated, - correct=int( - result.accuracy * result.samples_evaluated - ), - errors=0, - total_input_tokens=0, - total_output_tokens=result.total_tokens, - ) - feedback = self.llm_optimizer.analyze_trial( - config, - synth, - per_benchmark=result.per_benchmark, - ) - result.structured_feedback = feedback - result.analysis = feedback.summary_text - else: - result.analysis = "" - - # Record - history.append(result) - optimization_run.trials.append(result) - - # Recompute Pareto frontier - optimization_run.pareto_frontier = compute_pareto_frontier( - history, optimization_run.objectives, - ) - frontier_ids = {t.trial_id for t in optimization_run.pareto_frontier} - - # Persist trial - if self.store is not None: - self.store.save_trial(run_id, result) - - # Track best - if result.accuracy > best_accuracy: - best_accuracy = result.accuracy - optimization_run.best_trial = result - trials_without_improvement = 0 - else: - trials_without_improvement += 1 - - # Progress callback - if progress_callback is not None: - progress_callback(trial_num, self.max_trials) - - # Early stopping - if trials_without_improvement >= self.early_stop_patience: - LOGGER.info( - "Early stopping after %d trials without improvement.", - self.early_stop_patience, - ) - break - - # Propose next (unless this was the last trial) - if trial_num < self.max_trials: - # Decide proposal strategy - target_pillar = "" - if result.structured_feedback: - target_pillar = result.structured_feedback.target_pillar - - if ( - trial_num % 5 == 0 - and len(optimization_run.pareto_frontier) >= 2 - ): - # Merge frontier members periodically - candidates = optimization_run.pareto_frontier[:3] - config = self.llm_optimizer.propose_merge( - candidates, history, frontier_ids=frontier_ids, - ) - elif target_pillar and trial_num > 2: - # Targeted mutation on the suggested pillar - config = self.llm_optimizer.propose_targeted( - history, - result.config, - target_pillar, - frontier_ids=frontier_ids, - ) - else: - config = self.llm_optimizer.propose_next( - history, frontier_ids=frontier_ids, - ) - - optimization_run.status = "completed" - - if self.store is not None: - self.store.save_run(optimization_run) - - return optimization_run - - def export_best_recipe( - self, run: OptimizationRun, path: Path - ) -> Path: - """Export the best trial's config as a TOML recipe file. - - Args: - run: A completed :class:`OptimizationRun`. - path: Destination path for the TOML file. - - Returns: - The *path* written to. - - Raises: - ValueError: If there is no best trial in the run. - """ - if run.best_trial is None: - raise ValueError("No best trial to export.") - - recipe_data = self._trial_to_recipe_dict(run.best_trial) - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - - if tomli_w is not None: - with open(path, "wb") as fh: - tomli_w.dump(recipe_data, fh) - else: - # Fallback: write TOML manually - self._write_toml_fallback(recipe_data, path) - - run.best_recipe_path = str(path) - return path - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - @staticmethod - def _trial_to_recipe_dict(trial: TrialResult) -> Dict[str, Any]: - """Convert a TrialResult into a Recipe-style TOML dict.""" - params = trial.config.params - recipe: Dict[str, Any] = { - "recipe": { - "name": f"optimized-{trial.trial_id}", - "description": ( - f"Auto-optimized config (accuracy={trial.accuracy:.4f})" - ), - "version": "1.0.0", - }, - } - - # Intelligence section - intel: Dict[str, Any] = {} - if "intelligence.model" in params: - intel["model"] = params["intelligence.model"] - if "intelligence.temperature" in params: - intel["temperature"] = params["intelligence.temperature"] - if "intelligence.quantization" in params: - intel["quantization"] = params["intelligence.quantization"] - if "intelligence.system_prompt" in params: - intel["system_prompt"] = params["intelligence.system_prompt"] - if "intelligence.max_tokens" in params: - intel["max_tokens"] = params["intelligence.max_tokens"] - if "intelligence.top_p" in params: - intel["top_p"] = params["intelligence.top_p"] - if intel: - recipe["intelligence"] = intel - - # Engine section - engine: Dict[str, Any] = {} - if "engine.backend" in params: - engine["key"] = params["engine.backend"] - if engine: - recipe["engine"] = engine - - # Agent section - agent: Dict[str, Any] = {} - if "agent.type" in params: - agent["type"] = params["agent.type"] - if "agent.max_turns" in params: - agent["max_turns"] = params["agent.max_turns"] - if "agent.system_prompt" in params: - agent["system_prompt"] = params["agent.system_prompt"] - if "tools.tool_set" in params: - agent["tools"] = params["tools.tool_set"] - if agent: - recipe["agent"] = agent - - # Learning section - learning: Dict[str, Any] = {} - if "learning.routing_policy" in params: - learning["routing"] = params["learning.routing_policy"] - if "learning.agent_policy" in params: - learning["agent"] = params["learning.agent_policy"] - if learning: - recipe["learning"] = learning - - return recipe - - @staticmethod - def _write_toml_fallback( - data: Dict[str, Any], path: Path - ) -> None: - """Write a simple nested dict as TOML without tomli_w.""" - lines: List[str] = [] - for section, values in data.items(): - if not isinstance(values, dict): - continue - lines.append(f"[{section}]") - for key, val in values.items(): - if isinstance(val, str): - lines.append(f'{key} = "{val}"') - elif isinstance(val, bool): - lines.append(f"{key} = {'true' if val else 'false'}") - elif isinstance(val, (int, float)): - lines.append(f"{key} = {val}") - elif isinstance(val, list): - items = ", ".join( - f'"{v}"' if isinstance(v, str) else str(v) - for v in val - ) - lines.append(f"{key} = [{items}]") - else: - lines.append(f'{key} = "{val}"') - lines.append("") - path.write_text("\n".join(lines), encoding="utf-8") - - -__all__ = ["OptimizationEngine", "compute_pareto_frontier"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.optimizer import * # noqa: F401,F403 +from openjarvis.learning.optimize.optimizer import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/personal/__init__.py b/src/openjarvis/optimize/personal/__init__.py index f22ff8e5..5b4a0623 100644 --- a/src/openjarvis/optimize/personal/__init__.py +++ b/src/openjarvis/optimize/personal/__init__.py @@ -1,17 +1,3 @@ -"""Personal benchmark system -- synthesize benchmarks from interaction traces.""" - -from openjarvis.optimize.personal.dataset import PersonalBenchmarkDataset -from openjarvis.optimize.personal.scorer import PersonalBenchmarkScorer -from openjarvis.optimize.personal.synthesizer import ( - PersonalBenchmark, - PersonalBenchmarkSample, - PersonalBenchmarkSynthesizer, -) - -__all__ = [ - "PersonalBenchmark", - "PersonalBenchmarkSample", - "PersonalBenchmarkSynthesizer", - "PersonalBenchmarkDataset", - "PersonalBenchmarkScorer", -] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.personal import * # noqa: F401,F403 +from openjarvis.learning.optimize.personal import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/personal/dataset.py b/src/openjarvis/optimize/personal/dataset.py index efe6e540..bde2f478 100644 --- a/src/openjarvis/optimize/personal/dataset.py +++ b/src/openjarvis/optimize/personal/dataset.py @@ -1,54 +1,3 @@ -"""DatasetProvider adapter for personal benchmarks.""" - -from __future__ import annotations - -from typing import Iterable, List, Optional - -from openjarvis.evals.core.dataset import DatasetProvider -from openjarvis.evals.core.types import EvalRecord -from openjarvis.optimize.personal.synthesizer import PersonalBenchmark - - -class PersonalBenchmarkDataset(DatasetProvider): - """Wraps a PersonalBenchmark as a DatasetProvider for EvalRunner.""" - - dataset_id: str = "personal" - dataset_name: str = "Personal Benchmark" - - def __init__(self, benchmark: PersonalBenchmark) -> None: - self._benchmark = benchmark - self._records: List[EvalRecord] = [] - - def load( - self, - *, - max_samples: Optional[int] = None, - split: Optional[str] = None, - seed: Optional[int] = None, - ) -> None: - """Convert :class:`PersonalBenchmarkSample` instances to :class:`EvalRecord`.""" - samples = self._benchmark.samples - if max_samples is not None: - samples = samples[:max_samples] - self._records = [ - EvalRecord( - record_id=s.trace_id, - problem=s.query, - reference=s.reference_answer, - category=s.category, - subject=s.agent or "general", - metadata=s.metadata, - ) - for s in samples - ] - - def iter_records(self) -> Iterable[EvalRecord]: - """Iterate over loaded records.""" - return iter(self._records) - - def size(self) -> int: - """Return the number of loaded records.""" - return len(self._records) - - -__all__ = ["PersonalBenchmarkDataset"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.personal.dataset import * # noqa: F401,F403 +from openjarvis.learning.optimize.personal.dataset import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/personal/scorer.py b/src/openjarvis/optimize/personal/scorer.py index 8ed0bec1..58737002 100644 --- a/src/openjarvis/optimize/personal/scorer.py +++ b/src/openjarvis/optimize/personal/scorer.py @@ -1,47 +1,3 @@ -"""LLM-judge scorer for personal benchmarks.""" - -from __future__ import annotations - -from typing import Any, Dict, Optional, Tuple - -from openjarvis.evals.core.backend import InferenceBackend -from openjarvis.evals.core.scorer import LLMJudgeScorer -from openjarvis.evals.core.types import EvalRecord - - -class PersonalBenchmarkScorer(LLMJudgeScorer): - """Judges a candidate response against the best-known response from traces.""" - - scorer_id: str = "personal_judge" - - def __init__(self, judge_backend: InferenceBackend, judge_model: str) -> None: - super().__init__(judge_backend, judge_model) - - def score( - self, record: EvalRecord, model_answer: str, - ) -> Tuple[Optional[bool], Dict[str, Any]]: - """Compare *model_answer* against *record.reference* using the judge LLM. - - Returns ``(is_correct, metadata)`` where *is_correct* indicates whether - the candidate answer is at least as good as the reference. - """ - prompt = ( - "Compare these two answers to the query.\n\n" - f"Query: {record.problem}\n\n" - "Reference answer (known good):\n" - f"{record.reference}\n\n" - "Candidate answer:\n" - f"{model_answer}\n\n" - "Is the candidate answer at least as good as the reference? " - 'Respond with exactly "YES" or "NO" on the first line, ' - "then explain your reasoning." - ) - response = self._ask_judge( - prompt, system="You are an impartial quality judge.", - ) - first_line = response.strip().split("\n")[0].strip().upper() - is_correct = first_line.startswith("YES") - return is_correct, {"judge_response": response} - - -__all__ = ["PersonalBenchmarkScorer"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.personal.scorer import * # noqa: F401,F403 +from openjarvis.learning.optimize.personal.scorer import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/personal/synthesizer.py b/src/openjarvis/optimize/personal/synthesizer.py index ecd0f289..0fcc5718 100644 --- a/src/openjarvis/optimize/personal/synthesizer.py +++ b/src/openjarvis/optimize/personal/synthesizer.py @@ -1,119 +1,3 @@ -"""Synthesize personal benchmarks from interaction traces.""" - -from __future__ import annotations - -import time -from collections import defaultdict -from dataclasses import dataclass, field -from typing import Any, Dict, List - -from openjarvis.traces.store import TraceStore - - -@dataclass(slots=True) -class PersonalBenchmarkSample: - """A single sample in a personal benchmark.""" - - trace_id: str - query: str - reference_answer: str # best known answer from traces - agent: str = "" - category: str = "chat" - feedback_score: float = 0.0 - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass(slots=True) -class PersonalBenchmark: - """A synthesized benchmark from user interaction traces.""" - - workflow_id: str - samples: List[PersonalBenchmarkSample] = field(default_factory=list) - created_at: float = 0.0 - - -def _query_class_key(agent: str, query: str) -> str: - """Compute a grouping key from agent name and query prefix.""" - prefix = query[:50].strip().lower() - return f"{agent}::{prefix}" - - -def _infer_category(agent: str) -> str: - """Heuristic to map an agent name to an eval category.""" - agent_lower = agent.lower() - if any(tok in agent_lower for tok in ("react", "openhands", "orchestrator")): - return "agentic" - if any(tok in agent_lower for tok in ("rag", "memory", "retriev")): - return "rag" - if any(tok in agent_lower for tok in ("reason", "math", "code")): - return "reasoning" - return "chat" - - -class PersonalBenchmarkSynthesizer: - """Mines interaction traces into a reusable personal benchmark.""" - - def __init__(self, trace_store: TraceStore) -> None: - self._store = trace_store - - def synthesize( - self, - workflow_id: str = "default", - min_feedback: float = 0.7, - max_samples: int = 100, - ) -> PersonalBenchmark: - """Build a personal benchmark from high-quality traces. - - 1. Query traces that have feedback >= *min_feedback*. - 2. Group by query class (agent + first 50 chars of query). - 3. For each class, pick the trace with the highest feedback as reference. - 4. Return a :class:`PersonalBenchmark` capped at *max_samples*. - """ - # Fetch a large pool of traces (limit high enough to cover most stores) - all_traces = self._store.list_traces(limit=10_000) - - # Filter to traces with sufficient feedback - qualified = [ - t - for t in all_traces - if t.feedback is not None and t.feedback >= min_feedback - ] - - # Group by query class - groups: Dict[str, list] = defaultdict(list) - for trace in qualified: - key = _query_class_key(trace.agent, trace.query) - groups[key].append(trace) - - # Pick best trace per class - samples: List[PersonalBenchmarkSample] = [] - for _key, traces in groups.items(): - best = max(traces, key=lambda t: t.feedback or 0.0) - samples.append( - PersonalBenchmarkSample( - trace_id=best.trace_id, - query=best.query, - reference_answer=best.result, - agent=best.agent, - category=_infer_category(best.agent), - feedback_score=best.feedback or 0.0, - metadata=best.metadata, - ), - ) - - # Sort deterministically (highest feedback first) and cap - samples.sort(key=lambda s: (-s.feedback_score, s.trace_id)) - samples = samples[:max_samples] - - return PersonalBenchmark( - workflow_id=workflow_id, - samples=samples, - created_at=time.time(), - ) - - -__all__ = [ - "PersonalBenchmark", - "PersonalBenchmarkSample", - "PersonalBenchmarkSynthesizer", -] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.personal.synthesizer import * # noqa: F401,F403 +from openjarvis.learning.optimize.personal.synthesizer import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/search_space.py b/src/openjarvis/optimize/search_space.py index 1827600c..3f28b613 100644 --- a/src/openjarvis/optimize/search_space.py +++ b/src/openjarvis/optimize/search_space.py @@ -1,193 +1,3 @@ -"""Search space builder and default search space for configuration optimization.""" - -from __future__ import annotations - -from typing import Any, Dict, List - -from openjarvis.optimize.types import SearchDimension, SearchSpace - - -def build_search_space(config: Dict[str, Any]) -> SearchSpace: - """Build a SearchSpace from a TOML-style config dict. - - Expected format:: - - { - "optimize": { - "search": [ - { - "name": "agent.type", - "type": "categorical", - "values": ["orchestrator", "native_react"], - "description": "Agent architecture", - }, - { - "name": "intelligence.temperature", - "type": "continuous", - "low": 0.0, - "high": 1.0, - "description": "Generation temperature", - }, - ], - "fixed": {"engine": "ollama", "model": "qwen3:8b"}, - "constraints": { - "rules": ["SimpleAgent should only have max_turns = 1"], - }, - } - } - """ - opt = config.get("optimize", {}) - search_entries: List[Dict[str, Any]] = opt.get("search", []) - fixed: Dict[str, Any] = dict(opt.get("fixed", {})) - constraints_sec = opt.get("constraints", {}) - constraints: List[str] = list(constraints_sec.get("rules", [])) - - dimensions: List[SearchDimension] = [] - for entry in search_entries: - # Infer pillar from the first segment of the dotted name - name = entry.get("name", "") - pillar = name.split(".")[0] if "." in name else "" - - dimensions.append( - SearchDimension( - name=name, - dim_type=entry.get("type", "categorical"), - values=list(entry.get("values", [])), - low=entry.get("low"), - high=entry.get("high"), - description=entry.get("description", ""), - pillar=pillar, - ) - ) - - return SearchSpace( - dimensions=dimensions, - fixed=fixed, - constraints=constraints, - ) - - -# --------------------------------------------------------------------------- -# Default search space covering all 5 pillars -# --------------------------------------------------------------------------- - -DEFAULT_SEARCH_SPACE = SearchSpace( - dimensions=[ - # Intelligence pillar - SearchDimension( - name="intelligence.model", - dim_type="categorical", - values=[ - "qwen3:8b", - "qwen3:4b", - "qwen3:1.7b", - "llama3.1:8b", - "llama3.1:70b", - "gemma2:9b", - "mistral:7b", - "deepseek-r1:8b", - ], - description="The LLM model to use for generation", - pillar="intelligence", - ), - SearchDimension( - name="intelligence.temperature", - dim_type="continuous", - low=0.0, - high=1.0, - description="Generation temperature (0 = deterministic, 1 = creative)", - pillar="intelligence", - ), - SearchDimension( - name="intelligence.max_tokens", - dim_type="integer", - low=256, - high=8192, - description="Maximum tokens to generate per response", - pillar="intelligence", - ), - SearchDimension( - name="intelligence.top_p", - dim_type="continuous", - low=0.0, - high=1.0, - description="Nucleus sampling probability threshold", - pillar="intelligence", - ), - SearchDimension( - name="intelligence.system_prompt", - dim_type="text", - description="System prompt to guide model behavior", - pillar="intelligence", - ), - # Engine pillar - SearchDimension( - name="engine.backend", - dim_type="categorical", - values=[ - "ollama", "vllm", "sglang", "llamacpp", - "mlx", "lmstudio", "exo", "nexa", - "uzu", "apple_fm", - ], - description="Inference engine backend", - pillar="engine", - ), - # Agent pillar - SearchDimension( - name="agent.type", - dim_type="categorical", - values=["simple", "orchestrator", "native_react", "native_openhands"], - description="Agent architecture to use", - pillar="agent", - ), - SearchDimension( - name="agent.max_turns", - dim_type="integer", - low=1, - high=30, - description="Maximum number of agent reasoning turns", - pillar="agent", - ), - # Tools pillar - SearchDimension( - name="tools.tool_set", - dim_type="subset", - values=[ - "calculator", - "think", - "file_read", - "file_write", - "web_search", - "code_interpreter", - "llm", - "shell_exec", - "apply_patch", - "http_request", - "database_query", - ], - description="Set of tools available to the agent", - pillar="tools", - ), - # Learning pillar - SearchDimension( - name="learning.routing_policy", - dim_type="categorical", - values=["heuristic", "grpo", "bandit", "learned"], - description="Router policy for model/agent selection", - pillar="learning", - ), - ], - fixed={}, - constraints=[ - "SimpleAgent (agent.type='simple') should only have max_turns = 1", - "agent.max_turns must be >= 1", - "intelligence.temperature and intelligence.top_p " - "should not both be at extreme values", - ], -) - - -__all__ = [ - "build_search_space", - "DEFAULT_SEARCH_SPACE", -] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.search_space import * # noqa: F401,F403 +from openjarvis.learning.optimize.search_space import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/store.py b/src/openjarvis/optimize/store.py index 6721bf67..6f734980 100644 --- a/src/openjarvis/optimize/store.py +++ b/src/openjarvis/optimize/store.py @@ -1,497 +1,3 @@ -"""SQLite-backed storage for optimization runs and trials.""" - -from __future__ import annotations - -import json -import sqlite3 -import time -from pathlib import Path -from typing import Any, Dict, List, Optional, Union - -from openjarvis.optimize.types import ( - BenchmarkScore, - OptimizationRun, - SampleScore, - SearchSpace, - TrialConfig, - TrialFeedback, - TrialResult, -) - -_CREATE_RUNS = """\ -CREATE TABLE IF NOT EXISTS optimization_runs ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - run_id TEXT NOT NULL UNIQUE, - search_space TEXT NOT NULL DEFAULT '{}', - status TEXT NOT NULL DEFAULT 'running', - optimizer_model TEXT NOT NULL DEFAULT '', - benchmark TEXT NOT NULL DEFAULT '', - best_trial_id TEXT, - best_recipe_path TEXT, - created_at REAL NOT NULL DEFAULT 0.0, - updated_at REAL NOT NULL DEFAULT 0.0 -); -""" - -_CREATE_TRIALS = """\ -CREATE TABLE IF NOT EXISTS trial_results ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - trial_id TEXT NOT NULL, - run_id TEXT NOT NULL, - config TEXT NOT NULL DEFAULT '{}', - reasoning TEXT NOT NULL DEFAULT '', - accuracy REAL NOT NULL DEFAULT 0.0, - mean_latency_seconds REAL NOT NULL DEFAULT 0.0, - total_cost_usd REAL NOT NULL DEFAULT 0.0, - total_energy_joules REAL NOT NULL DEFAULT 0.0, - total_tokens INTEGER NOT NULL DEFAULT 0, - samples_evaluated INTEGER NOT NULL DEFAULT 0, - analysis TEXT NOT NULL DEFAULT '', - failure_modes TEXT NOT NULL DEFAULT '[]', - created_at REAL NOT NULL DEFAULT 0.0, - FOREIGN KEY (run_id) REFERENCES optimization_runs(run_id) -); -""" - -_INSERT_RUN = """\ -INSERT OR REPLACE INTO optimization_runs ( - run_id, search_space, status, optimizer_model, benchmark, - best_trial_id, best_recipe_path, created_at, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) -""" - -_INSERT_TRIAL = """\ -INSERT OR REPLACE INTO trial_results ( - trial_id, run_id, config, reasoning, accuracy, - mean_latency_seconds, total_cost_usd, total_energy_joules, - total_tokens, samples_evaluated, analysis, failure_modes, - created_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -""" - - -_MIGRATE_TRIALS = [ - "ALTER TABLE trial_results ADD COLUMN " - "sample_scores TEXT NOT NULL DEFAULT '[]'", - "ALTER TABLE trial_results ADD COLUMN " - "structured_feedback TEXT NOT NULL DEFAULT '{}'", - "ALTER TABLE trial_results ADD COLUMN " - "per_benchmark TEXT NOT NULL DEFAULT '[]'", -] - -_MIGRATE_RUNS = [ - "ALTER TABLE optimization_runs ADD COLUMN " - "pareto_frontier_ids TEXT NOT NULL DEFAULT '[]'", - "ALTER TABLE optimization_runs ADD COLUMN " - "benchmarks TEXT NOT NULL DEFAULT '[]'", -] - - -class OptimizationStore: - """SQLite-backed storage for optimization runs and trials.""" - - def __init__(self, db_path: Union[str, Path]) -> None: - self._db_path = str(db_path) - self._conn = sqlite3.connect(self._db_path) - self._conn.execute("PRAGMA journal_mode=WAL") - self._conn.execute(_CREATE_RUNS) - self._conn.execute(_CREATE_TRIALS) - self._conn.commit() - self._migrate() - - def _migrate(self) -> None: - """Add new columns to existing databases gracefully.""" - for stmt in _MIGRATE_TRIALS + _MIGRATE_RUNS: - try: - self._conn.execute(stmt) - except sqlite3.OperationalError: - pass # Column already exists - self._conn.commit() - - # ------------------------------------------------------------------ - # Runs - # ------------------------------------------------------------------ - - def save_run(self, run: OptimizationRun) -> None: - """Persist an optimization run (insert or update).""" - now = time.time() - search_space_json = self._search_space_to_json(run.search_space) - best_trial_id = run.best_trial.trial_id if run.best_trial else None - pareto_ids = json.dumps([t.trial_id for t in run.pareto_frontier]) - self._conn.execute( - _INSERT_RUN, - ( - run.run_id, - search_space_json, - run.status, - run.optimizer_model, - run.benchmark, - best_trial_id, - run.best_recipe_path, - now, - now, - ), - ) - benchmarks_json = json.dumps(run.benchmarks) - # Update pareto_frontier_ids and benchmarks separately - self._conn.execute( - "UPDATE optimization_runs SET pareto_frontier_ids = ?, " - "benchmarks = ? WHERE run_id = ?", - (pareto_ids, benchmarks_json, run.run_id), - ) - self._conn.commit() - - def get_run(self, run_id: str) -> Optional[OptimizationRun]: - """Retrieve an optimization run by id, or ``None``.""" - row = self._conn.execute( - "SELECT * FROM optimization_runs WHERE run_id = ?", - (run_id,), - ).fetchone() - if row is None: - return None - return self._row_to_run(row) - - def list_runs(self, limit: int = 50) -> List[Dict[str, Any]]: - """Return summary dicts of recent optimization runs.""" - rows = self._conn.execute( - "SELECT * FROM optimization_runs ORDER BY created_at DESC LIMIT ?", - (limit,), - ).fetchall() - result: List[Dict[str, Any]] = [] - for row in rows: - result.append( - { - "run_id": row[1], - "status": row[3], - "optimizer_model": row[4], - "benchmark": row[5], - "best_trial_id": row[6], - "best_recipe_path": row[7], - "created_at": row[8], - "updated_at": row[9], - } - ) - return result - - # ------------------------------------------------------------------ - # Trials - # ------------------------------------------------------------------ - - def save_trial(self, run_id: str, trial: TrialResult) -> None: - """Persist a single trial result.""" - now = time.time() - # Serialize sample_scores - scores_json = json.dumps([ - { - "record_id": s.record_id, - "is_correct": s.is_correct, - "score": s.score, - "latency_seconds": s.latency_seconds, - "prompt_tokens": s.prompt_tokens, - "completion_tokens": s.completion_tokens, - "cost_usd": s.cost_usd, - "error": s.error, - "ttft": s.ttft, - "energy_joules": s.energy_joules, - "power_watts": s.power_watts, - "gpu_utilization_pct": s.gpu_utilization_pct, - "throughput_tok_per_sec": s.throughput_tok_per_sec, - "mfu_pct": s.mfu_pct, - "mbu_pct": s.mbu_pct, - "ipw": s.ipw, - "ipj": s.ipj, - "energy_per_output_token_joules": s.energy_per_output_token_joules, - "throughput_per_watt": s.throughput_per_watt, - "mean_itl_ms": s.mean_itl_ms, - } - for s in trial.sample_scores - ]) - # Serialize structured_feedback - fb = trial.structured_feedback - fb_json = json.dumps({ - "summary_text": fb.summary_text, - "failure_patterns": fb.failure_patterns, - "pillar_ratings": fb.pillar_ratings, - "suggested_changes": fb.suggested_changes, - "target_pillar": fb.target_pillar, - }) if fb else "{}" - - self._conn.execute( - _INSERT_TRIAL, - ( - trial.trial_id, - run_id, - json.dumps(trial.config.params), - trial.config.reasoning, - trial.accuracy, - trial.mean_latency_seconds, - trial.total_cost_usd, - trial.total_energy_joules, - trial.total_tokens, - trial.samples_evaluated, - trial.analysis, - json.dumps(trial.failure_modes), - now, - ), - ) - # Serialize per_benchmark - pb_json = json.dumps([ - { - "benchmark": b.benchmark, - "accuracy": b.accuracy, - "mean_latency_seconds": b.mean_latency_seconds, - "total_cost_usd": b.total_cost_usd, - "total_energy_joules": b.total_energy_joules, - "total_tokens": b.total_tokens, - "samples_evaluated": b.samples_evaluated, - "errors": b.errors, - "weight": b.weight, - } - for b in trial.per_benchmark - ]) - - # Update new columns separately - self._conn.execute( - "UPDATE trial_results SET sample_scores = ?, " - "structured_feedback = ?, per_benchmark = ? " - "WHERE trial_id = ? AND run_id = ?", - (scores_json, fb_json, pb_json, trial.trial_id, run_id), - ) - self._conn.commit() - - def get_trials(self, run_id: str) -> List[TrialResult]: - """Retrieve all trial results for a given run.""" - rows = self._conn.execute( - "SELECT * FROM trial_results WHERE run_id = ? ORDER BY id", - (run_id,), - ).fetchall() - return [self._row_to_trial(r) for r in rows] - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ - - def close(self) -> None: - """Close the underlying SQLite connection.""" - self._conn.close() - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - @staticmethod - def _search_space_to_json(space: SearchSpace) -> str: - """Serialize a SearchSpace to JSON.""" - dims = [] - for d in space.dimensions: - dims.append( - { - "name": d.name, - "dim_type": d.dim_type, - "values": d.values, - "low": d.low, - "high": d.high, - "description": d.description, - "pillar": d.pillar, - } - ) - return json.dumps( - { - "dimensions": dims, - "fixed": space.fixed, - "constraints": space.constraints, - } - ) - - @staticmethod - def _json_to_search_space(raw: str) -> SearchSpace: - """Deserialize a SearchSpace from JSON.""" - from openjarvis.optimize.types import SearchDimension - - data = json.loads(raw) - dims = [] - for d in data.get("dimensions", []): - dims.append( - SearchDimension( - name=d.get("name", ""), - dim_type=d.get("dim_type", "categorical"), - values=d.get("values", []), - low=d.get("low"), - high=d.get("high"), - description=d.get("description", ""), - pillar=d.get("pillar", ""), - ) - ) - return SearchSpace( - dimensions=dims, - fixed=data.get("fixed", {}), - constraints=data.get("constraints", []), - ) - - def _row_to_run(self, row: tuple) -> OptimizationRun: - """Convert a database row to an OptimizationRun.""" - run_id = row[1] - search_space = self._json_to_search_space(row[2]) - status = row[3] - optimizer_model = row[4] - benchmark = row[5] - best_trial_id = row[6] - best_recipe_path = row[7] - - # Load trials for this run - trials = self.get_trials(run_id) - - # Find the best trial - best_trial: Optional[TrialResult] = None - if best_trial_id: - for t in trials: - if t.trial_id == best_trial_id: - best_trial = t - break - - # Reconstruct benchmarks list - benchmarks: List[str] = [] - if len(row) > 11: - try: - benchmarks = json.loads(row[11]) if row[11] else [] - except (json.JSONDecodeError, TypeError): - pass - - # Reconstruct pareto frontier from IDs - pareto_frontier: List[TrialResult] = [] - if len(row) > 10: - try: - frontier_ids = json.loads(row[10]) if row[10] else [] - trial_map = {t.trial_id: t for t in trials} - pareto_frontier = [ - trial_map[tid] for tid in frontier_ids if tid in trial_map - ] - except (json.JSONDecodeError, TypeError): - pass - - return OptimizationRun( - run_id=run_id, - search_space=search_space, - trials=trials, - best_trial=best_trial, - best_recipe_path=best_recipe_path, - status=status, - optimizer_model=optimizer_model, - benchmark=benchmark, - benchmarks=benchmarks, - pareto_frontier=pareto_frontier, - ) - - @staticmethod - def _row_to_trial(row: tuple) -> TrialResult: - """Convert a database row to a TrialResult.""" - trial_id = row[1] - # row[2] = run_id (not stored on TrialResult) - params = json.loads(row[3]) - reasoning = row[4] - accuracy = row[5] - mean_latency = row[6] - cost = row[7] - energy = row[8] - tokens = row[9] - samples = row[10] - analysis = row[11] - failure_modes = json.loads(row[12]) - # row[13] = created_at - - # New columns (may be absent in old DBs) - sample_scores: List[SampleScore] = [] - structured_feedback: Optional[TrialFeedback] = None - - if len(row) > 14: - try: - raw_scores = json.loads(row[14]) if row[14] else [] - sample_scores = [ - SampleScore( - record_id=s.get("record_id", ""), - is_correct=s.get("is_correct"), - score=s.get("score"), - latency_seconds=s.get("latency_seconds", 0.0), - prompt_tokens=s.get("prompt_tokens", 0), - completion_tokens=s.get("completion_tokens", 0), - cost_usd=s.get("cost_usd", 0.0), - error=s.get("error"), - ttft=s.get("ttft", 0.0), - energy_joules=s.get("energy_joules", 0.0), - power_watts=s.get("power_watts", 0.0), - gpu_utilization_pct=s.get("gpu_utilization_pct", 0.0), - throughput_tok_per_sec=s.get("throughput_tok_per_sec", 0.0), - mfu_pct=s.get("mfu_pct", 0.0), - mbu_pct=s.get("mbu_pct", 0.0), - ipw=s.get("ipw", 0.0), - ipj=s.get("ipj", 0.0), - energy_per_output_token_joules=s.get( - "energy_per_output_token_joules", 0.0, - ), - throughput_per_watt=s.get("throughput_per_watt", 0.0), - mean_itl_ms=s.get("mean_itl_ms", 0.0), - ) - for s in raw_scores - ] - except (json.JSONDecodeError, TypeError): - pass - - if len(row) > 15: - try: - raw_fb = json.loads(row[15]) if row[15] else {} - if raw_fb and raw_fb.get("summary_text", "") != "": - structured_feedback = TrialFeedback( - summary_text=raw_fb.get("summary_text", ""), - failure_patterns=raw_fb.get("failure_patterns", []), - pillar_ratings=raw_fb.get("pillar_ratings", {}), - suggested_changes=raw_fb.get("suggested_changes", []), - target_pillar=raw_fb.get("target_pillar", ""), - ) - except (json.JSONDecodeError, TypeError): - pass - - # per_benchmark column - per_benchmark: List[BenchmarkScore] = [] - if len(row) > 16: - try: - raw_pb = json.loads(row[16]) if row[16] else [] - per_benchmark = [ - BenchmarkScore( - benchmark=b.get("benchmark", ""), - accuracy=b.get("accuracy", 0.0), - mean_latency_seconds=b.get("mean_latency_seconds", 0.0), - total_cost_usd=b.get("total_cost_usd", 0.0), - total_energy_joules=b.get("total_energy_joules", 0.0), - total_tokens=b.get("total_tokens", 0), - samples_evaluated=b.get("samples_evaluated", 0), - errors=b.get("errors", 0), - weight=b.get("weight", 1.0), - ) - for b in raw_pb - ] - except (json.JSONDecodeError, TypeError): - pass - - config = TrialConfig( - trial_id=trial_id, - params=params, - reasoning=reasoning, - ) - return TrialResult( - trial_id=trial_id, - config=config, - accuracy=accuracy, - mean_latency_seconds=mean_latency, - total_cost_usd=cost, - total_energy_joules=energy, - total_tokens=tokens, - samples_evaluated=samples, - analysis=analysis, - failure_modes=failure_modes, - sample_scores=sample_scores, - structured_feedback=structured_feedback, - per_benchmark=per_benchmark, - ) - - -__all__ = ["OptimizationStore"] +"""Backward-compatibility shim -- optimize.store moved to learning.optimize.store.""" +from openjarvis.learning.optimize.store import * # noqa: F401,F403 +from openjarvis.learning.optimize.store import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/trial_runner.py b/src/openjarvis/optimize/trial_runner.py index d07d5289..f1bd5abc 100644 --- a/src/openjarvis/optimize/trial_runner.py +++ b/src/openjarvis/optimize/trial_runner.py @@ -1,392 +1,3 @@ -"""TrialRunner -- evaluates a proposed config against a benchmark.""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from pathlib import Path -from typing import Any, List, Optional - -from openjarvis.evals.core.types import RunConfig, RunSummary -from openjarvis.optimize.types import ( - BenchmarkScore, - SampleScore, - TrialConfig, - TrialResult, -) - -LOGGER = logging.getLogger(__name__) - - -@dataclass(slots=True) -class BenchmarkSpec: - """Specification for one benchmark in a multi-benchmark optimization.""" - - benchmark: str - max_samples: int = 200 - weight: float = 1.0 - - -class TrialRunner: - """Evaluates a proposed config against a benchmark. - - Bridges the optimization types (:class:`TrialConfig`) to the eval - framework (:class:`EvalRunner`) so the optimizer can score candidate - configurations end-to-end. - """ - - def __init__( - self, - benchmark: str, - max_samples: int = 50, - judge_model: str = "gpt-5-mini-2025-08-07", - output_dir: str = "results/optimize/", - ) -> None: - self.benchmark = benchmark - self.max_samples = max_samples - self.judge_model = judge_model - self.output_dir = output_dir - - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - def run_trial(self, trial: TrialConfig) -> TrialResult: - """Run *trial* against the configured benchmark and return a result. - - Steps: - 1. Convert ``trial`` to a :class:`Recipe` and extract params. - 2. Build a :class:`RunConfig` from recipe + benchmark settings. - 3. Lazily import eval-framework registries to resolve the - benchmark -> dataset + scorer, and build the backend. - 4. Execute via ``EvalRunner.run()`` -> :class:`RunSummary`. - 5. Map the summary into a :class:`TrialResult`. - """ - recipe = trial.to_recipe() - run_config = self._build_run_config(trial, recipe) - - # Lazy imports so the optimize package stays lightweight - from openjarvis.evals.cli import ( - _build_backend, - _build_dataset, - _build_judge_backend, - _build_scorer, - ) - from openjarvis.evals.core.runner import EvalRunner - - dataset = _build_dataset(self.benchmark) - backend = _build_backend( - run_config.backend, - run_config.engine_key, - run_config.agent_name or "orchestrator", - run_config.tools, - ) - judge_backend = _build_judge_backend(run_config.judge_model) - scorer = _build_scorer( - self.benchmark, judge_backend, run_config.judge_model, - ) - - try: - eval_runner = EvalRunner( - run_config, dataset, backend, scorer, - ) - summary: RunSummary = eval_runner.run() - eval_results = eval_runner.results - finally: - backend.close() - judge_backend.close() - - return self._summary_to_result(trial, summary, eval_results=eval_results) - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _build_run_config(self, trial: TrialConfig, recipe: Any) -> RunConfig: - """Map recipe fields into a :class:`RunConfig`.""" - model = recipe.model or "default" - backend_name = "jarvis-direct" - if recipe.agent_type is not None: - backend_name = "jarvis-agent" - - model_slug = model.replace("/", "-").replace(":", "-") - output_path = str( - Path(self.output_dir) / f"{trial.trial_id}_{model_slug}.jsonl", - ) - - max_tokens = recipe.max_tokens if recipe.max_tokens is not None else 2048 - - return RunConfig( - benchmark=self.benchmark, - backend=backend_name, - model=model, - max_samples=self.max_samples, - temperature=recipe.temperature if recipe.temperature is not None else 0.0, - max_tokens=max_tokens, - judge_model=self.judge_model, - engine_key=recipe.engine_key, - agent_name=recipe.agent_type, - tools=list(recipe.tools) if recipe.tools else [], - output_path=output_path, - system_prompt=recipe.system_prompt or "", - ) - - @staticmethod - def _summary_to_result( - trial: TrialConfig, - summary: RunSummary, - eval_results: Optional[List[Any]] = None, - ) -> TrialResult: - """Convert a :class:`RunSummary` to a :class:`TrialResult`.""" - total_tokens = summary.total_input_tokens + summary.total_output_tokens - - failure_modes: List[str] = [] - if summary.errors > 0: - failure_modes.append(f"{summary.errors} evaluation errors") - - sample_scores: List[SampleScore] = [] - if eval_results: - for er in eval_results: - sample_scores.append( - SampleScore( - record_id=er.record_id, - is_correct=er.is_correct, - score=er.score, - latency_seconds=er.latency_seconds, - prompt_tokens=er.prompt_tokens, - completion_tokens=er.completion_tokens, - cost_usd=er.cost_usd, - error=er.error, - ttft=er.ttft, - energy_joules=er.energy_joules, - power_watts=er.power_watts, - gpu_utilization_pct=er.gpu_utilization_pct, - throughput_tok_per_sec=er.throughput_tok_per_sec, - mfu_pct=er.mfu_pct, - mbu_pct=er.mbu_pct, - ipw=er.ipw, - ipj=er.ipj, - energy_per_output_token_joules=er.energy_per_output_token_joules, - throughput_per_watt=er.throughput_per_watt, - mean_itl_ms=er.mean_itl_ms, - ) - ) - - return TrialResult( - trial_id=trial.trial_id, - config=trial, - accuracy=summary.accuracy, - mean_latency_seconds=summary.mean_latency_seconds, - total_cost_usd=summary.total_cost_usd, - total_energy_joules=summary.total_energy_joules, - total_tokens=total_tokens, - samples_evaluated=summary.total_samples, - failure_modes=failure_modes, - summary=summary, - sample_scores=sample_scores, - ) - - -class MultiBenchTrialRunner: - """Evaluates a proposed config across multiple benchmarks. - - Delegates to :class:`TrialRunner` per benchmark, then aggregates - results into a single composite :class:`TrialResult` with weighted - metrics and per-benchmark breakdowns. - """ - - def __init__( - self, - benchmark_specs: List[BenchmarkSpec], - judge_model: str = "gpt-5-mini-2025-08-07", - output_dir: str = "results/optimize/", - ) -> None: - self.benchmark_specs = benchmark_specs - self.judge_model = judge_model - self.output_dir = output_dir - - def run_trial(self, trial: TrialConfig) -> TrialResult: - """Run *trial* against all benchmarks and return a composite result.""" - per_benchmark: List[BenchmarkScore] = [] - - for spec in self.benchmark_specs: - if spec.benchmark == "terminalbench-native": - score = self._run_terminalbench_native(trial, spec) - else: - runner = TrialRunner( - benchmark=spec.benchmark, - max_samples=spec.max_samples, - judge_model=self.judge_model, - output_dir=self.output_dir, - ) - result = runner.run_trial(trial) - score = BenchmarkScore( - benchmark=spec.benchmark, - accuracy=result.accuracy, - mean_latency_seconds=result.mean_latency_seconds, - total_cost_usd=result.total_cost_usd, - total_energy_joules=result.total_energy_joules, - total_tokens=result.total_tokens, - samples_evaluated=result.samples_evaluated, - errors=len([s for s in result.sample_scores if s.error]), - weight=spec.weight, - summary=result.summary, - sample_scores=result.sample_scores, - ) - per_benchmark.append(score) - - return self._aggregate(trial, per_benchmark) - - def _run_terminalbench_native( - self, trial: TrialConfig, spec: BenchmarkSpec, - ) -> BenchmarkScore: - """Run terminal-bench natively via Harness with Docker execution.""" - import time - - from terminal_bench import BenchmarkResults, Harness - - recipe = trial.to_recipe() - model_name = recipe.model or "default" - - # For vLLM-served models, use openai/ prefix for litellm - if recipe.engine_key == "vllm": - litellm_model = f"openai/{model_name}" - api_base = "http://localhost:8000/v1" - else: - litellm_model = model_name - api_base = "http://localhost:8000/v1" - - temperature = recipe.temperature if recipe.temperature is not None else 0.2 - - output_path = Path(self.output_dir) / f"terminalbench/{trial.trial_id}" - output_path.mkdir(parents=True, exist_ok=True) - - harness_kwargs = { - "output_path": output_path, - "run_id": trial.trial_id, - "dataset_name": "terminal-bench-core", - "dataset_version": "0.1.1", - "model_name": litellm_model, - "agent_import_path": ( - "openjarvis.evals.backends.tb_agent" - ":OpenJarvisTerminalBenchAgent" - ), - "agent_kwargs": { - "model_name": litellm_model, - "api_base": api_base, - "temperature": temperature, - }, - "n_concurrent_trials": 4, - "cleanup": True, - } - - if spec.max_samples and spec.max_samples < 200: - harness_kwargs["n_tasks"] = spec.max_samples - - LOGGER.info( - "Running terminal-bench native: model=%s, max_tasks=%s", - litellm_model, spec.max_samples, - ) - - t0 = time.monotonic() - try: - harness = Harness(**harness_kwargs) - tb_results: BenchmarkResults = harness.run() - except Exception as e: - LOGGER.error("terminal-bench harness failed: %s", e) - return BenchmarkScore( - benchmark="terminalbench-native", - accuracy=0.0, - mean_latency_seconds=0.0, - samples_evaluated=0, - errors=1, - weight=spec.weight, - ) - elapsed = time.monotonic() - t0 - - # Extract results - total_tasks = len(tb_results.results) - resolved = sum(1 for r in tb_results.results if r.is_resolved) - accuracy = resolved / total_tasks if total_tasks > 0 else 0.0 - mean_latency = elapsed / total_tasks if total_tasks > 0 else 0.0 - - total_input = sum(r.total_input_tokens or 0 for r in tb_results.results) - total_output = sum(r.total_output_tokens or 0 for r in tb_results.results) - - # Build sample scores - sample_scores = [] - for r in tb_results.results: - sample_scores.append( - SampleScore( - record_id=f"terminalbench-native-{r.task_id}", - is_correct=bool(r.is_resolved), - score=1.0 if r.is_resolved else 0.0, - prompt_tokens=r.total_input_tokens or 0, - completion_tokens=r.total_output_tokens or 0, - ) - ) - - LOGGER.info( - "terminal-bench native: %d/%d resolved (%.1f%%), %.1fs total", - resolved, total_tasks, accuracy * 100, elapsed, - ) - - return BenchmarkScore( - benchmark="terminalbench-native", - accuracy=accuracy, - mean_latency_seconds=mean_latency, - total_tokens=total_input + total_output, - samples_evaluated=total_tasks, - errors=sum( - 1 for r in tb_results.results - if r.failure_mode.value not in ("none", "unset") - ), - weight=spec.weight, - sample_scores=sample_scores, - ) - - @staticmethod - def _aggregate( - trial: TrialConfig, - per_benchmark: List[BenchmarkScore], - ) -> TrialResult: - """Compute weighted-aggregate metrics from per-benchmark scores.""" - total_weight = sum(b.weight for b in per_benchmark) or 1.0 - accuracy = sum(b.accuracy * b.weight for b in per_benchmark) / total_weight - - # Weighted mean latency by samples evaluated - total_samples = sum(b.samples_evaluated for b in per_benchmark) or 1 - mean_latency = ( - sum(b.mean_latency_seconds * b.samples_evaluated for b in per_benchmark) - / total_samples - ) - - # Sums across benchmarks - total_cost = sum(b.total_cost_usd for b in per_benchmark) - total_energy = sum(b.total_energy_joules for b in per_benchmark) - total_tokens = sum(b.total_tokens for b in per_benchmark) - - # Merge all sample scores - all_scores: List[SampleScore] = [] - failure_modes: List[str] = [] - for b in per_benchmark: - all_scores.extend(b.sample_scores) - if b.errors > 0: - failure_modes.append(f"{b.benchmark}: {b.errors} errors") - - return TrialResult( - trial_id=trial.trial_id, - config=trial, - accuracy=accuracy, - mean_latency_seconds=mean_latency, - total_cost_usd=total_cost, - total_energy_joules=total_energy, - total_tokens=total_tokens, - samples_evaluated=total_samples, - failure_modes=failure_modes, - sample_scores=all_scores, - per_benchmark=per_benchmark, - ) - - -__all__ = ["BenchmarkSpec", "MultiBenchTrialRunner", "TrialRunner"] +"""Backward-compat shim: moved to learning.optimize.""" +from openjarvis.learning.optimize.trial_runner import * # noqa: F401,F403 +from openjarvis.learning.optimize.trial_runner import __all__ # noqa: F401 diff --git a/src/openjarvis/optimize/types.py b/src/openjarvis/optimize/types.py index ce1eef26..6b96442a 100644 --- a/src/openjarvis/optimize/types.py +++ b/src/openjarvis/optimize/types.py @@ -1,253 +1,6 @@ -"""Core data types for the optimization framework.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional - -from openjarvis.evals.core.types import RunSummary -from openjarvis.recipes.loader import Recipe - - -@dataclass(slots=True) -class SearchDimension: - """One tunable dimension in the config space.""" - - name: str # e.g. "agent.type", "intelligence.temperature" - dim_type: str # "categorical", "continuous", "integer", "subset", "text" - # categorical/subset: explicit options - values: List[Any] = field(default_factory=list) - low: Optional[float] = None # continuous/integer lower bound - high: Optional[float] = None # continuous/integer upper bound - description: str = "" # human-readable explanation for the LLM optimizer - pillar: str = "" # intelligence | engine | agent | tools | learning - - -@dataclass(slots=True) -class SearchSpace: - """The full space of configs the optimizer can propose.""" - - dimensions: List[SearchDimension] = field(default_factory=list) - fixed: Dict[str, Any] = field(default_factory=dict) # params NOT being optimized - constraints: List[str] = field(default_factory=list) # natural language constraints - - def to_prompt_description(self) -> str: - """Render search space as structured text for the LLM optimizer.""" - lines: List[str] = [] - lines.append("# Search Space") - lines.append("") - - # Group dimensions by pillar - by_pillar: Dict[str, List[SearchDimension]] = {} - for dim in self.dimensions: - key = dim.pillar or "other" - by_pillar.setdefault(key, []).append(dim) - - for pillar, dims in sorted(by_pillar.items()): - lines.append(f"## {pillar.title()}") - for dim in dims: - lines.append(f"- **{dim.name}** ({dim.dim_type})") - if dim.description: - lines.append(f" Description: {dim.description}") - if dim.dim_type in ("categorical", "subset"): - lines.append(f" Options: {dim.values}") - elif dim.dim_type in ("continuous", "integer"): - lines.append(f" Range: [{dim.low}, {dim.high}]") - elif dim.dim_type == "text": - lines.append(" Free-form text") - lines.append("") - - if self.fixed: - lines.append("## Fixed Parameters") - for k, v in sorted(self.fixed.items()): - lines.append(f"- {k} = {v}") - lines.append("") - - if self.constraints: - lines.append("## Constraints") - for c in self.constraints: - lines.append(f"- {c}") - lines.append("") - - return "\n".join(lines) - - -# Mapping from dotted param names to Recipe constructor fields. -_PARAM_TO_RECIPE: Dict[str, str] = { - "intelligence.model": "model", - "intelligence.temperature": "temperature", - "intelligence.max_tokens": "max_tokens", - "intelligence.quantization": "quantization", - "engine.backend": "engine_key", - "agent.type": "agent_type", - "agent.max_turns": "max_turns", - "agent.system_prompt": "system_prompt", - "intelligence.system_prompt": "system_prompt", - "tools.tool_set": "tools", - "learning.routing_policy": "routing_policy", - "learning.agent_policy": "agent_policy", -} - - -@dataclass(slots=True) -class BenchmarkScore: - """Per-benchmark metrics from a multi-benchmark evaluation trial.""" - - benchmark: str - accuracy: float = 0.0 - mean_latency_seconds: float = 0.0 - total_cost_usd: float = 0.0 - total_energy_joules: float = 0.0 - total_tokens: int = 0 - samples_evaluated: int = 0 - errors: int = 0 - weight: float = 1.0 - summary: Optional[Any] = None # RunSummary - sample_scores: List["SampleScore"] = field(default_factory=list) - - -@dataclass(slots=True) -class SampleScore: - """Per-sample metrics from an evaluation trial.""" - - record_id: str - is_correct: Optional[bool] = None - score: Optional[float] = None - latency_seconds: float = 0.0 - prompt_tokens: int = 0 - completion_tokens: int = 0 - cost_usd: float = 0.0 - error: Optional[str] = None - ttft: float = 0.0 - energy_joules: float = 0.0 - power_watts: float = 0.0 - gpu_utilization_pct: float = 0.0 - throughput_tok_per_sec: float = 0.0 - mfu_pct: float = 0.0 - mbu_pct: float = 0.0 - ipw: float = 0.0 - ipj: float = 0.0 - energy_per_output_token_joules: float = 0.0 - throughput_per_watt: float = 0.0 - mean_itl_ms: float = 0.0 - - -@dataclass(slots=True) -class TrialFeedback: - """Structured feedback from trial analysis.""" - - summary_text: str = "" - failure_patterns: List[str] = field(default_factory=list) - pillar_ratings: Dict[str, str] = field(default_factory=dict) - suggested_changes: List[str] = field(default_factory=list) - target_pillar: str = "" - - -@dataclass(slots=True) -class ObjectiveSpec: - """A single optimization objective.""" - - metric: str - direction: str # "maximize" or "minimize" - weight: float = 1.0 - - -DEFAULT_OBJECTIVES = [ - ObjectiveSpec("accuracy", "maximize"), - ObjectiveSpec("mean_latency_seconds", "minimize"), - ObjectiveSpec("total_cost_usd", "minimize"), -] - -ALL_OBJECTIVES = [ - ObjectiveSpec("accuracy", "maximize"), - ObjectiveSpec("mean_latency_seconds", "minimize"), - ObjectiveSpec("total_cost_usd", "minimize"), - ObjectiveSpec("total_energy_joules", "minimize"), - ObjectiveSpec("avg_power_watts", "minimize"), - ObjectiveSpec("throughput_tok_per_sec", "maximize"), - ObjectiveSpec("mfu_pct", "maximize"), - ObjectiveSpec("mbu_pct", "maximize"), - ObjectiveSpec("ipw", "maximize"), - ObjectiveSpec("ipj", "maximize"), - ObjectiveSpec("energy_per_output_token", "minimize"), - ObjectiveSpec("throughput_per_watt", "maximize"), - ObjectiveSpec("ttft", "minimize"), - ObjectiveSpec("mean_itl_ms", "minimize"), -] - - -@dataclass(slots=True) -class TrialConfig: - """A single candidate configuration proposed by the optimizer.""" - - trial_id: str - params: Dict[str, Any] = field(default_factory=dict) # dotted keys -> values - reasoning: str = "" # optimizer's explanation - - def to_recipe(self) -> Recipe: - """Map params back to Recipe fields.""" - kwargs: Dict[str, Any] = {} - for dotted_key, value in self.params.items(): - recipe_field = _PARAM_TO_RECIPE.get(dotted_key) - if recipe_field is not None: - kwargs[recipe_field] = value - - return Recipe( - name=f"trial-{self.trial_id}", - **kwargs, - ) - - -@dataclass(slots=True) -class TrialResult: - """Result of evaluating a trial, with both scalar and textual feedback.""" - - trial_id: str - config: TrialConfig - accuracy: float = 0.0 - mean_latency_seconds: float = 0.0 - total_cost_usd: float = 0.0 - total_energy_joules: float = 0.0 - total_tokens: int = 0 - samples_evaluated: int = 0 - analysis: str = "" - failure_modes: List[str] = field(default_factory=list) - per_sample_feedback: List[Dict[str, Any]] = field(default_factory=list) - summary: Optional[RunSummary] = None - sample_scores: List[SampleScore] = field(default_factory=list) - structured_feedback: Optional[TrialFeedback] = None - per_benchmark: List[BenchmarkScore] = field(default_factory=list) - - -@dataclass(slots=True) -class OptimizationRun: - """Complete optimization session.""" - - run_id: str - search_space: SearchSpace - trials: List[TrialResult] = field(default_factory=list) - best_trial: Optional[TrialResult] = None - best_recipe_path: Optional[str] = None - status: str = "running" # running | completed | failed - optimizer_model: str = "" - benchmark: str = "" - benchmarks: List[str] = field(default_factory=list) - pareto_frontier: List[TrialResult] = field(default_factory=list) - objectives: List[ObjectiveSpec] = field( - default_factory=lambda: list(DEFAULT_OBJECTIVES), - ) - - -__all__ = [ - "ALL_OBJECTIVES", - "BenchmarkScore", - "DEFAULT_OBJECTIVES", - "ObjectiveSpec", - "OptimizationRun", - "SampleScore", - "SearchDimension", - "SearchSpace", - "TrialConfig", - "TrialFeedback", - "TrialResult", -] +"""Backward-compatibility shim -- optimize.types moved to learning.optimize.types.""" +from openjarvis.learning.optimize.types import * # noqa: F401,F403 +from openjarvis.learning.optimize.types import ( + _PARAM_TO_RECIPE, # noqa: F401 + __all__, # noqa: F401 +) diff --git a/src/openjarvis/server/api_routes.py b/src/openjarvis/server/api_routes.py index 00c281fb..f6e04b52 100644 --- a/src/openjarvis/server/api_routes.py +++ b/src/openjarvis/server/api_routes.py @@ -722,7 +722,7 @@ async def list_optimize_runs(request: Request): """List optimization runs.""" try: from openjarvis.core.config import DEFAULT_CONFIG_DIR - from openjarvis.optimize.store import OptimizationStore + from openjarvis.learning.optimize.store import OptimizationStore db_path = DEFAULT_CONFIG_DIR / "optimize.db" if not db_path.exists(): @@ -741,7 +741,7 @@ async def get_optimize_run(run_id: str, request: Request): """Get optimization run details.""" try: from openjarvis.core.config import DEFAULT_CONFIG_DIR - from openjarvis.optimize.store import OptimizationStore + from openjarvis.learning.optimize.store import OptimizationStore db_path = DEFAULT_CONFIG_DIR / "optimize.db" if not db_path.exists(): diff --git a/tests/test_rust_bridge.py b/tests/core/test_rust_bridge.py similarity index 100% rename from tests/test_rust_bridge.py rename to tests/core/test_rust_bridge.py diff --git a/tests/test_docker.py b/tests/deployment/test_docker.py similarity index 96% rename from tests/test_docker.py rename to tests/deployment/test_docker.py index 703f7031..ccc3c96f 100644 --- a/tests/test_docker.py +++ b/tests/deployment/test_docker.py @@ -4,7 +4,7 @@ from __future__ import annotations from pathlib import Path -ROOT = Path(__file__).resolve().parent.parent +ROOT = Path(__file__).resolve().parent.parent.parent DOCKER_DIR = ROOT / "deploy" / "docker" diff --git a/tests/evals/test_deepplanning.py b/tests/evals/test_deepplanning.py new file mode 100644 index 00000000..d4db7903 --- /dev/null +++ b/tests/evals/test_deepplanning.py @@ -0,0 +1,87 @@ +"""Tests for DeepPlanning benchmark.""" + +from unittest.mock import MagicMock + +from openjarvis.evals.core.types import EvalRecord +from openjarvis.evals.datasets.deepplanning import DeepPlanningDataset +from openjarvis.evals.scorers.deepplanning_scorer import DeepPlanningScorer + + +def _mock_backend() -> MagicMock: + backend = MagicMock() + backend.generate.return_value = "CORRECT" + return backend + + +class TestDeepPlanningDataset: + def test_instantiation(self) -> None: + ds = DeepPlanningDataset() + assert ds.dataset_id == "deepplanning" + assert ds.dataset_name == "DeepPlanning" + + def test_has_required_methods(self) -> None: + ds = DeepPlanningDataset() + assert hasattr(ds, "load") + assert hasattr(ds, "iter_records") + assert hasattr(ds, "size") + + +class TestDeepPlanningScorer: + def test_instantiation(self) -> None: + s = DeepPlanningScorer(_mock_backend(), "test-model") + assert s.scorer_id == "deepplanning" + + def test_correct_plan(self) -> None: + s = DeepPlanningScorer(_mock_backend(), "test-model") + record = EvalRecord( + record_id="dp-1", + problem="## Task (travel)\nPlan a trip to Paris", + reference="Day 1: Flight to Paris...", + category="agentic", + subject="travel", + metadata={"task_type": "travel"}, + ) + is_correct, meta = s.score(record, "Day 1: Fly to Paris, visit Eiffel Tower") + assert is_correct is True + assert meta["match_type"] == "llm_judge" + + def test_incorrect_plan(self) -> None: + backend = MagicMock() + backend.generate.return_value = "INCORRECT - Missing budget constraint" + s = DeepPlanningScorer(backend, "test-model") + record = EvalRecord( + record_id="dp-2", + problem="## Task (shopping)\nBuild a cart", + reference="Cart: item A, item B, total $50", + category="agentic", + subject="shopping", + metadata={"task_type": "shopping"}, + ) + is_correct, meta = s.score(record, "Cart: item C, total $100") + assert is_correct is False + + def test_empty_response(self) -> None: + s = DeepPlanningScorer(_mock_backend(), "test-model") + record = EvalRecord( + record_id="dp-3", problem="task", + reference="answer", category="agentic", + ) + is_correct, meta = s.score(record, "") + assert is_correct is False + assert meta["reason"] == "empty_response" + + +class TestDeepPlanningCLI: + def test_in_benchmarks(self) -> None: + from openjarvis.evals.cli import BENCHMARKS + assert "deepplanning" in BENCHMARKS + + def test_build_dataset(self) -> None: + from openjarvis.evals.cli import _build_dataset + ds = _build_dataset("deepplanning") + assert ds.dataset_id == "deepplanning" + + def test_build_scorer(self) -> None: + from openjarvis.evals.cli import _build_scorer + s = _build_scorer("deepplanning", _mock_backend(), "test-model") + assert s.scorer_id == "deepplanning" diff --git a/tests/evals/test_paperarena.py b/tests/evals/test_paperarena.py new file mode 100644 index 00000000..022aeb9c --- /dev/null +++ b/tests/evals/test_paperarena.py @@ -0,0 +1,115 @@ +"""Tests for PaperArena benchmark.""" + +from unittest.mock import MagicMock + +from openjarvis.evals.core.types import EvalRecord +from openjarvis.evals.datasets.paperarena import PaperArenaDataset +from openjarvis.evals.scorers.paperarena_judge import PaperArenaScorer + + +def _mock_backend() -> MagicMock: + backend = MagicMock() + backend.generate.return_value = "CORRECT" + return backend + + +class TestPaperArenaDataset: + def test_instantiation(self) -> None: + ds = PaperArenaDataset() + assert ds.dataset_id == "paperarena" + assert ds.dataset_name == "PaperArena" + + def test_has_required_methods(self) -> None: + ds = PaperArenaDataset() + assert hasattr(ds, "load") + assert hasattr(ds, "iter_records") + assert hasattr(ds, "size") + + +class TestPaperArenaScorer: + def test_instantiation(self) -> None: + s = PaperArenaScorer(_mock_backend(), "test-model") + assert s.scorer_id == "paperarena" + + def test_mc_correct(self) -> None: + s = PaperArenaScorer(_mock_backend(), "test-model") + record = EvalRecord( + record_id="pa-1", + problem="## Question\nWhat is X?\nOptions:\n A) foo\n B) bar", + reference="A", + category="agentic", + subject="easy_mc", + metadata={"question_type": "MC"}, + ) + is_correct, meta = s.score(record, "The answer is A") + assert is_correct is True + assert meta["match_type"] == "exact_letter" + + def test_mc_wrong(self) -> None: + s = PaperArenaScorer(_mock_backend(), "test-model") + record = EvalRecord( + record_id="pa-2", + problem="Question", + reference="B", + category="agentic", + subject="medium_mc", + metadata={"question_type": "MC"}, + ) + is_correct, meta = s.score(record, "The answer is C") + assert is_correct is False + assert meta["candidate_letter"] == "C" + + def test_open_answer_judge(self) -> None: + s = PaperArenaScorer(_mock_backend(), "test-model") + record = EvalRecord( + record_id="pa-3", + problem="## Question\nExplain X", + reference="X is a method for...", + category="agentic", + subject="hard_oa", + metadata={"question_type": "OA"}, + ) + is_correct, meta = s.score(record, "X is a technique used to...") + assert is_correct is True + assert meta["match_type"] == "llm_judge" + + def test_closed_answer_judge(self) -> None: + backend = MagicMock() + backend.generate.return_value = "INCORRECT" + s = PaperArenaScorer(backend, "test-model") + record = EvalRecord( + record_id="pa-4", + problem="## Question\nWhat is the value?", + reference="42.5", + category="agentic", + subject="medium_ca", + metadata={"question_type": "CA"}, + ) + is_correct, meta = s.score(record, "The value is 100") + assert is_correct is False + + def test_empty_response(self) -> None: + s = PaperArenaScorer(_mock_backend(), "test-model") + record = EvalRecord( + record_id="pa-5", problem="q", + reference="a", category="agentic", + ) + is_correct, meta = s.score(record, "") + assert is_correct is False + assert meta["reason"] == "empty_response" + + +class TestPaperArenaCLI: + def test_in_benchmarks(self) -> None: + from openjarvis.evals.cli import BENCHMARKS + assert "paperarena" in BENCHMARKS + + def test_build_dataset(self) -> None: + from openjarvis.evals.cli import _build_dataset + ds = _build_dataset("paperarena") + assert ds.dataset_id == "paperarena" + + def test_build_scorer(self) -> None: + from openjarvis.evals.cli import _build_scorer + s = _build_scorer("paperarena", _mock_backend(), "test-model") + assert s.scorer_id == "paperarena" diff --git a/tests/test_integration.py b/tests/integration/test_integration.py similarity index 100% rename from tests/test_integration.py rename to tests/integration/test_integration.py diff --git a/tests/test_integration_extended.py b/tests/integration/test_integration_extended.py similarity index 100% rename from tests/test_integration_extended.py rename to tests/integration/test_integration_extended.py diff --git a/tests/test_feedback_collector.py b/tests/learning/test_feedback_collector.py similarity index 100% rename from tests/test_feedback_collector.py rename to tests/learning/test_feedback_collector.py diff --git a/tests/test_llm_optimizer.py b/tests/learning/test_llm_optimizer.py similarity index 100% rename from tests/test_llm_optimizer.py rename to tests/learning/test_llm_optimizer.py diff --git a/tests/test_multi_bench_runner.py b/tests/learning/test_multi_bench_runner.py similarity index 100% rename from tests/test_multi_bench_runner.py rename to tests/learning/test_multi_bench_runner.py diff --git a/tests/test_optimize_cmd.py b/tests/learning/test_optimize_cmd.py similarity index 100% rename from tests/test_optimize_cmd.py rename to tests/learning/test_optimize_cmd.py diff --git a/tests/test_optimize_store.py b/tests/learning/test_optimize_store.py similarity index 100% rename from tests/test_optimize_store.py rename to tests/learning/test_optimize_store.py diff --git a/tests/test_optimize_types.py b/tests/learning/test_optimize_types.py similarity index 100% rename from tests/test_optimize_types.py rename to tests/learning/test_optimize_types.py diff --git a/tests/test_optimizer_engine.py b/tests/learning/test_optimizer_engine.py similarity index 100% rename from tests/test_optimizer_engine.py rename to tests/learning/test_optimizer_engine.py diff --git a/tests/test_pareto.py b/tests/learning/test_pareto.py similarity index 100% rename from tests/test_pareto.py rename to tests/learning/test_pareto.py diff --git a/tests/test_personal_synthesizer.py b/tests/learning/test_personal_synthesizer.py similarity index 100% rename from tests/test_personal_synthesizer.py rename to tests/learning/test_personal_synthesizer.py diff --git a/tests/test_search_space.py b/tests/learning/test_search_space.py similarity index 100% rename from tests/test_search_space.py rename to tests/learning/test_search_space.py diff --git a/tests/test_system_learning.py b/tests/learning/test_system_learning.py similarity index 100% rename from tests/test_system_learning.py rename to tests/learning/test_system_learning.py diff --git a/tests/test_trace_judge.py b/tests/learning/test_trace_judge.py similarity index 100% rename from tests/test_trace_judge.py rename to tests/learning/test_trace_judge.py diff --git a/tests/test_trial_runner.py b/tests/learning/test_trial_runner.py similarity index 100% rename from tests/test_trial_runner.py rename to tests/learning/test_trial_runner.py diff --git a/tests/test_system.py b/tests/sdk/test_system.py similarity index 100% rename from tests/test_system.py rename to tests/sdk/test_system.py