admin/OpenJarvis
1
0
Fork 0
mirror of https://github.com/open-jarvis/OpenJarvis.git synced 2026-07-31 03:12:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expand test suite to 1031 tests: new agents, tools, MCP layer, model catalog

Browse Source 
Add ReAct and OpenHands agents, WebSearch and CodeInterpreter tools,
full MCP protocol layer (server/client/transport), Gemini cloud engine
support, 12 new model specs (4 local MoE + 8 cloud), trace system,
and comprehensive test coverage across all dimensions (hardware, engine,
memory, agents, tools, MCP, integration).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jon Saad-Falcon
2026-02-21 04:59:11 +00:00
co-authored by Claude Opus 4.6
parent 2faf58f3f7
commit 990d7d8a79
60 changed files with 8844 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+23 -11
View File
@@ -4,13 +4,13 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
## Project Status
OpenJarvis Phase 5 (v1.0) is complete — Python SDK (`Jarvis` class with `ask()`, `ask_full()`, `MemoryHandle`), OpenClaw agent infrastructure (protocol, transport ABC with HTTP/subprocess, plugin skeleton), benchmarking framework (`jarvis bench run` with latency/throughput benchmarks, `BenchmarkRegistry`), Docker deployment (Dockerfile, Dockerfile.gpu, docker-compose.yml, systemd/launchd), and full documentation updates. ~520 tests pass (8 skipped for optional deps).
OpenJarvis is a research framework for studying on-device AI systems. Phase 5 (v1.0) complete, Phase 6 (trace system + learning) in progress. Core abstractions: Intelligence, Engine, Agentic Logic, Memory — with trace-driven learning as a cross-cutting concern. Python SDK (`Jarvis` class), OpenClaw agent infrastructure, benchmarking framework, Docker deployment all ready. ~576 tests pass (8 skipped for optional deps).
## Build & Development Commands
```bash
uv sync --extra dev # Install deps + dev tools
uv run pytest tests/ -v # Run ~520 tests (8 skipped if optional deps missing)
uv run pytest tests/ -v # Run ~576 tests (8 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)
@@ -66,15 +66,19 @@ j.close() # Release resources
## Architecture
OpenJarvis is a modular AI assistant backend organized around **five composable pillars**, each with a clear ABC interface and a decorator-based registry for runtime discovery.
OpenJarvis is a research framework for on-device AI organized around **four core abstractions** plus a cross-cutting learning system, each with a clear ABC interface and a decorator-based registry for runtime discovery.
### Five Pillars
### Four Core Abstractions
1. **Intelligence** (`src/openjarvis/intelligence/`) — Model management and query routing. `ModelRegistry` maps model keys to `ModelSpec`. Heuristic router selects model based on query characteristics.
2. **Learning** (`src/openjarvis/learning/`) — Router policy that determines which model handles a query. `RouterPolicyRegistry` enables pluggable policies. Implementations: `HeuristicRouter` (6 priority rules, registered as "heuristic"), `GRPORouterPolicy` (stub for training, registered as "grpo"). `HeuristicRewardFunction` scores inference results on latency/cost/efficiency.
3. **Memory** (`src/openjarvis/memory/`) — Persistent searchable storage. Backends: SQLite/FTS5 (default), FAISS, ColBERTv2, BM25, Hybrid (RRF fusion). All implement `MemoryBackend` ABC with `store()`, `retrieve()`, `delete()`, `clear()`.
4. **Agents** (`src/openjarvis/agents/`) — Multi-turn reasoning and tool use. `SimpleAgent` (single-turn, no tools), `OrchestratorAgent` (multi-turn tool-calling loop with `ToolExecutor`), `CustomAgent` (template for user-defined agents), `OpenClawAgent` (HTTP/subprocess transport to OpenClaw Pi agent). All implement `BaseAgent` ABC with `run()`.
5. **Inference Engine** (`src/openjarvis/engine/`) — LLM runtime management. Backends: vLLM, SGLang, Ollama, llama.cpp, MLX. All implement `InferenceEngine` ABC with `generate()`, `stream()`, `list_models()`, `health()`. Engines extract and pass through `tool_calls` in OpenAI format.
1. **Intelligence** (`src/openjarvis/intelligence/`) — The local LM. Model management and query routing. `ModelRegistry` maps model keys to `ModelSpec`. Heuristic router selects model based on query characteristics.
2. **Engine** (`src/openjarvis/engine/`) — The inference runtime. Backends: vLLM, SGLang, Ollama, llama.cpp, MLX. All implement `InferenceEngine` ABC with `generate()`, `stream()`, `list_models()`, `health()`. Engines extract and pass through `tool_calls` in OpenAI format.
3. **Agentic Logic** (`src/openjarvis/agents/`) — Pluggable logic for handling queries, making tool/API calls, managing memory. `SimpleAgent` (single-turn, no tools), `OrchestratorAgent` (multi-turn tool-calling loop with `ToolExecutor`), `CustomAgent` (template for user-defined agents), `OpenClawAgent` (HTTP/subprocess transport). All implement `BaseAgent` ABC with `run()`. Can be static or learned from traces.
4. **Memory** (`src/openjarvis/memory/`) — Persistent searchable storage. Backends: SQLite/FTS5 (default), FAISS, ColBERTv2, BM25, Hybrid (RRF fusion). All implement `MemoryBackend` ABC with `store()`, `retrieve()`, `delete()`, `clear()`.
### Cross-cutting: Learning & Traces
- **Traces** (`src/openjarvis/traces/`) — Full interaction-level recording. Every agent interaction produces a `Trace` capturing the sequence of `TraceStep`s (route, retrieve, generate, tool_call, respond) with timing, inputs, outputs, and outcomes. `TraceStore` persists to SQLite. `TraceCollector` wraps any `BaseAgent` to record traces automatically. `TraceAnalyzer` provides aggregated stats for the learning system.
- **Learning** (`src/openjarvis/learning/`) — Router policy that determines which model handles a query. `RouterPolicyRegistry` enables pluggable policies. Implementations: `HeuristicRouter` (6 priority rules, registered as "heuristic"), `TraceDrivenPolicy` (learns from trace outcomes, registered as "learned"), `GRPORouterPolicy` (stub for RL training, registered as "grpo"). `HeuristicRewardFunction` scores inference results on latency/cost/efficiency.
### Python SDK (`src/openjarvis/sdk.py`)
@@ -112,6 +116,13 @@ OpenJarvis is a modular AI assistant backend organized around **five composable
- `GET /health` — health check
- Pydantic request/response models matching OpenAI API format
### Trace System (`src/openjarvis/traces/`)
- `store.py` — `TraceStore` writes complete `Trace` objects (with `TraceStep` lists) to SQLite. Supports filtering by agent, model, outcome, time range. Subscribes to `TRACE_COMPLETE` events on EventBus.
- `collector.py` — `TraceCollector` wraps any `BaseAgent` and records a `Trace` for every `run()`. Subscribes to EventBus events (inference, tool, memory) during agent execution, converting them to `TraceStep` objects. Automatically persists traces and publishes `TRACE_COMPLETE`.
- `analyzer.py` — `TraceAnalyzer` read-only query layer: `summary()`, `per_route_stats()`, `per_tool_stats()`, `traces_for_query_type()`, `export_traces()`. Time-range filtering. Provides inputs for the learning system.
- Dataclasses: `RouteStats`, `ToolStats`, `TraceSummary`
### Telemetry (`src/openjarvis/telemetry/`)
- `store.py` — `TelemetryStore` writes records to SQLite via EventBus subscription (append-only)
@@ -122,7 +133,7 @@ OpenJarvis is a modular AI assistant backend organized around **five composable
### Core Module (`src/openjarvis/core/`)
- `registry.py` — `RegistryBase[T]` generic base class adapted from IPW. Typed subclasses: `ModelRegistry`, `EngineRegistry`, `MemoryRegistry`, `AgentRegistry`, `ToolRegistry`, `RouterPolicyRegistry`, `BenchmarkRegistry`.
- `types.py` — `Message`, `Conversation`, `ModelSpec`, `ToolResult`, `TelemetryRecord`.
- `types.py` — `Message`, `Conversation`, `ModelSpec`, `ToolResult`, `TelemetryRecord`, `StepType`, `TraceStep`, `Trace`.
- `config.py` — `JarvisConfig` dataclass hierarchy with TOML loader. Includes `LearningConfig` (default_policy, reward_weights). User config lives at `~/.openjarvis/config.toml`. Hardware auto-detection populates defaults.
- `events.py` — Pub/sub event bus for inter-pillar telemetry (synchronous dispatch).
@@ -136,7 +147,7 @@ OpenJarvis is a modular AI assistant backend organized around **five composable
### Query Flow
User query &rarr; Agentic Logic (determine tools/memory needs) &rarr; Memory retrieval &rarr; Context injection with source attribution &rarr; Learning/Router selects model (via RouterPolicyRegistry) &rarr; Inference Engine generates response &rarr; Telemetry recorded to SQLite.
User query &rarr; Agentic Logic (determine tools/memory needs) &rarr; Memory retrieval &rarr; Context injection with source attribution &rarr; Learning/Router selects model (via RouterPolicyRegistry, heuristic or trace-driven) &rarr; Inference Engine generates response &rarr; Trace recorded to SQLite (full interaction sequence) &rarr; Telemetry recorded &rarr; Learning policies update from accumulated traces.
### API Surface
@@ -161,3 +172,4 @@ OpenAI-compatible server via `jarvis serve`: `POST /v1/chat/completions` and `GE
| v0.4 | Phase 3 | Agents, tool system, OpenAI-compatible API server |
| v0.5 | Phase 4 | Learning implementations, telemetry aggregation, `--router` CLI, `jarvis telemetry` |
| v1.0 | Phase 5 | SDK, OpenClaw infrastructure, benchmarks, Docker, documentation |
| v1.1 | Phase 6 | Trace system, trace-driven learning, pluggable agent architectures |
NOTES.md
+38 -7
View File
@@ -4,14 +4,14 @@ Living document tracking implementation progress, testing state, lessons learned
---
## Current State (2026-02-16)
## Current State (2026-02-21)
- **Version:** 1.0.0
- **All 6 roadmap phases complete** (Phase 0 through Phase 5)
- **Tests:** 520 passed, 8 skipped, 0 failures
- **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:** 72 Python files in `src/openjarvis/`
- **Test files:** 74 Python files in `tests/`
- **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
@@ -40,6 +40,7 @@ Living document tracking implementation progress, testing state, lessons learned
| 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 |
---
@@ -57,7 +58,8 @@ src/openjarvis/
│ ├── config.py # JarvisConfig dataclass hierarchy, TOML loader
│ └── events.py # EventBus pub/sub (synchronous)
├── intelligence/ # ModelRegistry, HeuristicRouter, model catalog
├── learning/ # RouterPolicyRegistry, HeuristicRouter policy, GRPO stub
├── 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
@@ -459,3 +461,32 @@ python -c "from openjarvis import Jarvis; print(Jarvis)"
- 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**
README.md
+17 -12
View File
@@ -1,20 +1,23 @@
# OpenJarvis
**Your AI stack, your rules.**
**Programming abstractions for on-device AI.**
A modular, pluggable AI assistant backend. Compose your own stack across five pillars — Intelligence, Learning, Memory, Agents, and Inference — then swap any piece without touching the rest.
OpenJarvis defines the abstractions needed to study and build AI systems that run entirely on local hardware. It provides four composable pillars — Intelligence, Engine, Agentic Logic, and Memory — with a trace-driven learning system that improves over time.
> **Status: v1.0** — All five pillars implemented. SDK, benchmarks, OpenClaw infrastructure, and Docker deployment ready.
> **Status: v1.0+** — All pillars implemented. Trace system, trace-driven learning, SDK, benchmarks, and Docker deployment ready. 576 tests passing.
## What is this?
OpenJarvis lets you build a personal AI assistant from composable parts:
Local AI is a new computing paradigm: intelligence as a *resource you own*, not a *service you rent*. Existing frameworks (LangChain, DSPy, CrewAI) assume cloud-class models and infinite compute. OpenJarvis provides the abstractions for building AI systems against local hardware constraints.
- **Intelligence** — multi-model management with automatic routing (Qwen3, GPT OSS, Kimi-K2.5, Claude, GPT-5, Gemini)
**Four core abstractions:**
- **Intelligence** — the local LM being run (Qwen3 8B, GPT OSS 120B, Kimi 2.5, etc.) with multi-model management and automatic routing
- **Engine** — the local inference engine (Ollama, SGLang, vLLM, llama.cpp, MLX) with hardware-aware selection
- **Agentic Logic** — pluggable logic for handling queries, making tool/API calls, managing memory. Can be static (rules, ReAct) or learned from traces
- **Memory** — persistent, searchable storage with multiple backends (SQLite, FAISS, ColBERTv2, BM25, hybrid)
- **Agents** — pluggable reasoning and tool use (OpenClaw Pi agent, simple, orchestrator, custom)
- **Inference** — hardware-aware engine selection (vLLM, SGLang, Ollama, llama.cpp, MLX)
- **Learning** — router that improves over time (heuristic now, learned later)
**Cross-cutting: Learning** — every interaction generates a trace. The system learns better routing, tool selection, and memory strategies from accumulated trace data.
## Quick Start — Python SDK
@@ -65,12 +68,14 @@ curl http://localhost:8000/health
src/openjarvis/
├── core/ # Registry, types, config, event bus
├── intelligence/ # Model management, routing
├── engine/ # Inference engine wrappers (Ollama, vLLM, SGLang, llama.cpp, MLX)
├── agents/ # Pluggable agent implementations + tool system
├── memory/ # Storage backends (SQLite, FAISS, ColBERT, BM25, hybrid)
├── agents/ # Agent implementations + tool system + OpenClaw
├── engine/ # Inference engine wrappers
├── learning/ # Router policy (heuristic, GRPO stub)
├── traces/ # Full interaction traces — store, collector, analyzer
├── learning/ # Router policies (heuristic, trace-driven, GRPO stub)
├── telemetry/ # Per-inference telemetry store + aggregator
├── tools/ # Built-in tools (calculator, think, retrieval, LLM, file read)
├── bench/ # Benchmarking framework (latency, throughput)
├── telemetry/ # Telemetry store + aggregator
├── server/ # OpenAI-compatible API server
├── cli/ # CLI entry points
└── sdk.py # Python SDK (Jarvis class)
ROADMAP.md
+29
View File
@@ -333,6 +333,34 @@ OpenJarvis/
---
## 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 |
@@ -343,3 +371,4 @@ OpenJarvis/
| **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 |
VISION.md
+24 -12
View File
@@ -1,10 +1,20 @@
# OpenJarvis
**Your AI stack, your rules.**
**Programming abstractions for on-device AI.**
OpenJarvis is a modular, pluggable AI assistant backend. Instead of locking you into one model, one memory system, or one inference engine, OpenJarvis lets you compose your own stack across five pillars — then swap any piece without touching the rest.
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 developers who want full control and researchers who need reproducible, measurable AI systems.
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.
---
@@ -61,16 +71,17 @@ OpenJarvis is organized around five composable pillars. Each pillar defines a cl
### 2. Learning Approach (Router Policy)
**What it does:** Determines *which* model handles a given query. V1 is heuristic; future versions will learn from usage data.
**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.
**What's pluggable:** Routing policy, reward functions, training pipeline, trace analyzers.
**V1 (Placeholder):**
- Heuristic routing based on query characteristics (length, complexity keywords, domain detection)
- Rule-based fallback chains
- All telemetry logged to SQLite for future training
**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 (Post-V1):**
**Future:**
- Learned router via GRPO (Group Relative Policy Optimization)
- Preference learning from user feedback
- Continual fine-tuning on accumulated trajectories
@@ -188,9 +199,10 @@ User query
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
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. **Telemetry** records timing, token counts, energy (if available), and cost
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
---
pyproject.toml
+14
View File
@@ -30,6 +30,12 @@ inference-cloud = [
"openai>=1.30",
"anthropic>=0.30",
]
inference-google = [
"google-genai>=1.0",
]
tools-search = [
"tavily-python>=0.3",
]
memory-faiss = [
"faiss-cpu>=1.7",
"sentence-transformers>=2.2",
@@ -58,6 +64,14 @@ packages = ["src/openjarvis"]
[tool.pytest.ini_options]
testpaths = ["tests"]
markers = [
"live: requires running inference engine",
"cloud: requires cloud API keys",
"nvidia: requires NVIDIA GPU",
"amd: requires AMD GPU",
"apple: requires Apple Silicon",
"slow: long-running test",
]
[tool.ruff]
target-version = "py310"
src/openjarvis/agents/__init__.py
+10
View File
@@ -25,4 +25,14 @@ try:
except ImportError:
pass
try:
import openjarvis.agents.react # noqa: F401
except ImportError:
pass
try:
import openjarvis.agents.openhands # noqa: F401
except ImportError:
pass
__all__ = ["AgentContext", "AgentResult", "BaseAgent"]
src/openjarvis/agents/openhands.py
+188
View File
@@ -0,0 +1,188 @@
"""OpenHandsAgent -- code-execution-centric agent."""
from __future__ import annotations
import json as _json
import re
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult
from openjarvis.engine._stubs import InferenceEngine
from openjarvis.tools._stubs import BaseTool, ToolExecutor
OPENHANDS_SYSTEM_PROMPT = (
"You are a CodeAct agent. To solve tasks, "
"write Python code enclosed in ```python blocks.\n"
"The code will be executed and you will see the output. "
"You can iterate.\n"
"If no code is needed, respond directly with your answer.\n"
"Available tools: {tool_names}"
)
@AgentRegistry.register("openhands")
class OpenHandsAgent(BaseAgent):
"""OpenHands CodeAct agent -- generates and executes Python code."""
agent_id = "openhands"
def __init__(
self,
engine: InferenceEngine,
model: str,
*,
tools: Optional[List[BaseTool]] = None,
bus: Optional[EventBus] = None,
max_turns: int = 15,
temperature: float = 0.7,
max_tokens: int = 2048,
) -> None:
self._engine = engine
self._model = model
self._tools = tools or []
self._executor = ToolExecutor(self._tools, bus=bus)
self._bus = bus
self._max_turns = max_turns
self._temperature = temperature
self._max_tokens = max_tokens
def _extract_code(self, text: str) -> str | None:
"""Extract Python code from markdown code blocks."""
match = re.search(r"```python\n(.*?)```", text, re.DOTALL)
if match:
return match.group(1).strip()
return None
def _extract_tool_call(self, text: str) -> tuple[str, str] | None:
"""Extract tool call from structured output."""
action_match = re.search(r"Action:\s*(.+)", text)
input_match = re.search(
r"Action Input:\s*(.+?)(?=\n\n|\Z)", text, re.DOTALL
)
if action_match:
return (
action_match.group(1).strip(),
input_match.group(1).strip() if input_match else "{}",
)
return None
def run(
self,
input: str,
context: Optional[AgentContext] = None,
**kwargs: Any,
) -> AgentResult:
bus = self._bus
if bus:
bus.publish(
EventType.AGENT_TURN_START,
{"agent": self.agent_id, "input": input},
)
tool_names = (
", ".join(t.spec.name for t in self._tools) if self._tools else "none"
)
system_prompt = OPENHANDS_SYSTEM_PROMPT.format(tool_names=tool_names)
messages: list[Message] = [Message(role=Role.SYSTEM, content=system_prompt)]
if context and context.conversation.messages:
messages.extend(context.conversation.messages)
messages.append(Message(role=Role.USER, content=input))
all_tool_results: list[ToolResult] = []
turns = 0
last_content = ""
for _turn in range(self._max_turns):
turns += 1
if bus:
bus.publish(
EventType.INFERENCE_START,
{"model": self._model, "turn": turns},
)
result = self._engine.generate(
messages,
model=self._model,
temperature=self._temperature,
max_tokens=self._max_tokens,
)
if bus:
bus.publish(
EventType.INFERENCE_END,
{"model": self._model, "turn": turns},
)
content = result.get("content", "")
last_content = content
# Try to extract code
code = self._extract_code(content)
if code:
messages.append(Message(role=Role.ASSISTANT, content=content))
# Execute via code_interpreter tool if available
tool_call = ToolCall(
id=f"code_{turns}",
name="code_interpreter",
arguments=_json.dumps({"code": code}),
)
tool_result = self._executor.execute(tool_call)
all_tool_results.append(tool_result)
observation = f"Output:\n{tool_result.content}"
messages.append(Message(role=Role.USER, content=observation))
continue
# Try tool call
tool_info = self._extract_tool_call(content)
if tool_info:
action, action_input = tool_info
messages.append(Message(role=Role.ASSISTANT, content=content))
tool_call = ToolCall(
id=f"tool_{turns}", name=action, arguments=action_input
)
tool_result = self._executor.execute(tool_call)
all_tool_results.append(tool_result)
observation = f"Result: {tool_result.content}"
messages.append(Message(role=Role.USER, content=observation))
continue
# No code or tool call -- this is the final answer
if bus:
bus.publish(
EventType.AGENT_TURN_END,
{"agent": self.agent_id, "turns": turns},
)
return AgentResult(
content=content, tool_results=all_tool_results, turns=turns
)
# Max turns
if bus:
bus.publish(
EventType.AGENT_TURN_END,
{
"agent": self.agent_id,
"turns": turns,
"max_turns_exceeded": True,
},
)
return AgentResult(
content=last_content or "Maximum turns reached.",
tool_results=all_tool_results,
turns=turns,
metadata={"max_turns_exceeded": True},
)
__all__ = ["OpenHandsAgent"]
src/openjarvis/agents/react.py
+196
View File
@@ -0,0 +1,196 @@
"""ReActAgent -- Thought-Action-Observation loop agent."""
from __future__ import annotations
import re
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult
from openjarvis.engine._stubs import InferenceEngine
from openjarvis.tools._stubs import BaseTool, ToolExecutor
REACT_SYSTEM_PROMPT = """\
You are a ReAct agent. For each step, respond with exactly one of:
1. To think and act:
Thought: <your reasoning>
Action: <tool_name>
Action Input: <json arguments>
2. To give a final answer:
Thought: <your reasoning>
Final Answer: <your answer>
Available tools: {tool_names}"""
@AgentRegistry.register("react")
class ReActAgent(BaseAgent):
"""ReAct agent: Thought -> Action -> Observation loop."""
agent_id = "react"
def __init__(
self,
engine: InferenceEngine,
model: str,
*,
tools: Optional[List[BaseTool]] = None,
bus: Optional[EventBus] = None,
max_turns: int = 10,
temperature: float = 0.7,
max_tokens: int = 1024,
) -> None:
self._engine = engine
self._model = model
self._tools = tools or []
self._executor = ToolExecutor(self._tools, bus=bus)
self._bus = bus
self._max_turns = max_turns
self._temperature = temperature
self._max_tokens = max_tokens
def _parse_response(self, text: str) -> dict:
"""Parse ReAct structured output."""
result = {"thought": "", "action": "", "action_input": "", "final_answer": ""}
# Extract Thought
thought_match = re.search(
r"Thought:\s*(.+?)(?=\nAction:|\nFinal Answer:|\Z)", text, re.DOTALL
)
if thought_match:
result["thought"] = thought_match.group(1).strip()
# Check for Final Answer
final_match = re.search(r"Final Answer:\s*(.+)", text, re.DOTALL)
if final_match:
result["final_answer"] = final_match.group(1).strip()
return result
# Extract Action and Action Input
action_match = re.search(r"Action:\s*(.+)", text)
if action_match:
result["action"] = action_match.group(1).strip()
input_match = re.search(
r"Action Input:\s*(.+?)(?=\n\n|\nThought:|\Z)", text, re.DOTALL
)
if input_match:
result["action_input"] = input_match.group(1).strip()
return result
def run(
self,
input: str,
context: Optional[AgentContext] = None,
**kwargs: Any,
) -> AgentResult:
bus = self._bus
if bus:
bus.publish(
EventType.AGENT_TURN_START,
{"agent": self.agent_id, "input": input},
)
# Build system prompt with available tools
tool_names = (
", ".join(t.spec.name for t in self._tools) if self._tools else "none"
)
system_prompt = REACT_SYSTEM_PROMPT.format(tool_names=tool_names)
messages: list[Message] = [Message(role=Role.SYSTEM, content=system_prompt)]
if context and context.conversation.messages:
messages.extend(context.conversation.messages)
messages.append(Message(role=Role.USER, content=input))
all_tool_results: list[ToolResult] = []
turns = 0
for _turn in range(self._max_turns):
turns += 1
if bus:
bus.publish(
EventType.INFERENCE_START,
{"model": self._model, "turn": turns},
)
result = self._engine.generate(
messages,
model=self._model,
temperature=self._temperature,
max_tokens=self._max_tokens,
)
if bus:
bus.publish(
EventType.INFERENCE_END,
{"model": self._model, "turn": turns},
)
content = result.get("content", "")
parsed = self._parse_response(content)
# Final answer?
if parsed["final_answer"]:
if bus:
bus.publish(
EventType.AGENT_TURN_END,
{"agent": self.agent_id, "turns": turns},
)
return AgentResult(
content=parsed["final_answer"],
tool_results=all_tool_results,
turns=turns,
)
# No action? Treat content as final answer
if not parsed["action"]:
if bus:
bus.publish(
EventType.AGENT_TURN_END,
{"agent": self.agent_id, "turns": turns},
)
return AgentResult(
content=content, tool_results=all_tool_results, turns=turns
)
# Execute action
messages.append(Message(role=Role.ASSISTANT, content=content))
tool_call = ToolCall(
id=f"react_{turns}",
name=parsed["action"],
arguments=parsed["action_input"] or "{}",
)
tool_result = self._executor.execute(tool_call)
all_tool_results.append(tool_result)
observation = f"Observation: {tool_result.content}"
messages.append(Message(role=Role.USER, content=observation))
# Max turns exceeded
if bus:
bus.publish(
EventType.AGENT_TURN_END,
{
"agent": self.agent_id,
"turns": turns,
"max_turns_exceeded": True,
},
)
return AgentResult(
content="Maximum turns reached without a final answer.",
tool_results=all_tool_results,
turns=turns,
metadata={"max_turns_exceeded": True},
)
__all__ = ["ReActAgent"]
src/openjarvis/core/events.py
+2
View File
@@ -30,6 +30,8 @@ class EventType(str, Enum):
AGENT_TURN_START = "agent_turn_start"
AGENT_TURN_END = "agent_turn_end"
TELEMETRY_RECORD = "telemetry_record"
TRACE_STEP = "trace_step"
TRACE_COMPLETE = "trace_complete"
@dataclass(slots=True)
src/openjarvis/core/types.py
+72
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Sequence # noqa: I001
@@ -32,6 +33,16 @@ class Quantization(str, Enum):
GGUF_Q8 = "gguf_q8"
class StepType(str, Enum):
"""Types of steps within an agent trace."""
ROUTE = "route"
RETRIEVE = "retrieve"
GENERATE = "generate"
TOOL_CALL = "tool_call"
RESPOND = "respond"
# ---------------------------------------------------------------------------
# Message types
# ---------------------------------------------------------------------------
@@ -130,13 +141,74 @@ class TelemetryRecord:
metadata: Dict[str, Any] = field(default_factory=dict)
# ---------------------------------------------------------------------------
# Trace types — full interaction-level recording
# ---------------------------------------------------------------------------
def _trace_id() -> str:
return uuid.uuid4().hex[:16]
@dataclass(slots=True)
class TraceStep:
"""A single step within an agent trace.
Each step records what the agent did (route, retrieve, generate,
tool_call, respond), its inputs and outputs, and timing.
"""
step_type: StepType
timestamp: float
duration_seconds: float = 0.0
input: Dict[str, Any] = field(default_factory=dict)
output: Dict[str, Any] = field(default_factory=dict)
metadata: Dict[str, Any] = field(default_factory=dict)
@dataclass(slots=True)
class Trace:
"""Complete trace of an agent handling a query.
A trace captures the full sequence of steps an agent took to handle a
query — which model was selected, what memory was retrieved, which tools
were called, and the final response. Traces are the primary input to the
learning system: by analyzing which decisions led to good outcomes, the
system can improve routing, tool selection, and memory strategies.
"""
trace_id: str = field(default_factory=_trace_id)
query: str = ""
agent: str = ""
model: str = ""
engine: str = ""
steps: List[TraceStep] = field(default_factory=list)
result: str = ""
outcome: Optional[str] = None # None=unknown, "success", "failure"
feedback: Optional[float] = None # user quality score [0, 1]
started_at: float = 0.0
ended_at: float = 0.0
total_tokens: int = 0
total_latency_seconds: float = 0.0
metadata: Dict[str, Any] = field(default_factory=dict)
def add_step(self, step: TraceStep) -> None:
"""Append a step and update running totals."""
self.steps.append(step)
self.total_latency_seconds += step.duration_seconds
self.total_tokens += step.output.get("tokens", 0)
__all__ = [
"Conversation",
"Message",
"ModelSpec",
"Quantization",
"Role",
"StepType",
"TelemetryRecord",
"ToolCall",
"ToolResult",
"Trace",
"TraceStep",
]
src/openjarvis/engine/cloud.py
+147 -4
View File
@@ -1,4 +1,4 @@
"""Cloud inference engine — OpenAI and Anthropic API backends."""
"""Cloud inference engine — OpenAI, Anthropic, and Google API backends."""
from __future__ import annotations
@@ -19,18 +19,35 @@ PRICING: Dict[str, tuple[float, float]] = {
"gpt-4o": (2.50, 10.00),
"gpt-4o-mini": (0.15, 0.60),
"gpt-5": (10.00, 30.00),
"gpt-5-mini": (0.25, 2.00),
"o3-mini": (1.10, 4.40),
"claude-sonnet-4-20250514": (3.00, 15.00),
"claude-opus-4-20250514": (15.00, 75.00),
"claude-haiku-3-5-20241022": (0.80, 4.00),
"claude-opus-4-6": (5.00, 25.00),
"claude-sonnet-4-6": (3.00, 15.00),
"claude-haiku-4-5": (1.00, 5.00),
"gemini-2.5-pro": (1.25, 10.00),
"gemini-2.5-flash": (0.30, 2.50),
"gemini-3-pro": (2.00, 12.00),
"gemini-3-flash": (0.50, 3.00),
}
# Well-known model IDs per provider
_OPENAI_MODELS = ["gpt-4o", "gpt-4o-mini", "gpt-5", "o3-mini"]
_OPENAI_MODELS = ["gpt-4o", "gpt-4o-mini", "gpt-5", "gpt-5-mini", "o3-mini"]
_ANTHROPIC_MODELS = [
"claude-sonnet-4-20250514",
"claude-opus-4-20250514",
"claude-haiku-3-5-20241022",
"claude-opus-4-6",
"claude-sonnet-4-6",
"claude-haiku-4-5",
]
_GOOGLE_MODELS = [
"gemini-2.5-pro",
"gemini-2.5-flash",
"gemini-3-pro",
"gemini-3-flash",
]
@@ -38,6 +55,10 @@ def _is_anthropic_model(model: str) -> bool:
return "claude" in model.lower()
def _is_google_model(model: str) -> bool:
return "gemini" in model.lower()
def estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> float:
"""Estimate USD cost based on the hardcoded pricing table."""
# Try exact match first, then prefix match
@@ -56,13 +77,14 @@ def estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> flo
@EngineRegistry.register("cloud")
class CloudEngine(InferenceEngine):
"""Cloud inference via OpenAI and Anthropic SDKs."""
"""Cloud inference via OpenAI, Anthropic, and Google SDKs."""
engine_id = "cloud"
def __init__(self) -> None:
self._openai_client: Any = None
self._anthropic_client: Any = None
self._google_client: Any = None
self._init_clients()
def _init_clients(self) -> None:
@@ -78,6 +100,16 @@ class CloudEngine(InferenceEngine):
self._anthropic_client = anthropic.Anthropic()
except ImportError:
pass
gemini_key = (
os.environ.get("GEMINI_API_KEY")
or os.environ.get("GOOGLE_API_KEY")
)
if gemini_key:
try:
from google import genai
self._google_client = genai.Client(api_key=gemini_key)
except ImportError:
pass
def _generate_openai(
self,
@@ -164,6 +196,66 @@ class CloudEngine(InferenceEngine):
"cost_usd": estimate_cost(model, prompt_tokens, completion_tokens),
}
def _generate_google(
self,
messages: Sequence[Message],
*,
model: str,
temperature: float,
max_tokens: int,
**kwargs: Any,
) -> Dict[str, Any]:
if self._google_client is None:
raise EngineConnectionError(
"Google client not available — set "
"GEMINI_API_KEY or GOOGLE_API_KEY and install "
"openjarvis[inference-google]"
)
# Build contents from messages
system_text = ""
contents: List[Dict[str, Any]] = []
for m in messages:
if m.role.value == "system":
system_text = m.content
elif m.role.value == "assistant":
contents.append({"role": "model", "parts": [{"text": m.content}]})
else:
contents.append({"role": "user", "parts": [{"text": m.content}]})
from google.genai import types as genai_types
config = genai_types.GenerateContentConfig(
temperature=temperature,
max_output_tokens=max_tokens,
)
if system_text:
config.system_instruction = system_text
resp = self._google_client.models.generate_content(
model=model,
contents=contents,
config=config,
)
content = resp.text or ""
um = resp.usage_metadata
prompt_tokens = (
getattr(um, "prompt_token_count", 0) if um else 0
)
completion_tokens = (
getattr(um, "candidates_token_count", 0) if um else 0
)
return {
"content": content,
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
"model": model,
"finish_reason": "stop",
"cost_usd": estimate_cost(model, prompt_tokens, completion_tokens),
}
def generate(
self,
messages: Sequence[Message],
@@ -181,6 +273,8 @@ class CloudEngine(InferenceEngine):
)
if _is_anthropic_model(model):
return self._generate_anthropic(messages, **kw)
if _is_google_model(model):
return self._generate_google(messages, **kw)
return self._generate_openai(messages, **kw)
async def stream(
@@ -203,6 +297,11 @@ class CloudEngine(InferenceEngine):
messages, **kw
):
yield token
elif _is_google_model(model):
async for token in self._stream_google(
messages, **kw
):
yield token
else:
async for token in self._stream_openai(
messages, **kw
@@ -263,16 +362,60 @@ class CloudEngine(InferenceEngine):
for text in stream.text_stream:
yield text
async def _stream_google(
self,
messages: Sequence[Message],
*,
model: str,
temperature: float,
max_tokens: int,
**kwargs: Any,
) -> AsyncIterator[str]:
if self._google_client is None:
raise EngineConnectionError("Google client not available")
system_text = ""
contents: List[Dict[str, Any]] = []
for m in messages:
if m.role.value == "system":
system_text = m.content
elif m.role.value == "assistant":
contents.append({"role": "model", "parts": [{"text": m.content}]})
else:
contents.append({"role": "user", "parts": [{"text": m.content}]})
from google.genai import types as genai_types
config = genai_types.GenerateContentConfig(
temperature=temperature,
max_output_tokens=max_tokens,
)
if system_text:
config.system_instruction = system_text
for chunk in self._google_client.models.generate_content_stream(
model=model,
contents=contents,
config=config,
):
if chunk.text:
yield chunk.text
def list_models(self) -> List[str]:
models: List[str] = []
if self._openai_client is not None:
models.extend(_OPENAI_MODELS)
if self._anthropic_client is not None:
models.extend(_ANTHROPIC_MODELS)
if self._google_client is not None:
models.extend(_GOOGLE_MODELS)
return models
def health(self) -> bool:
return self._openai_client is not None or self._anthropic_client is not None
return (
self._openai_client is not None
or self._anthropic_client is not None
or self._google_client is not None
)
__all__ = ["CloudEngine", "PRICING", "estimate_cost"]
src/openjarvis/intelligence/model_catalog.py
+140 -2
View File
@@ -8,13 +8,17 @@ from openjarvis.core.registry import ModelRegistry
from openjarvis.core.types import ModelSpec
BUILTIN_MODELS: List[ModelSpec] = [
# -----------------------------------------------------------------------
# Local models — Dense
# -----------------------------------------------------------------------
ModelSpec(
model_id="qwen3:8b",
name="Qwen3 8B",
parameter_count_b=8.0,
parameter_count_b=8.2,
context_length=32768,
supported_engines=("ollama", "vllm", "llamacpp"),
supported_engines=("vllm", "ollama", "llamacpp", "sglang"),
provider="alibaba",
metadata={"architecture": "dense"},
),
ModelSpec(
model_id="qwen3:32b",
@@ -24,6 +28,7 @@ BUILTIN_MODELS: List[ModelSpec] = [
min_vram_gb=20.0,
supported_engines=("ollama", "vllm"),
provider="alibaba",
metadata={"architecture": "dense"},
),
ModelSpec(
model_id="llama3.3:70b",
@@ -33,6 +38,7 @@ BUILTIN_MODELS: List[ModelSpec] = [
min_vram_gb=40.0,
supported_engines=("ollama", "vllm"),
provider="meta",
metadata={"architecture": "dense"},
),
ModelSpec(
model_id="llama3.2:3b",
@@ -41,6 +47,7 @@ BUILTIN_MODELS: List[ModelSpec] = [
context_length=131072,
supported_engines=("ollama", "vllm", "llamacpp"),
provider="meta",
metadata={"architecture": "dense"},
),
ModelSpec(
model_id="deepseek-coder-v2:16b",
@@ -49,6 +56,7 @@ BUILTIN_MODELS: List[ModelSpec] = [
context_length=131072,
supported_engines=("ollama", "vllm"),
provider="deepseek",
metadata={"architecture": "dense"},
),
ModelSpec(
model_id="mistral:7b",
@@ -57,7 +65,47 @@ BUILTIN_MODELS: List[ModelSpec] = [
context_length=32768,
supported_engines=("ollama", "vllm", "llamacpp"),
provider="mistral",
metadata={"architecture": "dense"},
),
# -----------------------------------------------------------------------
# Local models — Mixture of Experts (MoE)
# -----------------------------------------------------------------------
ModelSpec(
model_id="gpt-oss:120b",
name="GPT-OSS 120B",
parameter_count_b=117.0,
active_parameter_count_b=5.1,
context_length=131072,
min_vram_gb=12.0,
supported_engines=("vllm", "ollama"),
provider="open-source",
metadata={"architecture": "moe"},
),
ModelSpec(
model_id="glm-4.7-flash",
name="GLM 4.7 Flash",
parameter_count_b=30.0,
active_parameter_count_b=3.0,
context_length=131072,
min_vram_gb=8.0,
supported_engines=("vllm", "sglang"),
provider="zhipu",
metadata={"architecture": "moe"},
),
ModelSpec(
model_id="trinity-mini",
name="Trinity Mini",
parameter_count_b=26.0,
active_parameter_count_b=3.0,
context_length=128000,
min_vram_gb=8.0,
supported_engines=("vllm", "llamacpp"),
provider="trinity",
metadata={"architecture": "moe"},
),
# -----------------------------------------------------------------------
# Cloud models — OpenAI
# -----------------------------------------------------------------------
ModelSpec(
model_id="gpt-4o",
name="GPT-4o",
@@ -66,6 +114,7 @@ BUILTIN_MODELS: List[ModelSpec] = [
supported_engines=("cloud",),
provider="openai",
requires_api_key=True,
metadata={"pricing_input": 2.50, "pricing_output": 10.00},
),
ModelSpec(
model_id="gpt-4o-mini",
@@ -75,7 +124,21 @@ BUILTIN_MODELS: List[ModelSpec] = [
supported_engines=("cloud",),
provider="openai",
requires_api_key=True,
metadata={"pricing_input": 0.15, "pricing_output": 0.60},
),
ModelSpec(
model_id="gpt-5-mini",
name="GPT-5 Mini",
parameter_count_b=0.0,
context_length=400000,
supported_engines=("cloud",),
provider="openai",
requires_api_key=True,
metadata={"pricing_input": 0.25, "pricing_output": 2.00},
),
# -----------------------------------------------------------------------
# Cloud models — Anthropic
# -----------------------------------------------------------------------
ModelSpec(
model_id="claude-sonnet-4-20250514",
name="Claude Sonnet 4",
@@ -84,6 +147,7 @@ BUILTIN_MODELS: List[ModelSpec] = [
supported_engines=("cloud",),
provider="anthropic",
requires_api_key=True,
metadata={"pricing_input": 3.00, "pricing_output": 15.00},
),
ModelSpec(
model_id="claude-opus-4-20250514",
@@ -93,6 +157,80 @@ BUILTIN_MODELS: List[ModelSpec] = [
supported_engines=("cloud",),
provider="anthropic",
requires_api_key=True,
metadata={"pricing_input": 15.00, "pricing_output": 75.00},
),
ModelSpec(
model_id="claude-opus-4-6",
name="Claude Opus 4.6",
parameter_count_b=0.0,
context_length=200000,
supported_engines=("cloud",),
provider="anthropic",
requires_api_key=True,
metadata={"pricing_input": 5.00, "pricing_output": 25.00},
),
ModelSpec(
model_id="claude-sonnet-4-6",
name="Claude Sonnet 4.6",
parameter_count_b=0.0,
context_length=200000,
supported_engines=("cloud",),
provider="anthropic",
requires_api_key=True,
metadata={"pricing_input": 3.00, "pricing_output": 15.00},
),
ModelSpec(
model_id="claude-haiku-4-5",
name="Claude Haiku 4.5",
parameter_count_b=0.0,
context_length=200000,
supported_engines=("cloud",),
provider="anthropic",
requires_api_key=True,
metadata={"pricing_input": 1.00, "pricing_output": 5.00},
),
# -----------------------------------------------------------------------
# Cloud models — Google
# -----------------------------------------------------------------------
ModelSpec(
model_id="gemini-2.5-pro",
name="Gemini 2.5 Pro",
parameter_count_b=0.0,
context_length=1000000,
supported_engines=("cloud",),
provider="google",
requires_api_key=True,
metadata={"pricing_input": 1.25, "pricing_output": 10.00},
),
ModelSpec(
model_id="gemini-2.5-flash",
name="Gemini 2.5 Flash",
parameter_count_b=0.0,
context_length=1000000,
supported_engines=("cloud",),
provider="google",
requires_api_key=True,
metadata={"pricing_input": 0.30, "pricing_output": 2.50},
),
ModelSpec(
model_id="gemini-3-pro",
name="Gemini 3 Pro",
parameter_count_b=0.0,
context_length=1000000,
supported_engines=("cloud",),
provider="google",
requires_api_key=True,
metadata={"pricing_input": 2.00, "pricing_output": 12.00},
),
ModelSpec(
model_id="gemini-3-flash",
name="Gemini 3 Flash",
parameter_count_b=0.0,
context_length=1000000,
supported_engines=("cloud",),
provider="google",
requires_api_key=True,
metadata={"pricing_input": 0.50, "pricing_output": 3.00},
),
]
src/openjarvis/learning/__init__.py
+6
View File
@@ -26,6 +26,12 @@ def ensure_registered() -> None:
except ImportError:
pass
from openjarvis.learning.trace_policy import (
ensure_registered as _reg_trace,
)
_reg_trace()
__all__ = [
"HeuristicRewardFunction",
src/openjarvis/learning/trace_policy.py
+237
View File
@@ -0,0 +1,237 @@
"""Trace-driven router policy — learns from interaction history.
Unlike the ``HeuristicRouter`` which uses static rules, this policy
learns from accumulated traces which model/agent/tool combinations
produce the best outcomes for different query types. It maintains a
lightweight mapping of (query_class → model) that is updated
periodically from the ``TraceAnalyzer``.
"""
from __future__ import annotations
import re
from typing import Any, Dict, List, Optional
from openjarvis.core.registry import RouterPolicyRegistry
from openjarvis.learning._stubs import RouterPolicy, RoutingContext
from openjarvis.traces.analyzer import TraceAnalyzer
# Query classification for grouping traces
_CODE_RE = re.compile(
r"```|`[^`]+`|\bdef\s|\bclass\s|\bimport\s|\bfunction\s",
re.IGNORECASE,
)
_MATH_RE = re.compile(
r"\bsolve\b|\bintegral\b|\bequation\b|\bcalculate\b|\bcompute\b",
re.IGNORECASE,
)
def classify_query(query: str) -> str:
"""Classify a query into a broad category for routing."""
if _CODE_RE.search(query):
return "code"
if _MATH_RE.search(query):
return "math"
if len(query) < 50:
return "short"
if len(query) > 500:
return "long"
return "general"
class TraceDrivenPolicy(RouterPolicy):
"""Router policy that learns from historical traces.
Maintains a mapping of ``query_class → best_model`` derived from
trace outcomes. Falls back to the provided default when no trace
data is available for a query class.
The policy is updated by calling :meth:`update_from_traces`, which
reads the ``TraceAnalyzer`` and recomputes the mapping.
"""
def __init__(
self,
analyzer: Optional[TraceAnalyzer] = None,
*,
available_models: Optional[List[str]] = None,
default_model: str = "",
fallback_model: str = "",
) -> None:
self._analyzer = analyzer
self._available = available_models or []
self._default = default_model
self._fallback = fallback_model
# Learned mapping: query_class → model key
self._policy_map: Dict[str, str] = {}
# Track confidence: query_class → sample count
self._confidence: Dict[str, int] = {}
# Minimum samples before trusting learned policy
self.min_samples: int = 5
@property
def policy_map(self) -> Dict[str, str]:
"""Current learned routing decisions (read-only copy)."""
return dict(self._policy_map)
def select_model(self, context: RoutingContext) -> str:
"""Select the best model based on learned policy or fallback."""
query_class = classify_query(context.query)
# Use learned policy if we have enough confidence
if (
query_class in self._policy_map
and self._confidence.get(query_class, 0) >= self.min_samples
):
model = self._policy_map[query_class]
if not self._available or model in self._available:
return model
# Fallback chain
avail = self._available
if self._default and (not avail or self._default in avail):
return self._default
if self._fallback and (not avail or self._fallback in avail):
return self._fallback
if self._available:
return self._available[0]
return self._default or ""
def update_from_traces(
self,
*,
since: Optional[float] = None,
until: Optional[float] = None,
) -> Dict[str, Any]:
"""Recompute the policy map from trace history.
Returns a summary of what changed for logging/debugging.
"""
if self._analyzer is None:
return {"error": "no analyzer configured"}
traces = self._analyzer._store.list_traces(
since=since, until=until, limit=10_000
)
if not traces:
return {"updated": False, "reason": "no traces"}
# Group traces by query class
groups: Dict[str, list] = {}
for t in traces:
qclass = classify_query(t.query)
groups.setdefault(qclass, []).append(t)
old_map = dict(self._policy_map)
changes: Dict[str, Dict[str, str]] = {}
for qclass, class_traces in groups.items():
# Score each model for this query class
model_scores: Dict[str, _ModelScore] = {}
for t in class_traces:
if not t.model:
continue
if t.model not in model_scores:
model_scores[t.model] = _ModelScore()
score = model_scores[t.model]
score.count += 1
score.total_latency += t.total_latency_seconds
if t.outcome == "success":
score.successes += 1
if t.feedback is not None:
score.feedback_sum += t.feedback
score.feedback_count += 1
if not model_scores:
continue
# Pick the best model: weighted score of success_rate and feedback
best_model = max(
model_scores.items(),
key=lambda kv: kv[1].composite_score(),
)[0]
self._policy_map[qclass] = best_model
self._confidence[qclass] = sum(s.count for s in model_scores.values())
if old_map.get(qclass) != best_model:
changes[qclass] = {
"old": old_map.get(qclass, ""),
"new": best_model,
}
return {
"updated": True,
"query_classes": len(groups),
"total_traces": len(traces),
"changes": changes,
}
def observe(
self,
query: str,
model: str,
outcome: Optional[str],
feedback: Optional[float],
) -> None:
"""Record a single observation for online (incremental) updates.
This is a lighter-weight alternative to :meth:`update_from_traces`
for use cases where you want to update the policy after every
interaction rather than in batch.
"""
qclass = classify_query(query)
current_count = self._confidence.get(qclass, 0)
# Simple exponential moving average for online update
if qclass not in self._policy_map:
self._policy_map[qclass] = model
self._confidence[qclass] = 1
return
self._confidence[qclass] = current_count + 1
# Only switch models if the new model shows clearly better outcomes
if outcome == "success" and feedback is not None and feedback > 0.7:
# Weight new evidence against existing policy
if current_count < self.min_samples:
self._policy_map[qclass] = model
class _ModelScore:
"""Accumulator for per-model scoring during policy update."""
__slots__ = (
"count", "successes", "total_latency",
"feedback_sum", "feedback_count",
)
def __init__(self) -> None:
self.count = 0
self.successes = 0
self.total_latency = 0.0
self.feedback_sum = 0.0
self.feedback_count = 0
def composite_score(self) -> float:
"""Weighted score combining success rate and feedback."""
sr = self.successes / self.count if self.count else 0.0
fb = (
self.feedback_sum / self.feedback_count
if self.feedback_count else 0.5
)
# 60% success rate + 40% feedback
return 0.6 * sr + 0.4 * fb
def ensure_registered() -> None:
"""Register TraceDrivenPolicy if not already present."""
if not RouterPolicyRegistry.contains("learned"):
RouterPolicyRegistry.register_value("learned", TraceDrivenPolicy)
ensure_registered()
__all__ = ["TraceDrivenPolicy", "classify_query", "ensure_registered"]
src/openjarvis/mcp/__init__.py
+24
View File
@@ -0,0 +1,24 @@
"""MCP (Model Context Protocol) layer for OpenJarvis."""
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.protocol import MCPError, MCPNotification, MCPRequest, MCPResponse
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import (
InProcessTransport,
MCPTransport,
SSETransport,
StdioTransport,
)
__all__ = [
"MCPClient",
"MCPError",
"MCPNotification",
"MCPRequest",
"MCPResponse",
"MCPServer",
"MCPTransport",
"InProcessTransport",
"SSETransport",
"StdioTransport",
]
src/openjarvis/mcp/client.py
+91
View File
@@ -0,0 +1,91 @@
"""MCP Client — connects to MCP servers and discovers/calls tools."""
from __future__ import annotations
import itertools
from typing import Any, Dict, List
from openjarvis.mcp.protocol import MCPError, MCPRequest, MCPResponse
from openjarvis.mcp.transport import MCPTransport
from openjarvis.tools._stubs import ToolSpec
class MCPClient:
"""Client that communicates with an MCP server via a transport.
Parameters
----------
transport:
The transport layer to use for communication.
"""
def __init__(self, transport: MCPTransport) -> None:
self._transport = transport
self._initialized = False
self._capabilities: Dict[str, Any] = {}
self._id_counter = itertools.count(1)
def _next_id(self) -> int:
return next(self._id_counter)
def _send(self, method: str, params: Dict[str, Any] | None = None) -> MCPResponse:
"""Send a request and check for errors."""
request = MCPRequest(
method=method,
params=params or {},
id=self._next_id(),
)
response = self._transport.send(request)
if response.error is not None:
raise MCPError(
code=response.error.get("code", -1),
message=response.error.get("message", "Unknown error"),
data=response.error.get("data"),
)
return response
def initialize(self) -> Dict[str, Any]:
"""Perform the MCP initialize handshake.
Returns the server capabilities.
"""
response = self._send("initialize")
self._initialized = True
self._capabilities = response.result.get("capabilities", {})
return response.result
def list_tools(self) -> List[ToolSpec]:
"""Discover available tools from the server.
Returns a list of ``ToolSpec`` objects.
"""
response = self._send("tools/list")
tools = response.result.get("tools", [])
return [
ToolSpec(
name=t["name"],
description=t.get("description", ""),
parameters=t.get("inputSchema", {}),
)
for t in tools
]
def call_tool(
self, name: str, arguments: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
"""Call a tool on the server.
Returns the result dictionary with ``content`` and ``isError`` fields.
"""
response = self._send(
"tools/call",
{"name": name, "arguments": arguments or {}},
)
return response.result
def close(self) -> None:
"""Close the transport connection."""
self._transport.close()
__all__ = ["MCPClient"]
src/openjarvis/mcp/protocol.py
+134
View File
@@ -0,0 +1,134 @@
"""MCP JSON-RPC 2.0 protocol message types."""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import Any, Dict, Optional
# Error codes per JSON-RPC 2.0 / MCP spec
PARSE_ERROR = -32700
INVALID_REQUEST = -32600
METHOD_NOT_FOUND = -32601
INVALID_PARAMS = -32602
INTERNAL_ERROR = -32603
@dataclass
class MCPRequest:
"""JSON-RPC 2.0 request message."""
method: str
params: Dict[str, Any] = field(default_factory=dict)
id: int | str = 0
jsonrpc: str = "2.0"
def to_json(self) -> str:
"""Serialize to JSON string."""
return json.dumps(
{
"jsonrpc": self.jsonrpc,
"id": self.id,
"method": self.method,
"params": self.params,
}
)
@classmethod
def from_json(cls, data: str) -> MCPRequest:
"""Deserialize from JSON string."""
parsed = json.loads(data)
return cls(
method=parsed["method"],
params=parsed.get("params", {}),
id=parsed.get("id", 0),
jsonrpc=parsed.get("jsonrpc", "2.0"),
)
@dataclass
class MCPResponse:
"""JSON-RPC 2.0 response message."""
result: Any = None
error: Optional[Dict[str, Any]] = None
id: int | str = 0
jsonrpc: str = "2.0"
def to_json(self) -> str:
"""Serialize to JSON string."""
obj: Dict[str, Any] = {"jsonrpc": self.jsonrpc, "id": self.id}
if self.error is not None:
obj["error"] = self.error
else:
obj["result"] = self.result
return json.dumps(obj)
@classmethod
def from_json(cls, data: str) -> MCPResponse:
"""Deserialize from JSON string."""
parsed = json.loads(data)
return cls(
result=parsed.get("result"),
error=parsed.get("error"),
id=parsed.get("id", 0),
jsonrpc=parsed.get("jsonrpc", "2.0"),
)
@classmethod
def error_response(
cls,
id: int | str,
code: int,
message: str,
data: Any = None,
) -> MCPResponse:
"""Create an error response."""
error: Dict[str, Any] = {"code": code, "message": message}
if data is not None:
error["data"] = data
return cls(error=error, id=id)
@dataclass
class MCPNotification:
"""JSON-RPC 2.0 notification (no id, no response expected)."""
method: str
params: Dict[str, Any] = field(default_factory=dict)
jsonrpc: str = "2.0"
def to_json(self) -> str:
"""Serialize to JSON string."""
return json.dumps(
{
"jsonrpc": self.jsonrpc,
"method": self.method,
"params": self.params,
}
)
@dataclass
class MCPError(Exception):
"""MCP protocol error with JSON-RPC error code."""
code: int
message: str
data: Any = None
def __str__(self) -> str:
return f"MCPError({self.code}): {self.message}"
__all__ = [
"INTERNAL_ERROR",
"INVALID_PARAMS",
"INVALID_REQUEST",
"MCPError",
"MCPNotification",
"MCPRequest",
"MCPResponse",
"METHOD_NOT_FOUND",
"PARSE_ERROR",
]
src/openjarvis/mcp/server.py
+125
View File
@@ -0,0 +1,125 @@
"""MCP Server — wraps OpenJarvis tools as MCP-discoverable tools."""
from __future__ import annotations
from typing import Dict, List
from openjarvis.core.types import ToolCall
from openjarvis.mcp.protocol import (
INTERNAL_ERROR,
INVALID_PARAMS,
METHOD_NOT_FOUND,
MCPRequest,
MCPResponse,
)
from openjarvis.tools._stubs import BaseTool, ToolExecutor
class MCPServer:
"""MCP server that exposes OpenJarvis tools via JSON-RPC.
Parameters
----------
tools:
List of ``BaseTool`` instances to expose.
"""
SERVER_NAME = "openjarvis"
SERVER_VERSION = "1.0.0"
PROTOCOL_VERSION = "2024-11-05"
def __init__(self, tools: List[BaseTool]) -> None:
self._tools: Dict[str, BaseTool] = {t.spec.name: t for t in tools}
self._executor = ToolExecutor(tools)
def handle(self, request: MCPRequest) -> MCPResponse:
"""Dispatch an MCP request and return a response."""
if request.method == "initialize":
return self._handle_initialize(request)
elif request.method == "tools/list":
return self._handle_tools_list(request)
elif request.method == "tools/call":
return self._handle_tools_call(request)
else:
return MCPResponse.error_response(
request.id,
METHOD_NOT_FOUND,
f"Unknown method: {request.method}",
)
def _handle_initialize(self, req: MCPRequest) -> MCPResponse:
"""Handle the initialize handshake."""
return MCPResponse(
result={
"protocolVersion": self.PROTOCOL_VERSION,
"capabilities": {
"tools": {"listChanged": False},
},
"serverInfo": {
"name": self.SERVER_NAME,
"version": self.SERVER_VERSION,
},
},
id=req.id,
)
def _handle_tools_list(self, req: MCPRequest) -> MCPResponse:
"""Handle tools/list — return specs for all registered tools."""
tool_list = []
for tool in self._tools.values():
s = tool.spec
tool_list.append(
{
"name": s.name,
"description": s.description,
"inputSchema": s.parameters,
}
)
return MCPResponse(result={"tools": tool_list}, id=req.id)
def _handle_tools_call(self, req: MCPRequest) -> MCPResponse:
"""Handle tools/call — execute a tool and return the result."""
tool_name = req.params.get("name")
arguments = req.params.get("arguments", {})
if not tool_name:
return MCPResponse.error_response(
req.id,
INVALID_PARAMS,
"Missing required parameter: name",
)
if tool_name not in self._tools:
return MCPResponse.error_response(
req.id,
INVALID_PARAMS,
f"Unknown tool: {tool_name}",
)
try:
import json
tool_call = ToolCall(
id=f"mcp-{req.id}",
name=tool_name,
arguments=json.dumps(arguments),
)
result = self._executor.execute(tool_call)
return MCPResponse(
result={
"content": [
{"type": "text", "text": result.content},
],
"isError": not result.success,
},
id=req.id,
)
except Exception as exc:
return MCPResponse.error_response(
req.id,
INTERNAL_ERROR,
f"Tool execution error: {exc}",
)
__all__ = ["MCPServer"]
src/openjarvis/mcp/transport.py
+122
View File
@@ -0,0 +1,122 @@
"""MCP transport implementations."""
from __future__ import annotations
import json
import subprocess
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, List, Optional
from openjarvis.mcp.protocol import MCPRequest, MCPResponse
if TYPE_CHECKING:
from openjarvis.mcp.server import MCPServer
class MCPTransport(ABC):
"""Abstract transport layer for MCP communication."""
@abstractmethod
def send(self, request: MCPRequest) -> MCPResponse:
"""Send a request and return the response."""
@abstractmethod
def close(self) -> None:
"""Release transport resources."""
class InProcessTransport(MCPTransport):
"""Direct in-process transport for testing.
Routes requests directly to an ``MCPServer`` instance without
serialization overhead.
"""
def __init__(self, server: MCPServer) -> None:
self._server = server
def send(self, request: MCPRequest) -> MCPResponse:
"""Dispatch request directly to the server."""
return self._server.handle(request)
def close(self) -> None:
"""No resources to release."""
class StdioTransport(MCPTransport):
"""JSON-RPC over stdin/stdout subprocess transport.
Launches a subprocess and communicates via JSON lines on
stdin/stdout.
"""
def __init__(self, command: List[str]) -> None:
self._command = command
self._process: Optional[subprocess.Popen[str]] = None
self._start()
def _start(self) -> None:
"""Start the subprocess."""
self._process = subprocess.Popen(
self._command,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
)
def send(self, request: MCPRequest) -> MCPResponse:
"""Write request as JSON line, read response line."""
proc = self._process
if proc is None or proc.stdin is None or proc.stdout is None:
raise RuntimeError("Transport process is not running")
line = request.to_json() + "\n"
proc.stdin.write(line)
proc.stdin.flush()
response_line = proc.stdout.readline()
if not response_line:
raise RuntimeError("No response from subprocess")
return MCPResponse.from_json(response_line.strip())
def close(self) -> None:
"""Terminate the subprocess."""
if self._process is not None:
self._process.terminate()
self._process.wait(timeout=5)
self._process = None
class SSETransport(MCPTransport):
"""JSON-RPC over HTTP with Server-Sent Events.
Sends requests via HTTP POST and reads SSE responses.
"""
def __init__(self, url: str) -> None:
self._url = url
def send(self, request: MCPRequest) -> MCPResponse:
"""Send request via HTTP POST."""
import httpx
response = httpx.post(
self._url,
json=json.loads(request.to_json()),
headers={"Content-Type": "application/json"},
timeout=30.0,
)
response.raise_for_status()
return MCPResponse.from_json(response.text)
def close(self) -> None:
"""No persistent connection to close."""
__all__ = [
"InProcessTransport",
"MCPTransport",
"SSETransport",
"StdioTransport",
]
src/openjarvis/tools/__init__.py
+10
View File
@@ -32,4 +32,14 @@ try:
except ImportError:
pass
try:
import openjarvis.tools.web_search # noqa: F401
except ImportError:
pass
try:
import openjarvis.tools.code_interpreter # noqa: F401
except ImportError:
pass
__all__ = ["BaseTool", "ToolExecutor", "ToolSpec"]
src/openjarvis/tools/code_interpreter.py
+111
View File
@@ -0,0 +1,111 @@
"""Code interpreter tool — safe Python code execution in subprocess."""
from __future__ import annotations
import subprocess
import sys
from typing import Any
from openjarvis.core.registry import ToolRegistry
from openjarvis.core.types import ToolResult
from openjarvis.tools._stubs import BaseTool, ToolSpec
# Dangerous patterns to block
_BLOCKED_PATTERNS = [
"os.system",
"os.popen",
"subprocess.",
"shutil.rmtree",
"os.remove",
"os.unlink",
"os.rmdir",
"__import__",
"eval(",
"exec(",
"compile(",
"open(",
]
@ToolRegistry.register("code_interpreter")
class CodeInterpreterTool(BaseTool):
"""Execute Python code in an isolated subprocess."""
tool_id = "code_interpreter"
def __init__(self, timeout: int = 30, max_output: int = 10000):
self._timeout = timeout
self._max_output = max_output
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="code_interpreter",
description=(
"Execute Python code and return the output."
" Code runs in an isolated subprocess."
),
parameters={
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Python code to execute.",
},
},
"required": ["code"],
},
category="code",
)
def execute(self, **params: Any) -> ToolResult:
code = params.get("code", "")
if not code:
return ToolResult(
tool_name="code_interpreter",
content="No code provided.",
success=False,
)
# Security check
for pattern in _BLOCKED_PATTERNS:
if pattern in code:
return ToolResult(
tool_name="code_interpreter",
content=f"Blocked: code contains prohibited pattern '{pattern}'",
success=False,
)
try:
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
timeout=self._timeout,
)
output = result.stdout
if result.stderr:
output += ("\n" if output else "") + result.stderr
if len(output) > self._max_output:
output = output[: self._max_output] + "\n... (output truncated)"
return ToolResult(
tool_name="code_interpreter",
content=output or "(no output)",
success=result.returncode == 0,
metadata={"returncode": result.returncode},
)
except subprocess.TimeoutExpired:
return ToolResult(
tool_name="code_interpreter",
content=f"Execution timed out after {self._timeout} seconds.",
success=False,
)
except Exception as exc:
return ToolResult(
tool_name="code_interpreter",
content=f"Execution error: {exc}",
success=False,
)
__all__ = ["CodeInterpreterTool"]
src/openjarvis/tools/web_search.py
+95
View File
@@ -0,0 +1,95 @@
"""Web search tool — Tavily API with DuckDuckGo fallback."""
from __future__ import annotations
import os
from typing import Any
from openjarvis.core.registry import ToolRegistry
from openjarvis.core.types import ToolResult
from openjarvis.tools._stubs import BaseTool, ToolSpec
@ToolRegistry.register("web_search")
class WebSearchTool(BaseTool):
"""Search the web via Tavily API."""
tool_id = "web_search"
def __init__(self, api_key: str | None = None, max_results: int = 5):
self._api_key = api_key or os.environ.get("TAVILY_API_KEY")
self._max_results = max_results
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="web_search",
description=(
"Search the web for current information."
" Returns relevant search results."
),
parameters={
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query."},
"max_results": {
"type": "integer",
"description": "Maximum results to return.",
},
},
"required": ["query"],
},
category="search",
metadata={"requires_api_key": "TAVILY_API_KEY"},
)
def execute(self, **params: Any) -> ToolResult:
query = params.get("query", "")
if not query:
return ToolResult(
tool_name="web_search",
content="No query provided.",
success=False,
)
if not self._api_key:
return ToolResult(
tool_name="web_search",
content="No API key configured. Set TAVILY_API_KEY.",
success=False,
)
max_results = params.get("max_results", self._max_results)
try:
from tavily import TavilyClient
client = TavilyClient(api_key=self._api_key)
response = client.search(query, max_results=max_results)
results = response.get("results", [])
formatted = "\n\n".join(
f"**{r.get('title', 'Untitled')}**\n"
f"{r.get('url', '')}\n{r.get('content', '')}"
for r in results
)
return ToolResult(
tool_name="web_search",
content=formatted or "No results found.",
success=True,
metadata={"num_results": len(results)},
)
except ImportError:
return ToolResult(
tool_name="web_search",
content=(
"tavily-python not installed."
" Install with: pip install tavily-python"
),
success=False,
)
except Exception as exc:
return ToolResult(
tool_name="web_search",
content=f"Search error: {exc}",
success=False,
)
__all__ = ["WebSearchTool"]
src/openjarvis/traces/__init__.py
+15
View File
@@ -0,0 +1,15 @@
"""Trace system — full interaction-level recording and analysis.
The trace system captures the complete sequence of steps an agent takes to
handle a query. Unlike telemetry (which records per-inference metrics), traces
record the *decision-making process*: which model was selected, what memory was
retrieved, which tools were called, and the final response.
Traces are the primary input to the learning system.
"""
from openjarvis.traces.analyzer import TraceAnalyzer
from openjarvis.traces.collector import TraceCollector
from openjarvis.traces.store import TraceStore
__all__ = ["TraceAnalyzer", "TraceCollector", "TraceStore"]
src/openjarvis/traces/analyzer.py
+246
View File
@@ -0,0 +1,246 @@
"""TraceAnalyzer — read-only query layer over stored traces.
Provides aggregated statistics that the learning system uses to update
routing policies, tool selection strategies, and memory configuration.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from openjarvis.core.types import StepType, Trace, TraceStep
from openjarvis.traces.store import TraceStore
@dataclass(slots=True)
class RouteStats:
"""Aggregated statistics for a specific routing decision (model+agent)."""
model: str
agent: str
count: int = 0
avg_latency: float = 0.0
avg_tokens: float = 0.0
success_rate: float = 0.0
avg_feedback: Optional[float] = None
@dataclass(slots=True)
class ToolStats:
"""Aggregated statistics for a specific tool."""
tool_name: str
call_count: int = 0
avg_latency: float = 0.0
success_rate: float = 0.0
@dataclass(slots=True)
class TraceSummary:
"""Overall summary statistics across all traces."""
total_traces: int = 0
total_steps: int = 0
avg_steps_per_trace: float = 0.0
avg_latency: float = 0.0
avg_tokens: float = 0.0
success_rate: float = 0.0
step_type_distribution: Dict[str, int] = field(default_factory=dict)
class TraceAnalyzer:
"""Read-only query layer over a :class:`TraceStore`.
Computes aggregated statistics from stored traces, providing the
inputs that the learning system needs to update routing policies.
"""
def __init__(self, store: TraceStore) -> None:
self._store = store
def summary(
self,
*,
since: Optional[float] = None,
until: Optional[float] = None,
) -> TraceSummary:
"""Compute an overall summary of all traces in the time range."""
traces = self._store.list_traces(since=since, until=until, limit=10_000)
if not traces:
return TraceSummary()
total_steps = sum(len(t.steps) for t in traces)
evaluated = [t for t in traces if t.outcome is not None]
successes = [t for t in evaluated if t.outcome == "success"]
step_dist: Dict[str, int] = {}
for t in traces:
for s in t.steps:
key = _step_type_str(s)
step_dist[key] = step_dist.get(key, 0) + 1
return TraceSummary(
total_traces=len(traces),
total_steps=total_steps,
avg_steps_per_trace=total_steps / len(traces) if traces else 0.0,
avg_latency=_avg([t.total_latency_seconds for t in traces]),
avg_tokens=_avg([float(t.total_tokens) for t in traces]),
success_rate=len(successes) / len(evaluated) if evaluated else 0.0,
step_type_distribution=step_dist,
)
def per_route_stats(
self,
*,
since: Optional[float] = None,
until: Optional[float] = None,
) -> List[RouteStats]:
"""Compute stats grouped by (model, agent) routing decisions."""
traces = self._store.list_traces(since=since, until=until, limit=10_000)
groups: Dict[tuple, list[Trace]] = {}
for t in traces:
key = (t.model, t.agent)
groups.setdefault(key, []).append(t)
results = []
for (model, agent), group in sorted(groups.items()):
evaluated = [t for t in group if t.outcome is not None]
successes = [t for t in evaluated if t.outcome == "success"]
feedbacks = [t.feedback for t in group if t.feedback is not None]
results.append(
RouteStats(
model=model,
agent=agent,
count=len(group),
avg_latency=_avg([t.total_latency_seconds for t in group]),
avg_tokens=_avg([float(t.total_tokens) for t in group]),
success_rate=len(successes) / len(evaluated) if evaluated else 0.0,
avg_feedback=_avg(feedbacks) if feedbacks else None,
)
)
return results
def per_tool_stats(
self,
*,
since: Optional[float] = None,
until: Optional[float] = None,
) -> List[ToolStats]:
"""Compute stats grouped by tool name."""
traces = self._store.list_traces(since=since, until=until, limit=10_000)
tools: Dict[str, Dict[str, Any]] = {}
for t in traces:
for s in t.steps:
stype = _step_type_str(s)
if stype != "tool_call":
continue
name = s.input.get("tool", "unknown")
if name not in tools:
tools[name] = {"count": 0, "latencies": [], "successes": 0}
tools[name]["count"] += 1
tools[name]["latencies"].append(s.duration_seconds)
if s.output.get("success"):
tools[name]["successes"] += 1
return [
ToolStats(
tool_name=name,
call_count=data["count"],
avg_latency=_avg(data["latencies"]),
success_rate=(
data["successes"] / data["count"]
if data["count"] else 0.0
),
)
for name, data in sorted(tools.items())
]
def traces_for_query_type(
self,
*,
has_code: bool = False,
min_length: Optional[int] = None,
max_length: Optional[int] = None,
since: Optional[float] = None,
until: Optional[float] = None,
) -> List[Trace]:
"""Retrieve traces matching query characteristics.
Useful for the learning system to find traces similar to a new
query and learn which routing decisions worked best.
"""
traces = self._store.list_traces(since=since, until=until, limit=10_000)
filtered = []
for t in traces:
if has_code and not _looks_like_code(t.query):
continue
if min_length is not None and len(t.query) < min_length:
continue
if max_length is not None and len(t.query) > max_length:
continue
filtered.append(t)
return filtered
def export_traces(
self,
*,
since: Optional[float] = None,
until: Optional[float] = None,
limit: int = 1000,
) -> List[Dict[str, Any]]:
"""Export traces as plain dicts (for JSON serialization)."""
traces = self._store.list_traces(since=since, until=until, limit=limit)
return [_trace_to_dict(t) for t in traces]
# -- helpers -------------------------------------------------------------------
def _avg(values: list[float]) -> float:
return sum(values) / len(values) if values else 0.0
def _step_type_str(step: TraceStep) -> str:
st = step.step_type
return st.value if isinstance(st, StepType) else st
def _looks_like_code(text: str) -> bool:
indicators = [
"def ", "class ", "import ",
"function ", "const ", "var ", "```",
]
return any(ind in text for ind in indicators)
def _trace_to_dict(trace: Trace) -> Dict[str, Any]:
return {
"trace_id": trace.trace_id,
"query": trace.query,
"agent": trace.agent,
"model": trace.model,
"engine": trace.engine,
"result": trace.result,
"outcome": trace.outcome,
"feedback": trace.feedback,
"started_at": trace.started_at,
"ended_at": trace.ended_at,
"total_tokens": trace.total_tokens,
"total_latency_seconds": trace.total_latency_seconds,
"metadata": trace.metadata,
"steps": [
{
"step_type": _step_type_str(s),
"timestamp": s.timestamp,
"duration_seconds": s.duration_seconds,
"input": s.input,
"output": s.output,
"metadata": s.metadata,
}
for s in trace.steps
],
}
__all__ = ["RouteStats", "ToolStats", "TraceAnalyzer", "TraceSummary"]
src/openjarvis/traces/collector.py
+179
View File
@@ -0,0 +1,179 @@
"""TraceCollector — wraps any BaseAgent to record interaction traces."""
from __future__ import annotations
import time
from typing import Any, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.types import StepType, Trace, TraceStep
from openjarvis.traces.store import TraceStore
class TraceCollector:
"""Wraps a ``BaseAgent`` and records a :class:`Trace` for every ``run()``.
The collector subscribes to the ``EventBus`` to capture inference, tool,
and memory events emitted during agent execution, converting them into
``TraceStep`` objects. When the agent finishes, the complete ``Trace``
is persisted to the ``TraceStore`` and published on the bus.
Usage::
agent = OrchestratorAgent(engine, model, tools=tools, bus=bus)
collector = TraceCollector(agent, store=trace_store, bus=bus)
result = collector.run("What is 2+2?")
# Trace is automatically saved to trace_store
"""
def __init__(
self,
agent: BaseAgent,
*,
store: Optional[TraceStore] = None,
bus: Optional[EventBus] = None,
) -> None:
self._agent = agent
self._store = store
self._bus = bus
self._current_steps: list[TraceStep] = []
self._current_model: str = ""
self._current_engine: str = ""
def run(
self,
input: str,
context: Optional[AgentContext] = None,
**kwargs: Any,
) -> AgentResult:
"""Execute the wrapped agent and record a trace."""
self._current_steps = []
self._current_model = ""
self._current_engine = ""
# Subscribe to events for trace collection
unsubs = self._subscribe()
started_at = time.time()
try:
result = self._agent.run(input, context=context, **kwargs)
finally:
self._unsubscribe(unsubs)
ended_at = time.time()
# Add final respond step
self._current_steps.append(
TraceStep(
step_type=StepType.RESPOND,
timestamp=ended_at,
duration_seconds=0.0,
output={"content": result.content, "turns": result.turns},
)
)
# Build and persist the trace
trace = Trace(
query=input,
agent=getattr(self._agent, "agent_id", "unknown"),
model=self._current_model,
engine=self._current_engine,
steps=list(self._current_steps),
result=result.content,
started_at=started_at,
ended_at=ended_at,
)
# Recompute totals from steps
for step in trace.steps:
trace.total_latency_seconds += step.duration_seconds
trace.total_tokens += step.output.get("tokens", 0)
if self._store is not None:
self._store.save(trace)
if self._bus is not None:
self._bus.publish(EventType.TRACE_COMPLETE, {"trace": trace})
return result
@property
def last_trace(self) -> Optional[Trace]:
"""Return the trace from the most recent ``run()``, if available."""
if not self._current_steps:
return None
# Reconstruct from saved steps (steps cleared on next run)
return None # Use TraceStore.get() for retrieval after run
# -- event handlers --------------------------------------------------------
def _subscribe(self) -> list[tuple]:
if self._bus is None:
return []
handlers = [
(EventType.INFERENCE_START, self._on_inference_start),
(EventType.INFERENCE_END, self._on_inference_end),
(EventType.TOOL_CALL_START, self._on_tool_start),
(EventType.TOOL_CALL_END, self._on_tool_end),
(EventType.MEMORY_RETRIEVE, self._on_memory_retrieve),
]
for evt_type, handler in handlers:
self._bus.subscribe(evt_type, handler)
return handlers
def _unsubscribe(self, handlers: list[tuple]) -> None:
if self._bus is None:
return
for evt_type, handler in handlers:
self._bus.unsubscribe(evt_type, handler)
def _on_inference_start(self, event: Any) -> None:
self._current_model = event.data.get("model", self._current_model)
self._current_engine = event.data.get("engine", self._current_engine)
self._inference_start_time = event.timestamp
def _on_inference_end(self, event: Any) -> None:
start = getattr(self, "_inference_start_time", event.timestamp)
self._current_steps.append(
TraceStep(
step_type=StepType.GENERATE,
timestamp=start,
duration_seconds=event.timestamp - start,
input={"model": self._current_model},
output={
"tokens": event.data.get("total_tokens", 0),
},
metadata={"engine": self._current_engine},
)
)
def _on_tool_start(self, event: Any) -> None:
self._tool_start_time = event.timestamp
def _on_tool_end(self, event: Any) -> None:
start = getattr(self, "_tool_start_time", event.timestamp)
self._current_steps.append(
TraceStep(
step_type=StepType.TOOL_CALL,
timestamp=start,
duration_seconds=event.data.get("latency", event.timestamp - start),
input={"tool": event.data.get("tool", "")},
output={"success": event.data.get("success", False)},
)
)
def _on_memory_retrieve(self, event: Any) -> None:
self._current_steps.append(
TraceStep(
step_type=StepType.RETRIEVE,
timestamp=event.timestamp,
duration_seconds=event.data.get("latency", 0.0),
input={"query": event.data.get("query", "")},
output={
"num_results": event.data.get("num_results", 0),
},
)
)
__all__ = ["TraceCollector"]
src/openjarvis/traces/store.py
+214
View File
@@ -0,0 +1,214 @@
"""SQLite-backed trace storage."""
from __future__ import annotations
import json
import sqlite3
from pathlib import Path
from typing import Any, List, Optional
from openjarvis.core.events import Event, EventBus, EventType
from openjarvis.core.types import StepType, Trace, TraceStep
_CREATE_TRACES = """\
CREATE TABLE IF NOT EXISTS traces (
id INTEGER PRIMARY KEY AUTOINCREMENT,
trace_id TEXT NOT NULL UNIQUE,
query TEXT NOT NULL DEFAULT '',
agent TEXT NOT NULL DEFAULT '',
model TEXT NOT NULL DEFAULT '',
engine TEXT NOT NULL DEFAULT '',
result TEXT NOT NULL DEFAULT '',
outcome TEXT,
feedback REAL,
started_at REAL NOT NULL DEFAULT 0.0,
ended_at REAL NOT NULL DEFAULT 0.0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_latency_seconds REAL NOT NULL DEFAULT 0.0,
metadata TEXT NOT NULL DEFAULT '{}'
);
"""
_CREATE_STEPS = """\
CREATE TABLE IF NOT EXISTS trace_steps (
id INTEGER PRIMARY KEY AUTOINCREMENT,
trace_id TEXT NOT NULL,
step_index INTEGER NOT NULL,
step_type TEXT NOT NULL,
timestamp REAL NOT NULL DEFAULT 0.0,
duration_seconds REAL NOT NULL DEFAULT 0.0,
input TEXT NOT NULL DEFAULT '{}',
output TEXT NOT NULL DEFAULT '{}',
metadata TEXT NOT NULL DEFAULT '{}',
FOREIGN KEY (trace_id) REFERENCES traces(trace_id)
);
"""
_INSERT_TRACE = """\
INSERT INTO traces (
trace_id, query, agent, model, engine, result,
outcome, feedback, started_at, ended_at,
total_tokens, total_latency_seconds, metadata
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"""
_INSERT_STEP = """\
INSERT INTO trace_steps (
trace_id, step_index, step_type, timestamp,
duration_seconds, input, output, metadata
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
"""
class TraceStore:
"""Append-only SQLite store for interaction traces."""
def __init__(self, db_path: 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_TRACES)
self._conn.execute(_CREATE_STEPS)
self._conn.commit()
def save(self, trace: Trace) -> None:
"""Persist a complete trace with all its steps."""
self._conn.execute(
_INSERT_TRACE,
(
trace.trace_id,
trace.query,
trace.agent,
trace.model,
trace.engine,
trace.result,
trace.outcome,
trace.feedback,
trace.started_at,
trace.ended_at,
trace.total_tokens,
trace.total_latency_seconds,
json.dumps(trace.metadata),
),
)
for idx, step in enumerate(trace.steps):
self._conn.execute(
_INSERT_STEP,
(
trace.trace_id,
idx,
step.step_type.value
if isinstance(step.step_type, StepType)
else step.step_type,
step.timestamp,
step.duration_seconds,
json.dumps(step.input),
json.dumps(step.output),
json.dumps(step.metadata),
),
)
self._conn.commit()
def get(self, trace_id: str) -> Optional[Trace]:
"""Retrieve a trace by id, or ``None`` if not found."""
row = self._conn.execute(
"SELECT * FROM traces WHERE trace_id = ?", (trace_id,)
).fetchone()
if row is None:
return None
return self._row_to_trace(row)
def list_traces(
self,
*,
agent: Optional[str] = None,
model: Optional[str] = None,
outcome: Optional[str] = None,
since: Optional[float] = None,
until: Optional[float] = None,
limit: int = 100,
) -> List[Trace]:
"""Query traces with optional filters."""
clauses: List[str] = []
params: List[Any] = []
if agent is not None:
clauses.append("agent = ?")
params.append(agent)
if model is not None:
clauses.append("model = ?")
params.append(model)
if outcome is not None:
clauses.append("outcome = ?")
params.append(outcome)
if since is not None:
clauses.append("started_at >= ?")
params.append(since)
if until is not None:
clauses.append("started_at <= ?")
params.append(until)
where = " AND ".join(clauses) if clauses else "1=1"
sql = f"SELECT * FROM traces WHERE {where} ORDER BY started_at DESC LIMIT ?"
params.append(limit)
rows = self._conn.execute(sql, params).fetchall()
return [self._row_to_trace(r) for r in rows]
def count(self) -> int:
"""Return the total number of stored traces."""
row = self._conn.execute("SELECT COUNT(*) FROM traces").fetchone()
return row[0] if row else 0
def subscribe_to_bus(self, bus: EventBus) -> None:
"""Subscribe to ``TRACE_COMPLETE`` events on *bus*."""
bus.subscribe(EventType.TRACE_COMPLETE, self._on_event)
def _on_event(self, event: Event) -> None:
trace = event.data.get("trace")
if isinstance(trace, Trace):
self.save(trace)
def close(self) -> None:
"""Close the underlying SQLite connection."""
self._conn.close()
# -- internal helpers ------------------------------------------------------
def _row_to_trace(self, row: tuple) -> Trace:
"""Convert a traces table row + its steps into a Trace object."""
trace_id = row[1]
step_rows = self._conn.execute(
"SELECT * FROM trace_steps WHERE trace_id = ? ORDER BY step_index",
(trace_id,),
).fetchall()
steps = [
TraceStep(
step_type=StepType(sr[3]),
timestamp=sr[4],
duration_seconds=sr[5],
input=json.loads(sr[6]),
output=json.loads(sr[7]),
metadata=json.loads(sr[8]),
)
for sr in step_rows
]
return Trace(
trace_id=trace_id,
query=row[2],
agent=row[3],
model=row[4],
engine=row[5],
result=row[6],
outcome=row[7],
feedback=row[8],
started_at=row[9],
ended_at=row[10],
total_tokens=row[11],
total_latency_seconds=row[12],
metadata=json.loads(row[13]),
steps=steps,
)
def _fetchall(self, sql: str = "SELECT * FROM traces") -> list:
return self._conn.execute(sql).fetchall()
__all__ = ["TraceStore"]
tests/agents/test_agent_routing_matrix.py
+183
View File
@@ -0,0 +1,183 @@
"""Cross-product tests: agent x engine x model."""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
from openjarvis.agents._stubs import AgentResult
from openjarvis.agents.openhands import OpenHandsAgent
from openjarvis.agents.orchestrator import OrchestratorAgent
from openjarvis.agents.react import ReActAgent
from openjarvis.agents.simple import SimpleAgent
from openjarvis.core.events import EventBus, EventType
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_mock_engine(response="Hello!"):
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = {
"content": response,
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
"model": "test-model",
"finish_reason": "stop",
}
return engine
# ReAct needs a structured response to produce a result on first turn
_REACT_RESPONSE = "Thought: Simple query.\nFinal Answer: Hello!"
_OPENHANDS_RESPONSE = "Hello!"
_SIMPLE_RESPONSE = "Hello!"
_ORCHESTRATOR_RESPONSE = "Hello!"
AGENT_FACTORIES = {
"simple": lambda e, m, bus: SimpleAgent(e, m, bus=bus),
"orchestrator": lambda e, m, bus: OrchestratorAgent(e, m, bus=bus),
"react": lambda e, m, bus: ReActAgent(e, m, bus=bus),
"openhands": lambda e, m, bus: OpenHandsAgent(e, m, bus=bus),
}
AGENT_RESPONSES = {
"simple": _SIMPLE_RESPONSE,
"orchestrator": _ORCHESTRATOR_RESPONSE,
"react": _REACT_RESPONSE,
"openhands": _OPENHANDS_RESPONSE,
}
# ---------------------------------------------------------------------------
# Common agent tests (parametrized)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "react", "openhands"])
class TestAgentCommon:
def test_runs_with_mock_engine(self, agent_key):
"""Each agent type can run with a mock engine."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
bus = EventBus(record_history=True)
agent = AGENT_FACTORIES[agent_key](engine, "test-model", bus)
result = agent.run("Hello")
assert isinstance(result, AgentResult)
def test_returns_valid_result(self, agent_key):
"""Result has correct structure."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
bus = EventBus(record_history=True)
agent = AGENT_FACTORIES[agent_key](engine, "test-model", bus)
result = agent.run("Hello")
assert hasattr(result, "content")
assert hasattr(result, "turns")
assert hasattr(result, "tool_results")
assert hasattr(result, "metadata")
def test_returns_nonempty_content(self, agent_key):
"""Result has non-empty content."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
bus = EventBus(record_history=True)
agent = AGENT_FACTORIES[agent_key](engine, "test-model", bus)
result = agent.run("Hello")
assert result.content != ""
def test_emits_events(self, agent_key):
"""Each agent emits at least AGENT_TURN_START and inference events."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
bus = EventBus(record_history=True)
agent = AGENT_FACTORIES[agent_key](engine, "test-model", bus)
agent.run("Hello")
event_types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in event_types
def test_handles_empty_input(self, agent_key):
"""Agent handles empty string input without crashing."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
bus = EventBus(record_history=True)
agent = AGENT_FACTORIES[agent_key](engine, "test-model", bus)
result = agent.run("")
assert isinstance(result, AgentResult)
# ---------------------------------------------------------------------------
# Agent identity tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"agent_key,expected_id",
[
("simple", "simple"),
("orchestrator", "orchestrator"),
("react", "react"),
("openhands", "openhands"),
],
)
def test_agent_id(agent_key, expected_id):
engine = _make_mock_engine()
agent = AGENT_FACTORIES[agent_key](engine, "test-model", None)
assert agent.agent_id == expected_id
# ---------------------------------------------------------------------------
# Model parametrization tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "react", "openhands"])
@pytest.mark.parametrize("model", ["qwen3:8b", "llama3:70b", "gpt-oss:120b"])
def test_model_passthrough(agent_key, model):
"""Each agent passes the model name through to the engine."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
agent = AGENT_FACTORIES[agent_key](engine, model, None)
agent.run("Hello")
call_kwargs = engine.generate.call_args[1]
assert call_kwargs["model"] == model
# ---------------------------------------------------------------------------
# No-bus tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "react", "openhands"])
def test_no_bus(agent_key):
"""All agents work without an event bus."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
agent = AGENT_FACTORIES[agent_key](engine, "test-model", None)
result = agent.run("Hello")
assert isinstance(result, AgentResult)
assert result.content != ""
# ---------------------------------------------------------------------------
# Turns tracking
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "react", "openhands"])
def test_single_turn_count(agent_key):
"""All agents report at least 1 turn for a simple query."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
agent = AGENT_FACTORIES[agent_key](engine, "test-model", None)
result = agent.run("Hello")
assert result.turns >= 1
# ---------------------------------------------------------------------------
# Tool results are empty for no-tool queries
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "react", "openhands"])
def test_no_tool_results_for_simple_query(agent_key):
"""When no tools are used, tool_results should be empty."""
engine = _make_mock_engine(AGENT_RESPONSES[agent_key])
agent = AGENT_FACTORIES[agent_key](engine, "test-model", None)
result = agent.run("Hello")
assert result.tool_results == []
tests/agents/test_openhands.py
+356
View File
@@ -0,0 +1,356 @@
"""Tests for OpenHands agent."""
from __future__ import annotations
from unittest.mock import MagicMock
from openjarvis.agents._stubs import AgentContext
from openjarvis.agents.openhands import OpenHandsAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Conversation, Message, Role, ToolResult
from openjarvis.tools._stubs import BaseTool, ToolSpec
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _CodeInterpreterStub(BaseTool):
"""Stub code_interpreter tool for testing."""
tool_id = "code_interpreter"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="code_interpreter",
description="Execute Python code.",
parameters={
"type": "object",
"properties": {"code": {"type": "string"}},
"required": ["code"],
},
)
def execute(self, **params) -> ToolResult:
code = params.get("code", "")
# Simple simulation: if it contains print(), capture the content
if "print(" in code:
import re
match = re.search(r"print\((.+?)\)", code)
if match:
try:
val = eval(match.group(1)) # noqa: S307
return ToolResult(
tool_name="code_interpreter",
content=str(val),
success=True,
)
except Exception:
pass
return ToolResult(
tool_name="code_interpreter",
content=f"Executed: {code}",
success=True,
)
class _CalculatorStub(BaseTool):
tool_id = "calculator"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="calculator",
description="Math calculator.",
parameters={
"type": "object",
"properties": {"expression": {"type": "string"}},
"required": ["expression"],
},
)
def execute(self, **params) -> ToolResult:
expr = params.get("expression", "0")
try:
val = eval(expr) # noqa: S307
except Exception as e:
return ToolResult(tool_name="calculator", content=str(e), success=False)
return ToolResult(tool_name="calculator", content=str(val), success=True)
def _engine_response(content, **extra):
base = {
"content": content,
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
"model": "test-model",
"finish_reason": "stop",
}
base.update(extra)
return base
# ---------------------------------------------------------------------------
# Registration tests
# ---------------------------------------------------------------------------
class TestOpenHandsRegistration:
def test_registration(self):
AgentRegistry.register_value("openhands", OpenHandsAgent)
assert AgentRegistry.contains("openhands")
def test_agent_id(self):
engine = MagicMock()
engine.engine_id = "mock"
agent = OpenHandsAgent(engine, "test-model")
assert agent.agent_id == "openhands"
# ---------------------------------------------------------------------------
# Agent execution tests
# ---------------------------------------------------------------------------
class TestOpenHandsAgent:
def test_simple_response(self):
"""No code -> direct answer."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("The answer is 42.")
bus = EventBus(record_history=True)
agent = OpenHandsAgent(engine, "test-model", bus=bus)
result = agent.run("What is the meaning of life?")
assert result.content == "The answer is 42."
assert result.turns == 1
assert result.tool_results == []
def test_code_generation_execution(self):
"""Turn 1: returns code block -> code_interpreter executed. Turn 2: final."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
"Let me calculate:\n```python\nprint(2+2)\n```"
),
_engine_response("The result is 4."),
]
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub()],
)
result = agent.run("What is 2+2?")
assert result.content == "The result is 4."
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].tool_name == "code_interpreter"
def test_multi_step_code(self):
"""Multiple code blocks across turns."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response("Step 1:\n```python\nprint(1+1)\n```"),
_engine_response("Step 2:\n```python\nprint(3*3)\n```"),
_engine_response("First was 2, second was 9."),
]
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub()],
)
result = agent.run("Two calculations")
assert result.turns == 3
assert len(result.tool_results) == 2
assert result.content == "First was 2, second was 9."
def test_max_turns(self):
"""Engine keeps generating code -> hits max_turns."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"More code:\n```python\nprint('hello')\n```"
)
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub()],
max_turns=3,
)
result = agent.run("Keep coding")
assert result.turns == 3
assert result.metadata.get("max_turns_exceeded") is True
def test_event_bus_emissions(self):
"""Verify AGENT_TURN_START, INFERENCE_START/END, AGENT_TURN_END events."""
bus = EventBus(record_history=True)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Direct answer.")
agent = OpenHandsAgent(engine, "test-model", bus=bus)
agent.run("Hello")
event_types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in event_types
assert EventType.INFERENCE_START in event_types
assert EventType.INFERENCE_END in event_types
assert EventType.AGENT_TURN_END in event_types
def test_event_bus_tool_events(self):
"""Tool execution should trigger TOOL_CALL_START/END events."""
bus = EventBus(record_history=True)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response("Code:\n```python\nprint(1)\n```"),
_engine_response("Done."),
]
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub()],
bus=bus,
)
agent.run("Run code")
event_types = [e.event_type for e in bus.history]
assert EventType.TOOL_CALL_START in event_types
assert EventType.TOOL_CALL_END in event_types
def test_context_passing(self):
"""Pass AgentContext with conversation history."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Hello!")
conv = Conversation()
conv.add(Message(role=Role.USER, content="Previous"))
conv.add(Message(role=Role.ASSISTANT, content="Previous response"))
ctx = AgentContext(conversation=conv)
agent = OpenHandsAgent(engine, "test-model")
agent.run("Hello", context=ctx)
call_args = engine.generate.call_args
messages = call_args[0][0]
# System prompt + 2 context + user input
assert len(messages) == 4
assert messages[0].role == Role.SYSTEM
assert messages[1].content == "Previous"
assert messages[3].content == "Hello"
def test_tool_fallback(self):
"""Non-code tool use via Action: syntax."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Action: calculator\nAction Input: {"expression": "7*6"}'
),
_engine_response("The answer is 42."),
]
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CalculatorStub()],
)
result = agent.run("What is 7 times 6?")
assert result.content == "The answer is 42."
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].tool_name == "calculator"
assert result.tool_results[0].content == "42"
def test_no_code_interpreter_tool(self):
"""Agent has code but no code_interpreter -> tool not found."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response("```python\nprint(1)\n```"),
_engine_response("Could not run code."),
]
# No tools at all
agent = OpenHandsAgent(engine, "test-model")
result = agent.run("Run code")
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].success is False
assert "Unknown tool" in result.tool_results[0].content
def test_no_bus_works(self):
"""Agent runs correctly without an event bus."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Works!")
agent = OpenHandsAgent(engine, "test-model")
result = agent.run("Hello")
assert result.content == "Works!"
def test_system_prompt_includes_tool_names(self):
"""System prompt should list available tool names."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Ok")
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub(), _CalculatorStub()],
)
agent.run("Hello")
call_args = engine.generate.call_args
messages = call_args[0][0]
system_msg = messages[0]
assert "code_interpreter" in system_msg.content
assert "calculator" in system_msg.content
def test_max_turns_content_preserved(self):
"""When max turns exceeded, last content should be preserved."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Still working:\n```python\nx = 1\n```"
)
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub()],
max_turns=2,
)
result = agent.run("Loop")
assert result.metadata.get("max_turns_exceeded") is True
# Should have the last content, not the fallback message
assert "Still working" in result.content
def test_observation_appended_to_messages(self):
"""Code output is sent back to the engine as observation."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response("```python\nprint(42)\n```"),
_engine_response("Got 42."),
]
agent = OpenHandsAgent(
engine, "test-model",
tools=[_CodeInterpreterStub()],
)
agent.run("Print 42")
second_call = engine.generate.call_args_list[1]
messages = second_call[0][0]
last_msg = messages[-1]
assert last_msg.role == Role.USER
assert "Output:" in last_msg.content
def test_event_data_agent_turn_start(self):
"""AGENT_TURN_START event data should include agent id and input."""
bus = EventBus(record_history=True)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Hi")
agent = OpenHandsAgent(engine, "test-model", bus=bus)
agent.run("test input")
start_events = [
e for e in bus.history
if e.event_type == EventType.AGENT_TURN_START
]
assert len(start_events) == 1
assert start_events[0].data["agent"] == "openhands"
assert start_events[0].data["input"] == "test input"
def test_empty_input(self):
"""Agent handles empty input."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Empty input received.")
agent = OpenHandsAgent(engine, "test-model")
result = agent.run("")
assert result.content == "Empty input received."
tests/agents/test_react.py
+476
View File
@@ -0,0 +1,476 @@
"""Tests for ReAct agent."""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
from openjarvis.agents._stubs import AgentContext
from openjarvis.agents.react import ReActAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Conversation, Message, Role, ToolResult
from openjarvis.tools._stubs import BaseTool, ToolSpec
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _CalculatorStub(BaseTool):
tool_id = "calculator"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="calculator",
description="Math calculator.",
parameters={
"type": "object",
"properties": {"expression": {"type": "string"}},
"required": ["expression"],
},
)
def execute(self, **params) -> ToolResult:
expr = params.get("expression", "0")
try:
val = eval(expr) # noqa: S307
except Exception as e:
return ToolResult(tool_name="calculator", content=str(e), success=False)
return ToolResult(tool_name="calculator", content=str(val), success=True)
class _ThinkStub(BaseTool):
tool_id = "think"
@property
def spec(self) -> ToolSpec:
return ToolSpec(
name="think",
description="Thinking tool.",
parameters={
"type": "object",
"properties": {"thought": {"type": "string"}},
},
)
def execute(self, **params) -> ToolResult:
return ToolResult(
tool_name="think",
content=params.get("thought", ""),
success=True,
)
def _engine_response(content, **extra):
"""Helper to build an engine response dict."""
base = {
"content": content,
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
"model": "test-model",
"finish_reason": "stop",
}
base.update(extra)
return base
# ---------------------------------------------------------------------------
# Registration tests
# ---------------------------------------------------------------------------
class TestReActRegistration:
def test_registration(self):
AgentRegistry.register_value("react", ReActAgent)
assert AgentRegistry.contains("react")
def test_agent_id(self):
engine = MagicMock()
engine.engine_id = "mock"
agent = ReActAgent(engine, "test-model")
assert agent.agent_id == "react"
# ---------------------------------------------------------------------------
# Parsing tests
# ---------------------------------------------------------------------------
class TestReActParsing:
def _parser(self):
engine = MagicMock()
engine.engine_id = "mock"
agent = ReActAgent(engine, "test-model")
return agent._parse_response
def test_parse_thought_action(self):
parse = self._parser()
text = (
'Thought: I need to calculate 2+2.\n'
'Action: calculator\n'
'Action Input: {"expression": "2+2"}'
)
result = parse(text)
assert result["thought"] == "I need to calculate 2+2."
assert result["action"] == "calculator"
assert "expression" in result["action_input"]
assert result["final_answer"] == ""
def test_parse_final_answer(self):
parse = self._parser()
text = "Thought: I know the answer.\nFinal Answer: 42"
result = parse(text)
assert result["thought"] == "I know the answer."
assert result["final_answer"] == "42"
assert result["action"] == ""
def test_parse_no_structure(self):
parse = self._parser()
text = "Just a plain response with no structure."
result = parse(text)
assert result["thought"] == ""
assert result["action"] == ""
assert result["final_answer"] == ""
def test_parse_multiline_thought(self):
parse = self._parser()
text = (
"Thought: First I need to think.\n"
"Then consider options.\n"
"Final Answer: done"
)
result = parse(text)
assert result["final_answer"] == "done"
def test_parse_action_without_input(self):
parse = self._parser()
text = "Thought: Let me check.\nAction: calculator"
result = parse(text)
assert result["action"] == "calculator"
assert result["action_input"] == ""
# ---------------------------------------------------------------------------
# Agent execution tests
# ---------------------------------------------------------------------------
class TestReActAgent:
def test_simple_no_tool_response(self):
"""Engine returns Final Answer on first call."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Simple greeting.\nFinal Answer: Hello!"
)
bus = EventBus(record_history=True)
agent = ReActAgent(engine, "test-model", bus=bus)
result = agent.run("Hello")
assert result.content == "Hello!"
assert result.turns == 1
assert result.tool_results == []
def test_thought_action_observation(self):
"""Turn 1: action, Turn 2: final answer."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: I need to calculate.\n'
'Action: calculator\n'
'Action Input: {"expression": "2+2"}'
),
_engine_response(
"Thought: The result is 4.\nFinal Answer: 4"
),
]
bus = EventBus(record_history=True)
agent = ReActAgent(engine, "test-model", tools=[_CalculatorStub()], bus=bus)
result = agent.run("What is 2+2?")
assert result.content == "4"
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].tool_name == "calculator"
assert result.tool_results[0].content == "4"
def test_calculator_tool_use(self):
"""Verify calculator tool execution produces correct result."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: Calculate.\nAction: calculator\n'
'Action Input: {"expression": "3*7"}'
),
_engine_response("Thought: Done.\nFinal Answer: 21"),
]
agent = ReActAgent(engine, "test-model", tools=[_CalculatorStub()])
result = agent.run("3 times 7")
assert result.tool_results[0].content == "21"
assert result.tool_results[0].success is True
def test_multi_tool_turns(self):
"""Three tool calls before final answer."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: Step 1.\nAction: calculator\n'
'Action Input: {"expression": "1+1"}'
),
_engine_response(
'Thought: Step 2.\nAction: calculator\n'
'Action Input: {"expression": "2+2"}'
),
_engine_response(
'Thought: Step 3.\nAction: think\n'
'Action Input: {"thought": "combining results"}'
),
_engine_response("Thought: All done.\nFinal Answer: Complete."),
]
agent = ReActAgent(
engine, "test-model",
tools=[_CalculatorStub(), _ThinkStub()],
)
result = agent.run("Multi step")
assert result.turns == 4
assert len(result.tool_results) == 3
assert result.content == "Complete."
def test_max_turns_exceeded(self):
"""Engine always returns actions -- hits max_turns."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
'Thought: Keep going.\nAction: calculator\n'
'Action Input: {"expression": "1+1"}'
)
agent = ReActAgent(
engine, "test-model",
tools=[_CalculatorStub()],
max_turns=3,
)
result = agent.run("Loop forever")
assert result.turns == 3
assert result.metadata.get("max_turns_exceeded") is True
assert result.content == "Maximum turns reached without a final answer."
def test_unknown_tool_error(self):
"""Action references nonexistent tool -- ToolResult with success=False."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: Use a tool.\nAction: nonexistent\n'
'Action Input: {}'
),
_engine_response(
"Thought: Error occurred.\n"
"Final Answer: Could not run tool."
),
]
agent = ReActAgent(engine, "test-model", tools=[_CalculatorStub()])
result = agent.run("Do something")
assert len(result.tool_results) == 1
assert result.tool_results[0].success is False
assert "Unknown tool" in result.tool_results[0].content
def test_event_bus_emissions(self):
"""Verify AGENT_TURN_START, INFERENCE_START/END, AGENT_TURN_END events."""
bus = EventBus(record_history=True)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Quick.\nFinal Answer: Done."
)
agent = ReActAgent(engine, "test-model", bus=bus)
agent.run("Hello")
event_types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in event_types
assert EventType.INFERENCE_START in event_types
assert EventType.INFERENCE_END in event_types
assert EventType.AGENT_TURN_END in event_types
def test_event_bus_tool_events(self):
"""Tool call should trigger TOOL_CALL_START and TOOL_CALL_END events."""
bus = EventBus(record_history=True)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: Calc.\nAction: calculator\n'
'Action Input: {"expression": "1+1"}'
),
_engine_response("Thought: Done.\nFinal Answer: 2"),
]
agent = ReActAgent(engine, "test-model", tools=[_CalculatorStub()], bus=bus)
agent.run("Calc")
event_types = [e.event_type for e in bus.history]
assert EventType.TOOL_CALL_START in event_types
assert EventType.TOOL_CALL_END in event_types
def test_context_passing(self):
"""Pass AgentContext with conversation history."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Simple.\nFinal Answer: Hi!"
)
conv = Conversation()
conv.add(Message(role=Role.USER, content="Previous message"))
conv.add(Message(role=Role.ASSISTANT, content="Previous response"))
ctx = AgentContext(conversation=conv)
agent = ReActAgent(engine, "test-model")
agent.run("Hello", context=ctx)
call_args = engine.generate.call_args
messages = call_args[0][0]
# System prompt + 2 context messages + user input
assert len(messages) == 4
assert messages[0].role == Role.SYSTEM
assert messages[1].role == Role.USER
assert messages[1].content == "Previous message"
assert messages[3].role == Role.USER
assert messages[3].content == "Hello"
def test_with_think_tool(self):
"""Use think tool for internal reasoning."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: Let me reason.\nAction: think\n'
'Action Input: {"thought": "The user wants a greeting"}'
),
_engine_response("Thought: Now I know.\nFinal Answer: Greetings!"),
]
agent = ReActAgent(engine, "test-model", tools=[_ThinkStub()])
result = agent.run("Say hi")
assert result.content == "Greetings!"
assert result.tool_results[0].tool_name == "think"
assert result.tool_results[0].content == "The user wants a greeting"
def test_no_bus_works(self):
"""Agent runs correctly without an event bus."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Easy.\nFinal Answer: Works!"
)
agent = ReActAgent(engine, "test-model")
result = agent.run("Hello")
assert result.content == "Works!"
def test_plain_response_no_structure(self):
"""If engine returns no ReAct structure, treat as final answer."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response("Just a plain answer.")
agent = ReActAgent(engine, "test-model")
result = agent.run("Hello")
assert result.content == "Just a plain answer."
assert result.turns == 1
def test_observation_appended_to_messages(self):
"""Observation from tool result is sent back to the engine."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.side_effect = [
_engine_response(
'Thought: Calc.\nAction: calculator\n'
'Action Input: {"expression": "5+5"}'
),
_engine_response("Thought: Got it.\nFinal Answer: 10"),
]
agent = ReActAgent(engine, "test-model", tools=[_CalculatorStub()])
agent.run("What is 5+5?")
# Check second call messages
second_call = engine.generate.call_args_list[1]
messages = second_call[0][0]
# Last message should be the observation
last_msg = messages[-1]
assert last_msg.role == Role.USER
assert "Observation:" in last_msg.content
assert "10" in last_msg.content
def test_system_prompt_includes_tool_names(self):
"""System prompt should list available tool names."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Done.\nFinal Answer: ok"
)
agent = ReActAgent(
engine, "test-model",
tools=[_CalculatorStub(), _ThinkStub()],
)
agent.run("Hello")
call_args = engine.generate.call_args
messages = call_args[0][0]
system_msg = messages[0]
assert "calculator" in system_msg.content
assert "think" in system_msg.content
def test_system_prompt_no_tools(self):
"""System prompt should say 'none' when no tools are available."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: No tools.\nFinal Answer: ok"
)
agent = ReActAgent(engine, "test-model")
agent.run("Hello")
call_args = engine.generate.call_args
messages = call_args[0][0]
assert "none" in messages[0].content
def test_max_turns_1(self):
"""With max_turns=1 and an action, should stop after 1 turn."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
'Thought: Go.\nAction: calculator\n'
'Action Input: {"expression": "1"}'
)
agent = ReActAgent(
engine, "test-model",
tools=[_CalculatorStub()],
max_turns=1,
)
result = agent.run("Calc")
assert result.turns == 1
assert result.metadata.get("max_turns_exceeded") is True
def test_event_data_agent_turn_start(self):
"""AGENT_TURN_START event data should include agent id and input."""
bus = EventBus(record_history=True)
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Quick.\nFinal Answer: Hi"
)
agent = ReActAgent(engine, "test-model", bus=bus)
agent.run("test input")
start_events = [
e for e in bus.history
if e.event_type == EventType.AGENT_TURN_START
]
assert len(start_events) == 1
assert start_events[0].data["agent"] == "react"
assert start_events[0].data["input"] == "test input"
@pytest.mark.parametrize("model", ["qwen3:8b", "gpt-oss:120b"])
def test_react_with_different_models(model):
"""ReAct agent works with different model names."""
engine = MagicMock()
engine.engine_id = "mock"
engine.generate.return_value = _engine_response(
"Thought: Responding.\nFinal Answer: Hello!"
)
agent = ReActAgent(engine, model)
result = agent.run("Hello")
assert result.content == "Hello!"
call_kwargs = engine.generate.call_args[1]
assert call_kwargs["model"] == model
tests/conftest.py
+205 -1
View File
@@ -2,9 +2,13 @@
from __future__ import annotations
import os
from unittest.mock import MagicMock
import pytest
from openjarvis.core.events import reset_event_bus
from openjarvis.core.config import GpuInfo, HardwareInfo
from openjarvis.core.events import EventBus, reset_event_bus
from openjarvis.core.registry import (
AgentRegistry,
BenchmarkRegistry,
@@ -27,3 +31,203 @@ def _clean_registries() -> None:
RouterPolicyRegistry.clear()
BenchmarkRegistry.clear()
reset_event_bus()
# ---------------------------------------------------------------------------
# Hardware fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def nvidia_gpu() -> GpuInfo:
"""NVIDIA A100 GPU fixture."""
return GpuInfo(vendor="nvidia", name="NVIDIA A100-SXM4-80GB", vram_gb=80.0, count=1)
@pytest.fixture
def nvidia_consumer_gpu() -> GpuInfo:
"""NVIDIA consumer GPU fixture."""
return GpuInfo(
vendor="nvidia", name="NVIDIA GeForce RTX 4090",
vram_gb=24.0, count=1,
)
@pytest.fixture
def nvidia_multi_gpu() -> GpuInfo:
"""NVIDIA multi-GPU fixture."""
return GpuInfo(vendor="nvidia", name="NVIDIA H100", vram_gb=80.0, count=4)
@pytest.fixture
def amd_gpu() -> GpuInfo:
"""AMD MI300X GPU fixture."""
return GpuInfo(vendor="amd", name="AMD Instinct MI300X", vram_gb=192.0, count=1)
@pytest.fixture
def apple_gpu() -> GpuInfo:
"""Apple Silicon GPU fixture."""
return GpuInfo(vendor="apple", name="Apple M4 Max", vram_gb=128.0, count=1)
@pytest.fixture
def hardware_nvidia(nvidia_gpu: GpuInfo) -> HardwareInfo:
"""Full NVIDIA hardware profile."""
return HardwareInfo(
platform="linux",
cpu_brand="AMD EPYC 7763",
cpu_count=64,
ram_gb=512.0,
gpu=nvidia_gpu,
)
@pytest.fixture
def hardware_nvidia_consumer(nvidia_consumer_gpu: GpuInfo) -> HardwareInfo:
"""Consumer NVIDIA hardware profile."""
return HardwareInfo(
platform="linux",
cpu_brand="Intel Core i9-14900K",
cpu_count=24,
ram_gb=64.0,
gpu=nvidia_consumer_gpu,
)
@pytest.fixture
def hardware_amd(amd_gpu: GpuInfo) -> HardwareInfo:
"""Full AMD hardware profile."""
return HardwareInfo(
platform="linux",
cpu_brand="AMD EPYC 9654",
cpu_count=96,
ram_gb=768.0,
gpu=amd_gpu,
)
@pytest.fixture
def hardware_apple(apple_gpu: GpuInfo) -> HardwareInfo:
"""Apple Silicon hardware profile."""
return HardwareInfo(
platform="darwin",
cpu_brand="Apple M4 Max",
cpu_count=16,
ram_gb=128.0,
gpu=apple_gpu,
)
@pytest.fixture
def hardware_cpu_only() -> HardwareInfo:
"""CPU-only hardware profile (no GPU)."""
return HardwareInfo(
platform="linux",
cpu_brand="Intel Xeon E5-2686 v4",
cpu_count=8,
ram_gb=32.0,
gpu=None,
)
# ---------------------------------------------------------------------------
# Engine availability fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def has_ollama() -> bool:
"""Check if Ollama is running locally."""
try:
import httpx
resp = httpx.get("http://localhost:11434/api/tags", timeout=2.0)
return resp.status_code == 200
except Exception:
return False
@pytest.fixture
def has_vllm() -> bool:
"""Check if vLLM is running locally."""
try:
import httpx
resp = httpx.get("http://localhost:8000/v1/models", timeout=2.0)
return resp.status_code == 200
except Exception:
return False
@pytest.fixture
def has_llamacpp() -> bool:
"""Check if llama.cpp server is running locally."""
try:
import httpx
resp = httpx.get("http://localhost:8080/v1/models", timeout=2.0)
return resp.status_code == 200
except Exception:
return False
# ---------------------------------------------------------------------------
# Cloud API key fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def has_openai_key() -> bool:
"""Check if OPENAI_API_KEY is set."""
return bool(os.environ.get("OPENAI_API_KEY"))
@pytest.fixture
def has_anthropic_key() -> bool:
"""Check if ANTHROPIC_API_KEY is set."""
return bool(os.environ.get("ANTHROPIC_API_KEY"))
@pytest.fixture
def has_gemini_key() -> bool:
"""Check if GEMINI_API_KEY or GOOGLE_API_KEY is set."""
return bool(os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY"))
# ---------------------------------------------------------------------------
# Mock engine factory
# ---------------------------------------------------------------------------
@pytest.fixture
def mock_engine():
"""Factory for mock InferenceEngine instances."""
def _factory(
engine_id: str = "mock",
model_response: str = "Hello!",
tool_calls: list | None = None,
models: list[str] | None = None,
) -> MagicMock:
engine = MagicMock()
engine.engine_id = engine_id
engine.health.return_value = True
engine.list_models.return_value = models or ["test-model"]
result = {
"content": model_response,
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
"model": "test-model",
"finish_reason": "stop",
}
if tool_calls:
result["tool_calls"] = tool_calls
result["finish_reason"] = "tool_calls"
engine.generate.return_value = result
return engine
return _factory
@pytest.fixture
def event_bus() -> EventBus:
"""Fresh EventBus with history recording enabled."""
return EventBus(record_history=True)
tests/engine/test_cloud_extended.py
+427
View File
@@ -0,0 +1,427 @@
"""Extended cloud engine tests -- Gemini support and updated models."""
from __future__ import annotations
from types import SimpleNamespace
from unittest import mock
import pytest
from openjarvis.core.registry import EngineRegistry
from openjarvis.core.types import Message, Role
from openjarvis.engine._base import EngineConnectionError
from openjarvis.engine.cloud import (
_ANTHROPIC_MODELS,
_GOOGLE_MODELS,
_OPENAI_MODELS,
PRICING,
CloudEngine,
_is_anthropic_model,
_is_google_model,
estimate_cost,
)
def _make_cloud_engine(monkeypatch: pytest.MonkeyPatch) -> CloudEngine:
"""Create a CloudEngine with all API keys cleared."""
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
if not EngineRegistry.contains("cloud"):
EngineRegistry.register_value("cloud", CloudEngine)
return CloudEngine()
def _fake_openai_response(
content: str = "Hello!",
model: str = "gpt-5-mini",
prompt_tokens: int = 10,
completion_tokens: int = 5,
tool_calls: list | None = None,
) -> SimpleNamespace:
usage = SimpleNamespace(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
message = SimpleNamespace(content=content, tool_calls=tool_calls)
choice = SimpleNamespace(message=message, finish_reason="stop")
return SimpleNamespace(choices=[choice], usage=usage, model=model)
def _fake_anthropic_response(
content: str = "Hello!",
model: str = "claude-opus-4-6",
input_tokens: int = 12,
output_tokens: int = 8,
) -> SimpleNamespace:
usage = SimpleNamespace(input_tokens=input_tokens, output_tokens=output_tokens)
text_block = SimpleNamespace(text=content)
return SimpleNamespace(
content=[text_block], usage=usage, model=model, stop_reason="end_turn"
)
def _fake_gemini_response(
content: str = "Hello!",
prompt_tokens: int = 15,
completion_tokens: int = 10,
) -> SimpleNamespace:
usage = SimpleNamespace(
prompt_token_count=prompt_tokens,
candidates_token_count=completion_tokens,
)
return SimpleNamespace(text=content, usage_metadata=usage)
# ---------------------------------------------------------------------------
# OpenAI tests
# ---------------------------------------------------------------------------
class TestCloudOpenAI:
def test_gpt_5_mini_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.chat.completions.create.return_value = _fake_openai_response(
content="I am GPT-5 Mini", model="gpt-5-mini"
)
engine._openai_client = fake_client
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="gpt-5-mini"
)
assert result["content"] == "I am GPT-5 Mini"
assert result["model"] == "gpt-5-mini"
assert result["usage"]["prompt_tokens"] == 10
def test_gpt_5_mini_cost_estimate(self) -> None:
cost = estimate_cost("gpt-5-mini", 1_000_000, 1_000_000)
# $0.25/M input + $2.00/M output = $2.25
assert cost == pytest.approx(2.25)
def test_gpt_5_mini_tool_calls(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_tool_call = SimpleNamespace(
id="call_xyz",
type="function",
function=SimpleNamespace(name="calc", arguments='{"x":1}'),
)
fake_resp = _fake_openai_response(content="", model="gpt-5-mini")
fake_resp.choices[0].message.tool_calls = [fake_tool_call]
fake_resp.choices[0].finish_reason = "tool_calls"
fake_client = mock.MagicMock()
fake_client.chat.completions.create.return_value = fake_resp
engine._openai_client = fake_client
result = engine.generate(
[Message(role=Role.USER, content="Calculate")], model="gpt-5-mini"
)
assert result["content"] == ""
# ---------------------------------------------------------------------------
# Anthropic tests
# ---------------------------------------------------------------------------
class TestCloudAnthropic:
def test_claude_opus_4_6_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.messages.create.return_value = _fake_anthropic_response(
content="I am Opus 4.6", model="claude-opus-4-6"
)
engine._anthropic_client = fake_client
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="claude-opus-4-6"
)
assert result["content"] == "I am Opus 4.6"
assert result["model"] == "claude-opus-4-6"
def test_claude_sonnet_4_6_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.messages.create.return_value = _fake_anthropic_response(
content="I am Sonnet 4.6", model="claude-sonnet-4-6"
)
engine._anthropic_client = fake_client
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="claude-sonnet-4-6"
)
assert result["content"] == "I am Sonnet 4.6"
def test_claude_haiku_4_5_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.messages.create.return_value = _fake_anthropic_response(
content="I am Haiku 4.5", model="claude-haiku-4-5"
)
engine._anthropic_client = fake_client
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="claude-haiku-4-5"
)
assert result["content"] == "I am Haiku 4.5"
def test_claude_cost_estimate(self) -> None:
# claude-opus-4-6: $5.00/M in, $25.00/M out
cost = estimate_cost("claude-opus-4-6", 1_000_000, 1_000_000)
assert cost == pytest.approx(30.00)
# claude-sonnet-4-6: $3.00/M in, $15.00/M out
cost = estimate_cost("claude-sonnet-4-6", 1_000_000, 1_000_000)
assert cost == pytest.approx(18.00)
# claude-haiku-4-5: $1.00/M in, $5.00/M out
cost = estimate_cost("claude-haiku-4-5", 1_000_000, 1_000_000)
assert cost == pytest.approx(6.00)
def test_anthropic_routing(self) -> None:
assert _is_anthropic_model("claude-opus-4-6") is True
assert _is_anthropic_model("claude-sonnet-4-6") is True
assert _is_anthropic_model("claude-haiku-4-5") is True
assert _is_anthropic_model("gpt-5-mini") is False
assert _is_anthropic_model("gemini-3-pro") is False
def test_anthropic_system_message(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.messages.create.return_value = _fake_anthropic_response(
content="With system", model="claude-opus-4-6"
)
engine._anthropic_client = fake_client
engine.generate(
[
Message(role=Role.SYSTEM, content="You are helpful"),
Message(role=Role.USER, content="Hi"),
],
model="claude-opus-4-6",
)
call_kwargs = fake_client.messages.create.call_args
assert call_kwargs.kwargs.get("system") == "You are helpful" or \
call_kwargs[1].get("system") == "You are helpful"
# ---------------------------------------------------------------------------
# Gemini tests
# ---------------------------------------------------------------------------
class TestCloudGemini:
def test_gemini_init_with_api_key(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("GEMINI_API_KEY", "test-gemini-key")
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
fake_genai = mock.MagicMock()
with mock.patch.dict("sys.modules", {
"google": mock.MagicMock(),
"google.genai": fake_genai,
}):
if not EngineRegistry.contains("cloud"):
EngineRegistry.register_value("cloud", CloudEngine)
engine = CloudEngine()
assert engine._google_client is not None
def test_gemini_2_5_pro_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.models.generate_content.return_value = _fake_gemini_response(
content="I am Gemini 2.5 Pro"
)
engine._google_client = fake_client
# Mock the genai types import
fake_config = mock.MagicMock()
fake_types = mock.MagicMock()
fake_types.GenerateContentConfig.return_value = fake_config
with mock.patch.dict("sys.modules", {
"google": mock.MagicMock(),
"google.genai": mock.MagicMock(),
"google.genai.types": fake_types,
}):
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="gemini-2.5-pro"
)
assert result["content"] == "I am Gemini 2.5 Pro"
assert result["model"] == "gemini-2.5-pro"
assert result["usage"]["prompt_tokens"] == 15
assert result["usage"]["completion_tokens"] == 10
def test_gemini_2_5_flash_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.models.generate_content.return_value = _fake_gemini_response(
content="I am Gemini 2.5 Flash"
)
engine._google_client = fake_client
fake_types = mock.MagicMock()
fake_types.GenerateContentConfig.return_value = mock.MagicMock()
with mock.patch.dict("sys.modules", {
"google": mock.MagicMock(),
"google.genai": mock.MagicMock(),
"google.genai.types": fake_types,
}):
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="gemini-2.5-flash"
)
assert result["content"] == "I am Gemini 2.5 Flash"
def test_gemini_3_pro_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.models.generate_content.return_value = _fake_gemini_response(
content="I am Gemini 3 Pro"
)
engine._google_client = fake_client
fake_types = mock.MagicMock()
fake_types.GenerateContentConfig.return_value = mock.MagicMock()
with mock.patch.dict("sys.modules", {
"google": mock.MagicMock(),
"google.genai": mock.MagicMock(),
"google.genai.types": fake_types,
}):
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="gemini-3-pro"
)
assert result["content"] == "I am Gemini 3 Pro"
def test_gemini_3_flash_generate(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
fake_client = mock.MagicMock()
fake_client.models.generate_content.return_value = _fake_gemini_response(
content="I am Gemini 3 Flash"
)
engine._google_client = fake_client
fake_types = mock.MagicMock()
fake_types.GenerateContentConfig.return_value = mock.MagicMock()
with mock.patch.dict("sys.modules", {
"google": mock.MagicMock(),
"google.genai": mock.MagicMock(),
"google.genai.types": fake_types,
}):
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model="gemini-3-flash"
)
assert result["content"] == "I am Gemini 3 Flash"
def test_gemini_cost_estimate(self) -> None:
# gemini-2.5-pro: $1.25/M in, $10.00/M out
cost = estimate_cost("gemini-2.5-pro", 1_000_000, 1_000_000)
assert cost == pytest.approx(11.25)
# gemini-2.5-flash: $0.30/M in, $2.50/M out
cost = estimate_cost("gemini-2.5-flash", 1_000_000, 1_000_000)
assert cost == pytest.approx(2.80)
# gemini-3-pro: $2.00/M in, $12.00/M out
cost = estimate_cost("gemini-3-pro", 1_000_000, 1_000_000)
assert cost == pytest.approx(14.00)
# gemini-3-flash: $0.50/M in, $3.00/M out
cost = estimate_cost("gemini-3-flash", 1_000_000, 1_000_000)
assert cost == pytest.approx(3.50)
def test_gemini_routing(self) -> None:
assert _is_google_model("gemini-2.5-pro") is True
assert _is_google_model("gemini-2.5-flash") is True
assert _is_google_model("gemini-3-pro") is True
assert _is_google_model("gemini-3-flash") is True
assert _is_google_model("gpt-5-mini") is False
assert _is_google_model("claude-opus-4-6") is False
def test_gemini_no_client_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
assert engine._google_client is None
fake_types = mock.MagicMock()
fake_types.GenerateContentConfig.return_value = mock.MagicMock()
with mock.patch.dict("sys.modules", {
"google": mock.MagicMock(),
"google.genai": mock.MagicMock(),
"google.genai.types": fake_types,
}):
with pytest.raises(
EngineConnectionError,
match="Google client not available",
):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="gemini-3-pro"
)
# ---------------------------------------------------------------------------
# Model discovery
# ---------------------------------------------------------------------------
class TestCloudModelDiscovery:
def test_list_models_includes_all(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""When all clients are set, all model lists are returned."""
engine = _make_cloud_engine(monkeypatch)
engine._openai_client = mock.MagicMock()
engine._anthropic_client = mock.MagicMock()
engine._google_client = mock.MagicMock()
models = engine.list_models()
for m in _OPENAI_MODELS:
assert m in models
for m in _ANTHROPIC_MODELS:
assert m in models
for m in _GOOGLE_MODELS:
assert m in models
def test_no_api_key_empty_list(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
assert engine.list_models() == []
def test_only_google_client(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
engine._google_client = mock.MagicMock()
models = engine.list_models()
assert set(models) == set(_GOOGLE_MODELS)
# No OpenAI or Anthropic models
for m in _OPENAI_MODELS:
assert m not in models
for m in _ANTHROPIC_MODELS:
assert m not in models
def test_health_with_google_client(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
engine._google_client = mock.MagicMock()
assert engine.health() is True
def test_health_no_clients(self, monkeypatch: pytest.MonkeyPatch) -> None:
engine = _make_cloud_engine(monkeypatch)
assert engine.health() is False
# ---------------------------------------------------------------------------
# Pricing completeness
# ---------------------------------------------------------------------------
class TestPricingTable:
def test_all_new_models_in_pricing(self) -> None:
expected = [
"gpt-5-mini",
"claude-opus-4-6", "claude-sonnet-4-6", "claude-haiku-4-5",
"gemini-2.5-pro", "gemini-2.5-flash", "gemini-3-pro", "gemini-3-flash",
]
for model_id in expected:
assert model_id in PRICING, f"{model_id} missing from PRICING dict"
def test_pricing_values_positive(self) -> None:
for model_id, (inp, out) in PRICING.items():
assert inp >= 0, f"{model_id} has negative input price"
assert out >= 0, f"{model_id} has negative output price"
def test_zero_tokens_zero_cost(self) -> None:
assert estimate_cost("gpt-5-mini", 0, 0) == 0.0
tests/engine/test_engine_model_matrix.py
+198
View File
@@ -0,0 +1,198 @@
"""Cross-product parametrized tests: engine x scenario."""
from __future__ import annotations
import httpx
import pytest
from openjarvis.core.registry import EngineRegistry
from openjarvis.core.types import Message, Role
from openjarvis.engine._base import EngineConnectionError
from openjarvis.engine.ollama import OllamaEngine
from openjarvis.engine.vllm import VLLMEngine
ENGINES_AND_HOSTS = [
("vllm", "http://testhost:8000"),
("ollama", "http://testhost:11434"),
]
MODELS = ["gpt-oss:120b", "qwen3:8b", "glm-4.7-flash", "trinity-mini"]
def _create_engine(engine_key: str, host: str):
"""Instantiate the right engine class for the given key."""
if engine_key == "vllm":
if not EngineRegistry.contains("vllm"):
EngineRegistry.register_value("vllm", VLLMEngine)
return VLLMEngine(host=host)
elif engine_key == "ollama":
if not EngineRegistry.contains("ollama"):
EngineRegistry.register_value("ollama", OllamaEngine)
return OllamaEngine(host=host)
else:
raise ValueError(f"Unknown engine: {engine_key}")
def _mock_simple_chat(respx_mock, engine_key: str, host: str, model: str):
"""Set up mock for a simple chat response."""
if engine_key == "vllm":
respx_mock.post(f"{host}/v1/chat/completions").mock(
return_value=httpx.Response(200, json={
"choices": [
{"message": {"content": "Hello!"}, "finish_reason": "stop"},
],
"usage": {
"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15,
},
"model": model,
})
)
elif engine_key == "ollama":
respx_mock.post(f"{host}/api/chat").mock(
return_value=httpx.Response(200, json={
"message": {"role": "assistant", "content": "Hello!"},
"model": model,
"prompt_eval_count": 10,
"eval_count": 5,
"done": True,
})
)
def _mock_tool_call(respx_mock, engine_key: str, host: str, model: str):
"""Set up mock for a tool-call response."""
if engine_key == "vllm":
respx_mock.post(f"{host}/v1/chat/completions").mock(
return_value=httpx.Response(200, json={
"choices": [{
"message": {
"content": "",
"tool_calls": [{
"id": "call_1",
"type": "function",
"function": {"name": "calculator", "arguments": '{"x":1}'},
}],
},
"finish_reason": "tool_calls",
}],
"usage": {
"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18,
},
"model": model,
})
)
elif engine_key == "ollama":
respx_mock.post(f"{host}/api/chat").mock(
return_value=httpx.Response(200, json={
"message": {
"content": "",
"tool_calls": [{
"function": {"name": "calculator", "arguments": '{"x":1}'},
}],
},
"model": model,
"prompt_eval_count": 10,
"eval_count": 8,
"done": True,
})
)
def _mock_error(respx_mock, engine_key: str, host: str):
"""Set up mock for connection error."""
if engine_key == "vllm":
respx_mock.post(f"{host}/v1/chat/completions").mock(
side_effect=httpx.ConnectError("refused")
)
elif engine_key == "ollama":
respx_mock.post(f"{host}/api/chat").mock(
side_effect=httpx.ConnectError("refused")
)
# ---------------------------------------------------------------------------
# Cross-product: engine x scenario
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("engine_key,host", ENGINES_AND_HOSTS)
class TestEngineScenarios:
def test_simple_chat(self, respx_mock, engine_key: str, host: str) -> None:
engine = _create_engine(engine_key, host)
_mock_simple_chat(respx_mock, engine_key, host, "qwen3:8b")
result = engine.generate(
[Message(role=Role.USER, content="Hello")], model="qwen3:8b"
)
assert result["content"] == "Hello!"
assert result["usage"]["prompt_tokens"] == 10
def test_tool_call(self, respx_mock, engine_key: str, host: str) -> None:
engine = _create_engine(engine_key, host)
_mock_tool_call(respx_mock, engine_key, host, "qwen3:8b")
result = engine.generate(
[Message(role=Role.USER, content="Calculate")],
model="qwen3:8b",
tools=[{"type": "function", "function": {"name": "calculator"}}],
)
assert "tool_calls" in result
assert result["tool_calls"][0]["name"] == "calculator"
def test_error_handling(self, respx_mock, engine_key: str, host: str) -> None:
engine = _create_engine(engine_key, host)
_mock_error(respx_mock, engine_key, host)
with pytest.raises(EngineConnectionError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="qwen3:8b"
)
# ---------------------------------------------------------------------------
# Cross-product: engine x model
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("engine_key,host", ENGINES_AND_HOSTS)
@pytest.mark.parametrize("model_id", MODELS)
class TestEngineModelMatrix:
def test_generate_with_model(
self, respx_mock, engine_key: str, host: str, model_id: str,
) -> None:
engine = _create_engine(engine_key, host)
_mock_simple_chat(respx_mock, engine_key, host, model_id)
result = engine.generate(
[Message(role=Role.USER, content="Hi")], model=model_id
)
assert result["content"] == "Hello!"
assert result["model"] == model_id
# ---------------------------------------------------------------------------
# Health checks across engines
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("engine_key,host", ENGINES_AND_HOSTS)
class TestEngineHealth:
def test_health_true(self, respx_mock, engine_key: str, host: str) -> None:
engine = _create_engine(engine_key, host)
if engine_key == "vllm":
respx_mock.get(f"{host}/v1/models").mock(
return_value=httpx.Response(200, json={"data": []})
)
elif engine_key == "ollama":
respx_mock.get(f"{host}/api/tags").mock(
return_value=httpx.Response(200, json={"models": []})
)
assert engine.health() is True
def test_health_false(self, respx_mock, engine_key: str, host: str) -> None:
engine = _create_engine(engine_key, host)
if engine_key == "vllm":
respx_mock.get(f"{host}/v1/models").mock(
side_effect=httpx.ConnectError("refused")
)
elif engine_key == "ollama":
respx_mock.get(f"{host}/api/tags").mock(
side_effect=httpx.ConnectError("refused")
)
assert engine.health() is False
tests/engine/test_llamacpp_models.py
+205
View File
@@ -0,0 +1,205 @@
"""Tests for llama.cpp engine with compatible models."""
from __future__ import annotations
import json
import httpx
import pytest
import respx
from openjarvis.core.registry import EngineRegistry
from openjarvis.core.types import Message, Role
from openjarvis.engine._base import EngineConnectionError
from openjarvis.engine.llamacpp import LlamaCppEngine
LLAMACPP_HOST = "http://testhost:8080"
# Only models with llamacpp in supported_engines
COMPATIBLE_MODELS = ["qwen3:8b", "trinity-mini"]
def _make_engine() -> LlamaCppEngine:
if not EngineRegistry.contains("llamacpp"):
EngineRegistry.register_value("llamacpp", LlamaCppEngine)
return LlamaCppEngine(host=LLAMACPP_HOST)
def _openai_response(
content: str = "Hello!",
model: str = "qwen3:8b",
prompt_tokens: int = 10,
completion_tokens: int = 5,
tool_calls: list | None = None,
finish_reason: str = "stop",
) -> dict:
"""Build an OpenAI-format response dict."""
message: dict = {"content": content}
if tool_calls:
message["tool_calls"] = tool_calls
finish_reason = "tool_calls"
return {
"choices": [{"message": message, "finish_reason": finish_reason}],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
"model": model,
}
# ---------------------------------------------------------------------------
# Generate tests (parametrized over compatible models)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("model_id", COMPATIBLE_MODELS)
class TestLlamaCppGenerate:
def test_generate_basic(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
respx_mock.post(f"{LLAMACPP_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(
200, json=_openai_response(content="Test reply", model=model_id)
)
)
result = engine.generate(
[Message(role=Role.USER, content="Hello")], model=model_id
)
assert result["content"] == "Test reply"
assert result["model"] == model_id
assert result["usage"]["total_tokens"] == 15
def test_generate_with_tools(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
tool_calls = [
{
"id": "call_abc",
"type": "function",
"function": {
"name": "calculator",
"arguments": '{"expression":"3*3"}',
},
}
]
respx_mock.post(f"{LLAMACPP_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(
200,
json=_openai_response(
content="", model=model_id, tool_calls=tool_calls,
),
)
)
result = engine.generate(
[Message(role=Role.USER, content="3*3")],
model=model_id,
tools=[{"type": "function", "function": {"name": "calculator"}}],
)
assert "tool_calls" in result
assert result["tool_calls"][0]["name"] == "calculator"
def test_generate_tool_fallback(self, respx_mock, model_id: str) -> None:
"""400 with tools in payload → retry without tools."""
engine = _make_engine()
call_count = 0
def handler(request):
nonlocal call_count
call_count += 1
body = json.loads(request.content)
if "tools" in body:
return httpx.Response(400, json={"error": "unsupported"})
return httpx.Response(
200, json=_openai_response(content="Fallback", model=model_id)
)
respx_mock.post(f"{LLAMACPP_HOST}/v1/chat/completions").mock(
side_effect=handler
)
result = engine.generate(
[Message(role=Role.USER, content="Hello")],
model=model_id,
tools=[{"type": "function", "function": {"name": "calc"}}],
)
assert result["content"] == "Fallback"
assert call_count == 2
def test_generate_streaming(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
sse = (
'data: {"choices":[{"delta":{"content":"Hi"}}]}\n'
'data: {"choices":[{"delta":{"content":" there"}}]}\n'
"data: [DONE]\n"
)
respx_mock.post(f"{LLAMACPP_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(200, text=sse)
)
async def collect():
tokens = []
async for tok in engine.stream(
[Message(role=Role.USER, content="Hi")], model=model_id
):
tokens.append(tok)
return tokens
import asyncio
tokens = asyncio.get_event_loop().run_until_complete(collect())
assert tokens == ["Hi", " there"]
# ---------------------------------------------------------------------------
# Model discovery & health
# ---------------------------------------------------------------------------
class TestLlamaCppModelDiscovery:
def test_list_models(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{LLAMACPP_HOST}/v1/models").mock(
return_value=httpx.Response(
200,
json={"data": [{"id": m} for m in COMPATIBLE_MODELS]},
)
)
assert engine.list_models() == COMPATIBLE_MODELS
def test_health_healthy(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{LLAMACPP_HOST}/v1/models").mock(
return_value=httpx.Response(200, json={"data": []})
)
assert engine.health() is True
def test_health_unhealthy(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{LLAMACPP_HOST}/v1/models").mock(
side_effect=httpx.ConnectError("refused")
)
assert engine.health() is False
# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------
class TestLlamaCppErrors:
def test_connection_refused(self) -> None:
engine = _make_engine()
with respx.mock:
respx.post(f"{LLAMACPP_HOST}/v1/chat/completions").mock(
side_effect=httpx.ConnectError("refused")
)
with pytest.raises(EngineConnectionError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="qwen3:8b"
)
def test_default_host_is_8080(self) -> None:
"""LlamaCppEngine defaults to port 8080."""
engine = LlamaCppEngine()
assert engine._host == "http://localhost:8080"
def test_engine_id(self) -> None:
engine = LlamaCppEngine()
assert engine.engine_id == "llamacpp"
tests/engine/test_ollama_models.py
+252
View File
@@ -0,0 +1,252 @@
"""Tests for Ollama engine with extended local model set."""
from __future__ import annotations
import json
import httpx
import pytest
import respx
from openjarvis.core.registry import EngineRegistry
from openjarvis.core.types import Message, Role
from openjarvis.engine._base import EngineConnectionError
from openjarvis.engine.ollama import OllamaEngine
OLLAMA_HOST = "http://testhost:11434"
NEW_MODELS = ["gpt-oss:120b", "qwen3:8b", "glm-4.7-flash", "trinity-mini"]
def _make_engine() -> OllamaEngine:
if not EngineRegistry.contains("ollama"):
EngineRegistry.register_value("ollama", OllamaEngine)
return OllamaEngine(host=OLLAMA_HOST)
def _ollama_response(
content: str = "Hello!",
model: str = "qwen3:8b",
prompt_eval_count: int = 10,
eval_count: int = 5,
tool_calls: list | None = None,
) -> dict:
"""Build an Ollama-format response dict."""
message: dict = {"role": "assistant", "content": content}
if tool_calls:
message["tool_calls"] = tool_calls
result: dict = {
"message": message,
"model": model,
"prompt_eval_count": prompt_eval_count,
"eval_count": eval_count,
"done": True,
}
return result
# ---------------------------------------------------------------------------
# Generate tests (parametrized over new models)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("model_id", NEW_MODELS)
class TestOllamaGenerate:
def test_generate_basic(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
respx_mock.post(f"{OLLAMA_HOST}/api/chat").mock(
return_value=httpx.Response(
200, json=_ollama_response(content="Test reply", model=model_id)
)
)
result = engine.generate(
[Message(role=Role.USER, content="Hello")], model=model_id
)
assert result["content"] == "Test reply"
assert result["model"] == model_id
assert result["usage"]["prompt_tokens"] == 10
assert result["usage"]["completion_tokens"] == 5
assert result["usage"]["total_tokens"] == 15
def test_generate_with_tools(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
tool_calls = [
{
"function": {
"name": "calculator",
"arguments": '{"expression":"2+2"}',
},
}
]
respx_mock.post(f"{OLLAMA_HOST}/api/chat").mock(
return_value=httpx.Response(
200,
json=_ollama_response(
content="", model=model_id, tool_calls=tool_calls,
),
)
)
result = engine.generate(
[Message(role=Role.USER, content="What is 2+2?")],
model=model_id,
tools=[{"type": "function", "function": {"name": "calculator"}}],
)
assert "tool_calls" in result
assert result["tool_calls"][0]["name"] == "calculator"
assert result["tool_calls"][0]["id"] == "call_0"
def test_generate_with_multiple_tool_calls(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
tool_calls = [
{"function": {"name": "tool_a", "arguments": "{}"}},
{"function": {"name": "tool_b", "arguments": "{}"}},
]
respx_mock.post(f"{OLLAMA_HOST}/api/chat").mock(
return_value=httpx.Response(
200,
json=_ollama_response(
content="", model=model_id,
tool_calls=tool_calls,
),
)
)
result = engine.generate(
[Message(role=Role.USER, content="Use tools")], model=model_id
)
assert len(result["tool_calls"]) == 2
assert result["tool_calls"][0]["name"] == "tool_a"
assert result["tool_calls"][1]["name"] == "tool_b"
def test_generate_streaming(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
lines = [
json.dumps({"message": {"content": "Hello"}, "done": False}),
json.dumps({"message": {"content": " world"}, "done": True}),
]
body = "\n".join(lines)
respx_mock.post(f"{OLLAMA_HOST}/api/chat").mock(
return_value=httpx.Response(200, text=body)
)
async def collect():
tokens = []
async for tok in engine.stream(
[Message(role=Role.USER, content="Hi")], model=model_id
):
tokens.append(tok)
return tokens
import asyncio
tokens = asyncio.get_event_loop().run_until_complete(collect())
assert "Hello" in tokens
# ---------------------------------------------------------------------------
# Model discovery
# ---------------------------------------------------------------------------
class TestOllamaModelDiscovery:
def test_list_models(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{OLLAMA_HOST}/api/tags").mock(
return_value=httpx.Response(
200,
json={"models": [{"name": m} for m in NEW_MODELS]},
)
)
models = engine.list_models()
assert models == NEW_MODELS
def test_list_models_empty(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{OLLAMA_HOST}/api/tags").mock(
return_value=httpx.Response(200, json={"models": []})
)
assert engine.list_models() == []
def test_list_models_connection_error(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{OLLAMA_HOST}/api/tags").mock(
side_effect=httpx.ConnectError("refused")
)
assert engine.list_models() == []
def test_health_healthy(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{OLLAMA_HOST}/api/tags").mock(
return_value=httpx.Response(200, json={"models": []})
)
assert engine.health() is True
def test_health_unhealthy(self) -> None:
engine = _make_engine()
with respx.mock:
respx.get(f"{OLLAMA_HOST}/api/tags").mock(
side_effect=httpx.ConnectError("refused")
)
assert engine.health() is False
# ---------------------------------------------------------------------------
# Error handling
# ---------------------------------------------------------------------------
class TestOllamaErrors:
def test_connection_refused(self) -> None:
engine = _make_engine()
with respx.mock:
respx.post(f"{OLLAMA_HOST}/api/chat").mock(
side_effect=httpx.ConnectError("refused")
)
with pytest.raises(EngineConnectionError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="qwen3:8b"
)
def test_timeout_raises_connection_error(self) -> None:
engine = _make_engine()
with respx.mock:
respx.post(f"{OLLAMA_HOST}/api/chat").mock(
side_effect=httpx.TimeoutException("timed out")
)
with pytest.raises(EngineConnectionError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="qwen3:8b"
)
def test_tools_payload_included(self, respx_mock) -> None:
"""Tools are included in the Ollama payload when provided."""
engine = _make_engine()
captured = {}
def capture(request):
captured["body"] = json.loads(request.content)
return httpx.Response(
200, json=_ollama_response(content="ok")
)
respx_mock.post(f"{OLLAMA_HOST}/api/chat").mock(side_effect=capture)
engine.generate(
[Message(role=Role.USER, content="Hello")],
model="qwen3:8b",
tools=[{"type": "function", "function": {"name": "calc"}}],
)
assert "tools" in captured["body"]
def test_no_tools_no_tools_key(self, respx_mock) -> None:
"""Without tools kwarg, payload has no tools key."""
engine = _make_engine()
captured = {}
def capture(request):
captured["body"] = json.loads(request.content)
return httpx.Response(
200, json=_ollama_response(content="ok")
)
respx_mock.post(f"{OLLAMA_HOST}/api/chat").mock(side_effect=capture)
engine.generate(
[Message(role=Role.USER, content="Hello")], model="qwen3:8b"
)
assert "tools" not in captured["body"]
tests/engine/test_vllm_models.py
+220
View File
@@ -0,0 +1,220 @@
"""Tests for vLLM engine with extended local model set."""
from __future__ import annotations
import json
import httpx
import pytest
import respx
from openjarvis.core.registry import EngineRegistry
from openjarvis.core.types import Message, Role
from openjarvis.engine._base import EngineConnectionError
from openjarvis.engine.vllm import VLLMEngine
VLLM_HOST = "http://testhost:8000"
NEW_MODELS = ["gpt-oss:120b", "qwen3:8b", "glm-4.7-flash", "trinity-mini"]
def _make_engine() -> VLLMEngine:
if not EngineRegistry.contains("vllm"):
EngineRegistry.register_value("vllm", VLLMEngine)
return VLLMEngine(host=VLLM_HOST)
def _openai_response(
content: str = "Hello!",
model: str = "qwen3:8b",
prompt_tokens: int = 10,
completion_tokens: int = 5,
tool_calls: list | None = None,
finish_reason: str = "stop",
) -> dict:
"""Build an OpenAI-format response dict."""
message: dict = {"content": content}
if tool_calls:
message["tool_calls"] = tool_calls
finish_reason = "tool_calls"
return {
"choices": [{"message": message, "finish_reason": finish_reason}],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
"model": model,
}
# ---------------------------------------------------------------------------
# Generate tests (parametrized over new models)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("model_id", NEW_MODELS)
class TestVLLMGenerate:
def test_generate_basic(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
respx_mock.post(f"{VLLM_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(
200, json=_openai_response(content="Test reply", model=model_id)
)
)
result = engine.generate(
[Message(role=Role.USER, content="Hello")], model=model_id
)
assert result["content"] == "Test reply"
assert result["model"] == model_id
assert result["usage"]["total_tokens"] == 15
def test_generate_with_tools(self, respx_mock, model_id: str) -> None:
engine = _make_engine()
tool_calls = [
{
"id": "call_123",
"type": "function",
"function": {
"name": "calculator",
"arguments": '{"expression":"2+2"}',
},
}
]
respx_mock.post(f"{VLLM_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(
200,
json=_openai_response(
content="", model=model_id, tool_calls=tool_calls,
),
)
)
result = engine.generate(
[Message(role=Role.USER, content="What is 2+2?")],
model=model_id,
tools=[{"type": "function", "function": {"name": "calculator"}}],
)
assert "tool_calls" in result
assert result["tool_calls"][0]["name"] == "calculator"
assert result["tool_calls"][0]["id"] == "call_123"
def test_generate_tool_fallback(self, respx_mock, model_id: str) -> None:
"""When tools cause a 400, engine retries without tools."""
engine = _make_engine()
call_count = 0
def handler(request):
nonlocal call_count
call_count += 1
body = json.loads(request.content)
if "tools" in body:
return httpx.Response(400, json={"error": "tools not supported"})
return httpx.Response(
200, json=_openai_response(content="Fallback reply", model=model_id)
)
respx_mock.post(f"{VLLM_HOST}/v1/chat/completions").mock(
side_effect=handler
)
result = engine.generate(
[Message(role=Role.USER, content="Hello")],
model=model_id,
tools=[{"type": "function", "function": {"name": "calc"}}],
)
assert result["content"] == "Fallback reply"
assert call_count == 2
def test_generate_streaming(self, respx_mock, model_id: str) -> None:
"""SSE stream yields content tokens."""
engine = _make_engine()
sse = (
'data: {"choices":[{"delta":{"content":"Hello"}}]}\n'
'data: {"choices":[{"delta":{"content":" world"}}]}\n'
"data: [DONE]\n"
)
respx_mock.post(f"{VLLM_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(200, text=sse)
)
async def collect():
tokens = []
async for tok in engine.stream(
[Message(role=Role.USER, content="Hi")], model=model_id
):
tokens.append(tok)
return tokens
import asyncio
tokens = asyncio.get_event_loop().run_until_complete(collect())
assert tokens == ["Hello", " world"]
# ---------------------------------------------------------------------------
# Model discovery
# ---------------------------------------------------------------------------
class TestVLLMModelDiscovery:
def test_list_models(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{VLLM_HOST}/v1/models").mock(
return_value=httpx.Response(
200,
json={"data": [{"id": m} for m in NEW_MODELS]},
)
)
models = engine.list_models()
assert models == NEW_MODELS
def test_health_check_healthy(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{VLLM_HOST}/v1/models").mock(
return_value=httpx.Response(200, json={"data": []})
)
assert engine.health() is True
def test_health_check_unhealthy(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.get(f"{VLLM_HOST}/v1/models").mock(
side_effect=httpx.ConnectError("refused")
)
assert engine.health() is False
# ---------------------------------------------------------------------------
# Error handling
# ---------------------------------------------------------------------------
class TestVLLMErrors:
def test_connection_refused(self) -> None:
"""No mock — ConnectError raises EngineConnectionError."""
engine = VLLMEngine(host="http://localhost:19999")
with respx.mock:
respx.post("http://localhost:19999/v1/chat/completions").mock(
side_effect=httpx.ConnectError("refused")
)
with pytest.raises(EngineConnectionError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="qwen3:8b"
)
def test_invalid_model_404(self, respx_mock) -> None:
engine = _make_engine()
respx_mock.post(f"{VLLM_HOST}/v1/chat/completions").mock(
return_value=httpx.Response(404, json={"error": "model not found"})
)
with pytest.raises(httpx.HTTPStatusError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="nonexistent"
)
def test_timeout_raises_connection_error(self) -> None:
engine = _make_engine()
with respx.mock:
respx.post(f"{VLLM_HOST}/v1/chat/completions").mock(
side_effect=httpx.TimeoutException("timed out")
)
with pytest.raises(EngineConnectionError):
engine.generate(
[Message(role=Role.USER, content="Hi")], model="qwen3:8b"
)
tests/hardware/__init__.py
View File
tests/hardware/test_amd.py
+113
View File
@@ -0,0 +1,113 @@
"""AMD-specific hardware tests."""
from __future__ import annotations
from unittest.mock import patch
import pytest
from openjarvis.core.config import (
GpuInfo,
HardwareInfo,
_detect_amd_gpu,
recommend_engine,
)
pytestmark = pytest.mark.amd
# ---------------------------------------------------------------------------
# Detection / rocm-smi parsing
# ---------------------------------------------------------------------------
class TestAMDDetection:
"""Tests for _detect_amd_gpu() against various rocm-smi outputs."""
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="AMD Instinct MI300X",
)
def test_rocm_smi_parsing(self, mock_run, mock_which):
gpu = _detect_amd_gpu()
assert gpu is not None
assert gpu.vendor == "amd"
assert "MI300X" in gpu.name
@patch("openjarvis.core.config.shutil.which", return_value=None)
def test_rocm_smi_not_found(self, mock_which):
assert _detect_amd_gpu() is None
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="AMD Instinct MI250X\nAMD Instinct MI250X",
)
def test_amd_gpu_model(self, mock_run, mock_which):
"""First line of rocm-smi output is used as the GPU name."""
gpu = _detect_amd_gpu()
assert gpu is not None
assert "MI250X" in gpu.name
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
@patch("openjarvis.core.config._run_cmd", return_value="")
def test_rocm_smi_empty_output(self, mock_run, mock_which):
"""Empty output from rocm-smi returns None."""
assert _detect_amd_gpu() is None
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="AMD Instinct MI300X",
)
def test_amd_vram(self, mock_run, mock_which):
"""AMD detection does not parse VRAM; defaults to 0.0."""
gpu = _detect_amd_gpu()
assert gpu is not None
assert gpu.vram_gb == 0.0
# ---------------------------------------------------------------------------
# Engine recommendation
# ---------------------------------------------------------------------------
class TestAMDEngineRecommendation:
"""Tests that AMD cards map to vllm."""
def test_mi300x_recommends_vllm(self):
hw = HardwareInfo(
platform="linux",
cpu_brand="AMD EPYC 9654",
cpu_count=96,
ram_gb=768.0,
gpu=GpuInfo(
vendor="amd", name="AMD Instinct MI300X",
vram_gb=192.0, count=1,
),
)
assert recommend_engine(hw) == "vllm"
def test_amd_generic_recommends_vllm(self):
hw = HardwareInfo(
platform="linux",
cpu_brand="AMD EPYC",
cpu_count=64,
ram_gb=256.0,
gpu=GpuInfo(vendor="amd", name="AMD GPU", vram_gb=0.0, count=1),
)
assert recommend_engine(hw) == "vllm"
def test_amd_multi_gpu_recommends_vllm(self):
hw = HardwareInfo(
platform="linux",
cpu_brand="AMD EPYC 9654",
cpu_count=128,
ram_gb=1024.0,
gpu=GpuInfo(
vendor="amd", name="AMD Instinct MI300X",
vram_gb=192.0, count=4,
),
)
assert recommend_engine(hw) == "vllm"
tests/hardware/test_apple.py
+140
View File
@@ -0,0 +1,140 @@
"""Apple Silicon hardware tests."""
from __future__ import annotations
from unittest.mock import patch
import pytest
from openjarvis.core.config import (
GpuInfo,
HardwareInfo,
_detect_apple_gpu,
recommend_engine,
)
pytestmark = pytest.mark.apple
# ---------------------------------------------------------------------------
# Detection / system_profiler parsing
# ---------------------------------------------------------------------------
class TestAppleDetection:
"""Tests for _detect_apple_gpu() against system_profiler outputs."""
@patch("openjarvis.core.config.platform.system", return_value="Darwin")
@patch(
"openjarvis.core.config._run_cmd",
return_value=(
"Graphics/Displays:\n"
"\n"
" Apple M4 Max:\n"
"\n"
" Chipset Model: Apple M4 Max\n"
" Type: GPU\n"
" Bus: Built-In\n"
" Total Number of Cores: 40\n"
),
)
def test_system_profiler_parsing(self, mock_run, mock_system):
gpu = _detect_apple_gpu()
assert gpu is not None
assert gpu.vendor == "apple"
assert "M4 Max" in gpu.name
@patch("openjarvis.core.config.platform.system", return_value="Linux")
def test_non_darwin_returns_none(self, mock_system):
"""On non-Darwin platforms, Apple GPU detection returns None."""
assert _detect_apple_gpu() is None
@patch("openjarvis.core.config.platform.system", return_value="Darwin")
@patch(
"openjarvis.core.config._run_cmd",
return_value=(
"Graphics/Displays:\n"
"\n"
" Intel UHD Graphics 630:\n"
"\n"
" Chipset Model: Intel UHD Graphics 630\n"
" Type: GPU\n"
),
)
def test_no_apple_in_output(self, mock_run, mock_system):
"""system_profiler output without 'Apple' returns None."""
assert _detect_apple_gpu() is None
@patch("openjarvis.core.config.platform.system", return_value="Darwin")
@patch(
"openjarvis.core.config._run_cmd",
return_value=(
"Graphics/Displays:\n"
"\n"
" Apple M2 Ultra:\n"
"\n"
" Chipset Model: Apple M2 Ultra\n"
" Type: GPU\n"
" Bus: Built-In\n"
),
)
def test_apple_chip_model_extraction(self, mock_run, mock_system):
"""'Chipset Model:' line is used to extract the chip name."""
gpu = _detect_apple_gpu()
assert gpu is not None
assert gpu.name == "Apple M2 Ultra"
@patch("openjarvis.core.config.platform.system", return_value="Darwin")
@patch(
"openjarvis.core.config._run_cmd",
return_value=(
"Graphics/Displays:\n"
" Apple Silicon\n"
" Type: GPU\n"
),
)
def test_apple_no_chipset_line_falls_back(self, mock_run, mock_system):
"""When no 'Chipset Model' line exists, falls back to 'Apple Silicon'."""
gpu = _detect_apple_gpu()
assert gpu is not None
assert gpu.vendor == "apple"
assert gpu.name == "Apple Silicon"
@patch("openjarvis.core.config.platform.system", return_value="Darwin")
@patch("openjarvis.core.config._run_cmd", return_value="")
def test_empty_profiler_output(self, mock_run, mock_system):
"""Empty system_profiler output returns None (no 'Apple' substring)."""
assert _detect_apple_gpu() is None
# ---------------------------------------------------------------------------
# Engine recommendation
# ---------------------------------------------------------------------------
class TestAppleEngineRecommendation:
"""Tests that Apple Silicon hardware maps to ollama."""
def test_m4_max_recommends_ollama(self):
hw = HardwareInfo(
platform="darwin",
cpu_brand="Apple M4 Max",
cpu_count=16,
ram_gb=128.0,
gpu=GpuInfo(vendor="apple", name="Apple M4 Max", vram_gb=128.0, count=1),
)
assert recommend_engine(hw) == "ollama"
def test_unified_memory(self):
"""On Apple Silicon, GPU VRAM equals system RAM (unified memory)."""
ram_gb = 192.0
gpu = GpuInfo(vendor="apple", name="Apple M4 Ultra", vram_gb=ram_gb, count=1)
hw = HardwareInfo(
platform="darwin",
cpu_brand="Apple M4 Ultra",
cpu_count=24,
ram_gb=ram_gb,
gpu=gpu,
)
assert hw.gpu.vram_gb == hw.ram_gb
assert recommend_engine(hw) == "ollama"
tests/hardware/test_hardware_profiles.py
+127
View File
@@ -0,0 +1,127 @@
"""Tests for hardware detection, GPU vendor identification,
and engine recommendation."""
from __future__ import annotations
from unittest.mock import patch
from openjarvis.core.config import (
GpuInfo,
_detect_amd_gpu,
_detect_apple_gpu,
_detect_nvidia_gpu,
recommend_engine,
)
# ---------------------------------------------------------------------------
# Hardware detection
# ---------------------------------------------------------------------------
class TestDetectHardware:
"""Tests for the top-level detect_hardware() function."""
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/nvidia-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="NVIDIA A100-SXM4-80GB, 81920, 1",
)
def test_detect_nvidia_gpu(self, mock_run, mock_which):
gpu = _detect_nvidia_gpu()
assert gpu is not None
assert gpu.vendor == "nvidia"
assert "A100" in gpu.name
assert gpu.vram_gb == 80.0
assert gpu.count == 1
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/rocm-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="AMD Instinct MI300X",
)
def test_detect_amd_gpu(self, mock_run, mock_which):
gpu = _detect_amd_gpu()
assert gpu is not None
assert gpu.vendor == "amd"
assert "MI300X" in gpu.name
@patch("openjarvis.core.config.platform.system", return_value="Darwin")
@patch(
"openjarvis.core.config._run_cmd",
return_value=(
"Graphics/Displays:\n"
" Apple M4 Max:\n"
" Chipset Model: Apple M4 Max\n"
" Type: GPU\n"
" Bus: Built-In\n"
),
)
def test_detect_apple_silicon(self, mock_run, mock_system):
gpu = _detect_apple_gpu()
assert gpu is not None
assert gpu.vendor == "apple"
assert "M4 Max" in gpu.name
@patch("openjarvis.core.config.shutil.which", return_value=None)
@patch("openjarvis.core.config.platform.system", return_value="Linux")
def test_detect_no_gpu(self, mock_system, mock_which):
"""All GPU detection methods return None when no GPU is present."""
assert _detect_nvidia_gpu() is None
assert _detect_amd_gpu() is None
assert _detect_apple_gpu() is None
# ---------------------------------------------------------------------------
# Engine recommendation
# ---------------------------------------------------------------------------
class TestRecommendEngine:
"""Tests for recommend_engine() logic across hardware profiles."""
def test_nvidia_datacenter_vllm(self, hardware_nvidia):
assert recommend_engine(hardware_nvidia) == "vllm"
def test_nvidia_consumer_ollama(self, hardware_nvidia_consumer):
assert recommend_engine(hardware_nvidia_consumer) == "ollama"
def test_amd_vllm(self, hardware_amd):
assert recommend_engine(hardware_amd) == "vllm"
def test_apple_ollama(self, hardware_apple):
assert recommend_engine(hardware_apple) == "ollama"
def test_cpu_only_llamacpp(self, hardware_cpu_only):
assert recommend_engine(hardware_cpu_only) == "llamacpp"
# ---------------------------------------------------------------------------
# Dataclass serialization / field access
# ---------------------------------------------------------------------------
class TestHardwareProfileSerialization:
"""Tests that dataclass fields are accessible and hold correct values."""
def test_gpu_info_fields(self, nvidia_gpu):
assert nvidia_gpu.vendor == "nvidia"
assert nvidia_gpu.name == "NVIDIA A100-SXM4-80GB"
assert nvidia_gpu.vram_gb == 80.0
assert nvidia_gpu.count == 1
def test_hardware_info_fields(self, hardware_nvidia):
assert hardware_nvidia.platform == "linux"
assert hardware_nvidia.cpu_brand == "AMD EPYC 7763"
assert hardware_nvidia.cpu_count == 64
assert hardware_nvidia.ram_gb == 512.0
assert hardware_nvidia.gpu is not None
def test_vram_sufficient_for_model(self):
"""Verify that VRAM capacity can be checked against a model requirement."""
gpu = GpuInfo(vendor="nvidia", name="NVIDIA A100", vram_gb=80.0, count=1)
model_requirement_gb = 40.0
assert gpu.vram_gb >= model_requirement_gb
small_gpu = GpuInfo(vendor="nvidia", name="RTX 3060", vram_gb=12.0, count=1)
large_model_gb = 70.0
assert small_gpu.vram_gb < large_model_gb
tests/hardware/test_nvidia.py
+150
View File
@@ -0,0 +1,150 @@
"""NVIDIA-specific hardware tests."""
from __future__ import annotations
from unittest.mock import patch
import pytest
from openjarvis.core.config import (
GpuInfo,
HardwareInfo,
_detect_nvidia_gpu,
recommend_engine,
)
pytestmark = pytest.mark.nvidia
# ---------------------------------------------------------------------------
# Detection / nvidia-smi parsing
# ---------------------------------------------------------------------------
class TestNVIDIADetection:
"""Tests for _detect_nvidia_gpu() against various nvidia-smi outputs."""
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/nvidia-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="NVIDIA A100-SXM4-80GB, 81920, 1",
)
def test_nvidia_smi_parsing(self, mock_run, mock_which):
gpu = _detect_nvidia_gpu()
assert gpu is not None
assert gpu.name == "NVIDIA A100-SXM4-80GB"
assert gpu.vram_gb == 80.0
assert gpu.count == 1
assert gpu.vendor == "nvidia"
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/nvidia-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value=(
"NVIDIA H100 80GB HBM3, 81920, 4\n"
"NVIDIA H100 80GB HBM3, 81920, 4\n"
"NVIDIA H100 80GB HBM3, 81920, 4\n"
"NVIDIA H100 80GB HBM3, 81920, 4"
),
)
def test_nvidia_smi_multi_gpu(self, mock_run, mock_which):
"""First line is parsed; count field captures GPU count."""
gpu = _detect_nvidia_gpu()
assert gpu is not None
assert gpu.count == 4
assert "H100" in gpu.name
@patch("openjarvis.core.config.shutil.which", return_value=None)
def test_nvidia_smi_not_found(self, mock_which):
assert _detect_nvidia_gpu() is None
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/nvidia-smi")
@patch("openjarvis.core.config._run_cmd", return_value="")
def test_nvidia_smi_error(self, mock_run, mock_which):
"""Empty output from nvidia-smi returns None."""
assert _detect_nvidia_gpu() is None
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/nvidia-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="NVIDIA GeForce RTX 4090, 24564, 1",
)
def test_vram_detection(self, mock_run, mock_which):
gpu = _detect_nvidia_gpu()
assert gpu is not None
# 24564 MB -> ~24.0 GB
assert gpu.vram_gb == pytest.approx(24.0, abs=0.1)
@patch("openjarvis.core.config.shutil.which", return_value="/usr/bin/nvidia-smi")
@patch(
"openjarvis.core.config._run_cmd",
return_value="NVIDIA A100-SXM4-80GB, 81920, 1",
)
def test_compute_capability(self, mock_run, mock_which):
"""compute_capability defaults to empty string when not parsed."""
gpu = _detect_nvidia_gpu()
assert gpu is not None
assert gpu.compute_capability == ""
# ---------------------------------------------------------------------------
# Engine recommendation
# ---------------------------------------------------------------------------
class TestNVIDIAEngineRecommendation:
"""Tests that NVIDIA cards map to the correct inference engine."""
def test_a100_recommends_vllm(self):
hw = HardwareInfo(
platform="linux",
cpu_brand="EPYC",
cpu_count=64,
ram_gb=512.0,
gpu=GpuInfo(
vendor="nvidia",
name="NVIDIA A100-SXM4-80GB",
vram_gb=80.0, count=1,
),
)
assert recommend_engine(hw) == "vllm"
def test_h100_recommends_vllm(self):
hw = HardwareInfo(
platform="linux",
cpu_brand="EPYC",
cpu_count=64,
ram_gb=512.0,
gpu=GpuInfo(
vendor="nvidia",
name="NVIDIA H100 80GB HBM3",
vram_gb=80.0, count=1,
),
)
assert recommend_engine(hw) == "vllm"
def test_rtx_4090_recommends_ollama(self):
hw = HardwareInfo(
platform="linux",
cpu_brand="i9-14900K",
cpu_count=24,
ram_gb=64.0,
gpu=GpuInfo(
vendor="nvidia",
name="NVIDIA GeForce RTX 4090",
vram_gb=24.0, count=1,
),
)
assert recommend_engine(hw) == "ollama"
def test_multi_gpu_config(self):
"""Multi-GPU datacenter setup still recommends vllm."""
hw = HardwareInfo(
platform="linux",
cpu_brand="EPYC",
cpu_count=128,
ram_gb=1024.0,
gpu=GpuInfo(vendor="nvidia", name="NVIDIA H100", vram_gb=80.0, count=8),
)
assert recommend_engine(hw) == "vllm"
assert hw.gpu.count == 8
tests/intelligence/test_model_catalog_extended.py
+239
View File
@@ -0,0 +1,239 @@
"""Tests for the extended model catalog -- all new ModelSpec entries."""
from __future__ import annotations
import pytest
from openjarvis.core.registry import ModelRegistry
from openjarvis.core.types import ModelSpec
from openjarvis.intelligence.model_catalog import (
BUILTIN_MODELS,
merge_discovered_models,
register_builtin_models,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _get_spec(model_id: str) -> ModelSpec:
"""Lookup a spec from the BUILTIN_MODELS list (not registry)."""
for spec in BUILTIN_MODELS:
if spec.model_id == model_id:
return spec
raise KeyError(model_id)
# ---------------------------------------------------------------------------
# Local models
# ---------------------------------------------------------------------------
class TestLocalModelSpecs:
"""Verify every new local model has correct fields."""
def test_gpt_oss_120b_model_id(self) -> None:
spec = _get_spec("gpt-oss:120b")
assert spec.model_id == "gpt-oss:120b"
def test_gpt_oss_120b_params(self) -> None:
spec = _get_spec("gpt-oss:120b")
assert spec.parameter_count_b == 117.0
assert spec.active_parameter_count_b == 5.1
def test_gpt_oss_120b_context(self) -> None:
spec = _get_spec("gpt-oss:120b")
assert spec.context_length == 131072
def test_gpt_oss_120b_engines(self) -> None:
spec = _get_spec("gpt-oss:120b")
assert "vllm" in spec.supported_engines
assert "ollama" in spec.supported_engines
def test_gpt_oss_120b_architecture(self) -> None:
spec = _get_spec("gpt-oss:120b")
assert spec.metadata["architecture"] == "moe"
def test_qwen3_8b_model_id(self) -> None:
spec = _get_spec("qwen3:8b")
assert spec.model_id == "qwen3:8b"
def test_qwen3_8b_params(self) -> None:
spec = _get_spec("qwen3:8b")
assert spec.parameter_count_b == 8.2
def test_qwen3_8b_context(self) -> None:
spec = _get_spec("qwen3:8b")
assert spec.context_length == 32768
def test_qwen3_8b_engines(self) -> None:
spec = _get_spec("qwen3:8b")
for e in ("vllm", "ollama", "llamacpp", "sglang"):
assert e in spec.supported_engines
def test_qwen3_8b_architecture(self) -> None:
spec = _get_spec("qwen3:8b")
assert spec.metadata["architecture"] == "dense"
def test_glm_47_flash_params(self) -> None:
spec = _get_spec("glm-4.7-flash")
assert spec.parameter_count_b == 30.0
assert spec.active_parameter_count_b == 3.0
def test_glm_47_flash_context(self) -> None:
spec = _get_spec("glm-4.7-flash")
assert spec.context_length == 131072
def test_glm_47_flash_engines(self) -> None:
spec = _get_spec("glm-4.7-flash")
assert "vllm" in spec.supported_engines
assert "sglang" in spec.supported_engines
def test_trinity_mini_params(self) -> None:
spec = _get_spec("trinity-mini")
assert spec.parameter_count_b == 26.0
assert spec.active_parameter_count_b == 3.0
def test_trinity_mini_context(self) -> None:
spec = _get_spec("trinity-mini")
assert spec.context_length == 128000
def test_trinity_mini_engines(self) -> None:
spec = _get_spec("trinity-mini")
assert "vllm" in spec.supported_engines
assert "llamacpp" in spec.supported_engines
# ---------------------------------------------------------------------------
# Cloud models
# ---------------------------------------------------------------------------
class TestCloudModelSpecs:
"""Verify every new cloud model has correct fields."""
def test_gpt_5_mini_provider(self) -> None:
spec = _get_spec("gpt-5-mini")
assert spec.provider == "openai"
assert spec.requires_api_key is True
def test_gpt_5_mini_context(self) -> None:
spec = _get_spec("gpt-5-mini")
assert spec.context_length == 400000
def test_gpt_5_mini_pricing(self) -> None:
spec = _get_spec("gpt-5-mini")
assert spec.metadata["pricing_input"] == pytest.approx(0.25)
assert spec.metadata["pricing_output"] == pytest.approx(2.00)
def test_claude_opus_4_6_provider(self) -> None:
spec = _get_spec("claude-opus-4-6")
assert spec.provider == "anthropic"
assert spec.requires_api_key is True
def test_claude_opus_4_6_context(self) -> None:
spec = _get_spec("claude-opus-4-6")
assert spec.context_length == 200000
def test_claude_sonnet_4_6_provider(self) -> None:
spec = _get_spec("claude-sonnet-4-6")
assert spec.provider == "anthropic"
assert spec.requires_api_key is True
def test_claude_haiku_4_5_provider(self) -> None:
spec = _get_spec("claude-haiku-4-5")
assert spec.provider == "anthropic"
assert spec.requires_api_key is True
def test_claude_haiku_4_5_pricing(self) -> None:
spec = _get_spec("claude-haiku-4-5")
assert spec.metadata["pricing_input"] == pytest.approx(1.00)
assert spec.metadata["pricing_output"] == pytest.approx(5.00)
def test_gemini_25_pro_provider(self) -> None:
spec = _get_spec("gemini-2.5-pro")
assert spec.provider == "google"
assert spec.requires_api_key is True
assert spec.context_length == 1000000
def test_gemini_25_flash_provider(self) -> None:
spec = _get_spec("gemini-2.5-flash")
assert spec.provider == "google"
assert spec.context_length == 1000000
def test_gemini_3_pro_provider(self) -> None:
spec = _get_spec("gemini-3-pro")
assert spec.provider == "google"
assert spec.requires_api_key is True
def test_gemini_3_flash_provider(self) -> None:
spec = _get_spec("gemini-3-flash")
assert spec.provider == "google"
assert spec.requires_api_key is True
def test_gemini_3_pro_pricing(self) -> None:
spec = _get_spec("gemini-3-pro")
assert spec.metadata["pricing_input"] == pytest.approx(2.00)
assert spec.metadata["pricing_output"] == pytest.approx(12.00)
# ---------------------------------------------------------------------------
# Discovery / invariant tests
# ---------------------------------------------------------------------------
class TestModelDiscovery:
def test_local_models_have_engine_compat(self) -> None:
"""Every local model has at least one supported engine."""
for spec in BUILTIN_MODELS:
if not spec.requires_api_key:
assert len(spec.supported_engines) >= 1, (
f"{spec.model_id} has no supported engines"
)
def test_cloud_models_require_api_key(self) -> None:
"""All cloud models have requires_api_key=True."""
cloud_ids = {
"gpt-4o", "gpt-4o-mini", "gpt-5-mini",
"claude-sonnet-4-20250514", "claude-opus-4-20250514",
"claude-opus-4-6", "claude-sonnet-4-6", "claude-haiku-4-5",
"gemini-2.5-pro", "gemini-2.5-flash", "gemini-3-pro", "gemini-3-flash",
}
for spec in BUILTIN_MODELS:
if spec.model_id in cloud_ids:
assert spec.requires_api_key is True, (
f"{spec.model_id} should require API key"
)
def test_moe_models_have_active_params(self) -> None:
"""MoE models have active_parameter_count_b set."""
moe_ids = {"gpt-oss:120b", "glm-4.7-flash", "trinity-mini"}
for spec in BUILTIN_MODELS:
if spec.model_id in moe_ids:
assert spec.active_parameter_count_b is not None
assert spec.active_parameter_count_b > 0
def test_all_models_have_context_length(self) -> None:
"""No model has zero or None context length."""
for spec in BUILTIN_MODELS:
assert spec.context_length > 0, (
f"{spec.model_id} has context_length={spec.context_length}"
)
def test_merge_discovered_preserves_new(self) -> None:
"""merge_discovered_models works for all new model IDs."""
register_builtin_models()
new_ids = [
"gpt-oss:120b", "glm-4.7-flash", "trinity-mini",
"gpt-5-mini", "claude-opus-4-6", "gemini-3-pro",
]
# Merging known IDs should not raise
merge_discovered_models("vllm", new_ids)
for mid in new_ids:
assert ModelRegistry.contains(mid)
def test_register_builtin_is_idempotent(self) -> None:
"""Calling register_builtin_models twice does not raise."""
register_builtin_models()
register_builtin_models() # should not raise
assert ModelRegistry.contains("qwen3:8b")
tests/intelligence/test_routing_models.py
+178
View File
@@ -0,0 +1,178 @@
"""Tests for router behavior with the extended model catalog."""
from __future__ import annotations
import pytest
from openjarvis.intelligence.model_catalog import register_builtin_models
from openjarvis.intelligence.router import HeuristicRouter, build_routing_context
from openjarvis.learning._stubs import RoutingContext
# New local model keys for testing
NEW_LOCAL_MODELS = [
"gpt-oss:120b", # 117B total, 5.1B active, MoE
"qwen3:8b", # 8.2B, dense
"glm-4.7-flash", # 30B total, 3.0B active, MoE
"trinity-mini", # 26B total, 3.0B active, MoE
]
# Cloud model keys
CLOUD_MODELS = [
"gpt-5-mini",
"claude-opus-4-6",
"claude-sonnet-4-6",
"claude-haiku-4-5",
"gemini-2.5-pro",
"gemini-3-flash",
]
def _setup_models() -> None:
"""Register builtin models needed for the tests."""
register_builtin_models()
class TestRouterWithNewModels:
"""Router behavior when using the new local models."""
def test_short_query_routes_to_smallest(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = RoutingContext(query="hi", query_length=2)
selected = router.select_model(ctx)
# qwen3:8b is 8.2B -- the smallest by parameter_count_b
assert selected == "qwen3:8b"
def test_code_query_routes_to_largest(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = RoutingContext(
query="def merge_sort(arr):", query_length=22, has_code=True
)
selected = router.select_model(ctx)
# No model with "code"/"coder" in name, falls to largest → gpt-oss:120b (117B)
assert selected == "gpt-oss:120b"
def test_code_query_with_coder_available(self) -> None:
_setup_models()
models = NEW_LOCAL_MODELS + ["deepseek-coder-v2:16b"]
router = HeuristicRouter(available_models=models)
ctx = RoutingContext(
query="import numpy as np", query_length=18, has_code=True
)
selected = router.select_model(ctx)
assert selected == "deepseek-coder-v2:16b"
def test_math_query_routes_to_largest(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = RoutingContext(
query="solve the integral of x^2 dx", query_length=29, has_math=True
)
selected = router.select_model(ctx)
assert selected == "gpt-oss:120b"
def test_long_context_routes_to_largest(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = RoutingContext(query="x" * 501, query_length=501)
selected = router.select_model(ctx)
assert selected == "gpt-oss:120b"
def test_high_urgency_routes_to_smallest(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = RoutingContext(
query="solve the integral of x^2", query_length=25,
has_math=True, urgency=0.9,
)
selected = router.select_model(ctx)
# High urgency overrides everything → smallest params → qwen3:8b (8.2B)
assert selected == "qwen3:8b"
def test_reasoning_query_routes_to_largest(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = build_routing_context(
"Please explain step by step how neural networks learn"
)
selected = router.select_model(ctx)
assert selected == "gpt-oss:120b"
def test_medium_query_uses_default(self) -> None:
_setup_models()
router = HeuristicRouter(
available_models=NEW_LOCAL_MODELS,
default_model="glm-4.7-flash",
)
# Medium length query, no code/math/reasoning → rule 6 default
ctx = RoutingContext(query="Tell me about the weather today", query_length=60)
selected = router.select_model(ctx)
assert selected == "glm-4.7-flash"
class TestRouterCloudFallback:
"""Router behavior when only cloud models are available."""
def test_no_local_falls_to_cloud(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=CLOUD_MODELS)
ctx = RoutingContext(query="hi", query_length=2)
selected = router.select_model(ctx)
# Should return one of the cloud models (smallest params)
assert selected in CLOUD_MODELS
def test_cloud_model_selection_with_math(self) -> None:
_setup_models()
router = HeuristicRouter(available_models=CLOUD_MODELS)
ctx = RoutingContext(query="solve x", query_length=7, has_math=True)
selected = router.select_model(ctx)
# All cloud models have parameter_count_b=0 so falls to first
assert selected in CLOUD_MODELS
def test_empty_models_returns_fallback(self) -> None:
router = HeuristicRouter(
available_models=[],
fallback_model="gpt-5-mini",
)
ctx = RoutingContext(query="hello", query_length=5)
assert router.select_model(ctx) == "gpt-5-mini"
class TestRouterParameterized:
"""Parametrized cross-product tests for model/query combinations."""
@pytest.mark.parametrize("query,expected_is_largest", [
("hi", False),
("solve the integral of sin(x)", True),
("def foo(): pass", True), # code → largest when no coder
("x" * 501, True),
])
def test_query_type_selects_expected_size(
self, query: str, expected_is_largest: bool,
) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = build_routing_context(query)
selected = router.select_model(ctx)
if expected_is_largest:
assert selected == "gpt-oss:120b"
else:
assert selected == "qwen3:8b"
@pytest.mark.parametrize("model_id", NEW_LOCAL_MODELS)
def test_single_model_always_returns_it(self, model_id: str) -> None:
_setup_models()
router = HeuristicRouter(available_models=[model_id])
ctx = RoutingContext(query="hello world", query_length=11)
assert router.select_model(ctx) == model_id
@pytest.mark.parametrize("urgency", [0.85, 0.9, 1.0])
def test_high_urgency_always_smallest(self, urgency: float) -> None:
_setup_models()
router = HeuristicRouter(available_models=NEW_LOCAL_MODELS)
ctx = RoutingContext(
query="complex reasoning task", query_length=23,
has_math=True, urgency=urgency,
)
assert router.select_model(ctx) == "qwen3:8b"
tests/learning/test_trace_policy.py
+203
View File
@@ -0,0 +1,203 @@
"""Tests for the trace-driven router policy."""
from __future__ import annotations
import time
from pathlib import Path
from openjarvis.core.types import StepType, Trace, TraceStep
from openjarvis.learning._stubs import RoutingContext
from openjarvis.learning.trace_policy import TraceDrivenPolicy, classify_query
from openjarvis.traces.analyzer import TraceAnalyzer
from openjarvis.traces.store import TraceStore
def _make_trace(
query: str = "test",
model: str = "qwen3:8b",
outcome: str | None = "success",
feedback: float | None = 0.8,
) -> Trace:
now = time.time()
return Trace(
query=query,
agent="orchestrator",
model=model,
engine="ollama",
result="result",
outcome=outcome,
feedback=feedback,
started_at=now,
ended_at=now + 0.5,
total_tokens=100,
total_latency_seconds=0.5,
steps=[
TraceStep(
step_type=StepType.GENERATE,
timestamp=now,
duration_seconds=0.5,
output={"tokens": 100},
),
],
)
class TestClassifyQuery:
def test_code(self) -> None:
assert classify_query("def hello(): pass") == "code"
assert classify_query("```python\nprint()```") == "code"
def test_math(self) -> None:
assert classify_query("solve this equation for x") == "math"
assert classify_query("compute the integral") == "math"
def test_short(self) -> None:
assert classify_query("hello") == "short"
assert classify_query("what time is it?") == "short"
def test_long(self) -> None:
assert classify_query("a" * 501) == "long"
def test_general(self) -> None:
q = "Tell me about the history of artificial intelligence research"
assert classify_query(q) == "general"
class TestTraceDrivenPolicy:
def test_fallback_no_traces(self) -> None:
policy = TraceDrivenPolicy(default_model="qwen3:8b")
ctx = RoutingContext(query="hello")
assert policy.select_model(ctx) == "qwen3:8b"
def test_fallback_chain(self) -> None:
policy = TraceDrivenPolicy(
default_model="missing",
fallback_model="llama3:8b",
available_models=["llama3:8b"],
)
ctx = RoutingContext(query="hello")
assert policy.select_model(ctx) == "llama3:8b"
def test_fallback_first_available(self) -> None:
policy = TraceDrivenPolicy(available_models=["modelA", "modelB"])
ctx = RoutingContext(query="hello")
assert policy.select_model(ctx) == "modelA"
def test_update_from_traces(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
# Create traces: code queries succeed more with codestral
for _ in range(6):
store.save(_make_trace(
query="def foo(): pass",
model="codestral",
outcome="success",
feedback=0.9,
))
for _ in range(6):
store.save(_make_trace(
query="def bar(): return 1",
model="qwen3:8b",
outcome="failure",
feedback=0.3,
))
analyzer = TraceAnalyzer(store)
policy = TraceDrivenPolicy(
analyzer=analyzer,
default_model="qwen3:8b",
)
policy.min_samples = 3
result = policy.update_from_traces()
assert result["updated"] is True
# Policy should now route code to codestral
ctx = RoutingContext(query="import os; def main(): pass")
assert policy.select_model(ctx) == "codestral"
store.close()
def test_policy_map_readable(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
for _ in range(5):
store.save(_make_trace(
query="hello", model="small-model",
outcome="success",
))
analyzer = TraceAnalyzer(store)
policy = TraceDrivenPolicy(analyzer=analyzer, default_model="default")
policy.min_samples = 3
policy.update_from_traces()
pmap = policy.policy_map
assert isinstance(pmap, dict)
assert "short" in pmap
assert pmap["short"] == "small-model"
store.close()
def test_respects_min_samples(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
# Only 2 traces — below threshold
store.save(_make_trace(query="hello", model="small", outcome="success"))
store.save(_make_trace(query="hi", model="small", outcome="success"))
analyzer = TraceAnalyzer(store)
policy = TraceDrivenPolicy(
analyzer=analyzer,
default_model="default",
)
policy.min_samples = 5
policy.update_from_traces()
ctx = RoutingContext(query="hey")
# Should fallback since not enough confidence
assert policy.select_model(ctx) == "default"
store.close()
def test_respects_available_models(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
for _ in range(10):
store.save(_make_trace(
query="hello", model="unavailable-model",
outcome="success",
))
analyzer = TraceAnalyzer(store)
policy = TraceDrivenPolicy(
analyzer=analyzer,
available_models=["qwen3:8b", "llama3:8b"],
default_model="qwen3:8b",
)
policy.min_samples = 3
policy.update_from_traces()
ctx = RoutingContext(query="hey")
# Should fallback since learned model not in available_models
assert policy.select_model(ctx) == "qwen3:8b"
store.close()
def test_observe_online(self) -> None:
policy = TraceDrivenPolicy(default_model="default")
policy.min_samples = 3
# First observation creates the entry
policy.observe("hello", "fast-model", "success", 0.9)
assert policy.policy_map.get("short") == "fast-model"
# Not enough samples yet for high confidence
ctx = RoutingContext(query="hi")
# Confidence is 1 < min_samples=3, so fallback
assert policy.select_model(ctx) == "default"
def test_update_no_analyzer(self) -> None:
policy = TraceDrivenPolicy()
result = policy.update_from_traces()
assert result["error"] == "no analyzer configured"
def test_update_empty_store(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
analyzer = TraceAnalyzer(store)
policy = TraceDrivenPolicy(analyzer=analyzer)
result = policy.update_from_traces()
assert result["updated"] is False
store.close()
tests/mcp/__init__.py
View File
tests/mcp/test_client.py
+102
View File
@@ -0,0 +1,102 @@
"""Tests for the MCP client."""
from __future__ import annotations
import pytest
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.protocol import MCPError
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import InProcessTransport
from openjarvis.tools._stubs import ToolSpec
from openjarvis.tools.calculator import CalculatorTool
from openjarvis.tools.think import ThinkTool
@pytest.fixture
def client():
"""MCP client connected via in-process transport."""
server = MCPServer([CalculatorTool(), ThinkTool()])
transport = InProcessTransport(server)
return MCPClient(transport)
class TestMCPClient:
def test_initialize_handshake(self, client):
result = client.initialize()
assert "protocolVersion" in result
assert "serverInfo" in result
assert result["serverInfo"]["name"] == "openjarvis"
assert client._initialized is True
def test_initialize_sets_capabilities(self, client):
client.initialize()
assert "tools" in client._capabilities
def test_list_tools(self, client):
tools = client.list_tools()
assert len(tools) == 2
assert all(isinstance(t, ToolSpec) for t in tools)
names = {t.name for t in tools}
assert "calculator" in names
assert "think" in names
def test_list_tools_have_descriptions(self, client):
tools = client.list_tools()
for t in tools:
assert t.description # non-empty
def test_list_tools_have_parameters(self, client):
tools = client.list_tools()
for t in tools:
assert "properties" in t.parameters
def test_call_tool_calculator(self, client):
result = client.call_tool("calculator", {"expression": "10 + 5"})
assert result["isError"] is False
assert "15" in result["content"][0]["text"]
def test_call_tool_think(self, client):
result = client.call_tool("think", {"thought": "Reasoning step."})
assert result["isError"] is False
assert "Reasoning step." in result["content"][0]["text"]
def test_call_tool_error(self, client):
result = client.call_tool("calculator", {"expression": "1/0"})
assert result["isError"] is True
assert "division by zero" in result["content"][0]["text"]
def test_call_unknown_tool_raises(self, client):
with pytest.raises(MCPError) as exc_info:
client.call_tool("nonexistent", {})
assert "Unknown tool" in str(exc_info.value)
def test_client_server_roundtrip(self, client):
"""Full lifecycle: initialize -> list -> call -> close."""
info = client.initialize()
assert "serverInfo" in info
tools = client.list_tools()
assert len(tools) >= 1
result = client.call_tool("calculator", {"expression": "7 * 8"})
assert "56" in result["content"][0]["text"]
client.close()
def test_close(self, client):
client.close()
# Close should not raise even if called multiple times
client.close()
def test_incremental_ids(self, client):
"""Each request should get a unique ID."""
id1 = client._next_id()
id2 = client._next_id()
assert id2 > id1
def test_call_tool_with_no_arguments(self, client):
"""Calling a tool with no arguments passes empty dict."""
result = client.call_tool("think")
# Think tool echoes empty thought
assert result["isError"] is False
tests/mcp/test_mcp_tools_matrix.py
+82
View File
@@ -0,0 +1,82 @@
"""Matrix tests — verify every built-in tool is discoverable and callable via MCP."""
from __future__ import annotations
import pytest
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import InProcessTransport
from openjarvis.tools.calculator import CalculatorTool
from openjarvis.tools.think import ThinkTool
# Tool configs: (tool_class, call_arguments, expected_substring)
_TOOL_CONFIGS = {
"calculator": (CalculatorTool, {"expression": "2+2"}, "4"),
"think": (ThinkTool, {"thought": "test thought"}, "test thought"),
}
def _make_client(tool_classes):
"""Create an MCP client with the given tool instances."""
tools = [cls() for cls in tool_classes]
server = MCPServer(tools)
transport = InProcessTransport(server)
return MCPClient(transport)
@pytest.mark.parametrize("tool_name", ["calculator", "think"])
class TestMCPToolsMatrix:
def test_tool_discoverable_via_mcp(self, tool_name):
tool_cls, _, _ = _TOOL_CONFIGS[tool_name]
client = _make_client([tool_cls])
tools = client.list_tools()
names = [t.name for t in tools]
assert tool_name in names
def test_tool_callable_via_mcp(self, tool_name):
tool_cls, arguments, expected = _TOOL_CONFIGS[tool_name]
client = _make_client([tool_cls])
result = client.call_tool(tool_name, arguments)
assert result["isError"] is False
assert expected in result["content"][0]["text"]
def test_tool_has_description(self, tool_name):
tool_cls, _, _ = _TOOL_CONFIGS[tool_name]
client = _make_client([tool_cls])
tools = client.list_tools()
tool_spec = next(t for t in tools if t.name == tool_name)
assert tool_spec.description # non-empty
def test_tool_has_input_schema(self, tool_name):
tool_cls, _, _ = _TOOL_CONFIGS[tool_name]
client = _make_client([tool_cls])
tools = client.list_tools()
tool_spec = next(t for t in tools if t.name == tool_name)
assert "properties" in tool_spec.parameters
def test_tool_result_format(self, tool_name):
tool_cls, arguments, _ = _TOOL_CONFIGS[tool_name]
client = _make_client([tool_cls])
result = client.call_tool(tool_name, arguments)
assert "content" in result
assert "isError" in result
assert isinstance(result["content"], list)
assert result["content"][0]["type"] == "text"
def test_tool_error_handling(self, tool_name):
"""Calling with bad args should not crash the server."""
tool_cls, _, _ = _TOOL_CONFIGS[tool_name]
client = _make_client([tool_cls])
# Call with empty arguments
result = client.call_tool(tool_name, {})
# Should return a response (not crash)
assert "content" in result
def test_tool_in_full_server(self, tool_name):
"""Tool should be discoverable when all tools are registered."""
all_classes = [cls for cls, _, _ in _TOOL_CONFIGS.values()]
client = _make_client(all_classes)
tools = client.list_tools()
names = [t.name for t in tools]
assert tool_name in names
tests/mcp/test_protocol.py
+164
View File
@@ -0,0 +1,164 @@
"""Tests for MCP protocol message types."""
from __future__ import annotations
import json
from openjarvis.mcp.protocol import (
INTERNAL_ERROR,
INVALID_PARAMS,
INVALID_REQUEST,
METHOD_NOT_FOUND,
PARSE_ERROR,
MCPError,
MCPNotification,
MCPRequest,
MCPResponse,
)
class TestMCPRequest:
def test_serialize_deserialize(self):
req = MCPRequest(method="tools/list", params={"cursor": None}, id=1)
data = req.to_json()
restored = MCPRequest.from_json(data)
assert restored.method == "tools/list"
assert restored.id == 1
assert restored.jsonrpc == "2.0"
def test_initialize_request(self):
req = MCPRequest(method="initialize", params={}, id=1)
parsed = json.loads(req.to_json())
assert parsed["method"] == "initialize"
assert parsed["jsonrpc"] == "2.0"
assert parsed["id"] == 1
def test_tools_list_request(self):
req = MCPRequest(method="tools/list", id=2)
parsed = json.loads(req.to_json())
assert parsed["method"] == "tools/list"
def test_tools_call_request(self):
req = MCPRequest(
method="tools/call",
params={"name": "calculator", "arguments": {"expression": "2+2"}},
id=3,
)
parsed = json.loads(req.to_json())
assert parsed["method"] == "tools/call"
assert parsed["params"]["name"] == "calculator"
assert parsed["params"]["arguments"]["expression"] == "2+2"
def test_default_params_empty(self):
req = MCPRequest(method="test")
assert req.params == {}
def test_string_id(self):
req = MCPRequest(method="test", id="abc-123")
data = req.to_json()
restored = MCPRequest.from_json(data)
assert restored.id == "abc-123"
def test_from_json_missing_params(self):
raw = json.dumps({"jsonrpc": "2.0", "method": "test", "id": 1})
req = MCPRequest.from_json(raw)
assert req.params == {}
class TestMCPResponse:
def test_serialize_deserialize_success(self):
resp = MCPResponse(result={"tools": []}, id=1)
data = resp.to_json()
restored = MCPResponse.from_json(data)
assert restored.result == {"tools": []}
assert restored.error is None
assert restored.id == 1
def test_serialize_deserialize_error(self):
resp = MCPResponse.error_response(1, METHOD_NOT_FOUND, "Not found")
data = resp.to_json()
restored = MCPResponse.from_json(data)
assert restored.error is not None
assert restored.error["code"] == METHOD_NOT_FOUND
assert restored.error["message"] == "Not found"
def test_error_response_factory(self):
resp = MCPResponse.error_response(42, INVALID_PARAMS, "Bad params")
assert resp.error["code"] == INVALID_PARAMS
assert resp.error["message"] == "Bad params"
assert resp.id == 42
assert resp.result is None
def test_error_response_with_data(self):
resp = MCPResponse.error_response(
1, INTERNAL_ERROR, "Oops", data={"detail": "stack"},
)
assert resp.error["data"] == {"detail": "stack"}
def test_success_response(self):
resp = MCPResponse(result={"value": 42}, id=5)
parsed = json.loads(resp.to_json())
assert "result" in parsed
assert "error" not in parsed
assert parsed["result"]["value"] == 42
def test_error_excludes_result(self):
resp = MCPResponse.error_response(1, PARSE_ERROR, "Parse error")
parsed = json.loads(resp.to_json())
assert "error" in parsed
assert "result" not in parsed
def test_jsonrpc_version(self):
resp = MCPResponse(result={}, id=1)
parsed = json.loads(resp.to_json())
assert parsed["jsonrpc"] == "2.0"
class TestMCPNotification:
def test_format(self):
notif = MCPNotification(method="notifications/initialized", params={})
parsed = json.loads(notif.to_json())
assert parsed["method"] == "notifications/initialized"
assert parsed["jsonrpc"] == "2.0"
def test_no_id_field(self):
notif = MCPNotification(method="test")
parsed = json.loads(notif.to_json())
assert "id" not in parsed
def test_with_params(self):
notif = MCPNotification(method="progress", params={"percent": 50})
parsed = json.loads(notif.to_json())
assert parsed["params"]["percent"] == 50
class TestMCPError:
def test_error_is_exception(self):
err = MCPError(code=METHOD_NOT_FOUND, message="Not found")
assert isinstance(err, Exception)
def test_error_str(self):
err = MCPError(code=-32601, message="Not found")
assert "-32601" in str(err)
assert "Not found" in str(err)
def test_error_with_data(self):
err = MCPError(code=INTERNAL_ERROR, message="Oops", data={"trace": "..."})
assert err.data == {"trace": "..."}
class TestErrorCodes:
def test_parse_error(self):
assert PARSE_ERROR == -32700
def test_invalid_request(self):
assert INVALID_REQUEST == -32600
def test_method_not_found(self):
assert METHOD_NOT_FOUND == -32601
def test_invalid_params(self):
assert INVALID_PARAMS == -32602
def test_internal_error(self):
assert INTERNAL_ERROR == -32603
tests/mcp/test_server.py
+194
View File
@@ -0,0 +1,194 @@
"""Tests for the MCP server."""
from __future__ import annotations
import pytest
from openjarvis.mcp.protocol import (
INVALID_PARAMS,
METHOD_NOT_FOUND,
MCPRequest,
)
from openjarvis.mcp.server import MCPServer
from openjarvis.tools.calculator import CalculatorTool
from openjarvis.tools.think import ThinkTool
@pytest.fixture
def server():
"""Create an MCP server with calculator and think tools."""
return MCPServer([CalculatorTool(), ThinkTool()])
class TestMCPServer:
def test_initialize(self, server):
req = MCPRequest(method="initialize", id=1)
resp = server.handle(req)
assert resp.error is None
assert resp.id == 1
result = resp.result
assert "protocolVersion" in result
assert "capabilities" in result
assert "serverInfo" in result
assert result["serverInfo"]["name"] == "openjarvis"
def test_initialize_capabilities(self, server):
req = MCPRequest(method="initialize", id=1)
resp = server.handle(req)
caps = resp.result["capabilities"]
assert "tools" in caps
def test_tools_list_all(self, server):
req = MCPRequest(method="tools/list", id=2)
resp = server.handle(req)
assert resp.error is None
tools = resp.result["tools"]
assert len(tools) == 2
names = {t["name"] for t in tools}
assert "calculator" in names
assert "think" in names
def test_tools_list_includes_spec(self, server):
req = MCPRequest(method="tools/list", id=2)
resp = server.handle(req)
tools = resp.result["tools"]
for tool in tools:
assert "name" in tool
assert "description" in tool
assert "inputSchema" in tool
def test_tools_list_input_schema_has_properties(self, server):
req = MCPRequest(method="tools/list", id=2)
resp = server.handle(req)
tools = resp.result["tools"]
for tool in tools:
schema = tool["inputSchema"]
assert "properties" in schema
def test_tools_call_calculator(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "calculator", "arguments": {"expression": "2+2"}},
id=3,
)
resp = server.handle(req)
assert resp.error is None
assert resp.result["isError"] is False
content = resp.result["content"]
assert len(content) == 1
assert content[0]["type"] == "text"
assert "4" in content[0]["text"]
def test_tools_call_calculator_complex(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "calculator", "arguments": {"expression": "3 * (4 + 5)"}},
id=4,
)
resp = server.handle(req)
assert resp.error is None
assert "27" in resp.result["content"][0]["text"]
def test_tools_call_think(self, server):
req = MCPRequest(
method="tools/call",
params={
"name": "think",
"arguments": {"thought": "Step 1: analyze the problem."},
},
id=5,
)
resp = server.handle(req)
assert resp.error is None
assert resp.result["isError"] is False
assert "Step 1" in resp.result["content"][0]["text"]
def test_tools_call_unknown_tool(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "nonexistent", "arguments": {}},
id=6,
)
resp = server.handle(req)
assert resp.error is not None
assert resp.error["code"] == INVALID_PARAMS
assert "Unknown tool" in resp.error["message"]
def test_tools_call_missing_name(self, server):
req = MCPRequest(
method="tools/call",
params={"arguments": {}},
id=7,
)
resp = server.handle(req)
assert resp.error is not None
assert resp.error["code"] == INVALID_PARAMS
assert "name" in resp.error["message"]
def test_tools_call_invalid_arguments(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "calculator", "arguments": {}},
id=8,
)
resp = server.handle(req)
# Calculator returns failure but not a protocol error
assert resp.error is None
assert resp.result["isError"] is True
def test_result_format_mcp_compliant(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "think", "arguments": {"thought": "hello"}},
id=9,
)
resp = server.handle(req)
result = resp.result
assert "content" in result
assert "isError" in result
assert isinstance(result["content"], list)
assert result["content"][0]["type"] == "text"
def test_unknown_method(self, server):
req = MCPRequest(method="resources/list", id=10)
resp = server.handle(req)
assert resp.error is not None
assert resp.error["code"] == METHOD_NOT_FOUND
assert "Unknown method" in resp.error["message"]
def test_response_preserves_request_id(self, server):
req = MCPRequest(method="tools/list", id=42)
resp = server.handle(req)
assert resp.id == 42
def test_response_preserves_string_id(self, server):
req = MCPRequest(method="tools/list", id="req-abc")
resp = server.handle(req)
assert resp.id == "req-abc"
def test_empty_tools_server(self):
empty_server = MCPServer([])
req = MCPRequest(method="tools/list", id=1)
resp = empty_server.handle(req)
assert resp.result["tools"] == []
def test_tools_call_with_no_arguments_key(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "calculator"},
id=11,
)
resp = server.handle(req)
# Should still execute (with empty arguments)
assert resp.error is None
def test_calculator_error_returns_is_error_true(self, server):
req = MCPRequest(
method="tools/call",
params={"name": "calculator", "arguments": {"expression": "1/0"}},
id=12,
)
resp = server.handle(req)
assert resp.error is None
assert resp.result["isError"] is True
assert "division by zero" in resp.result["content"][0]["text"]
tests/mcp/test_transport.py
+214
View File
@@ -0,0 +1,214 @@
"""Tests for MCP transport implementations."""
from __future__ import annotations
import json
import sys
import textwrap
from unittest.mock import MagicMock
import pytest
from openjarvis.mcp.protocol import MCPRequest
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import InProcessTransport, SSETransport, StdioTransport
from openjarvis.tools.calculator import CalculatorTool
from openjarvis.tools.think import ThinkTool
@pytest.fixture
def server():
"""MCP server with calculator and think tools."""
return MCPServer([CalculatorTool(), ThinkTool()])
class TestInProcessTransport:
def test_direct_call(self, server):
transport = InProcessTransport(server)
req = MCPRequest(method="initialize", id=1)
resp = transport.send(req)
assert resp.error is None
assert "serverInfo" in resp.result
def test_roundtrip_tools_list(self, server):
transport = InProcessTransport(server)
req = MCPRequest(method="tools/list", id=2)
resp = transport.send(req)
assert resp.error is None
tools = resp.result["tools"]
assert len(tools) == 2
def test_roundtrip_tools_call(self, server):
transport = InProcessTransport(server)
req = MCPRequest(
method="tools/call",
params={"name": "calculator", "arguments": {"expression": "5*5"}},
id=3,
)
resp = transport.send(req)
assert resp.error is None
assert "25" in resp.result["content"][0]["text"]
def test_multiple_calls(self, server):
transport = InProcessTransport(server)
for i in range(5):
req = MCPRequest(method="tools/list", id=i)
resp = transport.send(req)
assert resp.error is None
def test_close_is_noop(self, server):
transport = InProcessTransport(server)
transport.close() # Should not raise
def test_error_method(self, server):
transport = InProcessTransport(server)
req = MCPRequest(method="unknown/method", id=1)
resp = transport.send(req)
assert resp.error is not None
class TestStdioTransport:
def test_send_receive(self, tmp_path):
"""Use a simple Python echo script as the subprocess."""
script = tmp_path / "echo_server.py"
script.write_text(textwrap.dedent("""\
import sys
import json
for line in sys.stdin:
line = line.strip()
if not line:
continue
req = json.loads(line)
resp = {
"jsonrpc": "2.0",
"id": req.get("id", 0),
"result": {"echo": req.get("method", "")},
}
sys.stdout.write(json.dumps(resp) + "\\n")
sys.stdout.flush()
"""))
transport = StdioTransport([sys.executable, str(script)])
try:
req = MCPRequest(method="test/echo", id=1)
resp = transport.send(req)
assert resp.error is None
assert resp.result["echo"] == "test/echo"
assert resp.id == 1
finally:
transport.close()
def test_multiple_requests(self, tmp_path):
"""Send multiple requests to the subprocess."""
script = tmp_path / "echo_server.py"
script.write_text(textwrap.dedent("""\
import sys
import json
for line in sys.stdin:
line = line.strip()
if not line:
continue
req = json.loads(line)
resp = {
"jsonrpc": "2.0",
"id": req.get("id", 0),
"result": {"method": req.get("method", "")},
}
sys.stdout.write(json.dumps(resp) + "\\n")
sys.stdout.flush()
"""))
transport = StdioTransport([sys.executable, str(script)])
try:
for i in range(3):
req = MCPRequest(method=f"test/{i}", id=i)
resp = transport.send(req)
assert resp.result["method"] == f"test/{i}"
finally:
transport.close()
def test_close_terminates_process(self, tmp_path):
script = tmp_path / "sleep_server.py"
script.write_text(textwrap.dedent("""\
import sys
import time
time.sleep(300)
"""))
transport = StdioTransport([sys.executable, str(script)])
proc = transport._process
assert proc is not None
assert proc.poll() is None # still running
transport.close()
assert transport._process is None
def test_close_idempotent(self, tmp_path):
script = tmp_path / "sleep_server.py"
script.write_text("import time; time.sleep(300)")
transport = StdioTransport([sys.executable, str(script)])
transport.close()
transport.close() # Should not raise
class TestSSETransport:
def test_send_receive(self, monkeypatch):
"""Mock httpx to simulate HTTP response."""
mock_response = MagicMock()
mock_response.text = json.dumps(
{"jsonrpc": "2.0", "id": 1, "result": {"tools": []}}
)
mock_response.raise_for_status = MagicMock()
mock_httpx = MagicMock()
mock_httpx.post.return_value = mock_response
monkeypatch.setitem(sys.modules, "httpx", mock_httpx)
transport = SSETransport("http://localhost:8080/mcp")
req = MCPRequest(method="tools/list", id=1)
resp = transport.send(req)
assert resp.error is None
assert resp.result == {"tools": []}
def test_send_posts_json(self, monkeypatch):
"""Verify the HTTP POST includes correct headers and body."""
mock_response = MagicMock()
mock_response.text = json.dumps({"jsonrpc": "2.0", "id": 1, "result": {}})
mock_response.raise_for_status = MagicMock()
mock_httpx = MagicMock()
mock_httpx.post.return_value = mock_response
monkeypatch.setitem(sys.modules, "httpx", mock_httpx)
transport = SSETransport("http://localhost:8080/mcp")
req = MCPRequest(method="initialize", id=1)
transport.send(req)
call_args = mock_httpx.post.call_args
assert call_args[0][0] == "http://localhost:8080/mcp"
assert call_args[1]["headers"]["Content-Type"] == "application/json"
def test_close_is_noop(self):
transport = SSETransport("http://localhost:8080/mcp")
transport.close() # Should not raise
def test_error_response(self, monkeypatch):
"""Simulate server returning an error response."""
mock_response = MagicMock()
mock_response.text = json.dumps(
{
"jsonrpc": "2.0",
"id": 1,
"error": {"code": -32601, "message": "Not found"},
}
)
mock_response.raise_for_status = MagicMock()
mock_httpx = MagicMock()
mock_httpx.post.return_value = mock_response
monkeypatch.setitem(sys.modules, "httpx", mock_httpx)
transport = SSETransport("http://localhost:8080/mcp")
req = MCPRequest(method="unknown", id=1)
resp = transport.send(req)
assert resp.error is not None
assert resp.error["code"] == -32601
tests/memory/test_retrieval_quality.py
+192
View File
@@ -0,0 +1,192 @@
"""Retrieval quality tests with a fixed corpus across backends."""
from __future__ import annotations
import pytest
from openjarvis.core.registry import MemoryRegistry
from openjarvis.memory.sqlite import SQLiteMemory
# ---------------------------------------------------------------------------
# Shared corpus
# ---------------------------------------------------------------------------
CORPUS = [
("Machine learning automates statistical analysis of data", "ml.md"),
("Neural networks are inspired by biological brain structure", "nn.md"),
("Python is a popular programming language for data science", "python.md"),
("The capital of France is Paris, located on the Seine river", "geo.md"),
("Quantum computing uses qubits instead of classical bits", "quantum.md"),
]
def _build_corpus(backend):
"""Insert the fixed corpus into the given backend."""
for content, source in CORPUS:
backend.store(content, source=source)
# ---------------------------------------------------------------------------
# Backend factory helpers
# ---------------------------------------------------------------------------
def _make_sqlite(tmp_path):
if not MemoryRegistry.contains("sqlite"):
MemoryRegistry.register_value("sqlite", SQLiteMemory)
return SQLiteMemory(db_path=tmp_path / "quality_test.db")
def _make_bm25():
bm25_mod = pytest.importorskip("openjarvis.memory.bm25", exc_type=ImportError)
BM25Memory = bm25_mod.BM25Memory
if not MemoryRegistry.contains("bm25"):
MemoryRegistry.register_value("bm25", BM25Memory)
return BM25Memory()
def _make_backend(key, tmp_path):
if key == "sqlite":
return _make_sqlite(tmp_path)
elif key == "bm25":
return _make_bm25()
else:
pytest.skip(f"Unknown backend key: {key}")
# ---------------------------------------------------------------------------
# Parametrized quality tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("backend_key", ["sqlite", "bm25"])
class TestRetrievalQuality:
"""Retrieval quality assertions that should hold across backends."""
def test_exact_keyword_match(self, backend_key, tmp_path):
"""Querying 'Python programming' should return the Python document."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("Python programming")
assert len(results) >= 1
assert "Python" in results[0].content
def test_semantic_similarity(self, backend_key, tmp_path):
"""Querying 'data analysis' should return the ML or Python document."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("data analysis")
assert len(results) >= 1
# At least one result should mention 'data'
data_results = [r for r in results if "data" in r.content.lower()]
assert len(data_results) >= 1
def test_no_match_returns_empty_or_low(self, backend_key, tmp_path):
"""Querying unrelated terms should return no results."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("medieval castles architecture")
# For keyword-based backends, unrelated queries return nothing
assert len(results) == 0
def test_ranking_order(self, backend_key, tmp_path):
"""Most relevant document should rank first."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("quantum computing qubits")
assert len(results) >= 1
content = results[0].content.lower()
assert "quantum" in content or "qubits" in content
def test_top_k_limiting(self, backend_key, tmp_path):
"""top_k=1 should return at most 1 result."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("data", top_k=1)
assert len(results) <= 1
def test_source_preserved(self, backend_key, tmp_path):
"""Source file information should be preserved in retrieval results."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("France Paris")
assert len(results) >= 1
assert results[0].source == "geo.md"
def test_multiple_relevant_results(self, backend_key, tmp_path):
"""A broad query should return multiple matching documents."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
# Both ML and Python docs mention 'data'
results = backend.retrieve("data", top_k=5)
data_results = [r for r in results if "data" in r.content.lower()]
assert len(data_results) >= 1
def test_query_with_stopwords(self, backend_key, tmp_path):
"""Query containing common words should still match relevant docs."""
backend = _make_backend(backend_key, tmp_path)
_build_corpus(backend)
results = backend.retrieve("the capital of France")
assert len(results) >= 1
# The geography document should be among results
geo_found = any("France" in r.content for r in results)
assert geo_found
# ---------------------------------------------------------------------------
# SQLite-specific retrieval tests
# ---------------------------------------------------------------------------
class TestSQLiteRetrievalSpecifics:
"""Tests specific to the SQLite FTS5 backend behavior."""
def test_fts5_keyword_matching(self, tmp_path):
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve("neural networks brain")
assert len(results) >= 1
content = results[0].content.lower()
assert "neural" in content or "brain" in content
def test_fts5_partial_match(self, tmp_path):
"""FTS5 matches individual terms, not just full phrases."""
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve("programming language")
assert len(results) >= 1
assert "programming" in results[0].content.lower()
def test_score_nonzero_for_matches(self, tmp_path):
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve("machine learning")
assert len(results) >= 1
# FTS5 scores are converted to positive values
assert results[0].score > 0
def test_empty_query(self, tmp_path):
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve("")
assert results == []
def test_whitespace_only_query(self, tmp_path):
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve(" ")
assert results == []
def test_single_word_query(self, tmp_path):
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve("Python")
assert len(results) >= 1
assert "Python" in results[0].content
def test_retrieve_all_corpus_with_broad_query(self, tmp_path):
"""A term that does not appear in the corpus returns empty."""
backend = _make_sqlite(tmp_path)
_build_corpus(backend)
results = backend.retrieve("xylophone")
assert len(results) == 0
tests/memory/test_storage_suite.py
+169
View File
@@ -0,0 +1,169 @@
"""Parametrized storage tests across all memory backends."""
from __future__ import annotations
import pytest
from openjarvis.core.registry import MemoryRegistry
from openjarvis.memory.sqlite import SQLiteMemory
# ---------------------------------------------------------------------------
# Backend factory helpers
# ---------------------------------------------------------------------------
def _make_sqlite(tmp_path):
if not MemoryRegistry.contains("sqlite"):
MemoryRegistry.register_value("sqlite", SQLiteMemory)
return SQLiteMemory(db_path=tmp_path / "test.db")
def _make_bm25():
bm25_mod = pytest.importorskip("openjarvis.memory.bm25", exc_type=ImportError)
BM25Memory = bm25_mod.BM25Memory
if not MemoryRegistry.contains("bm25"):
MemoryRegistry.register_value("bm25", BM25Memory)
return BM25Memory()
def _make_backend(key, tmp_path):
"""Create a backend instance by key, skipping if dependencies are missing."""
if key == "sqlite":
return _make_sqlite(tmp_path)
elif key == "bm25":
return _make_bm25()
elif key == "faiss":
mod = pytest.importorskip("openjarvis.memory.faiss", exc_type=ImportError)
return mod.FAISSMemory(db_path=str(tmp_path / "faiss"))
elif key == "colbert":
mod = pytest.importorskip("openjarvis.memory.colbert", exc_type=ImportError)
return mod.ColBERTMemory(db_path=str(tmp_path / "colbert"))
elif key == "hybrid":
mod = pytest.importorskip("openjarvis.memory.hybrid", exc_type=ImportError)
sqlite = _make_sqlite(tmp_path)
bm25 = _make_bm25()
return mod.HybridMemory(backends=[sqlite, bm25])
else:
pytest.skip(f"Unknown backend key: {key}")
# ---------------------------------------------------------------------------
# Core backends (always available)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("backend_key", ["sqlite", "bm25"])
class TestStorageSuiteCore:
"""Storage operations that must pass for guaranteed-available backends."""
def test_store_and_retrieve(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
backend.store("Python is a programming language", source="wiki.md")
results = backend.retrieve("Python programming")
assert len(results) >= 1
assert "Python" in results[0].content
def test_store_multiple_documents(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
doc_ids = []
for i in range(10):
doc_id = backend.store(
f"Document number {i} about testing software",
source=f"doc{i}.md",
)
doc_ids.append(doc_id)
assert len(doc_ids) == 10
assert len(set(doc_ids)) == 10 # all unique
def test_retrieve_respects_top_k(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
for i in range(10):
backend.store(f"document {i} about testing software quality")
results = backend.retrieve("testing software", top_k=3)
assert len(results) <= 3
def test_delete_document(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
doc_id = backend.store("content to delete")
assert backend.delete(doc_id) is True
assert backend.delete(doc_id) is False # already deleted
def test_clear_all(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
backend.store("first document")
backend.store("second document")
backend.clear()
results = backend.retrieve("first")
assert len(results) == 0
def test_metadata_roundtrip(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
meta = {"author": "test_user", "version": 2}
backend.store(
"content with metadata fields",
source="paper.pdf",
metadata=meta,
)
results = backend.retrieve("content metadata")
assert len(results) >= 1
assert results[0].source == "paper.pdf"
assert results[0].metadata["author"] == "test_user"
assert results[0].metadata["version"] == 2
def test_empty_retrieve(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
results = backend.retrieve("anything at all")
assert results == []
def test_store_returns_doc_id(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
doc_id = backend.store("some content")
assert isinstance(doc_id, str)
assert len(doc_id) > 0
def test_duplicate_content(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
id1 = backend.store("identical content here")
id2 = backend.store("identical content here")
assert id1 != id2 # different IDs even for same content
def test_large_document(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
base = "This is a large test document about software engineering. "
large_content = base * 200
assert len(large_content) > 10000
doc_id = backend.store(large_content, source="large.txt")
assert isinstance(doc_id, str)
results = backend.retrieve("software engineering")
assert len(results) >= 1
# ---------------------------------------------------------------------------
# Optional backends (may need extra dependencies)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("backend_key", ["faiss", "colbert", "hybrid"])
class TestStorageSuiteOptional:
"""Same core operations for backends that require optional dependencies."""
def test_store_and_retrieve(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
backend.store("Python is a programming language", source="wiki.md")
results = backend.retrieve("Python programming")
assert len(results) >= 1
assert "Python" in results[0].content
def test_delete_document(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
doc_id = backend.store("content to delete")
assert backend.delete(doc_id) is True
assert backend.delete(doc_id) is False
def test_clear_all(self, backend_key, tmp_path):
backend = _make_backend(backend_key, tmp_path)
backend.store("first document")
backend.store("second document")
backend.clear()
results = backend.retrieve("first")
assert len(results) == 0
tests/test_integration_extended.py
+566
View File
@@ -0,0 +1,566 @@
"""Extended integration tests for new components (Phase 6+)."""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
from openjarvis.agents._stubs import AgentContext, AgentResult
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.registry import AgentRegistry, ToolRegistry
from openjarvis.core.types import (
Conversation,
Message,
Role,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _register_all():
"""Ensure agents and tools are registered."""
from openjarvis.agents.openhands import OpenHandsAgent
from openjarvis.agents.react import ReActAgent
from openjarvis.tools.calculator import CalculatorTool
from openjarvis.tools.think import ThinkTool
for key, cls in [
("react", ReActAgent),
("openhands", OpenHandsAgent),
]:
if not AgentRegistry.contains(key):
AgentRegistry.register_value(key, cls)
for key, cls in [
("calculator", CalculatorTool),
("think", ThinkTool),
]:
if not ToolRegistry.contains(key):
ToolRegistry.register_value(key, cls)
def _make_engine(responses):
"""Create a mock engine returning a sequence of responses."""
engine = MagicMock()
engine.engine_id = "mock"
engine.health.return_value = True
engine.list_models.return_value = ["test-model"]
if isinstance(responses, list):
engine.generate.side_effect = responses
else:
engine.generate.return_value = responses
return engine
def _simple_response(content, model="test-model"):
return {
"content": content,
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
},
"model": model,
"finish_reason": "stop",
}
# ---------------------------------------------------------------------------
# ReAct pipeline integration
# ---------------------------------------------------------------------------
class TestReActPipeline:
"""End-to-end: ReAct agent with calculator tool."""
def test_react_with_calculator_e2e(self):
_register_all()
from openjarvis.agents.react import ReActAgent
from openjarvis.tools.calculator import CalculatorTool
responses = [
_simple_response(
"Thought: I need to calculate 2+2.\n"
"Action: calculator\n"
'Action Input: {"expression":"2+2"}'
),
_simple_response(
"Thought: The result is 4.\n"
"Final Answer: 2+2 equals 4."
),
]
engine = _make_engine(responses)
bus = EventBus(record_history=True)
agent = ReActAgent(
engine, "test-model",
tools=[CalculatorTool()], bus=bus,
)
result = agent.run("What is 2+2?")
assert isinstance(result, AgentResult)
assert "4" in result.content
assert result.turns == 2
assert len(result.tool_results) == 1
assert result.tool_results[0].content == "4"
def test_react_with_think_tool(self):
_register_all()
from openjarvis.agents.react import ReActAgent
from openjarvis.tools.think import ThinkTool
responses = [
_simple_response(
"Thought: Let me reason about this.\n"
"Action: think\n"
'Action Input: {"thought":"Step 1: analyze"}'
),
_simple_response(
"Thought: I have my analysis.\n"
"Final Answer: The answer is clear."
),
]
engine = _make_engine(responses)
agent = ReActAgent(
engine, "test-model", tools=[ThinkTool()],
)
result = agent.run("Analyze this.")
assert result.turns == 2
assert result.tool_results[0].success is True
def test_react_direct_answer(self):
"""ReAct returns immediately when no tool use is needed."""
_register_all()
from openjarvis.agents.react import ReActAgent
engine = _make_engine(
_simple_response(
"Thought: This is simple.\n"
"Final Answer: Hello!"
)
)
agent = ReActAgent(engine, "test-model")
result = agent.run("Say hello")
assert result.content == "Hello!"
assert result.turns == 1
def test_react_event_chain(self):
"""Verify complete event chain through ReAct run."""
_register_all()
from openjarvis.agents.react import ReActAgent
engine = _make_engine(
_simple_response(
"Thought: done.\nFinal Answer: ok"
)
)
bus = EventBus(record_history=True)
agent = ReActAgent(
engine, "test-model", bus=bus,
)
agent.run("Test")
types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in types
assert EventType.INFERENCE_START in types
assert EventType.INFERENCE_END in types
assert EventType.AGENT_TURN_END in types
# ---------------------------------------------------------------------------
# OpenHands pipeline integration
# ---------------------------------------------------------------------------
class TestOpenHandsPipeline:
"""End-to-end: OpenHands agent with code execution."""
def test_openhands_code_execution_e2e(self):
_register_all()
from openjarvis.agents.openhands import OpenHandsAgent
from openjarvis.tools.code_interpreter import (
CodeInterpreterTool,
)
if not ToolRegistry.contains("code_interpreter"):
ToolRegistry.register_value(
"code_interpreter", CodeInterpreterTool,
)
responses = [
_simple_response(
"I'll calculate this:\n"
"```python\nprint(2 + 2)\n```"
),
_simple_response("The result is 4."),
]
engine = _make_engine(responses)
agent = OpenHandsAgent(
engine, "test-model",
tools=[CodeInterpreterTool()],
)
result = agent.run("What is 2+2?")
assert isinstance(result, AgentResult)
assert result.turns == 2
assert len(result.tool_results) == 1
# The code_interpreter actually runs print(2+2)
assert "4" in result.tool_results[0].content
def test_openhands_direct_answer(self):
"""OpenHands returns directly when no code is needed."""
_register_all()
from openjarvis.agents.openhands import OpenHandsAgent
engine = _make_engine(
_simple_response("Hello! How can I help?")
)
agent = OpenHandsAgent(engine, "test-model")
result = agent.run("Say hello")
assert result.content == "Hello! How can I help?"
assert result.turns == 1
def test_openhands_event_chain(self):
"""Verify event chain through OpenHands run."""
_register_all()
from openjarvis.agents.openhands import OpenHandsAgent
engine = _make_engine(
_simple_response("Direct answer.")
)
bus = EventBus(record_history=True)
agent = OpenHandsAgent(
engine, "test-model", bus=bus,
)
agent.run("Test")
types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in types
assert EventType.AGENT_TURN_END in types
# ---------------------------------------------------------------------------
# MCP integration
# ---------------------------------------------------------------------------
class TestMCPIntegration:
"""MCP server + client with real tools."""
def test_mcp_server_with_all_tools(self):
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import InProcessTransport
from openjarvis.tools.calculator import CalculatorTool
from openjarvis.tools.think import ThinkTool
tools = [CalculatorTool(), ThinkTool()]
server = MCPServer(tools)
transport = InProcessTransport(server)
client = MCPClient(transport)
# Initialize
caps = client.initialize()
assert "serverInfo" in caps
# List tools
specs = client.list_tools()
names = [s.name for s in specs]
assert "calculator" in names
assert "think" in names
# Call calculator
result = client.call_tool(
"calculator", {"expression": "10*5"},
)
assert result["content"][0]["text"] == "50"
assert result["isError"] is False
# Call think
result = client.call_tool(
"think", {"thought": "reasoning step"},
)
assert result["content"][0]["text"] == "reasoning step"
client.close()
def test_mcp_unknown_tool_error(self):
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.protocol import MCPError
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import InProcessTransport
from openjarvis.tools.calculator import CalculatorTool
server = MCPServer([CalculatorTool()])
client = MCPClient(InProcessTransport(server))
client.initialize()
with pytest.raises(MCPError):
client.call_tool("nonexistent", {})
client.close()
def test_mcp_roundtrip_lifecycle(self):
"""Full lifecycle: init -> list -> call -> result."""
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.server import MCPServer
from openjarvis.mcp.transport import InProcessTransport
from openjarvis.tools.calculator import CalculatorTool
server = MCPServer([CalculatorTool()])
client = MCPClient(InProcessTransport(server))
# 1. Initialize
caps = client.initialize()
assert caps["protocolVersion"] == "2024-11-05"
# 2. Discover tools
tools = client.list_tools()
assert len(tools) == 1
assert tools[0].name == "calculator"
# 3. Call tool
result = client.call_tool(
"calculator", {"expression": "7+3"},
)
assert result["content"][0]["text"] == "10"
# 4. Close
client.close()
# ---------------------------------------------------------------------------
# Cross-engine consistency
# ---------------------------------------------------------------------------
class TestCrossEngineConsistency:
"""Same query through different mock engine configs."""
def test_same_query_same_format(self):
"""All engines return the same result dict shape."""
_register_all()
from openjarvis.agents.react import ReActAgent
for engine_name in ["vllm", "ollama", "mock"]:
engine = _make_engine(
_simple_response(
"Thought: ok.\nFinal Answer: Result",
model="test-model",
)
)
engine.engine_id = engine_name
agent = ReActAgent(engine, "test-model")
result = agent.run("Test query")
assert isinstance(result, AgentResult)
assert result.content == "Result"
def test_tool_calls_across_engines(self):
"""Tool calling works regardless of engine mock."""
_register_all()
from openjarvis.agents.react import ReActAgent
from openjarvis.tools.calculator import CalculatorTool
for engine_name in ["vllm", "ollama"]:
responses = [
_simple_response(
"Thought: calc.\n"
"Action: calculator\n"
'Action Input: {"expression":"3*3"}'
),
_simple_response(
"Thought: got 9.\n"
"Final Answer: 9"
),
]
engine = _make_engine(responses)
engine.engine_id = engine_name
agent = ReActAgent(
engine, "test-model",
tools=[CalculatorTool()],
)
result = agent.run("What is 3*3?")
assert result.content == "9"
tr = result.tool_results[0]
assert tr.content == "9"
assert tr.success is True
# ---------------------------------------------------------------------------
# Memory pipeline integration
# ---------------------------------------------------------------------------
class TestMemoryPipeline:
"""Index and retrieve across available backends."""
def test_sqlite_index_and_retrieve(self, tmp_path):
from openjarvis.memory.sqlite import SQLiteMemory
backend = SQLiteMemory(db_path=str(tmp_path / "mem.db"))
backend.store(
"Machine learning uses data to learn patterns",
source="ml.md",
)
backend.store(
"Python is a versatile programming language",
source="py.md",
)
results = backend.retrieve("machine learning")
assert len(results) >= 1
assert "machine" in results[0].content.lower()
def test_bm25_index_and_retrieve(self, tmp_path):
try:
from openjarvis.memory.bm25 import BM25Memory
except ImportError:
pytest.skip("rank_bm25 not installed")
backend = BM25Memory(
db_path=str(tmp_path / "bm25.db"),
)
backend.store("Neural networks for NLP", source="a.md")
backend.store("Database indexing strategies", source="b.md")
results = backend.retrieve("neural NLP")
assert len(results) >= 1
# ---------------------------------------------------------------------------
# Model catalog integration
# ---------------------------------------------------------------------------
class TestModelCatalogIntegration:
"""All registered models have valid metadata."""
def test_all_models_routable(self):
"""Every model in catalog has required fields."""
from openjarvis.intelligence.model_catalog import (
BUILTIN_MODELS,
)
for spec in BUILTIN_MODELS:
assert spec.model_id, "model_id required"
assert spec.context_length > 0
if spec.requires_api_key:
assert spec.provider
def test_local_models_have_engine_compat(self):
"""Every local model has at least one engine."""
from openjarvis.intelligence.model_catalog import (
BUILTIN_MODELS,
)
local = [
s for s in BUILTIN_MODELS if not s.requires_api_key
]
for spec in local:
assert len(spec.supported_engines) >= 1, (
f"{spec.model_id} has no engines"
)
def test_cloud_models_require_api_key(self):
"""All cloud models require an API key."""
from openjarvis.intelligence.model_catalog import (
BUILTIN_MODELS,
)
cloud_ids = [
"gpt-5-mini",
"claude-opus-4-6",
"claude-sonnet-4-6",
"claude-haiku-4-5",
"gemini-2.5-pro",
"gemini-2.5-flash",
"gemini-3-pro",
"gemini-3-flash",
]
for mid in cloud_ids:
matches = [
s for s in BUILTIN_MODELS if s.model_id == mid
]
assert len(matches) == 1, f"Missing {mid}"
assert matches[0].requires_api_key is True
# ---------------------------------------------------------------------------
# Agent routing matrix (lightweight integration)
# ---------------------------------------------------------------------------
class TestAgentRoutingMatrix:
"""Agents run consistently across different configurations."""
@pytest.mark.parametrize(
"agent_key", ["react", "openhands"],
)
def test_agent_returns_valid_result(self, agent_key):
_register_all()
engine = _make_engine(
_simple_response(
"Thought: done.\nFinal Answer: ok"
if agent_key == "react"
else "The answer is ok"
)
)
agent_cls = AgentRegistry.get(agent_key)
agent = agent_cls(engine, "test-model")
result = agent.run("Test")
assert isinstance(result, AgentResult)
assert result.turns >= 1
assert len(result.content) > 0
@pytest.mark.parametrize(
"agent_key", ["react", "openhands"],
)
def test_agent_emits_events(self, agent_key):
_register_all()
engine = _make_engine(
_simple_response(
"Thought: done.\nFinal Answer: ok"
if agent_key == "react"
else "Direct answer"
)
)
bus = EventBus(record_history=True)
agent_cls = AgentRegistry.get(agent_key)
agent = agent_cls(engine, "test-model", bus=bus)
agent.run("Test")
types = [e.event_type for e in bus.history]
assert EventType.AGENT_TURN_START in types
assert EventType.AGENT_TURN_END in types
def test_context_passing(self):
"""Agents accept and use AgentContext."""
_register_all()
from openjarvis.agents.react import ReActAgent
engine = _make_engine(
_simple_response(
"Thought: I see the system message.\n"
"Final Answer: Got context."
)
)
conv = Conversation()
conv.add(Message(
role=Role.SYSTEM,
content="You are helpful.",
))
ctx = AgentContext(conversation=conv)
agent = ReActAgent(engine, "test-model")
result = agent.run("Hello", context=ctx)
assert result.content == "Got context."
# Verify engine received system message
call_args = engine.generate.call_args
msgs = call_args[0][0]
# First message is ReAct system prompt, then context
assert any(
m.content == "You are helpful." for m in msgs
)
tests/tools/test_code_interpreter.py
+132
View File
@@ -0,0 +1,132 @@
"""Tests for the code interpreter tool."""
from __future__ import annotations
from openjarvis.core.registry import ToolRegistry
from openjarvis.tools.code_interpreter import CodeInterpreterTool
class TestCodeInterpreterTool:
def test_spec_name_and_category(self):
tool = CodeInterpreterTool()
assert tool.spec.name == "code_interpreter"
assert tool.spec.category == "code"
def test_spec_parameters_require_code(self):
tool = CodeInterpreterTool()
assert "code" in tool.spec.parameters["properties"]
assert "code" in tool.spec.parameters["required"]
def test_execute_simple_code(self):
tool = CodeInterpreterTool()
result = tool.execute(code="print(2 + 2)")
assert result.success is True
assert "4" in result.content
def test_execute_with_imports(self):
tool = CodeInterpreterTool()
result = tool.execute(code="import math; print(round(math.pi, 5))")
assert result.success is True
assert "3.14159" in result.content
def test_execute_multiline(self):
tool = CodeInterpreterTool()
code = "x = 10\ny = 20\nprint(x + y)"
result = tool.execute(code=code)
assert result.success is True
assert "30" in result.content
def test_timeout_protection(self):
tool = CodeInterpreterTool(timeout=2)
result = tool.execute(code="import time; time.sleep(10)")
assert result.success is False
assert "timed out" in result.content
def test_syntax_error(self):
tool = CodeInterpreterTool()
result = tool.execute(code="def f(\n")
assert result.success is False
assert "SyntaxError" in result.content
def test_runtime_error(self):
tool = CodeInterpreterTool()
result = tool.execute(code="print(1/0)")
assert result.success is False
assert "ZeroDivisionError" in result.content
def test_dangerous_os_system_blocked(self):
tool = CodeInterpreterTool()
result = tool.execute(code="import os; os.system('ls')")
assert result.success is False
assert "Blocked" in result.content
assert "os.system" in result.content
def test_dangerous_subprocess_blocked(self):
tool = CodeInterpreterTool()
result = tool.execute(code="import subprocess; subprocess.run(['ls'])")
assert result.success is False
assert "Blocked" in result.content
def test_dangerous_eval_blocked(self):
tool = CodeInterpreterTool()
result = tool.execute(code="eval('2+2')")
assert result.success is False
assert "Blocked" in result.content
def test_dangerous_open_blocked(self):
tool = CodeInterpreterTool()
result = tool.execute(code="f = open('/etc/passwd')")
assert result.success is False
assert "Blocked" in result.content
def test_no_code_provided(self):
tool = CodeInterpreterTool()
result = tool.execute(code="")
assert result.success is False
assert "No code" in result.content
def test_no_code_param(self):
tool = CodeInterpreterTool()
result = tool.execute()
assert result.success is False
assert "No code" in result.content
def test_output_truncation(self):
tool = CodeInterpreterTool(max_output=50)
result = tool.execute(code="print('A' * 200)")
assert result.success is True
assert "truncated" in result.content
assert len(result.content) < 200
def test_to_openai_function(self):
tool = CodeInterpreterTool()
fn = tool.to_openai_function()
assert fn["type"] == "function"
assert fn["function"]["name"] == "code_interpreter"
assert "code" in fn["function"]["parameters"]["properties"]
def test_returncode_in_metadata(self):
tool = CodeInterpreterTool()
result = tool.execute(code="print('ok')")
assert result.success is True
assert result.metadata["returncode"] == 0
def test_returncode_nonzero_on_error(self):
tool = CodeInterpreterTool()
result = tool.execute(code="raise ValueError('boom')")
assert result.success is False
assert result.metadata["returncode"] != 0
def test_no_output_produces_placeholder(self):
tool = CodeInterpreterTool()
result = tool.execute(code="x = 42")
assert result.success is True
assert result.content == "(no output)"
def test_tool_id(self):
tool = CodeInterpreterTool()
assert tool.tool_id == "code_interpreter"
def test_registry_registration(self):
ToolRegistry.register_value("code_interpreter", CodeInterpreterTool)
assert ToolRegistry.contains("code_interpreter")
tests/tools/test_web_search.py
+143
View File
@@ -0,0 +1,143 @@
"""Tests for the web search tool."""
from __future__ import annotations
import sys
from unittest.mock import MagicMock, patch
from openjarvis.core.registry import ToolRegistry
from openjarvis.tools.web_search import WebSearchTool
class TestWebSearchTool:
def test_spec_name_and_category(self):
tool = WebSearchTool(api_key="test-key")
assert tool.spec.name == "web_search"
assert tool.spec.category == "search"
def test_spec_requires_api_key_metadata(self):
tool = WebSearchTool(api_key="test-key")
assert tool.spec.metadata["requires_api_key"] == "TAVILY_API_KEY"
def test_spec_parameters_require_query(self):
tool = WebSearchTool(api_key="test-key")
assert "query" in tool.spec.parameters["properties"]
assert "query" in tool.spec.parameters["required"]
def test_execute_no_query(self):
tool = WebSearchTool(api_key="test-key")
result = tool.execute(query="")
assert result.success is False
assert "No query" in result.content
def test_execute_no_query_param(self):
tool = WebSearchTool(api_key="test-key")
result = tool.execute()
assert result.success is False
assert "No query" in result.content
def test_execute_no_api_key(self):
tool = WebSearchTool(api_key=None)
# Clear env var to ensure no fallback
with patch.dict("os.environ", {}, clear=True):
tool._api_key = None
result = tool.execute(query="test query")
assert result.success is False
assert "No API key" in result.content
def test_execute_mocked_tavily(self, monkeypatch):
mock_client = MagicMock()
mock_client.search.return_value = {
"results": [
{
"title": "Result 1",
"url": "https://example.com/1",
"content": "Content about test.",
},
{
"title": "Result 2",
"url": "https://example.com/2",
"content": "More content.",
},
]
}
mock_tavily_module = MagicMock()
mock_tavily_module.TavilyClient.return_value = mock_client
monkeypatch.setitem(sys.modules, "tavily", mock_tavily_module)
tool = WebSearchTool(api_key="test-key")
result = tool.execute(query="test query")
assert result.success is True
assert "Result 1" in result.content
assert "Result 2" in result.content
assert result.metadata["num_results"] == 2
def test_execute_tavily_error(self, monkeypatch):
mock_client = MagicMock()
mock_client.search.side_effect = RuntimeError("API rate limit exceeded")
mock_tavily_module = MagicMock()
mock_tavily_module.TavilyClient.return_value = mock_client
monkeypatch.setitem(sys.modules, "tavily", mock_tavily_module)
tool = WebSearchTool(api_key="test-key")
result = tool.execute(query="test query")
assert result.success is False
assert "Search error" in result.content
def test_max_results_parameter(self, monkeypatch):
mock_client = MagicMock()
mock_client.search.return_value = {"results": []}
mock_tavily_module = MagicMock()
mock_tavily_module.TavilyClient.return_value = mock_client
monkeypatch.setitem(sys.modules, "tavily", mock_tavily_module)
tool = WebSearchTool(api_key="test-key", max_results=3)
tool.execute(query="test", max_results=7)
mock_client.search.assert_called_once_with("test", max_results=7)
def test_to_openai_function(self):
tool = WebSearchTool(api_key="test-key")
fn = tool.to_openai_function()
assert fn["type"] == "function"
assert fn["function"]["name"] == "web_search"
assert "query" in fn["function"]["parameters"]["properties"]
def test_execute_import_error(self, monkeypatch):
"""Simulate tavily-python not being installed."""
# Remove tavily from sys.modules if present, and make import fail
monkeypatch.delitem(sys.modules, "tavily", raising=False)
import builtins
original_import = builtins.__import__
def _mock_import(name, *args, **kwargs):
if name == "tavily":
raise ImportError("No module named 'tavily'")
return original_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", _mock_import)
tool = WebSearchTool(api_key="test-key")
result = tool.execute(query="test query")
assert result.success is False
assert "tavily-python not installed" in result.content
def test_empty_results(self, monkeypatch):
mock_client = MagicMock()
mock_client.search.return_value = {"results": []}
mock_tavily_module = MagicMock()
mock_tavily_module.TavilyClient.return_value = mock_client
monkeypatch.setitem(sys.modules, "tavily", mock_tavily_module)
tool = WebSearchTool(api_key="test-key")
result = tool.execute(query="obscure query")
assert result.success is True
assert result.content == "No results found."
def test_tool_id(self):
tool = WebSearchTool(api_key="test-key")
assert tool.tool_id == "web_search"
def test_registry_registration(self):
ToolRegistry.register_value("web_search", WebSearchTool)
assert ToolRegistry.contains("web_search")
tests/traces/__init__.py
View File
tests/traces/test_analyzer.py
+180
View File
@@ -0,0 +1,180 @@
"""Tests for the TraceAnalyzer."""
from __future__ import annotations
import time
from pathlib import Path
from openjarvis.core.types import StepType, Trace, TraceStep
from openjarvis.traces.analyzer import TraceAnalyzer
from openjarvis.traces.store import TraceStore
def _make_trace(
query: str = "test",
agent: str = "orchestrator",
model: str = "qwen3:8b",
outcome: str | None = None,
feedback: float | None = None,
latency: float = 1.0,
tokens: int = 100,
tool_name: str | None = None,
) -> Trace:
now = time.time()
steps = [
TraceStep(
step_type=StepType.GENERATE,
timestamp=now,
duration_seconds=latency * 0.8,
input={"model": model},
output={"tokens": tokens},
),
]
if tool_name:
steps.append(TraceStep(
step_type=StepType.TOOL_CALL,
timestamp=now + 0.1,
duration_seconds=latency * 0.2,
input={"tool": tool_name},
output={"success": True},
))
steps.append(TraceStep(
step_type=StepType.RESPOND,
timestamp=now + latency,
duration_seconds=0.0,
output={"content": "result"},
))
return Trace(
query=query,
agent=agent,
model=model,
engine="ollama",
result="result",
outcome=outcome,
feedback=feedback,
started_at=now,
ended_at=now + latency,
total_tokens=tokens,
total_latency_seconds=latency,
steps=steps,
)
class TestTraceAnalyzer:
def test_empty_summary(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
analyzer = TraceAnalyzer(store)
summary = analyzer.summary()
assert summary.total_traces == 0
assert summary.total_steps == 0
store.close()
def test_summary(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(outcome="success", latency=1.0, tokens=100))
store.save(_make_trace(outcome="success", latency=2.0, tokens=200))
store.save(_make_trace(outcome="failure", latency=0.5, tokens=50))
analyzer = TraceAnalyzer(store)
summary = analyzer.summary()
assert summary.total_traces == 3
assert summary.avg_latency > 0
assert summary.avg_tokens > 0
assert summary.success_rate == 2 / 3
assert "generate" in summary.step_type_distribution
assert "respond" in summary.step_type_distribution
store.close()
def test_per_route_stats(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(
model="qwen3:8b", agent="simple",
outcome="success", feedback=0.9,
))
store.save(_make_trace(
model="qwen3:8b", agent="simple",
outcome="success", feedback=0.8,
))
store.save(_make_trace(
model="llama3:70b", agent="orchestrator",
outcome="failure",
))
analyzer = TraceAnalyzer(store)
stats = analyzer.per_route_stats()
assert len(stats) == 2
qwen_stats = [s for s in stats if s.model == "qwen3:8b"][0]
assert qwen_stats.count == 2
assert qwen_stats.success_rate == 1.0
assert abs(qwen_stats.avg_feedback - 0.85) < 1e-9
llama_stats = [s for s in stats if s.model == "llama3:70b"][0]
assert llama_stats.count == 1
assert llama_stats.success_rate == 0.0
assert llama_stats.avg_feedback is None
store.close()
def test_per_tool_stats(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(tool_name="calculator"))
store.save(_make_trace(tool_name="calculator"))
store.save(_make_trace(tool_name="web_search"))
store.save(_make_trace()) # no tool
analyzer = TraceAnalyzer(store)
stats = analyzer.per_tool_stats()
assert len(stats) == 2
calc = [s for s in stats if s.tool_name == "calculator"][0]
assert calc.call_count == 2
assert calc.success_rate == 1.0
web = [s for s in stats if s.tool_name == "web_search"][0]
assert web.call_count == 1
store.close()
def test_per_route_stats_no_evaluated(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(outcome=None)) # unknown outcome
analyzer = TraceAnalyzer(store)
stats = analyzer.per_route_stats()
assert len(stats) == 1
assert stats[0].success_rate == 0.0 # no evaluated traces
store.close()
def test_export_traces(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(query="q1"))
store.save(_make_trace(query="q2"))
analyzer = TraceAnalyzer(store)
exported = analyzer.export_traces()
assert len(exported) == 2
assert all(isinstance(e, dict) for e in exported)
assert exported[0]["query"] in ("q1", "q2")
assert "steps" in exported[0]
assert len(exported[0]["steps"]) > 0
store.close()
def test_traces_for_query_type_code(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(query="def foo(): pass"))
store.save(_make_trace(query="what is the weather"))
analyzer = TraceAnalyzer(store)
code_traces = analyzer.traces_for_query_type(has_code=True)
assert len(code_traces) == 1
assert "def foo" in code_traces[0].query
store.close()
def test_traces_for_query_type_length(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(query="hi"))
store.save(_make_trace(query="a" * 200))
analyzer = TraceAnalyzer(store)
long_traces = analyzer.traces_for_query_type(min_length=100)
assert len(long_traces) == 1
store.close()
tests/traces/test_collector.py
+236
View File
@@ -0,0 +1,236 @@
"""Tests for the TraceCollector."""
from __future__ import annotations
import time
from pathlib import Path
from typing import Any, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.types import StepType
from openjarvis.traces.collector import TraceCollector
from openjarvis.traces.store import TraceStore
class _FakeAgent(BaseAgent):
"""Minimal agent that returns a fixed response."""
agent_id = "fake"
def __init__(
self,
response: str = "test response",
bus: Optional[EventBus] = None,
) -> None:
self._response = response
self._bus = bus
def run(
self, input: str, context: Optional[AgentContext] = None,
**kwargs: Any,
) -> AgentResult:
# Simulate an inference step via event bus
if self._bus:
self._bus.publish(EventType.INFERENCE_START, {
"model": "qwen3:8b",
"engine": "ollama",
})
self._bus.publish(EventType.INFERENCE_END, {
"total_tokens": 50,
})
return AgentResult(content=self._response, turns=1)
class _ToolAgent(BaseAgent):
"""Agent that simulates a tool call during execution."""
agent_id = "tool_agent"
def __init__(self, bus: EventBus) -> None:
self._bus = bus
def run(
self, input: str, context: Optional[AgentContext] = None,
**kwargs: Any,
) -> AgentResult:
# Simulate inference + tool call + inference
inf = {"model": "qwen3:8b", "engine": "ollama"}
self._bus.publish(EventType.INFERENCE_START, inf)
self._bus.publish(EventType.INFERENCE_END, {"total_tokens": 30})
self._bus.publish(EventType.TOOL_CALL_START, {
"tool": "calculator", "arguments": {"expr": "2+2"},
})
self._bus.publish(EventType.TOOL_CALL_END, {
"tool": "calculator", "success": True, "latency": 0.01,
})
self._bus.publish(EventType.INFERENCE_START, inf)
self._bus.publish(EventType.INFERENCE_END, {"total_tokens": 20})
return AgentResult(content="4", turns=2)
class TestTraceCollector:
def test_basic_collection(self, tmp_path: Path) -> None:
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(response="hello", bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
result = collector.run("say hello")
assert result.content == "hello"
assert store.count() == 1
traces = store.list_traces()
trace = traces[0]
assert trace.query == "say hello"
assert trace.agent == "fake"
assert trace.model == "qwen3:8b"
assert trace.engine == "ollama"
assert trace.result == "hello"
store.close()
def test_records_generate_steps(self, tmp_path: Path) -> None:
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
collector.run("test")
trace = store.list_traces()[0]
generate_steps = [s for s in trace.steps if s.step_type == StepType.GENERATE]
assert len(generate_steps) == 1
assert generate_steps[0].output.get("tokens") == 50
store.close()
def test_records_tool_steps(self, tmp_path: Path) -> None:
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _ToolAgent(bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
collector.run("What is 2+2?")
trace = store.list_traces()[0]
tool_steps = [s for s in trace.steps if s.step_type == StepType.TOOL_CALL]
assert len(tool_steps) == 1
assert tool_steps[0].input["tool"] == "calculator"
assert tool_steps[0].output["success"] is True
store.close()
def test_records_respond_step(self, tmp_path: Path) -> None:
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(response="final answer", bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
collector.run("test")
trace = store.list_traces()[0]
respond_steps = [s for s in trace.steps if s.step_type == StepType.RESPOND]
assert len(respond_steps) == 1
assert respond_steps[0].output["content"] == "final answer"
store.close()
def test_records_memory_retrieve(self, tmp_path: Path) -> None:
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
# Monkey-patch agent to emit memory event
original_run = agent.run
def run_with_memory(input, context=None, **kwargs):
bus.publish(EventType.MEMORY_RETRIEVE, {
"query": "meeting notes",
"num_results": 3,
"latency": 0.2,
})
return original_run(input, context=context, **kwargs)
agent.run = run_with_memory
collector.run("find my meeting notes")
trace = store.list_traces()[0]
retrieve_steps = [s for s in trace.steps if s.step_type == StepType.RETRIEVE]
assert len(retrieve_steps) == 1
assert retrieve_steps[0].input["query"] == "meeting notes"
store.close()
def test_publishes_trace_complete(self, tmp_path: Path) -> None:
bus = EventBus(record_history=True)
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
collector.run("test")
trace_events = [
e for e in bus.history
if e.event_type == EventType.TRACE_COMPLETE
]
assert len(trace_events) == 1
assert trace_events[0].data["trace"].query == "test"
store.close()
def test_no_store(self) -> None:
"""Collector works without a store (just collects, doesn't persist)."""
bus = EventBus()
agent = _FakeAgent(response="ok", bus=bus)
collector = TraceCollector(agent, bus=bus) # no store
result = collector.run("test")
assert result.content == "ok"
def test_no_bus(self, tmp_path: Path) -> None:
"""Collector works without a bus (no event-based step collection)."""
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(response="ok")
collector = TraceCollector(agent, store=store) # no bus
result = collector.run("test")
assert result.content == "ok"
assert store.count() == 1
# Only the RESPOND step (no events to capture)
trace = store.list_traces()[0]
assert len(trace.steps) == 1
assert trace.steps[0].step_type == StepType.RESPOND
store.close()
def test_timing(self, tmp_path: Path) -> None:
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
before = time.time()
collector.run("test")
after = time.time()
trace = store.list_traces()[0]
assert trace.started_at >= before
assert trace.ended_at <= after
assert trace.ended_at >= trace.started_at
store.close()
def test_unsubscribes_after_run(self, tmp_path: Path) -> None:
"""Events after run() completes should NOT affect the next trace."""
bus = EventBus()
store = TraceStore(tmp_path / "test.db")
agent = _FakeAgent(bus=bus)
collector = TraceCollector(agent, store=store, bus=bus)
collector.run("first")
# Emit events after run — should not affect stored trace
bus.publish(EventType.INFERENCE_START, {"model": "stray"})
bus.publish(EventType.INFERENCE_END, {"total_tokens": 999})
assert store.count() == 1
trace = store.list_traces()[0]
# No step with model="stray"
for s in trace.steps:
assert s.input.get("model") != "stray"
store.close()
tests/traces/test_store.py
+191
View File
@@ -0,0 +1,191 @@
"""Tests for the trace SQLite store."""
from __future__ import annotations
import time
from pathlib import Path
from openjarvis.core.events import EventBus, EventType
from openjarvis.core.types import StepType, Trace, TraceStep
from openjarvis.traces.store import TraceStore
def _make_trace(
query: str = "test query",
agent: str = "orchestrator",
model: str = "qwen3:8b",
engine: str = "ollama",
outcome: str | None = None,
feedback: float | None = None,
num_steps: int = 2,
) -> Trace:
"""Helper to create a trace with steps."""
now = time.time()
steps = []
for i in range(num_steps):
steps.append(TraceStep(
step_type=StepType.GENERATE if i % 2 == 0 else StepType.TOOL_CALL,
timestamp=now + i * 0.1,
duration_seconds=0.1 * (i + 1),
input={"model": model} if i % 2 == 0 else {"tool": "calculator"},
output={"tokens": 50} if i % 2 == 0 else {"success": True},
))
trace = Trace(
query=query,
agent=agent,
model=model,
engine=engine,
result="test result",
outcome=outcome,
feedback=feedback,
started_at=now,
ended_at=now + 1.0,
total_tokens=100,
total_latency_seconds=1.0,
steps=steps,
)
return trace
class TestTraceStore:
def test_creates_tables(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
assert store.count() == 0
store.close()
def test_save_and_get(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
trace = _make_trace()
store.save(trace)
retrieved = store.get(trace.trace_id)
assert retrieved is not None
assert retrieved.trace_id == trace.trace_id
assert retrieved.query == "test query"
assert retrieved.model == "qwen3:8b"
assert retrieved.engine == "ollama"
assert len(retrieved.steps) == 2
assert retrieved.steps[0].step_type == StepType.GENERATE
assert retrieved.steps[1].step_type == StepType.TOOL_CALL
store.close()
def test_get_nonexistent(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
assert store.get("nonexistent") is None
store.close()
def test_count(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(query="q1"))
store.save(_make_trace(query="q2"))
store.save(_make_trace(query="q3"))
assert store.count() == 3
store.close()
def test_list_traces_no_filter(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(query="q1"))
store.save(_make_trace(query="q2"))
traces = store.list_traces()
assert len(traces) == 2
store.close()
def test_list_traces_filter_agent(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(agent="simple"))
store.save(_make_trace(agent="orchestrator"))
store.save(_make_trace(agent="simple"))
traces = store.list_traces(agent="simple")
assert len(traces) == 2
assert all(t.agent == "simple" for t in traces)
store.close()
def test_list_traces_filter_model(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(model="qwen3:8b"))
store.save(_make_trace(model="llama3:70b"))
traces = store.list_traces(model="llama3:70b")
assert len(traces) == 1
assert traces[0].model == "llama3:70b"
store.close()
def test_list_traces_filter_outcome(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
store.save(_make_trace(outcome="success"))
store.save(_make_trace(outcome="failure"))
store.save(_make_trace(outcome="success"))
traces = store.list_traces(outcome="success")
assert len(traces) == 2
store.close()
def test_list_traces_time_range(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
now = time.time()
t1 = _make_trace(query="old")
t1.started_at = now - 3600 # 1 hour ago
t2 = _make_trace(query="recent")
t2.started_at = now
store.save(t1)
store.save(t2)
traces = store.list_traces(since=now - 60)
assert len(traces) == 1
assert traces[0].query == "recent"
store.close()
def test_list_traces_limit(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
for i in range(10):
store.save(_make_trace(query=f"q{i}"))
traces = store.list_traces(limit=3)
assert len(traces) == 3
store.close()
def test_bus_subscription(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
bus = EventBus()
store.subscribe_to_bus(bus)
trace = _make_trace()
bus.publish(EventType.TRACE_COMPLETE, {"trace": trace})
assert store.count() == 1
retrieved = store.get(trace.trace_id)
assert retrieved is not None
assert retrieved.query == trace.query
store.close()
def test_step_data_roundtrip(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
trace = _make_trace(num_steps=3)
store.save(trace)
retrieved = store.get(trace.trace_id)
assert retrieved is not None
assert len(retrieved.steps) == 3
for orig, retr in zip(trace.steps, retrieved.steps):
assert orig.step_type == retr.step_type
assert orig.input == retr.input
assert orig.output == retr.output
store.close()
def test_close_and_reopen(self, tmp_path: Path) -> None:
db_path = tmp_path / "test.db"
store = TraceStore(db_path)
store.save(_make_trace())
store.close()
store2 = TraceStore(db_path)
assert store2.count() == 1
store2.close()
def test_metadata_roundtrip(self, tmp_path: Path) -> None:
store = TraceStore(tmp_path / "test.db")
trace = _make_trace()
trace.metadata = {"key": "value", "nested": [1, 2, 3]}
store.save(trace)
retrieved = store.get(trace.trace_id)
assert retrieved is not None
assert retrieved.metadata["key"] == "value"
assert retrieved.metadata["nested"] == [1, 2, 3]
store.close()
tests/traces/test_types.py
+98
View File
@@ -0,0 +1,98 @@
"""Tests for Trace and TraceStep types."""
from __future__ import annotations
import time
from openjarvis.core.types import StepType, Trace, TraceStep
class TestTraceStep:
def test_create_step(self) -> None:
step = TraceStep(
step_type=StepType.GENERATE,
timestamp=time.time(),
duration_seconds=0.5,
input={"model": "qwen3:8b"},
output={"tokens": 100},
)
assert step.step_type == StepType.GENERATE
assert step.duration_seconds == 0.5
assert step.output["tokens"] == 100
def test_step_defaults(self) -> None:
step = TraceStep(step_type=StepType.ROUTE, timestamp=0.0)
assert step.duration_seconds == 0.0
assert step.input == {}
assert step.output == {}
assert step.metadata == {}
def test_all_step_types(self) -> None:
for st in StepType:
step = TraceStep(step_type=st, timestamp=0.0)
assert step.step_type == st
class TestTrace:
def test_create_trace(self) -> None:
trace = Trace(query="What is 2+2?", agent="orchestrator", model="qwen3:8b")
assert trace.query == "What is 2+2?"
assert trace.agent == "orchestrator"
assert len(trace.trace_id) == 16 # hex uuid
def test_trace_id_unique(self) -> None:
t1 = Trace()
t2 = Trace()
assert t1.trace_id != t2.trace_id
def test_add_step(self) -> None:
trace = Trace(query="test")
step = TraceStep(
step_type=StepType.GENERATE,
timestamp=time.time(),
duration_seconds=1.0,
output={"tokens": 50},
)
trace.add_step(step)
assert len(trace.steps) == 1
assert trace.total_latency_seconds == 1.0
assert trace.total_tokens == 50
def test_add_multiple_steps(self) -> None:
trace = Trace(query="test")
trace.add_step(TraceStep(
step_type=StepType.RETRIEVE,
timestamp=0.0,
duration_seconds=0.3,
output={"tokens": 0},
))
trace.add_step(TraceStep(
step_type=StepType.GENERATE,
timestamp=0.0,
duration_seconds=1.0,
output={"tokens": 100},
))
trace.add_step(TraceStep(
step_type=StepType.TOOL_CALL,
timestamp=0.0,
duration_seconds=0.5,
output={"tokens": 0},
))
assert len(trace.steps) == 3
assert trace.total_latency_seconds == 1.8
assert trace.total_tokens == 100
def test_trace_defaults(self) -> None:
trace = Trace()
assert trace.query == ""
assert trace.agent == ""
assert trace.model == ""
assert trace.outcome is None
assert trace.feedback is None
assert trace.steps == []
assert trace.total_tokens == 0
def test_trace_with_outcome(self) -> None:
trace = Trace(query="test", outcome="success", feedback=0.9)
assert trace.outcome == "success"
assert trace.feedback == 0.9