From 9ec402ab4a327df7d13243b16a7f62ab461fa4b8 Mon Sep 17 00:00:00 2001 From: Jon Saad-Falcon Date: Wed, 25 Feb 2026 01:58:18 +0000 Subject: [PATCH] Enrich agent tool awareness, normalize engine tool_calls, add telemetry and eval config - Add build_tool_descriptions() shared builder for enriched agent system prompts (NativeReAct, NativeOpenHands, RLM, Orchestrator structured mode) - Normalize tool_calls to flat {id, name, arguments} across CloudEngine (OpenAI/Anthropic/Google), LiteLLM, and Ollama - Add Anthropic tool_use extraction and input_schema conversion - Add Google function_call extraction - Make ReAct/OpenHands parsing case-insensitive - Add telemetry efficiency, GPU monitor, and vLLM metrics modules - Add TOML-based eval suite config system - Update documentation and changelog Co-Authored-By: Claude Opus 4.6 --- .gitignore | 5 + CLAUDE.md | 25 +- configs/openjarvis/config.toml | 97 +++++ docs/api/tools.md | 6 + docs/architecture/agents.md | 10 +- docs/architecture/engine.md | 22 + docs/development/changelog.md | 43 ++ docs/user-guide/agents.md | 4 +- docs/user-guide/tools.md | 41 ++ evals/backends/jarvis_agent.py | 16 +- evals/backends/jarvis_direct.py | 18 +- evals/cli.py | 53 ++- evals/configs/glm-4.7-flash-openhands.toml | 64 +++ evals/core/config.py | 6 +- evals/core/runner.py | 82 +++- evals/core/scorer.py | 2 +- evals/core/types.py | 41 +- evals/scorers/frames_judge.py | 2 +- evals/scorers/gaia_exact.py | 2 +- evals/scorers/supergpqa_mcq.py | 2 +- evals/scorers/wildchat_judge.py | 2 +- evals/tests/test_config.py | 8 +- evals/tests/test_types.py | 4 +- pyproject.toml | 2 + src/openjarvis/agents/native_openhands.py | 141 +++--- src/openjarvis/agents/native_react.py | 26 +- src/openjarvis/agents/openhands.py | 6 +- src/openjarvis/agents/orchestrator.py | 3 +- src/openjarvis/agents/react.py | 2 +- src/openjarvis/agents/rlm.py | 31 +- src/openjarvis/core/config.py | 4 + src/openjarvis/core/types.py | 6 + src/openjarvis/engine/_openai_compat.py | 3 + src/openjarvis/engine/_stubs.py | 3 + src/openjarvis/engine/cloud.py | 182 +++++++- src/openjarvis/engine/litellm.py | 9 +- src/openjarvis/engine/ollama.py | 8 + src/openjarvis/intelligence/model_catalog.py | 10 + .../learning/orchestrator/prompt_registry.py | 64 ++- src/openjarvis/system.py | 56 ++- src/openjarvis/telemetry/__init__.py | 23 + src/openjarvis/telemetry/aggregator.py | 28 +- src/openjarvis/telemetry/efficiency.py | 129 ++++++ src/openjarvis/telemetry/gpu_monitor.py | 313 +++++++++++++ .../telemetry/instrumented_engine.py | 95 +++- src/openjarvis/telemetry/store.py | 40 +- src/openjarvis/telemetry/vllm_metrics.py | 173 ++++++++ src/openjarvis/tools/_stubs.py | 64 ++- src/openjarvis/traces/collector.py | 16 +- tests/agents/test_agent_routing_matrix.py | 12 +- tests/agents/test_base_agent.py | 4 - tests/agents/test_native_openhands.py | 18 + tests/agents/test_native_react.py | 77 +++- tests/agents/test_orchestrator.py | 89 ++++ tests/agents/test_react.py | 2 +- tests/agents/test_rlm.py | 26 ++ tests/engine/test_cloud_extended.py | 194 +++++++++ tests/engine/test_litellm.py | 5 +- tests/telemetry/test_efficiency.py | 191 ++++++++ tests/telemetry/test_gpu_monitor.py | 410 ++++++++++++++++++ tests/telemetry/test_vllm_metrics.py | 201 +++++++++ tests/tools/test_tool_descriptions.py | 182 ++++++++ uv.lock | 29 +- 63 files changed, 3249 insertions(+), 183 deletions(-) create mode 100644 configs/openjarvis/config.toml create mode 100644 evals/configs/glm-4.7-flash-openhands.toml create mode 100644 src/openjarvis/telemetry/efficiency.py create mode 100644 src/openjarvis/telemetry/gpu_monitor.py create mode 100644 src/openjarvis/telemetry/vllm_metrics.py create mode 100644 tests/telemetry/test_efficiency.py create mode 100644 tests/telemetry/test_gpu_monitor.py create mode 100644 tests/telemetry/test_vllm_metrics.py create mode 100644 tests/tools/test_tool_descriptions.py diff --git a/.gitignore b/.gitignore index 3d26d3b2..aa0ad077 100644 --- a/.gitignore +++ b/.gitignore @@ -35,9 +35,14 @@ htmlcov/ .DS_Store Thumbs.db +# Secrets +.env +.env.* + # Project *.sqlite *.db +results/ # MkDocs build output site/ diff --git a/CLAUDE.md b/CLAUDE.md index 85a12826..916558d9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,19 +4,23 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co ## Project Status -OpenJarvis is a research framework for studying on-device AI systems. Phase 9 (Pillar-aligned config) complete. Five composable pillars: Intelligence, Engine, Agents, Tools (with storage + MCP), and Learning — with trace-driven learning as a cross-cutting concern. Python SDK (`Jarvis` class), composition layer (`SystemBuilder`/`JarvisSystem`), OpenClaw agent infrastructure, benchmarking framework, Docker deployment all ready. ~1780 tests pass (37 skipped for optional deps). +OpenJarvis is a research framework for studying on-device AI systems. Phase 10 (Agent restructuring) complete. Five composable pillars: Intelligence, Engine, Agents, Tools (with storage + MCP), and Learning — with trace-driven learning as a cross-cutting concern. Python SDK (`Jarvis` class), composition layer (`SystemBuilder`/`JarvisSystem`), OpenClaw agent infrastructure, benchmarking framework, Docker deployment all ready. Agent hierarchy refactored: `BaseAgent` (with shared helpers) → `ToolUsingAgent` (tool-using agents), `accepts_tools` introspection, real OpenHands SDK integration. ~1910 tests pass (36 skipped for optional deps). Eval framework with 4 benchmarks (SuperGPQA, GAIA, FRAMES, WildChat) and LLM-as-judge scoring via `gpt-5-mini-2025-08-07`. Tool system enriched with shared `build_tool_descriptions()` builder; engine tool_calls normalized across OpenAI, Anthropic, Google, and LiteLLM. ## Build & Development Commands ```bash uv sync --extra dev # Install deps + dev tools -uv run pytest tests/ -v # Run ~1780 tests (37 skipped if optional deps missing) +uv run pytest tests/ -v # Run ~1910 tests (36 skipped if optional deps missing) uv run ruff check src/ tests/ # Lint uv run jarvis --version # 1.0.0 uv run jarvis ask "Hello" # Query via discovered engine (direct mode) uv run jarvis ask --agent simple "Hello" # SimpleAgent route uv run jarvis ask --agent orchestrator "Hello" # OrchestratorAgent route uv run jarvis ask --agent orchestrator --tools calculator,think "What is 2+2?" +uv run jarvis ask --agent native_react --tools calculator "What is 2+2?" # NativeReActAgent +uv run jarvis ask --agent react "Hello" # Alias for native_react +uv run jarvis ask --agent native_openhands "Hello" # NativeOpenHandsAgent (CodeAct) +uv run jarvis ask --agent openhands "Hello" # Real OpenHands SDK (requires openhands-sdk) uv run jarvis ask --router heuristic "Hello" # Explicit heuristic policy uv run jarvis ask --no-context "Hello" # Query without memory context injection uv run jarvis model list # List models from running engines @@ -37,8 +41,20 @@ uv run jarvis bench run -b latency -o results.jsonl # Specific benchmark to fil uv run jarvis serve --port 8000 # OpenAI-compatible API server (requires openjarvis[server]) uv run jarvis --help # Show all subcommands uv run jarvis init --force # Detect hardware, write ~/.openjarvis/config.toml +# Eval framework +source .env # Load API keys before running evals +uv run python -m evals run -c evals/configs/glm-4.7-flash-openhands.toml -v # Run eval suite from TOML config +uv run python -m evals run -b supergpqa -m "qwen3:8b" -n 50 # Run single benchmark +uv run python -m evals summarize results/supergpqa_qwen3-8b.jsonl # Summarize results ``` +### Config File Conventions + +- **Runtime config (source of truth):** `configs/openjarvis/config.toml` — Pillar-aligned OpenJarvis config. Copied to `~/.openjarvis/config.toml` at runtime (which is where `load_config()` reads from). +- **Eval suite configs:** `evals/configs/*.toml` — TOML configs defining models x benchmarks matrices. +- **API keys:** `.env` file in project root (gitignored). Source with `source .env` before running evals or cloud operations. +- **Never save configs to `~/.openjarvis/` directly** — always maintain the canonical copy in `configs/openjarvis/` and copy/symlink to `~/.openjarvis/`. + ### Python SDK ```python @@ -75,7 +91,7 @@ OpenJarvis is a research framework for on-device AI organized around **five comp 1. **Intelligence** (`src/openjarvis/intelligence/`) — The model definition, catalog, and generation defaults. `ModelRegistry` maps model keys to `ModelSpec`. `IntelligenceConfig` holds model identity (default/fallback model, model_path, checkpoint_path, quantization, preferred_engine, provider) and generation defaults (temperature, max_tokens, top_p, top_k, repetition_penalty, stop_sequences). Model catalog (`model_catalog.py`) maintains `BUILTIN_MODELS` list with auto-discovery via `merge_discovered_models()`. Backward-compat shims in `intelligence/_stubs.py` and `intelligence/router.py` re-export from `learning/` for old import paths. 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. **Agents** (`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`), `ReActAgent` (Thought-Action-Observation loop), `OpenHandsAgent` (CodeAct-style), `CustomAgent` (template for user-defined agents), `OpenClawAgent` (HTTP/subprocess transport). All implement `BaseAgent` ABC with `run()`. Agents call `engine.generate()` directly — telemetry is handled by the `InstrumentedEngine` wrapper when enabled. +3. **Agents** (`src/openjarvis/agents/`) — Pluggable logic for handling queries, making tool/API calls, managing memory. Class hierarchy: `BaseAgent` ABC (with concrete helpers: `_emit_turn_start/end`, `_build_messages`, `_generate`, `_max_turns_result`, `_strip_think_tags`) → `ToolUsingAgent` (adds `tools`, `ToolExecutor`, `max_turns`). Agents: `SimpleAgent` (single-turn, no tools, extends `BaseAgent`), `OrchestratorAgent` (multi-turn tool-calling loop, extends `ToolUsingAgent`), `NativeReActAgent` (Thought-Action-Observation loop, extends `ToolUsingAgent`, registry key `"native_react"`, alias `"react"`), `NativeOpenHandsAgent` (CodeAct-style code execution, extends `ToolUsingAgent`, registry key `"native_openhands"`), `RLMAgent` (recursive LM with code gen, extends `ToolUsingAgent`), `OpenHandsAgent` (wraps real `openhands-sdk`, extends `BaseAgent`, registry key `"openhands"`, requires `openjarvis[openhands]` + Python 3.12+), `OpenClawAgent` (HTTP/subprocess transport). `accepts_tools` class attribute enables CLI/SDK to auto-detect tool-using agents. Backward-compat: `from openjarvis.agents.react import ReActAgent` still works (shim). Agents call `engine.generate()` directly — telemetry is handled by the `InstrumentedEngine` wrapper when enabled. 4. **Tools** (`src/openjarvis/tools/`) — All tools managed via MCP (Model Context Protocol). - **API tools**: `CalculatorTool`, `ThinkTool`, `FileReadTool`, `WebSearchTool`, `CodeInterpreterTool` — all implement `BaseTool` ABC - **Storage tools** (`tools/storage_tools.py`): `MemoryStoreTool`, `MemoryRetrieveTool`, `MemorySearchTool`, `MemoryIndexTool` — wrap `MemoryBackend` operations as MCP-discoverable tools @@ -194,7 +210,7 @@ OpenAI-compatible server via `jarvis serve`: `POST /v1/chat/completions`, `GET / - **Offline-first:** Cloud APIs are optional. All core functionality works without network. - **Hardware-aware:** Auto-detect GPU vendor/model/VRAM via `nvidia-smi`, `rocm-smi`, `system_profiler`, `/proc/cpuinfo`. Recommend engine accordingly. - **Telemetry opt-in:** When enabled, `InstrumentedEngine` wraps the inference engine to transparently record timing, tokens, energy, cost to SQLite via event bus. Agents call `engine.generate()` without awareness of telemetry. `TelemetryAggregator` provides read-only query/aggregation over stored records. -- **Backward-compat shims:** `memory/` re-exports from `tools/storage/`, `intelligence/_stubs.py` re-exports `RouterPolicy`/`QueryAnalyzer` from `learning/_stubs.py`, `intelligence/router.py` re-exports `HeuristicRouter`/`build_routing_context`/`DefaultQueryAnalyzer` from `learning/router.py`. Old import paths continue to work. Config backward-compat: `engine.ollama_host` → `engine.ollama.host`, `agent.default_tools` → `agent.tools`, `learning.default_policy` → `learning.routing.policy`, etc. TOML migration: `agent.temperature` → `intelligence.temperature`, `memory.context_injection` → `agent.context_from_memory`. +- **Backward-compat shims:** `memory/` re-exports from `tools/storage/`, `intelligence/_stubs.py` re-exports `RouterPolicy`/`QueryAnalyzer` from `learning/_stubs.py`, `intelligence/router.py` re-exports `HeuristicRouter`/`build_routing_context`/`DefaultQueryAnalyzer` from `learning/router.py`, `agents/react.py` re-exports `NativeReActAgent` as `ReActAgent` (old import path still works), registry alias `"react"` → `NativeReActAgent`. Old import paths continue to work. Config backward-compat: `engine.ollama_host` → `engine.ollama.host`, `agent.default_tools` → `agent.tools`, `learning.default_policy` → `learning.routing.policy`, etc. TOML migration: `agent.temperature` → `intelligence.temperature`, `memory.context_injection` → `agent.context_from_memory`. - **`ensure_registered()` pattern:** Benchmark and learning modules use lazy registration via `ensure_registered()` to survive registry clearing in tests. ## Development Phases @@ -211,3 +227,4 @@ OpenAI-compatible server via `jarvis serve`: `POST /v1/chat/completions`, `GET / | v1.2 | Phase 7 | 5-pillar restructuring: Intelligence ABCs, memory→tools/storage, MCP tool management, composition layer (SystemBuilder/JarvisSystem), InstrumentedEngine, structured learning (SFT/AgentAdvisor/ICL), config schema update | | v1.3 | Phase 8 | Intelligence = "The Model": routing moved to Learning, enriched IntelligenceConfig (model_path, checkpoint_path, quantization, preferred_engine, provider), intelligence-driven engine selection, backward-compat shims | | v1.4 | Phase 9 | Pillar-aligned config: generation params in Intelligence, nested engine/learning configs, agent objective/system_prompt/context_from_memory, structured learning sub-policies (routing/intelligence/agent/metrics), TOML migration layer | +| v1.5 | Phase 10 | Agent restructuring: BaseAgent helpers (`_emit_turn_start/end`, `_build_messages`, `_generate`, `_max_turns_result`), ToolUsingAgent intermediate base, `accepts_tools` introspection, NativeReActAgent/NativeOpenHandsAgent renames, real OpenHands SDK integration, CLI/SDK tool-passing bug fix, backward-compat shims | diff --git a/configs/openjarvis/config.toml b/configs/openjarvis/config.toml new file mode 100644 index 00000000..8eddb841 --- /dev/null +++ b/configs/openjarvis/config.toml @@ -0,0 +1,97 @@ +# OpenJarvis configuration — GLM-4.7-Flash eval on 8x A100-80GB + +# ═══════════════════════════════════════════════════════════════ +# PILLAR 1: Intelligence — The Model +# ═══════════════════════════════════════════════════════════════ +[intelligence] +default_model = "zai-org/GLM-4.7-Flash" # HuggingFace model ID +fallback_model = "glm-4.7-flash" # Catalog alias fallback +preferred_engine = "vllm" # MoE model, vLLM is best for A100s +provider = "local" # Running locally on this machine +quantization = "none" # Full precision, we have 640GB VRAM +# Generation defaults for eval +temperature = 0.0 # Deterministic for reproducibility +max_tokens = 2048 # Standard eval output length +top_p = 0.9 +top_k = 40 +repetition_penalty = 1.0 + +# ═══════════════════════════════════════════════════════════════ +# PILLAR 2: Agent — The Agentic Harness +# ═══════════════════════════════════════════════════════════════ +[agent] +default_agent = "native_openhands" # CodeAct-style agent +max_turns = 10 # Up to 10 tool-calling turns +tools = "code_interpreter,web_search,file_read,calculator,think" +objective = "Answer questions accurately using available tools" +context_from_memory = false # No memory injection during eval + +# ═══════════════════════════════════════════════════════════════ +# PILLAR 3: Tools — MCP Interface +# ═══════════════════════════════════════════════════════════════ +[tools.storage] +default_backend = "sqlite" +db_path = "~/.openjarvis/memory.db" + +[tools.mcp] +enabled = true + +# ═══════════════════════════════════════════════════════════════ +# PILLAR 4: Engine — The Inference Runtime +# ═══════════════════════════════════════════════════════════════ +[engine] +default = "vllm" + +[engine.vllm] +host = "http://localhost:8000" # Default vLLM port + +[engine.ollama] +host = "http://localhost:11434" + +[engine.sglang] +host = "http://localhost:30000" + +[engine.llamacpp] +host = "http://localhost:8080" + +# ═══════════════════════════════════════════════════════════════ +# PILLAR 5: Learning — Improvement Methodologies +# ═══════════════════════════════════════════════════════════════ +[learning] +enabled = false # No learning during eval + +[learning.routing] +policy = "heuristic" + +[learning.intelligence] +policy = "none" + +[learning.agent] +policy = "none" + +[learning.metrics] +accuracy_weight = 0.6 +latency_weight = 0.2 +cost_weight = 0.1 +efficiency_weight = 0.1 + +# ═══════════════════════════════════════════════════════════════ +# Supporting config +# ═══════════════════════════════════════════════════════════════ +[telemetry] +enabled = true +db_path = "~/.openjarvis/telemetry.db" +gpu_metrics = true +gpu_poll_interval_ms = 50 + +[traces] +enabled = true # Record traces for analysis +db_path = "~/.openjarvis/traces.db" + +[server] +host = "0.0.0.0" +port = 8000 +agent = "native_openhands" + +[security] +enabled = false # Disable for eval (no PII scanning overhead) diff --git a/docs/api/tools.md b/docs/api/tools.md index bac5cc9c..aa3a4a85 100644 --- a/docs/api/tools.md +++ b/docs/api/tools.md @@ -29,6 +29,12 @@ with JSON argument parsing, latency tracking, and event bus integration. show_source: true members_order: source +### build_tool_descriptions + +::: openjarvis.tools._stubs.build_tool_descriptions + options: + show_source: true + --- ## Built-in Tools diff --git a/docs/architecture/agents.md b/docs/architecture/agents.md index 67efb12f..6e252965 100644 --- a/docs/architecture/agents.md +++ b/docs/architecture/agents.md @@ -197,14 +197,14 @@ A ReAct (Reasoning + Acting) agent that implements a **Thought-Action-Observatio ```mermaid graph TD - Q["User Query"] --> SYS["Build system prompt
with tool names"] + Q["User Query"] --> SYS["Build system prompt
with tool descriptions"] SYS --> GEN["Generate response"] GEN --> PARSE["Parse ReAct output"] PARSE --> FINAL{"Final Answer?"} FINAL -->|Yes| DONE["Return answer"] FINAL -->|No| ACTION{"Has Action?"} ACTION -->|No| DONE2["Return content as-is"] - ACTION -->|Yes| EXEC["Execute tool
via ToolExecutor"] + ACTION -->|Yes| EXEC["Execute tool
via ToolExecutor
(case-insensitive)"] EXEC --> OBS["Append Observation"] OBS --> MAXCHECK{"Max turns
exceeded?"} MAXCHECK -->|No| GEN @@ -213,7 +213,7 @@ graph TD How it works: -1. Builds a system prompt listing available tool names +1. Builds a system prompt with enriched tool descriptions via `build_tool_descriptions()`. Parsing is case-insensitive. 2. Generates a response and parses the ReAct-structured output 3. If a `Final Answer:` is found, returns it 4. If an `Action:` is found, executes the tool and feeds the result back as an `Observation:` @@ -242,7 +242,7 @@ A CodeAct-style agent that generates and executes Python code. Extends `ToolUsin How it works: -1. Builds a detailed system prompt with tool descriptions and code execution instructions +1. Builds a detailed system prompt with enriched tool descriptions (via shared `build_tool_descriptions()` builder) and code execution instructions 2. Pre-fetches any URLs in the user input, inlining the content directly 3. For each turn: - Generates a response and strips `` tags @@ -287,7 +287,7 @@ graph TD How it works: -1. Creates a persistent REPL with `llm_query()` and `llm_batch()` callbacks +1. Creates a persistent REPL with `llm_query()` and `llm_batch()` callbacks. Tool descriptions are injected via the shared `build_tool_descriptions()` builder when tools are provided. 2. Injects context from `AgentContext` metadata or memory results into the REPL as a variable 3. Generates code and executes it in the REPL 4. If `FINAL(value)` or `FINAL_VAR("name")` is called, returns the final answer diff --git a/docs/architecture/engine.md b/docs/architecture/engine.md index f46c2bba..7710017f 100644 --- a/docs/architecture/engine.md +++ b/docs/architecture/engine.md @@ -80,6 +80,28 @@ When the model requests tool calls, they are extracted and passed through in Ope } ``` +### Multi-Provider Tool Call Extraction + +Engine backends normalize tool calls from different providers into the standard flat format used by agents: + +| Provider | Source Format | Extraction Logic | +|----------|-------------|-----------------| +| **OpenAI** | `choices[0].message.tool_calls[].function.{name, arguments}` | Direct extraction, add `id` from `tool_calls[].id` | +| **Anthropic** | `content[]` blocks with `type: "tool_use"` | Filter `tool_use` blocks, map `input` dict to JSON `arguments` | +| **Google** | `candidates[0].content.parts[]` with `function_call` | Extract `function_call.name` and `function_call.args`, serialize args to JSON | +| **LiteLLM** | Flat `{id, name, arguments}` dicts (proxy pre-normalizes) | Pass through directly | +| **Ollama** | `message.tool_calls[].function.{name, arguments}` | Extract from Ollama native format, serialize arguments dict to JSON | + +All providers produce the same output format consumed by agents: + +```python +{ + "tool_calls": [ + {"id": "call_abc", "name": "calculator", "arguments": "{\"expression\": \"2+2\"}"} + ] +} +``` + --- ## Backend Comparison diff --git a/docs/development/changelog.md b/docs/development/changelog.md index 931af5c4..50b7e611 100644 --- a/docs/development/changelog.md +++ b/docs/development/changelog.md @@ -4,6 +4,49 @@ All notable changes to OpenJarvis are documented in this file. --- +## Unreleased + +### Added + +- **`build_tool_descriptions()` shared builder** -- Single source of truth for + generating enriched tool descriptions in agent system prompts. Produces + Markdown sections with name, description, category, and parameter schemas. +- **Enriched agent prompts** -- `NativeReActAgent`, `NativeOpenHandsAgent`, + `RLMAgent`, and `OrchestratorAgent` (structured mode) now inject detailed + tool descriptions into their system prompts via the shared builder. +- **Case-insensitive parsing** -- ReAct (`Action:` / `Final Answer:`) and + Orchestrator structured-mode parsing (`TOOL:` / `FINAL_ANSWER:`) are now + case-insensitive. +- **Multi-provider tool_calls extraction** -- `CloudEngine` now extracts + `tool_calls` from Anthropic (`tool_use` content blocks) and Google + (`function_call` parts), normalizing to the flat `{id, name, arguments}` + format. `LiteLLM` engine handles the flat-format tool calls returned by + the LiteLLM proxy. +- **RLM tool awareness** -- `RLMAgent` injects an `## Available Tools` + section into its system prompt when tools are provided. +- **Orchestrator structured tool descriptions** -- Structured mode passes + `tools=self._tools` to `build_system_prompt()` for enriched descriptions. +- **Telemetry modules** -- `EfficiencyMetrics`, `GPUMonitor`, `VLLMMetrics` + for energy, GPU utilization, and vLLM server-side metrics collection. +- **Eval TOML config** -- TOML-based eval suite configuration system for + defining models x benchmarks matrices. + +### Changed + +- Agent prompt generation now uses `build_tool_descriptions()` instead of + inline tool name listing. +- `build_system_prompt()` in `prompt_registry.py` accepts an optional `tools` + parameter for enriched descriptions from `BaseTool` instances. +- ReAct and OpenHands regex patterns updated for case-insensitive matching. + +### Fixed + +- Engine `tool_calls` normalization -- Anthropic `tool_use` blocks and Google + `function_call` parts are now correctly extracted and converted to the + standard flat format used by agents. + +--- + ## v1.0.0 *Phase 5 -- SDK, Production Readiness, and Documentation* diff --git a/docs/user-guide/agents.md b/docs/user-guide/agents.md index f11aeb82..3ad8967a 100644 --- a/docs/user-guide/agents.md +++ b/docs/user-guide/agents.md @@ -142,7 +142,7 @@ The `NativeReActAgent` implements a **Thought-Action-Observation** loop followin **How it works:** -1. Builds a system prompt listing available tool names. +1. Builds a system prompt with enriched tool descriptions (names, parameter schemas, categories) via `build_tool_descriptions()`. Parsing is case-insensitive. 2. Generates a response and parses the ReAct-structured output. 3. If a `Final Answer:` is found, returns it. 4. If an `Action:` is found, executes the tool and feeds the result back as an `Observation:`. @@ -173,7 +173,7 @@ The `NativeOpenHandsAgent` is a CodeAct-style agent that generates and executes **How it works:** -1. Builds a detailed system prompt with tool descriptions and code execution instructions. +1. Builds a detailed system prompt with enriched tool descriptions (via shared `build_tool_descriptions()` builder) and code execution instructions. 2. Pre-fetches any URLs in the user input, inlining the content directly. 3. For each turn, generates a response and attempts to extract code blocks or tool calls. 4. Code is executed via `code_interpreter`; tool calls are dispatched via `ToolExecutor`. diff --git a/docs/user-guide/tools.md b/docs/user-guide/tools.md index 1df2a0a3..7713be09 100644 --- a/docs/user-guide/tools.md +++ b/docs/user-guide/tools.md @@ -135,6 +135,47 @@ If the tool name is unknown, a `ToolResult` with `success=False` is returned. If --- +## build_tool_descriptions() + +The `build_tool_descriptions()` function is the **single source of truth** for generating enriched tool descriptions used in agent system prompts. All text-based agents (`NativeReActAgent`, `NativeOpenHandsAgent`, `RLMAgent`, `OrchestratorAgent` in structured mode) use this function to produce Markdown-formatted tool sections. + +### Parameters + +| Parameter | Type | Default | Description | +|--------------------|------------------|---------|------------------------------------------------------| +| `tools` | `list[BaseTool]` | -- | List of tool instances to describe | +| `include_category` | `bool` | `True` | Whether to include the `Category:` line | +| `include_cost` | `bool` | `False` | Whether to include cost and latency estimate lines | + +### Usage + +```python +from openjarvis.tools._stubs import build_tool_descriptions +from openjarvis.tools.calculator import CalculatorTool +from openjarvis.tools.think import ThinkTool + +tools = [CalculatorTool(), ThinkTool()] +desc = build_tool_descriptions(tools) +print(desc) +# ### calculator +# Evaluate a mathematical expression safely. +# Category: math +# Parameters: +# - expression (string, required): Math expression to evaluate +# +# ### think +# Reasoning scratchpad ... +``` + +### Agents using this builder + +- **NativeReActAgent** -- Injects descriptions into the ReAct system prompt +- **NativeOpenHandsAgent** -- Injects descriptions into the CodeAct system prompt +- **RLMAgent** -- Adds an `## Available Tools` section to the REPL system prompt +- **OrchestratorAgent** (structured mode) -- Passes `tools=` to `build_system_prompt()` which delegates to this builder + +--- + ## Built-in Tools ### Calculator diff --git a/evals/backends/jarvis_agent.py b/evals/backends/jarvis_agent.py index 8cc86d16..06cd9a60 100644 --- a/evals/backends/jarvis_agent.py +++ b/evals/backends/jarvis_agent.py @@ -22,11 +22,15 @@ class JarvisAgentBackend(InferenceBackend): engine_key: Optional[str] = None, agent_name: str = "orchestrator", tools: Optional[List[str]] = None, + telemetry: bool = False, + gpu_metrics: bool = False, ) -> None: from openjarvis.system import SystemBuilder self._agent_name = agent_name self._tools = tools or [] + self._telemetry = telemetry + self._gpu_metrics = gpu_metrics builder = SystemBuilder() if engine_key: @@ -34,7 +38,11 @@ class JarvisAgentBackend(InferenceBackend): builder.agent(agent_name) if tools: builder.tools(tools) - self._system = builder.telemetry(False).traces(False).build() + # Propagate gpu_metrics to the runtime config so SystemBuilder + # creates a GpuMonitor when building the InstrumentedEngine. + if gpu_metrics: + builder._config.telemetry.gpu_metrics = True + self._system = builder.telemetry(telemetry).traces(telemetry).build() def generate( self, @@ -71,6 +79,7 @@ class JarvisAgentBackend(InferenceBackend): elapsed = time.monotonic() - t0 usage = result.get("usage", {}) + telemetry_data = result.get("_telemetry", {}) return { "content": result.get("content", ""), "usage": usage, @@ -79,6 +88,11 @@ class JarvisAgentBackend(InferenceBackend): "cost_usd": result.get("cost_usd", 0.0), "turns": result.get("turns", 1), "tool_results": result.get("tool_results", []), + "ttft": result.get("ttft", telemetry_data.get("ttft", 0.0)), + "energy_joules": telemetry_data.get("energy_joules", 0.0), + "power_watts": telemetry_data.get("power_watts", 0.0), + "gpu_utilization_pct": telemetry_data.get("gpu_utilization_pct", 0.0), + "throughput_tok_per_sec": telemetry_data.get("throughput_tok_per_sec", 0.0), } def close(self) -> None: diff --git a/evals/backends/jarvis_direct.py b/evals/backends/jarvis_direct.py index 69f6b5c6..d215f0fa 100644 --- a/evals/backends/jarvis_direct.py +++ b/evals/backends/jarvis_direct.py @@ -17,13 +17,21 @@ class JarvisDirectBackend(InferenceBackend): backend_id = "jarvis-direct" - def __init__(self, engine_key: Optional[str] = None) -> None: + def __init__( + self, + engine_key: Optional[str] = None, + telemetry: bool = False, + gpu_metrics: bool = False, + ) -> None: from openjarvis.system import SystemBuilder + self._telemetry = telemetry + self._gpu_metrics = gpu_metrics + builder = SystemBuilder() if engine_key: builder.engine(engine_key) - self._system = builder.telemetry(False).traces(False).build() + self._system = builder.telemetry(telemetry).traces(telemetry).build() def generate( self, @@ -64,12 +72,18 @@ class JarvisDirectBackend(InferenceBackend): elapsed = time.monotonic() - t0 usage = result.get("usage", {}) + telemetry_data = result.get("_telemetry", {}) return { "content": result.get("content", ""), "usage": usage, "model": result.get("model", model), "latency_seconds": elapsed, "cost_usd": result.get("cost_usd", 0.0), + "ttft": result.get("ttft", telemetry_data.get("ttft", 0.0)), + "energy_joules": telemetry_data.get("energy_joules", 0.0), + "power_watts": telemetry_data.get("power_watts", 0.0), + "gpu_utilization_pct": telemetry_data.get("gpu_utilization_pct", 0.0), + "throughput_tok_per_sec": telemetry_data.get("throughput_tok_per_sec", 0.0), } def close(self) -> None: diff --git a/evals/cli.py b/evals/cli.py index a9632dd4..fec0d1fa 100644 --- a/evals/cli.py +++ b/evals/cli.py @@ -33,7 +33,8 @@ def _setup_logging(verbose: bool) -> None: def _build_backend(backend_name: str, engine_key: Optional[str], - agent_name: str, tools: list[str]): + agent_name: str, tools: list[str], + telemetry: bool = False, gpu_metrics: bool = False): """Construct the appropriate backend.""" if backend_name == "jarvis-agent": from evals.backends.jarvis_agent import JarvisAgentBackend @@ -41,10 +42,16 @@ def _build_backend(backend_name: str, engine_key: Optional[str], engine_key=engine_key, agent_name=agent_name, tools=tools, + telemetry=telemetry, + gpu_metrics=gpu_metrics, ) else: from evals.backends.jarvis_direct import JarvisDirectBackend - return JarvisDirectBackend(engine_key=engine_key) + return JarvisDirectBackend( + engine_key=engine_key, + telemetry=telemetry, + gpu_metrics=gpu_metrics, + ) def _build_dataset(benchmark: str): @@ -107,6 +114,34 @@ def _print_summary(summary) -> None: for subj, stats in sorted(summary.per_subject.items()): click.echo(f" {subj}: {stats['accuracy']:.4f} " f"({int(stats['correct'])}/{int(stats['scored'])})") + # GPU telemetry stats + _stats_rows = [] + for label, stats_field in [ + ("Accuracy", "accuracy_stats"), + ("Latency (s)", "latency_stats"), + ("TTFT (s)", "ttft_stats"), + ("Energy (J)", "energy_stats"), + ("Power (W)", "power_stats"), + ("GPU Util (%)", "gpu_utilization_stats"), + ("Throughput (tok/s)", "throughput_stats"), + ("MFU (%)", "mfu_stats"), + ("MBU (%)", "mbu_stats"), + ("IPW", "ipw_stats"), + ("IPJ", "ipj_stats"), + ]: + ms = getattr(summary, stats_field, None) + if ms is not None: + _stats_rows.append((label, ms)) + if _stats_rows: + click.echo(f"\n{'Metric':20s} {'Mean':>10s} {'Median':>10s} " + f"{'Min':>10s} {'Max':>10s} {'Std':>10s}") + click.echo(f"{'-' * 20} {'-' * 10} {'-' * 10} " + f"{'-' * 10} {'-' * 10} {'-' * 10}") + for label, ms in _stats_rows: + click.echo(f"{label:20s} {ms.mean:10.4f} {ms.median:10.4f} " + f"{ms.min:10.4f} {ms.max:10.4f} {ms.std:10.4f}") + if getattr(summary, "total_energy_joules", 0.0) > 0: + click.echo(f"\nTotal Energy: {summary.total_energy_joules:.4f} J") click.echo(f"{'=' * 60}") @@ -119,6 +154,8 @@ def _run_single(config) -> object: config.engine_key, config.agent_name or "orchestrator", config.tools, + telemetry=getattr(config, "telemetry", False), + gpu_metrics=getattr(config, "gpu_metrics", False), ) dataset = _build_dataset(config.benchmark) judge_backend = _build_judge_backend(config.judge_model) @@ -200,7 +237,7 @@ def main(): help="Maximum samples to evaluate") @click.option("-w", "--max-workers", type=int, default=4, help="Parallel workers") -@click.option("--judge-model", default="gpt-4o", +@click.option("--judge-model", default="gpt-5-mini-2025-08-07", help="LLM judge model") @click.option("-o", "--output", "output_path", default=None, help="Output JSONL path") @@ -211,11 +248,15 @@ def main(): help="Generation temperature") @click.option("--max-tokens", type=int, default=2048, help="Max output tokens") +@click.option("--telemetry/--no-telemetry", default=False, + help="Enable telemetry collection during eval") +@click.option("--gpu-metrics/--no-gpu-metrics", default=False, + help="Enable GPU metrics collection") @click.option("-v", "--verbose", is_flag=True, help="Verbose logging") @click.pass_context def run(ctx, config_path, benchmark, backend, model, engine_key, agent_name, tools, max_samples, max_workers, judge_model, output_path, seed, - dataset_split, temperature, max_tokens, verbose): + dataset_split, temperature, max_tokens, telemetry, gpu_metrics, verbose): """Run a single benchmark evaluation, or a full suite from a TOML config.""" _setup_logging(verbose) @@ -255,6 +296,8 @@ def run(ctx, config_path, benchmark, backend, model, engine_key, agent_name, output_path=output_path, seed=seed, dataset_split=dataset_split, + telemetry=telemetry, + gpu_metrics=gpu_metrics, ) summary = _run_single(config) @@ -269,7 +312,7 @@ def run(ctx, config_path, benchmark, backend, model, engine_key, agent_name, help="Max samples per benchmark") @click.option("-w", "--max-workers", type=int, default=4, help="Parallel workers") -@click.option("--judge-model", default="gpt-4o", help="LLM judge model") +@click.option("--judge-model", default="gpt-5-mini-2025-08-07", help="LLM judge model") @click.option("--output-dir", default="results/", help="Output directory for results") @click.option("--seed", type=int, default=42, help="Random seed") diff --git a/evals/configs/glm-4.7-flash-openhands.toml b/evals/configs/glm-4.7-flash-openhands.toml new file mode 100644 index 00000000..20d700d3 --- /dev/null +++ b/evals/configs/glm-4.7-flash-openhands.toml @@ -0,0 +1,64 @@ +# GLM-4.7-Flash eval with NativeOpenHandsAgent on all 4 benchmarks. +# Machine: 8x A100-80GB, engine: vLLM + +[meta] +name = "glm-4.7-flash-openhands" +description = "Evaluate GLM-4.7-Flash via NativeOpenHandsAgent across all benchmark categories" + +[defaults] +temperature = 0.0 +max_tokens = 2048 + +[judge] +model = "gpt-5-mini-2025-08-07" +temperature = 0.0 +max_tokens = 1024 + +[run] +max_workers = 4 +output_dir = "results/glm-4.7-flash-openhands-v2/" +seed = 42 +telemetry = true +gpu_metrics = true + +# --- Model Under Test --- + +[[models]] +name = "zai-org/GLM-4.7-Flash" +engine = "vllm" + +# --- Benchmarks --- + +# Reasoning (MCQ): SuperGPQA — LLM judge extracts answer letter +[[benchmarks]] +name = "supergpqa" +backend = "jarvis-agent" +agent = "native_openhands" +tools = ["code_interpreter", "calculator", "think"] +max_samples = 200 +split = "train" + +# Agentic: GAIA — requires file reading, multi-step reasoning +[[benchmarks]] +name = "gaia" +backend = "jarvis-agent" +agent = "native_openhands" +tools = ["code_interpreter", "web_search", "file_read", "calculator", "think"] +max_samples = 50 + +# RAG: FRAMES — multi-hop factual retrieval +[[benchmarks]] +name = "frames" +backend = "jarvis-agent" +agent = "native_openhands" +tools = ["code_interpreter", "web_search", "calculator", "think"] +max_samples = 100 + +# Chat: WildChat — conversation quality +[[benchmarks]] +name = "wildchat" +backend = "jarvis-agent" +agent = "native_openhands" +tools = ["code_interpreter", "think"] +max_samples = 150 +temperature = 0.7 diff --git a/evals/core/config.py b/evals/core/config.py index e9cd7321..d3eed182 100644 --- a/evals/core/config.py +++ b/evals/core/config.py @@ -78,7 +78,7 @@ def load_eval_config(path: str | Path) -> EvalSuiteConfig: # Parse [judge] judge_raw = raw.get("judge", {}) judge = JudgeConfig( - model=judge_raw.get("model", "gpt-4o"), + model=judge_raw.get("model", "gpt-5-mini-2025-08-07"), provider=judge_raw.get("provider"), temperature=float(judge_raw.get("temperature", 0.0)), max_tokens=int(judge_raw.get("max_tokens", 1024)), @@ -90,6 +90,8 @@ def load_eval_config(path: str | Path) -> EvalSuiteConfig: max_workers=int(run_raw.get("max_workers", 4)), output_dir=run_raw.get("output_dir", "results/"), seed=int(run_raw.get("seed", 42)), + telemetry=bool(run_raw.get("telemetry", False)), + gpu_metrics=bool(run_raw.get("gpu_metrics", False)), ) # Parse [[models]] @@ -212,6 +214,8 @@ def expand_suite(suite: EvalSuiteConfig) -> List[RunConfig]: output_path=output_path, seed=suite.run.seed, dataset_split=bench.split, + telemetry=suite.run.telemetry, + gpu_metrics=suite.run.gpu_metrics, )) return configs diff --git a/evals/core/runner.py b/evals/core/runner.py index 96e54dcb..0c5faba5 100644 --- a/evals/core/runner.py +++ b/evals/core/runner.py @@ -4,6 +4,7 @@ from __future__ import annotations import json import logging +import statistics import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed @@ -13,7 +14,7 @@ from typing import Any, Dict, List, Optional from evals.core.backend import InferenceBackend from evals.core.dataset import DatasetProvider from evals.core.scorer import Scorer -from evals.core.types import EvalRecord, EvalResult, RunConfig, RunSummary +from evals.core.types import EvalRecord, EvalResult, MetricStats, RunConfig, RunSummary LOGGER = logging.getLogger(__name__) @@ -111,6 +112,11 @@ class EvalRunner: completion_tokens=usage.get("completion_tokens", 0), cost_usd=cost, scoring_metadata=scoring_meta, + ttft=full.get("ttft", 0.0), + energy_joules=full.get("energy_joules", 0.0), + power_watts=full.get("power_watts", 0.0), + gpu_utilization_pct=full.get("gpu_utilization_pct", 0.0), + throughput_tok_per_sec=full.get("throughput_tok_per_sec", 0.0), ) except Exception as exc: LOGGER.error("Error processing %s: %s", record.record_id, exc) @@ -138,6 +144,15 @@ class EvalRunner: "cost_usd": result.cost_usd, "error": result.error, "scoring_metadata": result.scoring_metadata, + "ttft": result.ttft, + "energy_joules": result.energy_joules, + "power_watts": result.power_watts, + "gpu_utilization_pct": result.gpu_utilization_pct, + "throughput_tok_per_sec": result.throughput_tok_per_sec, + "mfu_pct": result.mfu_pct, + "mbu_pct": result.mbu_pct, + "ipw": result.ipw, + "ipj": result.ipj, } self._output_file.write(json.dumps(record_dict) + "\n") self._output_file.flush() @@ -195,6 +210,21 @@ class EvalRunner: accuracy = len(correct) / len(scored) if scored else 0.0 + # Compute MetricStats for each metric + accuracy_vals = [1.0 if r.is_correct else 0.0 for r in scored] + latency_vals = [r.latency_seconds for r in results if r.latency_seconds > 0] + ttft_vals = [r.ttft for r in results if r.ttft > 0] + energy_vals = [r.energy_joules for r in results if r.energy_joules > 0] + power_vals = [r.power_watts for r in results if r.power_watts > 0] + gpu_util_vals = [r.gpu_utilization_pct for r in results if r.gpu_utilization_pct > 0] + throughput_vals = [r.throughput_tok_per_sec for r in results if r.throughput_tok_per_sec > 0] + mfu_vals = [r.mfu_pct for r in results if r.mfu_pct > 0] + mbu_vals = [r.mbu_pct for r in results if r.mbu_pct > 0] + ipw_vals = [r.ipw for r in results if r.ipw > 0] + ipj_vals = [r.ipj for r in results if r.ipj > 0] + + total_energy = sum(r.energy_joules for r in results) + return RunSummary( benchmark=cfg.benchmark, category=category, @@ -210,9 +240,47 @@ class EvalRunner: per_subject=per_subject, started_at=started_at, ended_at=ended_at, + accuracy_stats=_metric_stats(accuracy_vals), + latency_stats=_metric_stats(latency_vals), + ttft_stats=_metric_stats(ttft_vals), + energy_stats=_metric_stats(energy_vals), + power_stats=_metric_stats(power_vals), + gpu_utilization_stats=_metric_stats(gpu_util_vals), + throughput_stats=_metric_stats(throughput_vals), + mfu_stats=_metric_stats(mfu_vals), + mbu_stats=_metric_stats(mbu_vals), + ipw_stats=_metric_stats(ipw_vals), + ipj_stats=_metric_stats(ipj_vals), + total_energy_joules=round(total_energy, 6), ) +def _metric_stats(values: List[float]) -> Optional[MetricStats]: + """Compute MetricStats from a list of float values.""" + if not values: + return None + return MetricStats( + mean=statistics.mean(values), + median=statistics.median(values), + min=min(values), + max=max(values), + std=statistics.stdev(values) if len(values) > 1 else 0.0, + ) + + +def _metric_stats_to_dict(ms: Optional[MetricStats]) -> Optional[Dict[str, float]]: + """Convert MetricStats to a JSON-serializable dict.""" + if ms is None: + return None + return { + "mean": ms.mean, + "median": ms.median, + "min": ms.min, + "max": ms.max, + "std": ms.std, + } + + def _summary_to_dict(s: RunSummary) -> Dict[str, Any]: """Convert a RunSummary to a JSON-serializable dict.""" return { @@ -230,6 +298,18 @@ def _summary_to_dict(s: RunSummary) -> Dict[str, Any]: "per_subject": s.per_subject, "started_at": s.started_at, "ended_at": s.ended_at, + "accuracy_stats": _metric_stats_to_dict(s.accuracy_stats), + "latency_stats": _metric_stats_to_dict(s.latency_stats), + "ttft_stats": _metric_stats_to_dict(s.ttft_stats), + "energy_stats": _metric_stats_to_dict(s.energy_stats), + "power_stats": _metric_stats_to_dict(s.power_stats), + "gpu_utilization_stats": _metric_stats_to_dict(s.gpu_utilization_stats), + "throughput_stats": _metric_stats_to_dict(s.throughput_stats), + "mfu_stats": _metric_stats_to_dict(s.mfu_stats), + "mbu_stats": _metric_stats_to_dict(s.mbu_stats), + "ipw_stats": _metric_stats_to_dict(s.ipw_stats), + "ipj_stats": _metric_stats_to_dict(s.ipj_stats), + "total_energy_joules": s.total_energy_joules, } diff --git a/evals/core/scorer.py b/evals/core/scorer.py index 88f51ca5..f466cde7 100644 --- a/evals/core/scorer.py +++ b/evals/core/scorer.py @@ -38,7 +38,7 @@ class LLMJudgeScorer(Scorer): *, system: str = "", temperature: float = 0.0, - max_tokens: int = 1024, + max_tokens: int = 2048, ) -> str: """Send a prompt to the judge LLM and return the response text.""" return self._judge_backend.generate( diff --git a/evals/core/types.py b/evals/core/types.py index e822a345..4e80ff98 100644 --- a/evals/core/types.py +++ b/evals/core/types.py @@ -32,6 +32,15 @@ class EvalResult: cost_usd: float = 0.0 error: Optional[str] = None scoring_metadata: Dict[str, Any] = field(default_factory=dict) + ttft: float = 0.0 + energy_joules: float = 0.0 + power_watts: float = 0.0 + gpu_utilization_pct: float = 0.0 + throughput_tok_per_sec: float = 0.0 + mfu_pct: float = 0.0 + mbu_pct: float = 0.0 + ipw: float = 0.0 # Intelligence Per Watt + ipj: float = 0.0 # Intelligence Per Joule @dataclass(slots=True) @@ -45,13 +54,26 @@ class RunConfig: max_workers: int = 4 temperature: float = 0.0 max_tokens: int = 2048 - judge_model: str = "gpt-4o" + judge_model: str = "gpt-5-mini-2025-08-07" engine_key: Optional[str] = None agent_name: Optional[str] = None tools: List[str] = field(default_factory=list) output_path: Optional[str] = None seed: int = 42 dataset_split: Optional[str] = None + telemetry: bool = False + gpu_metrics: bool = False + + +@dataclass(slots=True) +class MetricStats: + """Descriptive statistics for a single metric across samples.""" + + mean: float = 0.0 + median: float = 0.0 + min: float = 0.0 + max: float = 0.0 + std: float = 0.0 @dataclass(slots=True) @@ -72,6 +94,18 @@ class RunSummary: per_subject: Dict[str, Dict[str, float]] = field(default_factory=dict) started_at: float = 0.0 ended_at: float = 0.0 + accuracy_stats: Optional[MetricStats] = None + latency_stats: Optional[MetricStats] = None + ttft_stats: Optional[MetricStats] = None + energy_stats: Optional[MetricStats] = None + power_stats: Optional[MetricStats] = None + gpu_utilization_stats: Optional[MetricStats] = None + throughput_stats: Optional[MetricStats] = None + mfu_stats: Optional[MetricStats] = None + mbu_stats: Optional[MetricStats] = None + ipw_stats: Optional[MetricStats] = None + ipj_stats: Optional[MetricStats] = None + total_energy_joules: float = 0.0 # --------------------------------------------------------------------------- @@ -99,7 +133,7 @@ class DefaultsConfig: class JudgeConfig: """Configuration for the LLM judge.""" - model: str = "gpt-4o" + model: str = "gpt-5-mini-2025-08-07" provider: Optional[str] = None temperature: float = 0.0 max_tokens: int = 1024 @@ -112,6 +146,8 @@ class ExecutionConfig: max_workers: int = 4 output_dir: str = "results/" seed: int = 42 + telemetry: bool = False + gpu_metrics: bool = False @dataclass(slots=True) @@ -155,6 +191,7 @@ class EvalSuiteConfig: __all__ = [ "EvalRecord", "EvalResult", + "MetricStats", "RunConfig", "RunSummary", "MetaConfig", diff --git a/evals/scorers/frames_judge.py b/evals/scorers/frames_judge.py index 966a3b64..96076f24 100644 --- a/evals/scorers/frames_judge.py +++ b/evals/scorers/frames_judge.py @@ -63,7 +63,7 @@ class FRAMESScorer(LLMJudgeScorer): ) try: - raw = self._ask_judge(prompt, temperature=0.0, max_tokens=1024) + raw = self._ask_judge(prompt, temperature=0.0, max_tokens=2048) structured_match = re.search( r"^correct:\s*(yes|no)", raw, re.MULTILINE | re.IGNORECASE, diff --git a/evals/scorers/gaia_exact.py b/evals/scorers/gaia_exact.py index e8095060..a52e1cc1 100644 --- a/evals/scorers/gaia_exact.py +++ b/evals/scorers/gaia_exact.py @@ -130,7 +130,7 @@ class GAIAScorer(LLMJudgeScorer): response=model_answer, ground_truth=reference, ) - raw = self._ask_judge(prompt, temperature=0.0, max_tokens=1024) + raw = self._ask_judge(prompt, temperature=0.0, max_tokens=2048) structured_match = re.search( r"^correct:\s*(yes|no)", raw, re.MULTILINE | re.IGNORECASE, diff --git a/evals/scorers/supergpqa_mcq.py b/evals/scorers/supergpqa_mcq.py index 97675a31..ce6f9648 100644 --- a/evals/scorers/supergpqa_mcq.py +++ b/evals/scorers/supergpqa_mcq.py @@ -50,7 +50,7 @@ class SuperGPQAScorer(LLMJudgeScorer): try: raw_response = self._ask_judge( user_prompt, system=system_prompt, - temperature=0.0, max_tokens=5, + temperature=0.0, max_tokens=2048, ) extracted = raw_response.strip().upper() diff --git a/evals/scorers/wildchat_judge.py b/evals/scorers/wildchat_judge.py index be6e9285..22923922 100644 --- a/evals/scorers/wildchat_judge.py +++ b/evals/scorers/wildchat_judge.py @@ -118,7 +118,7 @@ class WildChatScorer(LLMJudgeScorer): try: raw = self._ask_judge( prompt, system=SYSTEM_PROMPT, - temperature=0.0, max_tokens=1024, + temperature=0.0, max_tokens=2048, ) except Exception as exc: LOGGER.error("WildChat judge call failed: %s", exc) diff --git a/evals/tests/test_config.py b/evals/tests/test_config.py index 7319b123..5bb4912d 100644 --- a/evals/tests/test_config.py +++ b/evals/tests/test_config.py @@ -50,7 +50,7 @@ class TestDataclassDefaults: def test_judge_config_defaults(self): j = JudgeConfig() - assert j.model == "gpt-4o" + assert j.model == "gpt-5-mini-2025-08-07" assert j.provider is None assert j.temperature == 0.0 assert j.max_tokens == 1024 @@ -83,7 +83,7 @@ class TestDataclassDefaults: s = EvalSuiteConfig() assert s.meta.name == "" assert s.defaults.temperature == 0.0 - assert s.judge.model == "gpt-4o" + assert s.judge.model == "gpt-5-mini-2025-08-07" assert s.run.max_workers == 4 assert s.models == [] assert s.benchmarks == [] @@ -110,7 +110,7 @@ class TestLoadEvalConfig: assert suite.benchmarks[0].name == "supergpqa" # Defaults should be applied assert suite.defaults.temperature == 0.0 - assert suite.judge.model == "gpt-4o" + assert suite.judge.model == "gpt-5-mini-2025-08-07" def test_full_config(self, tmp_path): p = _write_toml(tmp_path, """\ @@ -280,7 +280,7 @@ class TestLoadEvalConfig: class TestExampleConfigs: - @pytest.fixture(params=["minimal.toml", "single-run.toml", "full-suite.toml"]) + @pytest.fixture(params=["minimal.toml", "single-run.toml", "full-suite.toml", "glm-4.7-flash-openhands.toml"]) def example_config(self, request): configs_dir = Path(__file__).resolve().parent.parent / "configs" return configs_dir / request.param diff --git a/evals/tests/test_types.py b/evals/tests/test_types.py index fda32629..e0d079eb 100644 --- a/evals/tests/test_types.py +++ b/evals/tests/test_types.py @@ -71,7 +71,7 @@ class TestRunConfig: assert c.max_workers == 4 assert c.temperature == 0.0 assert c.max_tokens == 2048 - assert c.judge_model == "gpt-4o" + assert c.judge_model == "gpt-5-mini-2025-08-07" assert c.seed == 42 assert c.tools == [] @@ -132,7 +132,7 @@ class TestDefaultsConfig: class TestJudgeConfig: def test_defaults(self): j = JudgeConfig() - assert j.model == "gpt-4o" + assert j.model == "gpt-5-mini-2025-08-07" assert j.provider is None assert j.temperature == 0.0 assert j.max_tokens == 1024 diff --git a/pyproject.toml b/pyproject.toml index eab8c4df..1ba93551 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ requires-python = ">=3.10" dependencies = [ "click>=8", "httpx>=0.27", + "pynvml>=13.0.1", "rich>=13", "tomli>=2.0; python_version < '3.11'", ] @@ -55,6 +56,7 @@ server = [ ] agents = [] openhands = ["openhands-sdk>=1.0; python_version >= '3.12'"] +gpu-metrics = ["pynvml>=12.0"] learning = [] orchestrator-training = ["torch>=2.0", "transformers>=4.40"] channel-telegram = ["python-telegram-bot>=21.0"] diff --git a/src/openjarvis/agents/native_openhands.py b/src/openjarvis/agents/native_openhands.py index f09c3c8a..171e99de 100644 --- a/src/openjarvis/agents/native_openhands.py +++ b/src/openjarvis/agents/native_openhands.py @@ -16,54 +16,37 @@ from openjarvis.core.events import EventBus 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 +from openjarvis.tools._stubs import BaseTool, build_tool_descriptions -OPENHANDS_SYSTEM_PROMPT = """\ -You are an AI assistant with access to tools. You MUST use tools when they would help answer the user's question. - -## How to use tools - -To call a tool, write on its own lines: - -Action: -Action Input: - -You will receive the result, then continue your response. - -## Available tools - -{tool_descriptions} - -## Important rules - -- When the user asks you to look up, search, fetch, or summarize a URL or topic, you MUST use web_search. Do NOT say you cannot browse the web. -- When the user provides a URL, pass the FULL URL (including https://) as the query to web_search. Do NOT rewrite URLs into search keywords. -- When the user asks a math question, use calculator. -- When the user asks to read a file, use file_read. -- You CAN write Python code in ```python blocks and it will be executed. Use this for computation, data processing, or when no specific tool fits. -- If no tool or code is needed, respond directly with your answer. -- Do NOT include tags or internal reasoning in your response. Respond directly.\ -""" - - -def _build_tool_descriptions(tools: list) -> str: - """Build detailed tool descriptions from ToolSpec objects.""" - if not tools: - return "No tools available." - lines = [] - for t in tools: - s = t.spec - params = s.parameters.get("properties", {}) - required = s.parameters.get("required", []) - param_parts = [] - for pname, pinfo in params.items(): - req_mark = " (required)" if pname in required else "" - param_parts.append( - f" - {pname}{req_mark}: {pinfo.get('description', pinfo.get('type', ''))}" - ) - param_str = "\n".join(param_parts) if param_parts else " (no parameters)" - lines.append(f"### {s.name}\n{s.description}\nParameters:\n{param_str}") - return "\n\n".join(lines) +OPENHANDS_SYSTEM_PROMPT = ( # noqa: E501 + "You are an AI assistant with access to tools. " + "You MUST use tools when they would help answer " + "the user's question.\n\n" + "## How to use tools\n\n" + "To call a tool, write on its own lines:\n\n" + "Action: \n" + "Action Input: \n\n" + "You will receive the result, then continue your " + "response.\n\n" + "## Available tools\n\n" + "{tool_descriptions}\n\n" + "## Important rules\n\n" + "- When the user asks you to look up, search, fetch, " + "or summarize a URL or topic, you MUST use web_search. " + "Do NOT say you cannot browse the web.\n" + "- When the user provides a URL, pass the FULL URL " + "(including https://) as the query to web_search. " + "Do NOT rewrite URLs into search keywords.\n" + "- When the user asks a math question, use calculator.\n" + "- When the user asks to read a file, use file_read.\n" + "- You CAN write Python code in ```python blocks and " + "it will be executed. Use this for computation, data " + "processing, or when no specific tool fits.\n" + "- If no tool or code is needed, respond directly " + "with your answer.\n" + "- Do NOT include tags or internal reasoning " + "in your response. Respond directly." +) @AgentRegistry.register("native_openhands") @@ -105,12 +88,20 @@ class NativeOpenHandsAgent(ToolUsingAgent): from openjarvis.tools.web_search import WebSearchTool content = WebSearchTool._fetch_url(url, max_chars=4000) - expanded = text.replace(url, f"\n\n--- Content from {url} ---\n{content}\n--- End of content ---\n") + header = f"\n\n--- Content from {url} ---\n" + footer = "\n--- End of content ---\n" + expanded = text.replace( + url, f"{header}{content}{footer}" + ) return expanded, True except Exception: return text, False - def _truncate_if_needed(self, messages: list[Message], max_prompt_tokens: int = 3000) -> list[Message]: + def _truncate_if_needed( + self, + messages: list[Message], + max_prompt_tokens: int = 3000, + ) -> list[Message]: """Truncate messages if estimated token count exceeds limit.""" total_chars = sum(len(m.content) for m in messages) estimated_tokens = total_chars // 4 @@ -126,7 +117,11 @@ class NativeOpenHandsAgent(ToolUsingAgent): truncated = original[: len(original) - excess_chars] messages[i] = Message( role=Role.USER, - content=truncated + "\n\n[Input truncated to fit context window]", + content=( + truncated + + "\n\n[Input truncated" + " to fit context window]" + ), ) break return messages @@ -146,9 +141,9 @@ class NativeOpenHandsAgent(ToolUsingAgent): 2. tool_name\\n$key=value (XML-style) """ # Format 1: Action / Action Input - action_match = re.search(r"Action:\s*(.+)", text) + action_match = re.search(r"Action:\s*(.+)", text, re.IGNORECASE) input_match = re.search( - r"Action Input:\s*(.+?)(?=\n\n|\Z)", text, re.DOTALL + r"Action Input:\s*(.+?)(?=\n\n|\Z)", text, re.DOTALL | re.IGNORECASE ) if action_match: return ( @@ -168,7 +163,8 @@ class NativeOpenHandsAgent(ToolUsingAgent): # Parse $key=value or value params into JSON params: dict[str, Any] = {} # $key=value format - for m in re.finditer(r"\$(\w+)=(.+?)(?=\$|\n<|\n") # value format for m in re.finditer(r"<(\w+)>(.*?)", raw_params, re.DOTALL): @@ -192,8 +188,10 @@ class NativeOpenHandsAgent(ToolUsingAgent): ) -> AgentResult: self._emit_turn_start(input) - tool_descriptions = _build_tool_descriptions(self._tools) - system_prompt = OPENHANDS_SYSTEM_PROMPT.format(tool_descriptions=tool_descriptions) + tool_descriptions = build_tool_descriptions(self._tools) + system_prompt = OPENHANDS_SYSTEM_PROMPT.format( + tool_descriptions=tool_descriptions, + ) # Pre-fetch any URLs in the input so the LLM gets the content directly input, url_expanded = self._expand_urls(input) @@ -201,7 +199,15 @@ class NativeOpenHandsAgent(ToolUsingAgent): # If URL content was inlined, skip the tool loop -- just summarize directly if url_expanded: direct_messages: list[Message] = [ - Message(role=Role.SYSTEM, content="You are a helpful assistant. Respond directly to the user's request using the provided content. Do NOT include tags."), + Message( + role=Role.SYSTEM, + content=( + "You are a helpful assistant. " + "Respond directly to the user's " + "request using the provided content." + " Do NOT include tags." + ), + ), Message(role=Role.USER, content=input), ] direct_messages = self._truncate_if_needed(direct_messages) @@ -212,13 +218,24 @@ class NativeOpenHandsAgent(ToolUsingAgent): return AgentResult(content=content, tool_results=[], turns=1) except Exception as exc: error_str = str(exc) - error_msg = ( - "The input is too long for the model's context window. Please try a shorter message." - if "400" in error_str - else f"The model returned an error: {error_str}" - ) + if "400" in error_str: + error_msg = ( + "The input is too long for the " + "model's context window. " + "Please try a shorter message." + ) + else: + error_msg = ( + "The model returned an error: " + + error_str + ) self._emit_turn_end(turns=1, error=True) - return AgentResult(content=error_msg, tool_results=[], turns=1, metadata={"error": True}) + return AgentResult( + content=error_msg, + tool_results=[], + turns=1, + metadata={"error": True}, + ) messages = self._build_messages(input, context, system_prompt=system_prompt) messages = self._truncate_if_needed(messages) diff --git a/src/openjarvis/agents/native_react.py b/src/openjarvis/agents/native_react.py index ebab4fd2..de27b4ec 100644 --- a/src/openjarvis/agents/native_react.py +++ b/src/openjarvis/agents/native_react.py @@ -14,7 +14,7 @@ from openjarvis.core.events import EventBus 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 +from openjarvis.tools._stubs import BaseTool, build_tool_descriptions REACT_SYSTEM_PROMPT = """\ You are a ReAct agent. For each step, respond with exactly one of: @@ -28,7 +28,7 @@ Action Input: Thought: Final Answer: -Available tools: {tool_names}""" +{tool_descriptions}""" @AgentRegistry.register("native_react") @@ -60,24 +60,30 @@ class NativeReActAgent(ToolUsingAgent): # Extract Thought thought_match = re.search( - r"Thought:\s*(.+?)(?=\nAction:|\nFinal Answer:|\Z)", text, re.DOTALL + r"Thought:\s*(.+?)(?=\nAction:|\nFinal Answer:|\Z)", + text, + re.DOTALL | re.IGNORECASE, ) 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) + final_match = re.search( + r"Final Answer:\s*(.+)", text, re.DOTALL | re.IGNORECASE + ) 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) + action_match = re.search(r"Action:\s*(.+)", text, re.IGNORECASE) 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 + r"Action Input:\s*(.+?)(?=\n\n|\nThought:|\Z)", + text, + re.DOTALL | re.IGNORECASE, ) if input_match: result["action_input"] = input_match.group(1).strip() @@ -92,11 +98,9 @@ class NativeReActAgent(ToolUsingAgent): ) -> AgentResult: self._emit_turn_start(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) + # Build system prompt with rich tool descriptions + tool_desc = build_tool_descriptions(self._tools) + system_prompt = REACT_SYSTEM_PROMPT.format(tool_descriptions=tool_desc) messages = self._build_messages(input, context, system_prompt=system_prompt) diff --git a/src/openjarvis/agents/openhands.py b/src/openjarvis/agents/openhands.py index 2273ca9c..7b7814cb 100644 --- a/src/openjarvis/agents/openhands.py +++ b/src/openjarvis/agents/openhands.py @@ -51,7 +51,11 @@ class OpenHandsAgent(BaseAgent): **kwargs: Any, ) -> AgentResult: try: - from openhands.sdk import Agent, Conversation, LLM # type: ignore[import-untyped] + from openhands.sdk import ( # type: ignore[import-untyped] + LLM, + Agent, + Conversation, + ) except ImportError: raise ImportError( "OpenHandsAgent requires the openhands-sdk package. " diff --git a/src/openjarvis/agents/orchestrator.py b/src/openjarvis/agents/orchestrator.py index f6685d5e..e4437da4 100644 --- a/src/openjarvis/agents/orchestrator.py +++ b/src/openjarvis/agents/orchestrator.py @@ -91,8 +91,7 @@ class OrchestratorAgent(ToolUsingAgent): from openjarvis.learning.orchestrator.prompt_registry import ( build_system_prompt, ) - tool_names = [t.spec.name for t in self._tools] - sys_prompt = build_system_prompt(tool_names) + sys_prompt = build_system_prompt(tools=self._tools) messages = self._build_messages(input, context, system_prompt=sys_prompt) diff --git a/src/openjarvis/agents/react.py b/src/openjarvis/agents/react.py index e25a2154..1cd478aa 100644 --- a/src/openjarvis/agents/react.py +++ b/src/openjarvis/agents/react.py @@ -1,6 +1,6 @@ """Backward-compat shim -- canonical location is agents.native_react.""" -from openjarvis.agents.native_react import NativeReActAgent as ReActAgent # noqa: F401 from openjarvis.agents.native_react import REACT_SYSTEM_PROMPT # noqa: F401 +from openjarvis.agents.native_react import NativeReActAgent as ReActAgent # noqa: F401 __all__ = ["ReActAgent", "REACT_SYSTEM_PROMPT"] diff --git a/src/openjarvis/agents/rlm.py b/src/openjarvis/agents/rlm.py index 7a363fd3..c85f5653 100644 --- a/src/openjarvis/agents/rlm.py +++ b/src/openjarvis/agents/rlm.py @@ -17,7 +17,7 @@ from openjarvis.core.events import EventBus 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 +from openjarvis.tools._stubs import BaseTool, build_tool_descriptions # --------------------------------------------------------------------------- # System prompt @@ -37,6 +37,7 @@ RLM_SYSTEM_PROMPT = ( "value of variable `var_name`.\n" "- `answer` dict — Set `answer[\"value\"] = ...` and " "`answer[\"ready\"] = True` to terminate.\n\n" + "{tool_section}" "## Available Modules\n\n" "json, re, math, collections, itertools, functools, " "textwrap, string, copy, datetime\n\n" @@ -112,7 +113,7 @@ class RLMAgent(ToolUsingAgent): self._sub_temperature = sub_temperature self._sub_max_tokens = sub_max_tokens self._max_output_chars = max_output_chars - self._system_prompt = system_prompt or RLM_SYSTEM_PROMPT + self._custom_system_prompt = system_prompt # ------------------------------------------------------------------ # Main run loop @@ -126,6 +127,30 @@ class RLMAgent(ToolUsingAgent): ) -> AgentResult: self._emit_turn_start(input) + # Build system prompt with tool section + if self._tools: + tool_section = ( + "## Available Tools\n\n" + "These tools are available to the sub-LM via " + "llm_query(). When writing prompts for llm_query(), " + "you can instruct it to use these tools:\n\n" + + build_tool_descriptions(self._tools) + + "\n\n" + ) + else: + tool_section = "" + + if self._custom_system_prompt: + system_prompt = self._custom_system_prompt + else: + try: + system_prompt = RLM_SYSTEM_PROMPT.format( + tool_section=tool_section, + ) + except KeyError: + # Custom system_prompt override without {tool_section} + system_prompt = RLM_SYSTEM_PROMPT + # Create REPL with sub-LM callbacks repl = RLMRepl( llm_query_fn=self._make_sub_query, @@ -140,7 +165,7 @@ class RLMAgent(ToolUsingAgent): # Build conversation messages = self._build_messages( - input, context, system_prompt=self._system_prompt, + input, context, system_prompt=system_prompt, ) all_tool_results: list[ToolResult] = [] diff --git a/src/openjarvis/core/config.py b/src/openjarvis/core/config.py index e05bba24..e581857f 100644 --- a/src/openjarvis/core/config.py +++ b/src/openjarvis/core/config.py @@ -480,6 +480,8 @@ class TelemetryConfig: enabled: bool = True db_path: str = str(DEFAULT_CONFIG_DIR / "telemetry.db") + gpu_metrics: bool = False + gpu_poll_interval_ms: int = 50 @dataclass(slots=True) @@ -871,6 +873,8 @@ policy = "heuristic" [telemetry] enabled = true +# gpu_metrics = false +# gpu_poll_interval_ms = 50 [traces] enabled = false diff --git a/src/openjarvis/core/types.py b/src/openjarvis/core/types.py index f0aa9285..01b76351 100644 --- a/src/openjarvis/core/types.py +++ b/src/openjarvis/core/types.py @@ -136,6 +136,12 @@ class TelemetryRecord: cost_usd: float = 0.0 energy_joules: float = 0.0 power_watts: float = 0.0 + gpu_utilization_pct: float = 0.0 + gpu_memory_used_gb: float = 0.0 + gpu_temperature_c: float = 0.0 + throughput_tok_per_sec: float = 0.0 + prefill_latency_seconds: float = 0.0 + decode_latency_seconds: float = 0.0 engine: str = "" agent: str = "" metadata: Dict[str, Any] = field(default_factory=dict) diff --git a/src/openjarvis/engine/_openai_compat.py b/src/openjarvis/engine/_openai_compat.py index fa27b40a..4e8358f8 100644 --- a/src/openjarvis/engine/_openai_compat.py +++ b/src/openjarvis/engine/_openai_compat.py @@ -138,5 +138,8 @@ class _OpenAICompatibleEngine(InferenceEngine): except Exception: return False + def close(self) -> None: + self._client.close() + __all__ = ["_OpenAICompatibleEngine"] diff --git a/src/openjarvis/engine/_stubs.py b/src/openjarvis/engine/_stubs.py index 68c2776d..30e89ad8 100644 --- a/src/openjarvis/engine/_stubs.py +++ b/src/openjarvis/engine/_stubs.py @@ -56,6 +56,9 @@ class InferenceEngine(ABC): def health(self) -> bool: """Return ``True`` when the engine is reachable and healthy.""" + def close(self) -> None: + """Release resources (HTTP clients, connections, threads, etc.).""" + def prepare(self, model: str) -> None: """Optional warm-up hook called before the first request.""" diff --git a/src/openjarvis/engine/cloud.py b/src/openjarvis/engine/cloud.py index c773d610..b52a3e83 100644 --- a/src/openjarvis/engine/cloud.py +++ b/src/openjarvis/engine/cloud.py @@ -2,7 +2,9 @@ from __future__ import annotations +import json import os +import time from collections.abc import AsyncIterator, Sequence from typing import Any, Dict, List @@ -59,6 +61,13 @@ def _is_google_model(model: str) -> bool: return "gemini" in model.lower() +def _is_openai_reasoning_model(model: str) -> bool: + """Check if model is an OpenAI reasoning model that restricts temperature.""" + m = model.lower() + # o1/o3 series and gpt-5-mini dated snapshots are reasoning models + return m.startswith(("o1", "o3")) or "gpt-5-mini-" in m + + 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 @@ -75,6 +84,36 @@ def estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> flo return input_cost + output_cost +def _convert_tools_to_anthropic( + openai_tools: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Convert OpenAI function-calling tools to Anthropic tool format.""" + result = [] + for tool in openai_tools: + func = tool.get("function", {}) + result.append({ + "name": func.get("name", ""), + "description": func.get("description", ""), + "input_schema": func.get("parameters", {}), + }) + return result + + +def _convert_tools_to_google( + openai_tools: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Convert OpenAI function-calling tools to Google function declarations.""" + declarations = [] + for tool in openai_tools: + func = tool.get("function", {}) + declarations.append({ + "name": func.get("name", ""), + "description": func.get("description", ""), + "parameters": func.get("parameters", {}), + }) + return declarations + + @EngineRegistry.register("cloud") class CloudEngine(InferenceEngine): """Cloud inference via OpenAI, Anthropic, and Google SDKs.""" @@ -126,18 +165,22 @@ class CloudEngine(InferenceEngine): "OPENAI_API_KEY and install " "openjarvis[inference-cloud]" ) - resp = self._openai_client.chat.completions.create( - model=model, - messages=messages_to_dicts(messages), - temperature=temperature, - max_tokens=max_tokens, + create_kwargs: Dict[str, Any] = { + "model": model, + "messages": messages_to_dicts(messages), + "max_completion_tokens": max_tokens, **kwargs, - ) + } + if not _is_openai_reasoning_model(model): + create_kwargs["temperature"] = temperature + t0 = time.monotonic() + resp = self._openai_client.chat.completions.create(**create_kwargs) + elapsed = time.monotonic() - t0 choice = resp.choices[0] usage = resp.usage prompt_tokens = usage.prompt_tokens if usage else 0 completion_tokens = usage.completion_tokens if usage else 0 - return { + result = { "content": choice.message.content or "", "usage": { "prompt_tokens": prompt_tokens, @@ -147,8 +190,22 @@ class CloudEngine(InferenceEngine): "model": resp.model, "finish_reason": choice.finish_reason or "stop", "cost_usd": estimate_cost(model, prompt_tokens, completion_tokens), + "ttft": elapsed, } + # Extract tool_calls if present + if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: + result["tool_calls"] = [ + { + "id": tc.id, + "name": tc.function.name, + "arguments": tc.function.arguments, + } + for tc in choice.message.tool_calls + ] + + return result + def _generate_anthropic( self, messages: Sequence[Message], @@ -180,11 +237,36 @@ class CloudEngine(InferenceEngine): } if system_text: create_kwargs["system"] = system_text + + # Convert and pass tools in Anthropic format + raw_tools = kwargs.pop("tools", None) + if raw_tools: + create_kwargs["tools"] = _convert_tools_to_anthropic(raw_tools) + + t0 = time.monotonic() resp = self._anthropic_client.messages.create(**create_kwargs) - content = resp.content[0].text if resp.content else "" + elapsed = time.monotonic() - t0 + + # Extract text and tool_use blocks from response content + content_parts: list[str] = [] + tool_calls: list[Dict[str, Any]] = [] + for block in resp.content: + if getattr(block, "type", None) == "tool_use": + tool_calls.append({ + "id": block.id, + "name": block.name, + "arguments": json.dumps(block.input) + if isinstance(block.input, dict) + else str(block.input), + }) + elif hasattr(block, "text"): + content_parts.append(block.text) + + content = "\n".join(content_parts) if content_parts else "" prompt_tokens = resp.usage.input_tokens if resp.usage else 0 completion_tokens = resp.usage.output_tokens if resp.usage else 0 - return { + + result: Dict[str, Any] = { "content": content, "usage": { "prompt_tokens": prompt_tokens, @@ -194,8 +276,14 @@ class CloudEngine(InferenceEngine): "model": resp.model, "finish_reason": resp.stop_reason or "stop", "cost_usd": estimate_cost(model, prompt_tokens, completion_tokens), + "ttft": elapsed, } + if tool_calls: + result["tool_calls"] = tool_calls + + return result + def _generate_google( self, messages: Sequence[Message], @@ -231,12 +319,49 @@ class CloudEngine(InferenceEngine): if system_text: config.system_instruction = system_text + # Convert and pass tools in Google format + raw_tools = kwargs.pop("tools", None) + if raw_tools: + declarations = _convert_tools_to_google(raw_tools) + config.tools = [{"function_declarations": declarations}] + + t0 = time.monotonic() resp = self._google_client.models.generate_content( model=model, contents=contents, config=config, ) - content = resp.text or "" + elapsed = time.monotonic() - t0 + + # Extract text and function_call parts from response + text_parts: list[str] = [] + tool_calls: list[Dict[str, Any]] = [] + candidates = getattr(resp, "candidates", None) + if candidates: + parts = getattr(candidates[0].content, "parts", []) + for part in parts: + if hasattr(part, "function_call") and part.function_call: + fc = part.function_call + fc_args = ( + dict(fc.args) if hasattr(fc.args, "items") else {} + ) + tool_calls.append({ + "id": f"google_{fc.name}", + "name": fc.name, + "arguments": json.dumps(fc_args), + }) + elif hasattr(part, "text") and part.text: + text_parts.append(part.text) + + # Guard against resp.text ValueError when only function_call parts + if text_parts: + content = "\n".join(text_parts) + else: + try: + content = resp.text or "" + except (ValueError, AttributeError): + content = "" + um = resp.usage_metadata prompt_tokens = ( getattr(um, "prompt_token_count", 0) if um else 0 @@ -244,7 +369,8 @@ class CloudEngine(InferenceEngine): completion_tokens = ( getattr(um, "candidates_token_count", 0) if um else 0 ) - return { + + result: Dict[str, Any] = { "content": content, "usage": { "prompt_tokens": prompt_tokens, @@ -254,8 +380,14 @@ class CloudEngine(InferenceEngine): "model": model, "finish_reason": "stop", "cost_usd": estimate_cost(model, prompt_tokens, completion_tokens), + "ttft": elapsed, } + if tool_calls: + result["tool_calls"] = tool_calls + + return result + def generate( self, messages: Sequence[Message], @@ -319,14 +451,16 @@ class CloudEngine(InferenceEngine): ) -> AsyncIterator[str]: if self._openai_client is None: raise EngineConnectionError("OpenAI client not available") - resp = self._openai_client.chat.completions.create( - model=model, - messages=messages_to_dicts(messages), - temperature=temperature, - max_tokens=max_tokens, - stream=True, + create_kwargs: Dict[str, Any] = { + "model": model, + "messages": messages_to_dicts(messages), + "max_completion_tokens": max_tokens, + "stream": True, **kwargs, - ) + } + if not _is_openai_reasoning_model(model): + create_kwargs["temperature"] = temperature + resp = self._openai_client.chat.completions.create(**create_kwargs) for chunk in resp: delta = chunk.choices[0].delta if chunk.choices else None if delta and delta.content: @@ -417,5 +551,17 @@ class CloudEngine(InferenceEngine): or self._google_client is not None ) + def close(self) -> None: + if self._openai_client is not None: + if hasattr(self._openai_client, "close"): + self._openai_client.close() + self._openai_client = None + if self._anthropic_client is not None: + if hasattr(self._anthropic_client, "close"): + self._anthropic_client.close() + self._anthropic_client = None + if self._google_client is not None: + self._google_client = None + __all__ = ["CloudEngine", "PRICING", "estimate_cost"] diff --git a/src/openjarvis/engine/litellm.py b/src/openjarvis/engine/litellm.py index 9fdfb693..bf8f98da 100644 --- a/src/openjarvis/engine/litellm.py +++ b/src/openjarvis/engine/litellm.py @@ -76,16 +76,13 @@ class LiteLLMEngine(InferenceEngine): "finish_reason": choice.finish_reason or "stop", } - # Extract tool_calls in OpenAI format + # Extract tool_calls in flat format (id, name, arguments) if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: result["tool_calls"] = [ { "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments, - }, + "name": tc.function.name, + "arguments": tc.function.arguments, } for tc in choice.message.tool_calls ] diff --git a/src/openjarvis/engine/ollama.py b/src/openjarvis/engine/ollama.py index bd7272c7..2c823aa9 100644 --- a/src/openjarvis/engine/ollama.py +++ b/src/openjarvis/engine/ollama.py @@ -71,6 +71,11 @@ class OllamaEngine(InferenceEngine): "model": data.get("model", model), "finish_reason": "stop", } + # Extract timing from Ollama response (nanoseconds → seconds) + result["ttft"] = data.get("prompt_eval_duration", 0) / 1e9 + result["engine_timing"] = {k: data[k] for k in + ("total_duration", "load_duration", "prompt_eval_duration", "eval_duration") + if k in data} # Extract tool calls if present raw_tool_calls = data.get("message", {}).get("tool_calls", []) if raw_tool_calls: @@ -135,5 +140,8 @@ class OllamaEngine(InferenceEngine): except Exception: return False + def close(self) -> None: + self._client.close() + __all__ = ["OllamaEngine"] diff --git a/src/openjarvis/intelligence/model_catalog.py b/src/openjarvis/intelligence/model_catalog.py index c37807f9..33c50c53 100644 --- a/src/openjarvis/intelligence/model_catalog.py +++ b/src/openjarvis/intelligence/model_catalog.py @@ -188,6 +188,16 @@ BUILTIN_MODELS: List[ModelSpec] = [ requires_api_key=True, metadata={"pricing_input": 0.25, "pricing_output": 2.00}, ), + ModelSpec( + model_id="gpt-5-mini-2025-08-07", + name="GPT-5 Mini (2025-08-07)", + 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 # ----------------------------------------------------------------------- diff --git a/src/openjarvis/learning/orchestrator/prompt_registry.py b/src/openjarvis/learning/orchestrator/prompt_registry.py index 68c4f251..5fac2af2 100644 --- a/src/openjarvis/learning/orchestrator/prompt_registry.py +++ b/src/openjarvis/learning/orchestrator/prompt_registry.py @@ -7,7 +7,10 @@ prompt template and tool descriptions used by the structured-mode from __future__ import annotations -from typing import Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional + +if TYPE_CHECKING: + from openjarvis.tools._stubs import BaseTool PROMPT_VERSION = "1.0" @@ -204,16 +207,61 @@ TOOL_DESCRIPTIONS: Dict[str, dict] = { } -def build_system_prompt(tool_names: Optional[List[str]] = None) -> str: +# Category labels for tool selection guide auto-generation +_CAT_LABELS: Dict[str, str] = { + "math": "MATH PROBLEMS", + "utility": "UTILITY / CODING TASKS", + "memory": "GENERAL Q&A / FACTUAL", + "llm": "REASONING/LOGIC", +} + + +def build_system_prompt( + tool_names: Optional[List[str]] = None, + *, + tools: Optional[List["BaseTool"]] = None, +) -> str: """Build the complete system prompt for the given tools. Args: tool_names: Tool names to include. If ``None``, uses all - tools from :data:`TOOL_DESCRIPTIONS`. + tools from :data:`TOOL_DESCRIPTIONS`. This path is kept for + backward compatibility with training pipelines. + tools: Optional list of ``BaseTool`` instances. When provided, + rich descriptions are auto-generated from ``ToolSpec``, + replacing the hardcoded :data:`TOOL_DESCRIPTIONS` lookup. + Unknown / MCP tools get full descriptions instead of + ``"Tool: {name}"``. Returns: Complete system prompt string. """ + # When BaseTool instances are provided, generate descriptions from spec + if tools is not None: + from openjarvis.tools._stubs import build_tool_descriptions + + desc_text = build_tool_descriptions(tools, include_cost=True) + + # Auto-generate tool selection guide by grouping tools by category + by_cat: Dict[str, List[str]] = {} + for t in tools: + cat = t.spec.category or "llm" + by_cat.setdefault(cat, []).append(t.spec.name) + + guide: list[str] = ["Choose tools based on task type:\n"] + for cat, names in by_cat.items(): + label = _CAT_LABELS.get(cat, cat.upper()) + guide.append(f"{label}:") + for n in names: + guide.append(f"- {n}") + guide.append("") + + return SYSTEM_PROMPT_TEMPLATE.format( + tools_description=desc_text, + tool_selection_guide="\n".join(guide), + ) + + # Backward-compat: tool_names-only path (used by training pipelines) if tool_names is None: tool_names = list(TOOL_DESCRIPTIONS) @@ -227,16 +275,16 @@ def build_system_prompt(tool_names: Optional[List[str]] = None) -> str: desc_lines.append(f"- {name}: {desc}") # Group tools by category - by_cat: Dict[str, List[str]] = {} + by_cat_names: Dict[str, List[str]] = {} for name in tool_names: cat = ( TOOL_DESCRIPTIONS[name]["category"] if name in TOOL_DESCRIPTIONS else "llm" ) - by_cat.setdefault(cat, []).append(name) + by_cat_names.setdefault(cat, []).append(name) - guide: list[str] = [ + guide = [ "Choose tools based on task type:\n", ] @@ -272,7 +320,7 @@ def build_system_prompt(tool_names: Optional[List[str]] = None) -> str: reasoning_lines.append( "- Step-by-step analysis -> think (organize thoughts first)" ) - llm_tools = by_cat.get("llm", []) + llm_tools = by_cat_names.get("llm", []) if llm_tools: reasoning_lines.append( f"- Complex reasoning -> {', '.join(llm_tools)}" @@ -286,7 +334,7 @@ def build_system_prompt(tool_names: Optional[List[str]] = None) -> str: general_lines: list[str] = [] if "web_search" in tool_names: general_lines.append("- Current events/recent info -> web_search") - memory_tools = by_cat.get("memory", []) + memory_tools = by_cat_names.get("memory", []) if memory_tools: general_lines.append( f"- Stored knowledge -> {', '.join(memory_tools)}" diff --git a/src/openjarvis/system.py b/src/openjarvis/system.py index 3e1efc76..f638dfe9 100644 --- a/src/openjarvis/system.py +++ b/src/openjarvis/system.py @@ -32,6 +32,7 @@ class JarvisSystem: telemetry_store: Optional[Any] = None trace_store: Optional[Any] = None trace_collector: Optional[Any] = None + gpu_monitor: Optional[Any] = None def ask( self, @@ -94,6 +95,7 @@ class JarvisSystem: ) -> Dict[str, Any]: """Run through an agent.""" from openjarvis.agents._stubs import AgentContext + from openjarvis.core.events import EventType from openjarvis.core.registry import AgentRegistry # Resolve agent @@ -134,8 +136,40 @@ class JarvisSystem: except TypeError: ag = agent_cls() + # Collect telemetry from all engine calls during agent run + telemetry_events: List[Dict[str, Any]] = [] + + def _on_inference_end(event: Any) -> None: + telemetry_events.append(event.data if hasattr(event, "data") else event) + + self.bus.subscribe(EventType.INFERENCE_END, _on_inference_end) + # Run - result = ag.run(query, context=ctx) + try: + result = ag.run(query, context=ctx) + finally: + self.bus.unsubscribe(EventType.INFERENCE_END, _on_inference_end) + + # Aggregate telemetry across all engine calls + _telemetry: Dict[str, Any] = {} + if telemetry_events: + total_energy = sum(e.get("energy_joules", 0.0) for e in telemetry_events) + total_latency = sum(e.get("latency", 0.0) for e in telemetry_events) + power_vals = [e.get("power_watts", 0.0) for e in telemetry_events if e.get("power_watts", 0.0) > 0] + util_vals = [e.get("gpu_utilization_pct", 0.0) for e in telemetry_events if e.get("gpu_utilization_pct", 0.0) > 0] + throughput_vals = [e.get("throughput_tok_per_sec", 0.0) for e in telemetry_events if e.get("throughput_tok_per_sec", 0.0) > 0] + _telemetry = { + "ttft": telemetry_events[0].get("ttft", 0.0), + "energy_joules": total_energy, + "power_watts": sum(power_vals) / len(power_vals) if power_vals else 0.0, + "gpu_utilization_pct": sum(util_vals) / len(util_vals) if util_vals else 0.0, + "throughput_tok_per_sec": sum(throughput_vals) / len(throughput_vals) if throughput_vals else 0.0, + "gpu_memory_used_gb": max((e.get("gpu_memory_used_gb", 0.0) for e in telemetry_events), default=0.0), + "gpu_temperature_c": max((e.get("gpu_temperature_c", 0.0) for e in telemetry_events), default=0.0), + "inference_calls": len(telemetry_events), + "total_inference_latency": total_latency, + } + return { "content": result.content, "usage": getattr(result, "usage", {}), @@ -150,6 +184,7 @@ class JarvisSystem: "turns": getattr(result, "turns", 1), "model": self.model, "engine": self.engine_key, + "_telemetry": _telemetry, } def _build_tools(self, tool_names: List[str]) -> List[BaseTool]: @@ -175,6 +210,10 @@ class JarvisSystem: def close(self) -> None: """Release resources.""" + if self.engine and hasattr(self.engine, "close"): + self.engine.close() + if self.gpu_monitor and hasattr(self.gpu_monitor, "close"): + self.gpu_monitor.close() if self.telemetry_store and hasattr(self.telemetry_store, "close"): self.telemetry_store.close() if self.trace_store and hasattr(self.trace_store, "close"): @@ -251,9 +290,21 @@ class SystemBuilder: self._telemetry if self._telemetry is not None else config.telemetry.enabled ) + gpu_monitor = None if telemetry_enabled: from openjarvis.telemetry.instrumented_engine import InstrumentedEngine - engine = InstrumentedEngine(engine, bus) + + if config.telemetry.gpu_metrics: + try: + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + if GpuMonitor.available(): + gpu_monitor = GpuMonitor( + poll_interval_ms=config.telemetry.gpu_poll_interval_ms, + ) + except ImportError: + pass + engine = InstrumentedEngine(engine, bus, gpu_monitor=gpu_monitor) # Apply security guardrails to engine engine = self._apply_security(config, engine, bus) @@ -295,6 +346,7 @@ class SystemBuilder: memory_backend=memory_backend, channel_backend=channel_backend, telemetry_store=telemetry_store, + gpu_monitor=gpu_monitor, ) def _resolve_engine(self, config: JarvisConfig): diff --git a/src/openjarvis/telemetry/__init__.py b/src/openjarvis/telemetry/__init__.py index 2290f8cb..e54d347d 100644 --- a/src/openjarvis/telemetry/__init__.py +++ b/src/openjarvis/telemetry/__init__.py @@ -11,11 +11,34 @@ from openjarvis.telemetry.aggregator import ( from openjarvis.telemetry.store import TelemetryStore from openjarvis.telemetry.wrapper import instrumented_generate +try: + from openjarvis.telemetry.gpu_monitor import GpuHardwareSpec, GpuMonitor, GpuSample, GpuSnapshot +except ImportError: + pass + +try: + from openjarvis.telemetry.efficiency import EfficiencyMetrics, compute_efficiency +except ImportError: + pass + +try: + from openjarvis.telemetry.vllm_metrics import VLLMMetrics, VLLMMetricsScraper +except ImportError: + pass + __all__ = [ "AggregatedStats", + "EfficiencyMetrics", "EngineStats", + "GpuHardwareSpec", + "GpuMonitor", + "GpuSample", + "GpuSnapshot", "ModelStats", "TelemetryAggregator", "TelemetryStore", + "VLLMMetrics", + "VLLMMetricsScraper", + "compute_efficiency", "instrumented_generate", ] diff --git a/src/openjarvis/telemetry/aggregator.py b/src/openjarvis/telemetry/aggregator.py index ec3700fb..a97244a8 100644 --- a/src/openjarvis/telemetry/aggregator.py +++ b/src/openjarvis/telemetry/aggregator.py @@ -20,6 +20,10 @@ class ModelStats: total_latency: float = 0.0 avg_latency: float = 0.0 total_cost: float = 0.0 + avg_ttft: float = 0.0 + total_energy_joules: float = 0.0 + avg_gpu_utilization_pct: float = 0.0 + avg_throughput_tok_per_sec: float = 0.0 @dataclass(slots=True) @@ -32,6 +36,10 @@ class EngineStats: total_latency: float = 0.0 avg_latency: float = 0.0 total_cost: float = 0.0 + avg_ttft: float = 0.0 + total_energy_joules: float = 0.0 + avg_gpu_utilization_pct: float = 0.0 + avg_throughput_tok_per_sec: float = 0.0 @dataclass(slots=True) @@ -87,7 +95,11 @@ class TelemetryAggregator: " SUM(completion_tokens) AS completion_tokens," " SUM(latency_seconds) AS total_latency," " AVG(latency_seconds) AS avg_latency," - " SUM(cost_usd) AS total_cost" + " SUM(cost_usd) AS total_cost," + " AVG(ttft) AS avg_ttft," + " SUM(energy_joules) AS total_energy_joules," + " AVG(gpu_utilization_pct) AS avg_gpu_utilization_pct," + " AVG(throughput_tok_per_sec) AS avg_throughput_tok_per_sec" f" FROM telemetry{where}" " GROUP BY model_id ORDER BY call_count DESC" ) @@ -102,6 +114,10 @@ class TelemetryAggregator: total_latency=r["total_latency"] or 0.0, avg_latency=r["avg_latency"] or 0.0, total_cost=r["total_cost"] or 0.0, + avg_ttft=r["avg_ttft"] or 0.0, + total_energy_joules=r["total_energy_joules"] or 0.0, + avg_gpu_utilization_pct=r["avg_gpu_utilization_pct"] or 0.0, + avg_throughput_tok_per_sec=r["avg_throughput_tok_per_sec"] or 0.0, ) for r in rows ] @@ -119,7 +135,11 @@ class TelemetryAggregator: " SUM(total_tokens) AS total_tokens," " SUM(latency_seconds) AS total_latency," " AVG(latency_seconds) AS avg_latency," - " SUM(cost_usd) AS total_cost" + " SUM(cost_usd) AS total_cost," + " AVG(ttft) AS avg_ttft," + " SUM(energy_joules) AS total_energy_joules," + " AVG(gpu_utilization_pct) AS avg_gpu_utilization_pct," + " AVG(throughput_tok_per_sec) AS avg_throughput_tok_per_sec" f" FROM telemetry{where}" " GROUP BY engine ORDER BY call_count DESC" ) @@ -132,6 +152,10 @@ class TelemetryAggregator: total_latency=r["total_latency"] or 0.0, avg_latency=r["avg_latency"] or 0.0, total_cost=r["total_cost"] or 0.0, + avg_ttft=r["avg_ttft"] or 0.0, + total_energy_joules=r["total_energy_joules"] or 0.0, + avg_gpu_utilization_pct=r["avg_gpu_utilization_pct"] or 0.0, + avg_throughput_tok_per_sec=r["avg_throughput_tok_per_sec"] or 0.0, ) for r in rows ] diff --git a/src/openjarvis/telemetry/efficiency.py b/src/openjarvis/telemetry/efficiency.py new file mode 100644 index 00000000..1b901ec4 --- /dev/null +++ b/src/openjarvis/telemetry/efficiency.py @@ -0,0 +1,129 @@ +"""MFU/MBU efficiency calculator for GPU inference telemetry. + +Computes Model FLOPs Utilization (MFU) and Model Bandwidth Utilization (MBU) +to quantify how efficiently a model uses available GPU compute and memory bandwidth. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class EfficiencyMetrics: + """Results of an MFU/MBU efficiency calculation.""" + + mfu_pct: float = 0.0 # Model FLOPs Utilization % + mbu_pct: float = 0.0 # Model Bandwidth Utilization % + actual_flops: float = 0.0 # Actual FLOPs achieved + peak_flops: float = 0.0 # Peak theoretical FLOPs + actual_bandwidth_gb_s: float = 0.0 # Actual memory bandwidth (GB/s) + peak_bandwidth_gb_s: float = 0.0 # Peak memory bandwidth (GB/s) + ipj: float = 0.0 # Intelligence Per Joule + + +def estimate_model_flops_per_token( + param_count_b: float, + active_params_b: float | None = None, +) -> float: + """Estimate FLOPs for one forward-pass token of a dense transformer. + + For dense models, FLOPs per token ≈ 2 * params. For MoE models, pass + ``active_params_b`` (the number of *active* parameters per token). + + Args: + param_count_b: Total parameter count in billions. + active_params_b: Active parameters per token in billions. If *None*, + defaults to ``param_count_b`` (dense model). + + Returns: + Estimated FLOPs per token. + """ + active = active_params_b if active_params_b is not None else param_count_b + return 2.0 * active * 1e9 + + +def estimate_model_bytes_per_token( + param_count_b: float, + bytes_per_param: float = 2.0, +) -> float: + """Estimate bytes of memory loaded per decode step. + + Args: + param_count_b: Total parameter count in billions. + bytes_per_param: Bytes per parameter (default 2.0 for FP16). + + Returns: + Bytes loaded per token. + """ + return param_count_b * 1e9 * bytes_per_param + + +def compute_efficiency( + param_count_b: float, + active_params_b: float | None, + gpu_peak_tflops: float, + gpu_peak_bandwidth_gb_s: float, + tokens_per_sec: float, + num_gpus: int = 1, + energy_joules: float = 0.0, + accuracy: float = 0.0, + bytes_per_param: float = 2.0, +) -> EfficiencyMetrics: + """Compute MFU, MBU, and derived efficiency metrics. + + Args: + param_count_b: Total parameter count in billions. + active_params_b: Active parameters per token in billions (*None* for dense). + gpu_peak_tflops: Peak theoretical TFLOPS per GPU (e.g. 312 for A100 SXM FP16). + gpu_peak_bandwidth_gb_s: Peak memory bandwidth per GPU in GB/s + (e.g. 2039 for A100 SXM). + tokens_per_sec: Measured generation throughput (tokens/second). + num_gpus: Number of GPUs used for inference. + energy_joules: Total energy consumed in joules (for IPJ calculation). + accuracy: Accuracy score in [0, 1] (for IPJ calculation). + bytes_per_param: Bytes per parameter (default 2.0 for FP16). + + Returns: + :class:`EfficiencyMetrics` with all computed values. + """ + flops_per_token = estimate_model_flops_per_token(param_count_b, active_params_b) + bytes_per_token = estimate_model_bytes_per_token(param_count_b, bytes_per_param) + + # Actual achieved rates + actual_flops = flops_per_token * tokens_per_sec + actual_bandwidth_bytes = bytes_per_token * tokens_per_sec + actual_bandwidth_gb_s = actual_bandwidth_bytes / 1e9 + + # Peak rates across all GPUs + peak_flops = gpu_peak_tflops * 1e12 * num_gpus + peak_bandwidth_gb_s = gpu_peak_bandwidth_gb_s * num_gpus + + # MFU and MBU + mfu_pct = (actual_flops / peak_flops * 100.0) if peak_flops > 0 else 0.0 + mbu_pct = ( + (actual_bandwidth_gb_s / peak_bandwidth_gb_s * 100.0) + if peak_bandwidth_gb_s > 0 + else 0.0 + ) + + # Intelligence Per Joule + ipj = (accuracy / energy_joules) if energy_joules > 0 else 0.0 + + return EfficiencyMetrics( + mfu_pct=mfu_pct, + mbu_pct=mbu_pct, + actual_flops=actual_flops, + peak_flops=peak_flops, + actual_bandwidth_gb_s=actual_bandwidth_gb_s, + peak_bandwidth_gb_s=peak_bandwidth_gb_s, + ipj=ipj, + ) + + +__all__ = [ + "EfficiencyMetrics", + "compute_efficiency", + "estimate_model_bytes_per_token", + "estimate_model_flops_per_token", +] diff --git a/src/openjarvis/telemetry/gpu_monitor.py b/src/openjarvis/telemetry/gpu_monitor.py new file mode 100644 index 00000000..7cc705e8 --- /dev/null +++ b/src/openjarvis/telemetry/gpu_monitor.py @@ -0,0 +1,313 @@ +"""GPU monitoring via pynvml — background poller for GPU metrics.""" + +from __future__ import annotations + +import threading +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Dict, Generator, List, Optional + +try: + import pynvml + + _PYNVML_AVAILABLE = True +except ImportError: + _PYNVML_AVAILABLE = False + + +# --------------------------------------------------------------------------- +# Hardware spec database +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class GpuHardwareSpec: + """Peak theoretical capabilities for a known GPU model.""" + + tflops_fp16: float + bandwidth_gb_s: float + tdp_watts: float + + +GPU_SPECS: Dict[str, GpuHardwareSpec] = { + "A100-SXM": GpuHardwareSpec(tflops_fp16=312, bandwidth_gb_s=2039, tdp_watts=400), + "A100-PCIE": GpuHardwareSpec(tflops_fp16=312, bandwidth_gb_s=2039, tdp_watts=300), + "H100-SXM": GpuHardwareSpec(tflops_fp16=990, bandwidth_gb_s=3350, tdp_watts=700), + "H100-PCIE": GpuHardwareSpec(tflops_fp16=756, bandwidth_gb_s=2000, tdp_watts=350), + "L40S": GpuHardwareSpec(tflops_fp16=366, bandwidth_gb_s=864, tdp_watts=350), + "A10": GpuHardwareSpec(tflops_fp16=125, bandwidth_gb_s=600, tdp_watts=150), + "RTX 4090": GpuHardwareSpec(tflops_fp16=165, bandwidth_gb_s=1008, tdp_watts=450), + "RTX 3090": GpuHardwareSpec(tflops_fp16=71, bandwidth_gb_s=936, tdp_watts=350), +} + + +def lookup_gpu_spec(name: str) -> Optional[GpuHardwareSpec]: + """Return the :class:`GpuHardwareSpec` for *name*, or ``None`` if unknown. + + Matches are case-insensitive substring lookups against the keys in + :data:`GPU_SPECS`. + """ + upper = name.upper() + for key, spec in GPU_SPECS.items(): + if key.upper() in upper: + return spec + return None + + +# --------------------------------------------------------------------------- +# Snapshot & aggregated sample +# --------------------------------------------------------------------------- + + +@dataclass +class GpuSnapshot: + """A single point-in-time reading from one GPU device.""" + + power_watts: float + utilization_pct: float + memory_used_gb: float + temperature_c: float + device_id: int = 0 + + +@dataclass +class GpuSample: + """Aggregated GPU metrics over an inference bracket.""" + + energy_joules: float = 0.0 + mean_power_watts: float = 0.0 + peak_power_watts: float = 0.0 + mean_utilization_pct: float = 0.0 + peak_utilization_pct: float = 0.0 + mean_memory_used_gb: float = 0.0 + peak_memory_used_gb: float = 0.0 + mean_temperature_c: float = 0.0 + peak_temperature_c: float = 0.0 + duration_seconds: float = 0.0 + num_snapshots: int = 0 + + +# --------------------------------------------------------------------------- +# Monitor +# --------------------------------------------------------------------------- + + +class GpuMonitor: + """Background GPU poller using pynvml. + + Usage:: + + mon = GpuMonitor(poll_interval_ms=50) + with mon.sample() as result: + # ... run inference ... + pass + print(result.energy_joules) + mon.close() + """ + + def __init__(self, poll_interval_ms: int = 50) -> None: + self._poll_interval_s = poll_interval_ms / 1000.0 + self._handles: List = [] + self._device_count = 0 + self._initialized = False + + if _PYNVML_AVAILABLE: + try: + pynvml.nvmlInit() + self._device_count = pynvml.nvmlDeviceGetCount() + self._handles = [ + pynvml.nvmlDeviceGetHandleByIndex(i) + for i in range(self._device_count) + ] + self._initialized = True + except Exception: + self._initialized = False + + @staticmethod + def available() -> bool: + """Return ``True`` if pynvml is importable and can be initialized.""" + if not _PYNVML_AVAILABLE: + return False + try: + pynvml.nvmlInit() + pynvml.nvmlShutdown() + return True + except Exception: + return False + + # -- polling thread internals --------------------------------------------- + + def _poll_once(self) -> List[GpuSnapshot]: + """Read current metrics from all GPU devices.""" + snapshots: List[GpuSnapshot] = [] + for idx, handle in enumerate(self._handles): + try: + power_mw = pynvml.nvmlDeviceGetPowerUsage(handle) + util = pynvml.nvmlDeviceGetUtilizationRates(handle) + mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) + temp = pynvml.nvmlDeviceGetTemperature( + handle, pynvml.NVML_TEMPERATURE_GPU + ) + snapshots.append( + GpuSnapshot( + power_watts=power_mw / 1000.0, + utilization_pct=float(util.gpu), + memory_used_gb=mem_info.used / (1024**3), + temperature_c=float(temp), + device_id=idx, + ) + ) + except Exception: + pass + return snapshots + + def _polling_loop( + self, + snapshots_out: List[List[GpuSnapshot]], + timestamps_out: List[float], + lock: threading.Lock, + stop_event: threading.Event, + ) -> None: + """Background thread: poll GPUs until *stop_event* is set.""" + while not stop_event.is_set(): + reading = self._poll_once() + if reading: + now = time.monotonic() + with lock: + snapshots_out.append(reading) + timestamps_out.append(now) + stop_event.wait(self._poll_interval_s) + + # -- aggregation ----------------------------------------------------------- + + @staticmethod + def _aggregate( + all_snapshots: List[List[GpuSnapshot]], + timestamps: List[float], + wall_duration: float, + ) -> GpuSample: + """Build a :class:`GpuSample` from collected snapshots. + + Energy is computed via trapezoidal integration of total power + (summed across all devices) over the timestamp series. + """ + if not all_snapshots: + return GpuSample(duration_seconds=wall_duration) + + # Flatten per-tick aggregates (sum power across devices per tick) + tick_powers: List[float] = [] + tick_utils: List[float] = [] + tick_mems: List[float] = [] + tick_temps: List[float] = [] + + for tick_snaps in all_snapshots: + total_power = sum(s.power_watts for s in tick_snaps) + mean_util = ( + sum(s.utilization_pct for s in tick_snaps) / len(tick_snaps) + ) + total_mem = sum(s.memory_used_gb for s in tick_snaps) + mean_temp = ( + sum(s.temperature_c for s in tick_snaps) / len(tick_snaps) + ) + tick_powers.append(total_power) + tick_utils.append(mean_util) + tick_mems.append(total_mem) + tick_temps.append(mean_temp) + + n = len(tick_powers) + + # Trapezoidal integration for energy + energy = 0.0 + for i in range(1, len(timestamps)): + dt = timestamps[i] - timestamps[i - 1] + energy += 0.5 * (tick_powers[i - 1] + tick_powers[i]) * dt + + return GpuSample( + energy_joules=energy, + mean_power_watts=sum(tick_powers) / n, + peak_power_watts=max(tick_powers), + mean_utilization_pct=sum(tick_utils) / n, + peak_utilization_pct=max(tick_utils), + mean_memory_used_gb=sum(tick_mems) / n, + peak_memory_used_gb=max(tick_mems), + mean_temperature_c=sum(tick_temps) / n, + peak_temperature_c=max(tick_temps), + duration_seconds=wall_duration, + num_snapshots=n, + ) + + # -- public API ----------------------------------------------------------- + + @contextmanager + def sample(self) -> Generator[GpuSample, None, None]: + """Context manager that polls GPUs during the block, then populates the sample. + + If pynvml is unavailable or no devices are found, yields an empty + :class:`GpuSample` without starting a background thread. + """ + result = GpuSample() + + if not self._initialized or self._device_count == 0: + t_start = time.monotonic() + yield result + result.duration_seconds = time.monotonic() - t_start + return + + snapshots: List[List[GpuSnapshot]] = [] + timestamps: List[float] = [] + lock = threading.Lock() + stop_event = threading.Event() + + thread = threading.Thread( + target=self._polling_loop, + args=(snapshots, timestamps, lock, stop_event), + daemon=True, + ) + + t_start = time.monotonic() + thread.start() + try: + yield result + finally: + stop_event.set() + thread.join(timeout=2.0) + wall = time.monotonic() - t_start + + with lock: + snap_copy = list(snapshots) + ts_copy = list(timestamps) + + aggregated = self._aggregate(snap_copy, ts_copy, wall) + + # Copy aggregated values into the yielded result object + result.energy_joules = aggregated.energy_joules + result.mean_power_watts = aggregated.mean_power_watts + result.peak_power_watts = aggregated.peak_power_watts + result.mean_utilization_pct = aggregated.mean_utilization_pct + result.peak_utilization_pct = aggregated.peak_utilization_pct + result.mean_memory_used_gb = aggregated.mean_memory_used_gb + result.peak_memory_used_gb = aggregated.peak_memory_used_gb + result.mean_temperature_c = aggregated.mean_temperature_c + result.peak_temperature_c = aggregated.peak_temperature_c + result.duration_seconds = aggregated.duration_seconds + result.num_snapshots = aggregated.num_snapshots + + def close(self) -> None: + """Shut down pynvml if it was initialized.""" + if self._initialized: + try: + pynvml.nvmlShutdown() + except Exception: + pass + self._initialized = False + + +__all__ = [ + "GpuHardwareSpec", + "GpuSnapshot", + "GpuSample", + "GpuMonitor", + "GPU_SPECS", + "lookup_gpu_spec", +] diff --git a/src/openjarvis/telemetry/instrumented_engine.py b/src/openjarvis/telemetry/instrumented_engine.py index 15d2d698..406038a5 100644 --- a/src/openjarvis/telemetry/instrumented_engine.py +++ b/src/openjarvis/telemetry/instrumented_engine.py @@ -3,11 +3,12 @@ from __future__ import annotations import time -from typing import Any, Dict, List, Sequence +from typing import Any, Dict, List, Optional, Sequence from openjarvis.core.events import EventBus, EventType from openjarvis.core.types import Message, TelemetryRecord from openjarvis.engine._stubs import InferenceEngine +from openjarvis.telemetry.gpu_monitor import GpuMonitor, GpuSample class InstrumentedEngine(InferenceEngine): @@ -20,9 +21,15 @@ class InstrumentedEngine(InferenceEngine): engine_id = "instrumented" - def __init__(self, engine: InferenceEngine, bus: EventBus) -> None: + def __init__( + self, + engine: InferenceEngine, + bus: EventBus, + gpu_monitor: Optional[Any] = None, + ) -> None: self._inner = engine self._bus = bus + self._gpu_monitor = gpu_monitor def generate( self, @@ -38,28 +45,93 @@ class InstrumentedEngine(InferenceEngine): "model": model, "message_count": len(messages), }) + gpu_sample: Optional[GpuSample] = None t0 = time.time() - result = self._inner.generate( - messages, model=model, temperature=temperature, - max_tokens=max_tokens, **kwargs, - ) + + if self._gpu_monitor is not None: + with self._gpu_monitor.sample() as gpu_sample: + result = self._inner.generate( + messages, model=model, temperature=temperature, + max_tokens=max_tokens, **kwargs, + ) + else: + result = self._inner.generate( + messages, model=model, temperature=temperature, + max_tokens=max_tokens, **kwargs, + ) + latency = time.time() - t0 usage = result.get("usage", {}) + completion_tokens = usage.get("completion_tokens", 0) + ttft = result.get("ttft", 0.0) + throughput = completion_tokens / latency if latency > 0 else 0.0 + + # GPU metrics from sample + energy_joules = 0.0 + power_watts = 0.0 + gpu_utilization_pct = 0.0 + gpu_memory_used_gb = 0.0 + gpu_temperature_c = 0.0 + prefill_latency = 0.0 + + if gpu_sample is not None: + energy_joules = gpu_sample.energy_joules + power_watts = gpu_sample.mean_power_watts + gpu_utilization_pct = gpu_sample.mean_utilization_pct + gpu_memory_used_gb = gpu_sample.peak_memory_used_gb + gpu_temperature_c = gpu_sample.mean_temperature_c + + if ttft > 0: + prefill_latency = ttft + record = TelemetryRecord( timestamp=t0, model_id=model, prompt_tokens=usage.get("prompt_tokens", 0), - completion_tokens=usage.get("completion_tokens", 0), + completion_tokens=completion_tokens, latency_seconds=latency, + ttft=ttft, + throughput_tok_per_sec=throughput, + energy_joules=energy_joules, + power_watts=power_watts, + gpu_utilization_pct=gpu_utilization_pct, + gpu_memory_used_gb=gpu_memory_used_gb, + gpu_temperature_c=gpu_temperature_c, + prefill_latency_seconds=prefill_latency, engine=getattr(self._inner, "engine_id", "unknown"), ) - self._bus.publish(EventType.INFERENCE_END, { - "model": model, "latency": latency, "usage": usage, - }) + event_data = { + "model": model, + "latency": latency, + "usage": usage, + "ttft": ttft, + "throughput_tok_per_sec": throughput, + "energy_joules": energy_joules, + "power_watts": power_watts, + "gpu_utilization_pct": gpu_utilization_pct, + "gpu_memory_used_gb": gpu_memory_used_gb, + "gpu_temperature_c": gpu_temperature_c, + "prefill_latency_seconds": prefill_latency, + } + + self._bus.publish(EventType.INFERENCE_END, event_data) self._bus.publish(EventType.TELEMETRY_RECORD, {"record": record}) + # Inject telemetry dict into result for downstream consumers (eval backend) + result["_telemetry"] = { + "latency": latency, + "ttft": ttft, + "throughput_tok_per_sec": throughput, + "energy_joules": energy_joules, + "power_watts": power_watts, + "gpu_utilization_pct": gpu_utilization_pct, + "gpu_memory_used_gb": gpu_memory_used_gb, + "gpu_temperature_c": gpu_temperature_c, + "prefill_latency_seconds": prefill_latency, + } + return result async def stream( @@ -92,5 +164,8 @@ class InstrumentedEngine(InferenceEngine): def health(self) -> bool: return self._inner.health() + def close(self) -> None: + self._inner.close() + __all__ = ["InstrumentedEngine"] diff --git a/src/openjarvis/telemetry/store.py b/src/openjarvis/telemetry/store.py index fa626d1d..0ff0a2a1 100644 --- a/src/openjarvis/telemetry/store.py +++ b/src/openjarvis/telemetry/store.py @@ -24,6 +24,12 @@ CREATE TABLE IF NOT EXISTS telemetry ( cost_usd REAL NOT NULL DEFAULT 0.0, energy_joules REAL NOT NULL DEFAULT 0.0, power_watts REAL NOT NULL DEFAULT 0.0, + gpu_utilization_pct REAL NOT NULL DEFAULT 0.0, + gpu_memory_used_gb REAL NOT NULL DEFAULT 0.0, + gpu_temperature_c REAL NOT NULL DEFAULT 0.0, + throughput_tok_per_sec REAL NOT NULL DEFAULT 0.0, + prefill_latency_seconds REAL NOT NULL DEFAULT 0.0, + decode_latency_seconds REAL NOT NULL DEFAULT 0.0, metadata TEXT NOT NULL DEFAULT '{}' ); """ @@ -32,10 +38,22 @@ _INSERT = """\ INSERT INTO telemetry ( timestamp, model_id, engine, agent, prompt_tokens, completion_tokens, total_tokens, - latency_seconds, ttft, cost_usd, energy_joules, power_watts, metadata -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + latency_seconds, ttft, cost_usd, energy_joules, power_watts, + gpu_utilization_pct, gpu_memory_used_gb, gpu_temperature_c, + throughput_tok_per_sec, prefill_latency_seconds, decode_latency_seconds, + metadata +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ +_MIGRATE_COLUMNS = [ + ("gpu_utilization_pct", "REAL NOT NULL DEFAULT 0.0"), + ("gpu_memory_used_gb", "REAL NOT NULL DEFAULT 0.0"), + ("gpu_temperature_c", "REAL NOT NULL DEFAULT 0.0"), + ("throughput_tok_per_sec", "REAL NOT NULL DEFAULT 0.0"), + ("prefill_latency_seconds", "REAL NOT NULL DEFAULT 0.0"), + ("decode_latency_seconds", "REAL NOT NULL DEFAULT 0.0"), +] + class TelemetryStore: """Append-only SQLite store for inference telemetry records.""" @@ -45,6 +63,18 @@ class TelemetryStore: self._conn = sqlite3.connect(self._db_path, check_same_thread=False) self._conn.execute(_CREATE_TABLE) self._conn.commit() + self._migrate_schema() + + def _migrate_schema(self) -> None: + """Add new columns to existing databases (idempotent).""" + for col_name, col_def in _MIGRATE_COLUMNS: + try: + self._conn.execute( + f"ALTER TABLE telemetry ADD COLUMN {col_name} {col_def}", + ) + except sqlite3.OperationalError: + pass # Column already exists + self._conn.commit() def record(self, rec: TelemetryRecord) -> None: """Persist a single telemetry record.""" @@ -63,6 +93,12 @@ class TelemetryStore: rec.cost_usd, rec.energy_joules, rec.power_watts, + rec.gpu_utilization_pct, + rec.gpu_memory_used_gb, + rec.gpu_temperature_c, + rec.throughput_tok_per_sec, + rec.prefill_latency_seconds, + rec.decode_latency_seconds, json.dumps(rec.metadata), ), ) diff --git a/src/openjarvis/telemetry/vllm_metrics.py b/src/openjarvis/telemetry/vllm_metrics.py new file mode 100644 index 00000000..ce3d04a3 --- /dev/null +++ b/src/openjarvis/telemetry/vllm_metrics.py @@ -0,0 +1,173 @@ +"""vLLM Prometheus metrics scraper — fetches and parses /metrics endpoint.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import List, Tuple + +import httpx + + +@dataclass +class VLLMMetrics: + """Parsed vLLM performance metrics.""" + + ttft_p50: float = 0.0 + ttft_p95: float = 0.0 + ttft_p99: float = 0.0 + gpu_cache_usage_pct: float = 0.0 + e2e_latency_p50: float = 0.0 + e2e_latency_p95: float = 0.0 + queue_depth: float = 0.0 + + +def _parse_histogram_buckets( + lines: List[str], metric_prefix: str +) -> Tuple[List[Tuple[float, float]], float, float]: + """Parse Prometheus histogram buckets, sum, and count for a metric. + + Returns (buckets, sum_value, count_value) where buckets is a sorted + list of (upper_bound, cumulative_count) pairs. + """ + buckets: List[Tuple[float, float]] = [] + sum_value = 0.0 + count_value = 0.0 + + for line in lines: + line = line.strip() + if line.startswith("#") or not line: + continue + + if line.startswith(f"{metric_prefix}_bucket{{"): + # Parse: metric_bucket{le="0.5"} 123 + try: + le_start = line.index('le="') + 4 + le_end = line.index('"', le_start) + le_str = line[le_start:le_end] + bound = float(le_str) if le_str != "+Inf" else math.inf + count_str = line.rsplit(None, 1)[-1] + buckets.append((bound, float(count_str))) + except (ValueError, IndexError): + continue + + elif line.startswith(f"{metric_prefix}_sum "): + try: + sum_value = float(line.split()[-1]) + except (ValueError, IndexError): + pass + + elif line.startswith(f"{metric_prefix}_count "): + try: + count_value = float(line.split()[-1]) + except (ValueError, IndexError): + pass + + buckets.sort(key=lambda x: x[0]) + return buckets, sum_value, count_value + + +def _percentile_from_buckets( + buckets: List[Tuple[float, float]], percentile: float +) -> float: + """Estimate a percentile from Prometheus histogram buckets. + + Uses linear interpolation between bucket boundaries. + """ + if not buckets: + return 0.0 + + total = buckets[-1][1] if buckets else 0.0 + if total == 0: + return 0.0 + + target = total * (percentile / 100.0) + + prev_bound = 0.0 + prev_count = 0.0 + for bound, count in buckets: + if bound == math.inf: + return prev_bound + if count >= target: + # Linear interpolation within this bucket + if count == prev_count: + return bound + fraction = (target - prev_count) / (count - prev_count) + return prev_bound + fraction * (bound - prev_bound) + prev_bound = bound + prev_count = count + + # Fallback: return last finite bound + for bound, _ in reversed(buckets): + if bound != math.inf: + return bound + return 0.0 + + +def _parse_gauge(lines: List[str], metric_name: str) -> float: + """Parse a Prometheus gauge value.""" + for line in lines: + line = line.strip() + if line.startswith("#") or not line: + continue + if line.startswith(f"{metric_name} "): + try: + return float(line.split()[-1]) + except (ValueError, IndexError): + pass + return 0.0 + + +class VLLMMetricsScraper: + """Scrapes vLLM's Prometheus /metrics endpoint.""" + + def __init__(self, host: str = "http://localhost:8000") -> None: + self._host = host.rstrip("/") + + def scrape(self) -> VLLMMetrics: + """Fetch and parse vLLM metrics. Returns zeroed metrics on error.""" + try: + resp = httpx.get(f"{self._host}/metrics", timeout=5.0) + resp.raise_for_status() + except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPStatusError): + return VLLMMetrics() + + return self._parse(resp.text) + + def _parse(self, text: str) -> VLLMMetrics: + """Parse Prometheus text format into VLLMMetrics.""" + lines = text.splitlines() + + # TTFT histogram + ttft_buckets, _, _ = _parse_histogram_buckets( + lines, "vllm:time_to_first_token_seconds" + ) + ttft_p50 = _percentile_from_buckets(ttft_buckets, 50) + ttft_p95 = _percentile_from_buckets(ttft_buckets, 95) + ttft_p99 = _percentile_from_buckets(ttft_buckets, 99) + + # GPU cache usage (gauge) + gpu_cache = _parse_gauge(lines, "vllm:gpu_cache_usage_perc") + + # E2E latency histogram + e2e_buckets, _, _ = _parse_histogram_buckets( + lines, "vllm:e2e_request_latency_seconds" + ) + e2e_p50 = _percentile_from_buckets(e2e_buckets, 50) + e2e_p95 = _percentile_from_buckets(e2e_buckets, 95) + + # Queue depth (gauge) — num_requests_waiting + queue_depth = _parse_gauge(lines, "vllm:num_requests_waiting") + + return VLLMMetrics( + ttft_p50=ttft_p50, + ttft_p95=ttft_p95, + ttft_p99=ttft_p99, + gpu_cache_usage_pct=gpu_cache, + e2e_latency_p50=e2e_p50, + e2e_latency_p95=e2e_p95, + queue_depth=queue_depth, + ) + + +__all__ = ["VLLMMetrics", "VLLMMetricsScraper"] diff --git a/src/openjarvis/tools/_stubs.py b/src/openjarvis/tools/_stubs.py index ea72b1dc..62a86472 100644 --- a/src/openjarvis/tools/_stubs.py +++ b/src/openjarvis/tools/_stubs.py @@ -185,4 +185,66 @@ class ToolExecutor: return [t.to_openai_function() for t in self._tools.values()] -__all__ = ["BaseTool", "ToolExecutor", "ToolSpec"] +def build_tool_descriptions( + tools: List[BaseTool], + *, + include_category: bool = True, + include_cost: bool = False, +) -> str: + """Build rich text descriptions from a list of tools. + + This is the single source of truth for all text-based agents that need + to describe available tools in their system prompts. + + Parameters + ---------- + tools: + List of tool instances. + include_category: + Whether to include the ``Category:`` line. + include_cost: + Whether to include ``Cost estimate:`` and ``Latency estimate:`` lines. + + Returns + ------- + str + Formatted multi-tool description, or ``"No tools available."`` if + *tools* is empty. + """ + if not tools: + return "No tools available." + + sections: list[str] = [] + for t in tools: + s = t.spec + lines = [f"### {s.name}", s.description] + + if include_category and s.category: + lines.append(f"Category: {s.category}") + + if include_cost: + if s.cost_estimate: + lines.append(f"Cost estimate: ${s.cost_estimate:.4f}") + if s.latency_estimate: + lines.append(f"Latency estimate: {s.latency_estimate:.1f}s") + + # Parameter descriptions + props = s.parameters.get("properties", {}) + required = set(s.parameters.get("required", [])) + if props: + lines.append("Parameters:") + for pname, pinfo in props.items(): + ptype = pinfo.get("type", "any") + req_mark = ", required" if pname in required else "" + desc = pinfo.get("description", "") + if desc: + lines.append(f" - {pname} ({ptype}{req_mark}): {desc}") + else: + lines.append(f" - {pname} ({ptype}{req_mark})") + + sections.append("\n".join(lines)) + + return "\n\n".join(sections) + + +__all__ = ["BaseTool", "ToolExecutor", "ToolSpec", "build_tool_descriptions"] diff --git a/src/openjarvis/traces/collector.py b/src/openjarvis/traces/collector.py index f692d166..ce7aa4c5 100644 --- a/src/openjarvis/traces/collector.py +++ b/src/openjarvis/traces/collector.py @@ -134,6 +134,8 @@ class TraceCollector: def _on_inference_end(self, event: Any) -> None: start = getattr(self, "_inference_start_time", event.timestamp) + data = event.data + usage = data.get("usage", {}) self._current_steps.append( TraceStep( step_type=StepType.GENERATE, @@ -141,9 +143,19 @@ class TraceCollector: duration_seconds=event.timestamp - start, input={"model": self._current_model}, output={ - "tokens": event.data.get("total_tokens", 0), + "prompt_tokens": usage.get("prompt_tokens", 0), + "completion_tokens": usage.get("completion_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + "tokens": usage.get("total_tokens", data.get("total_tokens", 0)), + }, + metadata={ + "engine": self._current_engine, + "ttft": data.get("ttft", 0.0), + "energy_joules": data.get("energy_joules", 0.0), + "power_watts": data.get("power_watts", 0.0), + "gpu_utilization_pct": data.get("gpu_utilization_pct", 0.0), + "throughput_tok_per_sec": data.get("throughput_tok_per_sec", 0.0), }, - metadata={"engine": self._current_engine}, ) ) diff --git a/tests/agents/test_agent_routing_matrix.py b/tests/agents/test_agent_routing_matrix.py index dbc2ef03..b32aa37a 100644 --- a/tests/agents/test_agent_routing_matrix.py +++ b/tests/agents/test_agent_routing_matrix.py @@ -51,13 +51,15 @@ AGENT_RESPONSES = { "native_openhands": _OPENHANDS_RESPONSE, } +_ALL_AGENTS = list(AGENT_RESPONSES) + # --------------------------------------------------------------------------- # Common agent tests (parametrized) # --------------------------------------------------------------------------- -@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "native_react", "native_openhands"]) +@pytest.mark.parametrize("agent_key", _ALL_AGENTS) class TestAgentCommon: def test_runs_with_mock_engine(self, agent_key): """Each agent type can run with a mock engine.""" @@ -129,7 +131,7 @@ def test_agent_id(agent_key, expected_id): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "native_react", "native_openhands"]) +@pytest.mark.parametrize("agent_key", _ALL_AGENTS) @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.""" @@ -145,7 +147,7 @@ def test_model_passthrough(agent_key, model): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "native_react", "native_openhands"]) +@pytest.mark.parametrize("agent_key", _ALL_AGENTS) def test_no_bus(agent_key): """All agents work without an event bus.""" engine = _make_mock_engine(AGENT_RESPONSES[agent_key]) @@ -160,7 +162,7 @@ def test_no_bus(agent_key): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "native_react", "native_openhands"]) +@pytest.mark.parametrize("agent_key", _ALL_AGENTS) 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]) @@ -174,7 +176,7 @@ def test_single_turn_count(agent_key): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("agent_key", ["simple", "orchestrator", "native_react", "native_openhands"]) +@pytest.mark.parametrize("agent_key", _ALL_AGENTS) 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]) diff --git a/tests/agents/test_base_agent.py b/tests/agents/test_base_agent.py index ef3ba606..a03197d2 100644 --- a/tests/agents/test_base_agent.py +++ b/tests/agents/test_base_agent.py @@ -2,11 +2,8 @@ from __future__ import annotations -from typing import Any, Optional from unittest.mock import MagicMock -import pytest - from openjarvis.agents._stubs import ( AgentContext, AgentResult, @@ -17,7 +14,6 @@ from openjarvis.core.events import EventBus, EventType from openjarvis.core.types import Conversation, Message, Role, ToolResult from openjarvis.tools._stubs import BaseTool, ToolSpec - # --------------------------------------------------------------------------- # Concrete subclass for testing # --------------------------------------------------------------------------- diff --git a/tests/agents/test_native_openhands.py b/tests/agents/test_native_openhands.py index aca4903b..8a37f880 100644 --- a/tests/agents/test_native_openhands.py +++ b/tests/agents/test_native_openhands.py @@ -295,6 +295,24 @@ class TestNativeOpenHandsAgent: assert "code_interpreter" in system_msg.content assert "calculator" in system_msg.content + def test_system_prompt_enriched_descriptions(self): + """System prompt has enriched descriptions with param schemas.""" + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.return_value = _engine_response("Ok") + agent = NativeOpenHandsAgent( + engine, "test-model", + tools=[_CodeInterpreterStub(), _CalculatorStub()], + ) + agent.run("Hello") + call_args = engine.generate.call_args + messages = call_args[0][0] + system_msg = messages[0].content + assert "### calculator" in system_msg + assert "### code_interpreter" in system_msg + assert "expression" in system_msg + assert "string" in system_msg + def test_max_turns_content_preserved(self): """When max turns exceeded, last content should be preserved.""" engine = MagicMock() diff --git a/tests/agents/test_native_react.py b/tests/agents/test_native_react.py index 814fcf3a..8a3797b1 100644 --- a/tests/agents/test_native_react.py +++ b/tests/agents/test_native_react.py @@ -154,6 +154,24 @@ class TestNativeReActParsing: assert result["action"] == "calculator" assert result["action_input"] == "" + def test_parse_case_insensitive_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"] + + def test_parse_case_insensitive_final_answer(self): + parse = self._parser() + text = "thought: I know the answer.\nfinal answer: 42" + result = parse(text) + assert result["final_answer"] == "42" + # --------------------------------------------------------------------------- # Agent execution tests @@ -190,7 +208,10 @@ class TestNativeReActAgent: ), ] bus = EventBus(record_history=True) - agent = NativeReActAgent(engine, "test-model", tools=[_CalculatorStub()], bus=bus) + agent = NativeReActAgent( + engine, "test-model", + tools=[_CalculatorStub()], bus=bus, + ) result = agent.run("What is 2+2?") assert result.content == "4" assert result.turns == 2 @@ -306,7 +327,10 @@ class TestNativeReActAgent: ), _engine_response("Thought: Done.\nFinal Answer: 2"), ] - agent = NativeReActAgent(engine, "test-model", tools=[_CalculatorStub()], bus=bus) + agent = NativeReActAgent( + 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 @@ -414,7 +438,7 @@ class TestNativeReActAgent: assert "think" in system_msg.content def test_system_prompt_no_tools(self): - """System prompt should say 'none' when no tools are available.""" + """System prompt should say 'No tools available.' when no tools.""" engine = MagicMock() engine.engine_id = "mock" engine.generate.return_value = _engine_response( @@ -424,7 +448,7 @@ class TestNativeReActAgent: agent.run("Hello") call_args = engine.generate.call_args messages = call_args[0][0] - assert "none" in messages[0].content + assert "No tools available." in messages[0].content def test_max_turns_1(self): """With max_turns=1 and an action, should stop after 1 turn.""" @@ -462,6 +486,51 @@ class TestNativeReActAgent: assert start_events[0].data["input"] == "test input" + def test_system_prompt_enriched_descriptions(self): + """System prompt should include parameter schemas, not just names.""" + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.return_value = _engine_response( + "Thought: Done.\nFinal Answer: ok" + ) + agent = NativeReActAgent( + engine, "test-model", + tools=[_CalculatorStub(), _ThinkStub()], + ) + agent.run("Hello") + call_args = engine.generate.call_args + messages = call_args[0][0] + system_content = messages[0].content + # Should contain tool name as header + assert "### calculator" in system_content + assert "### think" in system_content + # Should contain parameter info + assert "expression" in system_content + assert "string" in system_content + + def test_case_insensitive_execution(self): + """Agent handles lowercase action/thought/final answer from the LLM.""" + 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" + ), + ] + agent = NativeReActAgent( + engine, "test-model", + tools=[_CalculatorStub()], + ) + result = agent.run("What is 2+2?") + assert result.content == "4" + assert result.turns == 2 + + @pytest.mark.parametrize("model", ["qwen3:8b", "gpt-oss:120b"]) def test_native_react_with_different_models(model): """NativeReActAgent works with different model names.""" diff --git a/tests/agents/test_orchestrator.py b/tests/agents/test_orchestrator.py index 5f3dc8b4..7649ab78 100644 --- a/tests/agents/test_orchestrator.py +++ b/tests/agents/test_orchestrator.py @@ -440,3 +440,92 @@ class TestOrchestratorAgent: result = agent.run("Calc") # Should use the partial content if available assert result.content == "partial" + + +class TestOrchestratorStructuredMode: + """Tests for the structured (THOUGHT/TOOL/INPUT/FINAL_ANSWER) mode.""" + + def test_structured_mode_final_answer(self): + """Structured mode should parse FINAL_ANSWER: correctly.""" + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.return_value = { + "content": "THOUGHT: Easy question.\nFINAL_ANSWER: Paris", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + "model": "test-model", + "finish_reason": "stop", + } + agent = OrchestratorAgent( + engine, "test-model", mode="structured", + ) + result = agent.run("What is the capital of France?") + assert result.content == "Paris" + assert result.turns == 1 + assert result.tool_results == [] + + def test_structured_mode_tool_call(self): + """Parse TOOL:/INPUT:, execute tool, return final answer.""" + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.side_effect = [ + { + "content": ( + "THOUGHT: Need to calculate.\n" + 'TOOL: calculator\n' + 'INPUT: {"expression":"2+2"}' + ), + "usage": { + "prompt_tokens": 10, + "completion_tokens": 10, + "total_tokens": 20, + }, + "model": "test-model", + "finish_reason": "stop", + }, + { + "content": ( + "THOUGHT: Got 4.\n" + "FINAL_ANSWER: The answer is 4." + ), + "usage": { + "prompt_tokens": 20, + "completion_tokens": 10, + "total_tokens": 30, + }, + "model": "test-model", + "finish_reason": "stop", + }, + ] + agent = OrchestratorAgent( + engine, "test-model", + tools=[_CalculatorStub()], + mode="structured", + ) + result = agent.run("What is 2+2?") + assert result.content == "The answer is 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_structured_mode_enriched_descriptions(self): + """Structured mode system prompt should contain enriched tool descriptions.""" + engine = MagicMock() + engine.engine_id = "mock" + engine.generate.return_value = { + "content": "FINAL_ANSWER: ok", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + "model": "test-model", + "finish_reason": "stop", + } + agent = OrchestratorAgent( + engine, "test-model", + tools=[_CalculatorStub()], + mode="structured", + ) + agent.run("Hello") + call_args = engine.generate.call_args + messages = call_args[0][0] + system_msg = messages[0].content + assert "### calculator" in system_msg + assert "expression" in system_msg diff --git a/tests/agents/test_react.py b/tests/agents/test_react.py index 35e5af9c..871c1590 100644 --- a/tests/agents/test_react.py +++ b/tests/agents/test_react.py @@ -9,8 +9,8 @@ from __future__ import annotations from unittest.mock import MagicMock -from openjarvis.agents.react import ReActAgent from openjarvis.agents.native_react import NativeReActAgent +from openjarvis.agents.react import ReActAgent def _engine_response(content): diff --git a/tests/agents/test_rlm.py b/tests/agents/test_rlm.py index c5ad8d99..708cd950 100644 --- a/tests/agents/test_rlm.py +++ b/tests/agents/test_rlm.py @@ -367,6 +367,32 @@ class TestRLMBlockedCode: assert "Blocked" in result.tool_results[0].content +class TestRLMToolSectionInjection: + """Verify that tool descriptions are injected into the RLM system prompt.""" + + def test_system_prompt_includes_tool_section(self): + """Tools provided -> system prompt includes descriptions.""" + engine = _make_engine("Direct answer.") + agent = RLMAgent(engine, "test-model", tools=[_CalcStub()]) + agent.run("Hello") + call_args = engine.generate.call_args + messages = call_args[0][0] + system_msg = messages[0].content + assert "## Available Tools" in system_msg + assert "### calculator" in system_msg + assert "expression" in system_msg + + def test_system_prompt_no_tool_section_without_tools(self): + """No tools -> system prompt has no tool section.""" + engine = _make_engine("Direct answer.") + agent = RLMAgent(engine, "test-model") + agent.run("Hello") + call_args = engine.generate.call_args + messages = call_args[0][0] + system_msg = messages[0].content + assert "## Available Tools" not in system_msg + + class TestRLMReplResults: def test_repl_results_in_tool_results(self): engine = MagicMock() diff --git a/tests/engine/test_cloud_extended.py b/tests/engine/test_cloud_extended.py index 931c03b2..e63c8dda 100644 --- a/tests/engine/test_cloud_extended.py +++ b/tests/engine/test_cloud_extended.py @@ -120,6 +120,26 @@ class TestCloudOpenAI: [Message(role=Role.USER, content="Calculate")], model="gpt-5-mini" ) assert result["content"] == "" + # Verify flat tool_calls format + assert "tool_calls" in result + assert len(result["tool_calls"]) == 1 + tc = result["tool_calls"][0] + assert tc["id"] == "call_xyz" + assert tc["name"] == "calc" + assert tc["arguments"] == '{"x":1}' + + def test_no_tool_calls_when_absent(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="Just text", model="gpt-5-mini" + ) + engine._openai_client = fake_client + + result = engine.generate( + [Message(role=Role.USER, content="Hi")], model="gpt-5-mini" + ) + assert "tool_calls" not in result # --------------------------------------------------------------------------- @@ -188,6 +208,106 @@ class TestCloudAnthropic: assert _is_anthropic_model("gpt-5-mini") is False assert _is_anthropic_model("gemini-3-pro") is False + def test_anthropic_tool_use_extraction( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Anthropic tool_use blocks are extracted as flat tool_calls.""" + engine = _make_cloud_engine(monkeypatch) + fake_client = mock.MagicMock() + + # Build a response with a tool_use block + text_block = SimpleNamespace(type="text", text="Let me calculate.") + tool_block = SimpleNamespace( + type="tool_use", + id="toolu_123", + name="calculator", + input={"expression": "2+2"}, + ) + usage = SimpleNamespace(input_tokens=10, output_tokens=15) + fake_resp = SimpleNamespace( + content=[text_block, tool_block], + usage=usage, + model="claude-opus-4-6", + stop_reason="tool_use", + ) + fake_client.messages.create.return_value = fake_resp + engine._anthropic_client = fake_client + + openai_tools = [ + { + "type": "function", + "function": { + "name": "calculator", + "description": "Math", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + result = engine.generate( + [Message(role=Role.USER, content="What is 2+2?")], + model="claude-opus-4-6", + tools=openai_tools, + ) + assert result["content"] == "Let me calculate." + assert "tool_calls" in result + assert len(result["tool_calls"]) == 1 + tc = result["tool_calls"][0] + assert tc["id"] == "toolu_123" + assert tc["name"] == "calculator" + assert '"expression"' in tc["arguments"] + + def test_anthropic_tools_converted_to_input_schema( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Tools passed to Anthropic use input_schema format.""" + engine = _make_cloud_engine(monkeypatch) + fake_client = mock.MagicMock() + fake_client.messages.create.return_value = _fake_anthropic_response( + content="Ok", model="claude-opus-4-6" + ) + engine._anthropic_client = fake_client + + openai_tools = [ + { + "type": "function", + "function": { + "name": "calc", + "description": "Math", + "parameters": { + "type": "object", + "properties": {"x": {"type": "integer"}}, + }, + }, + } + ] + + engine.generate( + [Message(role=Role.USER, content="Hi")], + model="claude-opus-4-6", + tools=openai_tools, + ) + call_kwargs = fake_client.messages.create.call_args + passed_tools = call_kwargs.kwargs.get("tools") or call_kwargs[1].get("tools") + assert passed_tools is not None + assert passed_tools[0]["name"] == "calc" + assert "input_schema" in passed_tools[0] + + def test_anthropic_no_tool_calls_when_absent( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + engine = _make_cloud_engine(monkeypatch) + fake_client = mock.MagicMock() + fake_client.messages.create.return_value = _fake_anthropic_response( + content="Just text", 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 "tool_calls" not in result + def test_anthropic_system_message(self, monkeypatch: pytest.MonkeyPatch) -> None: engine = _make_cloud_engine(monkeypatch) fake_client = mock.MagicMock() @@ -313,6 +433,80 @@ class TestCloudGemini: ) assert result["content"] == "I am Gemini 3 Flash" + def test_gemini_function_call_extraction( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Google function_call parts are extracted as flat tool_calls.""" + engine = _make_cloud_engine(monkeypatch) + fake_client = mock.MagicMock() + + # Build a response with a function_call part + text_part = SimpleNamespace( + text="Let me calculate.", function_call=None + ) + fc = SimpleNamespace(name="calculator", args={"expression": "2+2"}) + fc_part = SimpleNamespace(text=None, function_call=fc) + content_obj = SimpleNamespace(parts=[text_part, fc_part]) + candidate = SimpleNamespace(content=content_obj) + usage = SimpleNamespace(prompt_token_count=10, candidates_token_count=8) + fake_resp = SimpleNamespace( + candidates=[candidate], usage_metadata=usage, text=None, + ) + fake_client.models.generate_content.return_value = fake_resp + engine._google_client = fake_client + + fake_types = mock.MagicMock() + fake_config = 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="What is 2+2?")], + model="gemini-3-pro", + tools=[ + { + "type": "function", + "function": { + "name": "calculator", + "description": "Math", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + assert result["content"] == "Let me calculate." + assert "tool_calls" in result + assert len(result["tool_calls"]) == 1 + tc = result["tool_calls"][0] + assert tc["name"] == "calculator" + assert '"expression"' in tc["arguments"] + + def test_gemini_no_tool_calls_when_absent( + 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="Just text" + ) + 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 "tool_calls" not in result + 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) diff --git a/tests/engine/test_litellm.py b/tests/engine/test_litellm.py index 78acae2a..b46b5b61 100644 --- a/tests/engine/test_litellm.py +++ b/tests/engine/test_litellm.py @@ -103,9 +103,8 @@ class TestLiteLLMEngineGenerate: assert len(result["tool_calls"]) == 1 tc = result["tool_calls"][0] assert tc["id"] == "call_123" - assert tc["type"] == "function" - assert tc["function"]["name"] == "calculator" - assert tc["function"]["arguments"] == '{"expression": "2+2"}' + assert tc["name"] == "calculator" + assert tc["arguments"] == '{"expression": "2+2"}' def test_generate_with_api_base(self) -> None: fake_usage = SimpleNamespace( diff --git a/tests/telemetry/test_efficiency.py b/tests/telemetry/test_efficiency.py new file mode 100644 index 00000000..ed6d0fbf --- /dev/null +++ b/tests/telemetry/test_efficiency.py @@ -0,0 +1,191 @@ +"""Tests for MFU/MBU efficiency calculator.""" + +from __future__ import annotations + +import pytest + +from openjarvis.telemetry.efficiency import ( + EfficiencyMetrics, + compute_efficiency, + estimate_model_bytes_per_token, + estimate_model_flops_per_token, +) + +# --------------------------------------------------------------------------- +# Constants: A100 SXM4 80GB specs +# --------------------------------------------------------------------------- +A100_PEAK_TFLOPS = 312.0 # FP16 Tensor Core +A100_PEAK_BW_GB_S = 2039.0 # HBM2e + + +# --------------------------------------------------------------------------- +# estimate_model_flops_per_token +# --------------------------------------------------------------------------- +class TestEstimateModelFlopsPerToken: + def test_7b_dense(self) -> None: + flops = estimate_model_flops_per_token(7.0) + assert flops == pytest.approx(14e9) + + def test_70b_dense(self) -> None: + flops = estimate_model_flops_per_token(70.0) + assert flops == pytest.approx(140e9) + + def test_moe_mixtral(self) -> None: + # Mixtral 8x7B: ~47B total, ~13B active per token + flops = estimate_model_flops_per_token(47.0, active_params_b=13.0) + assert flops == pytest.approx(26e9) + + def test_active_params_none_uses_total(self) -> None: + result = estimate_model_flops_per_token(7.0, None) + assert result == estimate_model_flops_per_token(7.0) + + +# --------------------------------------------------------------------------- +# estimate_model_bytes_per_token +# --------------------------------------------------------------------------- +class TestEstimateModelBytesPerToken: + def test_7b_fp16(self) -> None: + bpt = estimate_model_bytes_per_token(7.0) + assert bpt == pytest.approx(14e9) + + def test_7b_int8(self) -> None: + bpt = estimate_model_bytes_per_token(7.0, bytes_per_param=1.0) + assert bpt == pytest.approx(7e9) + + def test_70b_fp16(self) -> None: + bpt = estimate_model_bytes_per_token(70.0) + assert bpt == pytest.approx(140e9) + + +# --------------------------------------------------------------------------- +# compute_efficiency +# --------------------------------------------------------------------------- +class TestComputeEfficiency: + def test_7b_100tps_single_gpu(self) -> None: + """7B model, 100 tok/s, single A100.""" + m = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + ) + # actual_flops = 14e9 * 100 = 1.4e12 + assert m.actual_flops == pytest.approx(1.4e12) + # peak_flops = 312e12 + assert m.peak_flops == pytest.approx(312e12) + # MFU = 1.4e12 / 312e12 * 100 ≈ 0.4487% + expected_mfu = 1.4e12 / 312e12 * 100.0 + assert m.mfu_pct == pytest.approx(expected_mfu) + + # actual_bandwidth = 14e9 * 100 / 1e9 = 1400 GB/s + assert m.actual_bandwidth_gb_s == pytest.approx(1400.0) + # MBU = 1400 / 2039 * 100 ≈ 68.66% + expected_mbu = 1400.0 / 2039.0 * 100.0 + assert m.mbu_pct == pytest.approx(expected_mbu) + + def test_multi_gpu_scaling(self) -> None: + """2 GPUs should double peak, halving utilization at same throughput.""" + m1 = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + num_gpus=1, + ) + m2 = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + num_gpus=2, + ) + assert m2.mfu_pct == pytest.approx(m1.mfu_pct / 2.0) + assert m2.mbu_pct == pytest.approx(m1.mbu_pct / 2.0) + assert m2.peak_flops == pytest.approx(m1.peak_flops * 2.0) + assert m2.peak_bandwidth_gb_s == pytest.approx(m1.peak_bandwidth_gb_s * 2.0) + + def test_ipj_with_energy(self) -> None: + m = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + energy_joules=50.0, + accuracy=0.8, + ) + assert m.ipj == pytest.approx(0.8 / 50.0) + + def test_ipj_zero_energy(self) -> None: + m = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + energy_joules=0.0, + accuracy=0.8, + ) + assert m.ipj == 0.0 + + def test_zero_tokens_per_sec(self) -> None: + m = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=0.0, + ) + assert m.mfu_pct == 0.0 + assert m.mbu_pct == 0.0 + assert m.actual_flops == 0.0 + assert m.actual_bandwidth_gb_s == 0.0 + + def test_moe_efficiency(self) -> None: + """MoE model: FLOPs use active params, bandwidth uses total params.""" + m = compute_efficiency( + param_count_b=47.0, + active_params_b=13.0, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=50.0, + ) + # actual_flops based on active params: 2*13e9*50 = 1.3e12 + assert m.actual_flops == pytest.approx(1.3e12) + # bandwidth based on total params: 47e9*2*50/1e9 = 4700 GB/s + assert m.actual_bandwidth_gb_s == pytest.approx(4700.0) + + def test_returns_dataclass(self) -> None: + m = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + ) + assert isinstance(m, EfficiencyMetrics) + + def test_custom_bytes_per_param(self) -> None: + """INT4 quantization: 0.5 bytes per param.""" + m = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + bytes_per_param=0.5, + ) + # bandwidth = 7e9 * 0.5 * 100 / 1e9 = 350 GB/s + assert m.actual_bandwidth_gb_s == pytest.approx(350.0) + # MFU should be unchanged (FLOPs don't depend on bytes_per_param) + m_fp16 = compute_efficiency( + param_count_b=7.0, + active_params_b=None, + gpu_peak_tflops=A100_PEAK_TFLOPS, + gpu_peak_bandwidth_gb_s=A100_PEAK_BW_GB_S, + tokens_per_sec=100.0, + ) + assert m.mfu_pct == pytest.approx(m_fp16.mfu_pct) diff --git a/tests/telemetry/test_gpu_monitor.py b/tests/telemetry/test_gpu_monitor.py new file mode 100644 index 00000000..294895df --- /dev/null +++ b/tests/telemetry/test_gpu_monitor.py @@ -0,0 +1,410 @@ +"""Tests for GPU monitor -- mock pynvml (no real GPU required).""" + +from __future__ import annotations + +import sys +import time +import types +from dataclasses import dataclass +from unittest.mock import MagicMock, patch + +import pytest + +# --------------------------------------------------------------------------- +# Helpers: build a fake pynvml module +# --------------------------------------------------------------------------- + + +@dataclass +class _FakeUtilization: + gpu: int = 80 + memory: int = 50 + + +@dataclass +class _FakeMemInfo: + total: int = 24 * 1024**3 + used: int = 12 * 1024**3 + free: int = 12 * 1024**3 + + +def _make_fake_pynvml( + device_count: int = 1, power_mw: int = 300_000 +): + """Return a fake pynvml module object.""" + mod = types.ModuleType("pynvml") + mod.nvmlInit = MagicMock() + mod.nvmlShutdown = MagicMock() + mod.nvmlDeviceGetCount = MagicMock(return_value=device_count) + mod.nvmlDeviceGetHandleByIndex = MagicMock( + side_effect=lambda i: f"handle-{i}" + ) + mod.nvmlDeviceGetPowerUsage = MagicMock(return_value=power_mw) + mod.nvmlDeviceGetUtilizationRates = MagicMock( + return_value=_FakeUtilization() + ) + mod.nvmlDeviceGetMemoryInfo = MagicMock( + return_value=_FakeMemInfo() + ) + mod.nvmlDeviceGetTemperature = MagicMock(return_value=65) + mod.NVML_TEMPERATURE_GPU = 0 + return mod + + +def _snap( + power=300, util=80, mem=12, temp=65, dev=0 +): + """Shorthand to build a GpuSnapshot.""" + from openjarvis.telemetry.gpu_monitor import GpuSnapshot + + return GpuSnapshot( + power_watts=power, + utilization_pct=util, + memory_used_gb=mem, + temperature_c=temp, + device_id=dev, + ) + + +# --------------------------------------------------------------------------- +# Tests: GpuHardwareSpec lookup +# --------------------------------------------------------------------------- + + +class TestGpuHardwareSpec: + def test_lookup_exact_key(self): + from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec + + spec = lookup_gpu_spec("A100-SXM") + assert spec is not None + assert spec.tflops_fp16 == 312 + assert spec.bandwidth_gb_s == 2039 + assert spec.tdp_watts == 400 + + def test_lookup_substring_match(self): + from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec + + spec = lookup_gpu_spec("NVIDIA H100-SXM 80GB") + assert spec is not None + assert spec.tflops_fp16 == 990 + + def test_lookup_case_insensitive(self): + from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec + + spec = lookup_gpu_spec("nvidia rtx 4090") + assert spec is not None + assert spec.tflops_fp16 == 165 + + def test_lookup_unknown_returns_none(self): + from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec + + assert lookup_gpu_spec("SOME-UNKNOWN-GPU-9999") is None + + def test_all_specs_present(self): + from openjarvis.telemetry.gpu_monitor import GPU_SPECS + + expected = { + "A100-SXM", "A100-PCIE", + "H100-SXM", "H100-PCIE", + "L40S", "A10", + "RTX 4090", "RTX 3090", + } + assert set(GPU_SPECS.keys()) == expected + + def test_spec_is_frozen(self): + from openjarvis.telemetry.gpu_monitor import GPU_SPECS + + spec = GPU_SPECS["A100-SXM"] + with pytest.raises(AttributeError): + spec.tflops_fp16 = 999 + + +# --------------------------------------------------------------------------- +# Tests: energy integration math (trapezoidal rule) +# --------------------------------------------------------------------------- + + +class TestEnergyIntegration: + def test_constant_power(self): + """Constant 300W for 10 seconds = 3000 J.""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + snapshots = [ + [_snap(power=300)] + for _ in range(11) + ] + timestamps = [float(i) for i in range(11)] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=10.0 + ) + assert sample.energy_joules == pytest.approx(3000.0, rel=1e-6) + assert sample.mean_power_watts == pytest.approx(300.0, rel=1e-6) + assert sample.peak_power_watts == pytest.approx(300.0, rel=1e-6) + assert sample.duration_seconds == 10.0 + assert sample.num_snapshots == 11 + + def test_linear_ramp(self): + """Linear ramp 0W..400W over 4s => 800 J (trapezoidal).""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + snapshots = [ + [_snap(power=100.0 * i, util=50, mem=8, temp=60)] + for i in range(5) + ] + timestamps = [float(i) for i in range(5)] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=4.0 + ) + assert sample.energy_joules == pytest.approx(800.0, rel=1e-6) + assert sample.mean_power_watts == pytest.approx(200.0, rel=1e-6) + assert sample.peak_power_watts == pytest.approx(400.0) + + def test_empty_snapshots(self): + """No snapshots yields zeroed sample.""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + sample = GpuMonitor._aggregate([], [], wall_duration=5.0) + assert sample.energy_joules == 0.0 + assert sample.num_snapshots == 0 + assert sample.duration_seconds == 5.0 + + def test_single_snapshot(self): + """One snapshot: energy is zero (no interval).""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + snapshots = [ + [_snap(power=250, util=90, mem=16, temp=70)] + ] + timestamps = [0.0] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=0.05 + ) + assert sample.energy_joules == 0.0 + assert sample.num_snapshots == 1 + assert sample.mean_power_watts == pytest.approx(250.0) + + +# --------------------------------------------------------------------------- +# Tests: GpuSample aggregation +# --------------------------------------------------------------------------- + + +class TestGpuSampleAggregation: + def test_peak_values(self): + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + snapshots = [ + [_snap(power=200, util=60, mem=10, temp=55)], + [_snap(power=400, util=95, mem=20, temp=80)], + [_snap(power=300, util=70, mem=15, temp=65)], + ] + timestamps = [0.0, 1.0, 2.0] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=2.0 + ) + assert sample.peak_power_watts == pytest.approx(400.0) + assert sample.peak_utilization_pct == pytest.approx(95.0) + assert sample.peak_memory_used_gb == pytest.approx(20.0) + assert sample.peak_temperature_c == pytest.approx(80.0) + + def test_mean_values(self): + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + snapshots = [ + [_snap(power=100, util=40, mem=8, temp=50)], + [_snap(power=200, util=60, mem=12, temp=60)], + [_snap(power=300, util=80, mem=16, temp=70)], + ] + timestamps = [0.0, 1.0, 2.0] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=2.0 + ) + assert sample.mean_power_watts == pytest.approx(200.0) + assert sample.mean_utilization_pct == pytest.approx(60.0) + assert sample.mean_memory_used_gb == pytest.approx(12.0) + assert sample.mean_temperature_c == pytest.approx(60.0) + + +# --------------------------------------------------------------------------- +# Tests: Multi-GPU aggregation +# --------------------------------------------------------------------------- + + +class TestMultiGpu: + def test_multi_device_power_sum(self): + """Power summed across devices; util/temp averaged.""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + tick = [ + _snap(power=200, util=80, mem=10, temp=60, dev=0), + _snap(power=300, util=90, mem=12, temp=70, dev=1), + ] + snapshots = [tick, tick] + timestamps = [0.0, 1.0] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=1.0 + ) + assert sample.energy_joules == pytest.approx(500.0) + assert sample.mean_power_watts == pytest.approx(500.0) + assert sample.mean_utilization_pct == pytest.approx(85.0) + assert sample.mean_memory_used_gb == pytest.approx(22.0) + assert sample.mean_temperature_c == pytest.approx(65.0) + + def test_multi_device_varying_power(self): + """Multi-GPU with varying power over time.""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + snapshots = [ + [ + _snap(power=100, util=50, mem=8, temp=55, dev=0), + _snap(power=100, util=50, mem=8, temp=55, dev=1), + ], + [ + _snap(power=300, util=90, mem=16, temp=75, dev=0), + _snap(power=300, util=90, mem=16, temp=75, dev=1), + ], + ] + timestamps = [0.0, 2.0] + + sample = GpuMonitor._aggregate( + snapshots, timestamps, wall_duration=2.0 + ) + # Tick powers: 200, 600 => 0.5*(200+600)*2 = 800 + assert sample.energy_joules == pytest.approx(800.0) + assert sample.peak_power_watts == pytest.approx(600.0) + + +# --------------------------------------------------------------------------- +# Tests: context manager flow (with mocked pynvml) +# --------------------------------------------------------------------------- + + +class TestContextManager: + def test_sample_context_manager(self): + """sample() starts/stops polling and populates result.""" + fake_pynvml = _make_fake_pynvml( + device_count=1, power_mw=300_000 + ) + + import openjarvis.telemetry.gpu_monitor as mod + + orig_avail = mod._PYNVML_AVAILABLE + orig_pynvml = getattr(mod, "pynvml", None) + mod._PYNVML_AVAILABLE = True + mod.pynvml = fake_pynvml + try: + monitor = mod.GpuMonitor(poll_interval_ms=10) + monitor._initialized = True + monitor._device_count = 1 + monitor._handles = ["handle-0"] + + with monitor.sample() as result: + time.sleep(0.1) + + assert result.duration_seconds > 0 + assert result.num_snapshots > 0 + assert result.mean_power_watts > 0 + finally: + mod._PYNVML_AVAILABLE = orig_avail + if orig_pynvml is not None: + mod.pynvml = orig_pynvml + + def test_sample_no_gpu(self): + """sample() yields empty GpuSample when no GPU.""" + from openjarvis.telemetry.gpu_monitor import GpuMonitor + + monitor = GpuMonitor.__new__(GpuMonitor) + monitor._poll_interval_s = 0.05 + monitor._handles = [] + monitor._device_count = 0 + monitor._initialized = False + + with monitor.sample() as result: + pass + + assert result.num_snapshots == 0 + assert result.energy_joules == 0.0 + assert result.duration_seconds >= 0 + + +# --------------------------------------------------------------------------- +# Tests: available() +# --------------------------------------------------------------------------- + + +class TestAvailable: + def test_available_false_when_pynvml_missing(self): + """available() returns False when pynvml not importable.""" + import openjarvis.telemetry.gpu_monitor as mod + + orig = mod._PYNVML_AVAILABLE + mod._PYNVML_AVAILABLE = False + try: + assert mod.GpuMonitor.available() is False + finally: + mod._PYNVML_AVAILABLE = orig + + def test_available_true_with_fake_pynvml(self): + """available() returns True when pynvml can init.""" + fake_pynvml = _make_fake_pynvml() + + with patch.dict(sys.modules, {"pynvml": fake_pynvml}): + import openjarvis.telemetry.gpu_monitor as mod + + orig = mod._PYNVML_AVAILABLE + mod._PYNVML_AVAILABLE = True + mod.pynvml = fake_pynvml + try: + assert mod.GpuMonitor.available() is True + fake_pynvml.nvmlInit.assert_called() + fake_pynvml.nvmlShutdown.assert_called() + finally: + mod._PYNVML_AVAILABLE = orig + + def test_available_false_when_init_fails(self): + """available() returns False when nvmlInit raises.""" + fake_pynvml = _make_fake_pynvml() + fake_pynvml.nvmlInit.side_effect = RuntimeError("no driver") + + with patch.dict(sys.modules, {"pynvml": fake_pynvml}): + import openjarvis.telemetry.gpu_monitor as mod + + orig = mod._PYNVML_AVAILABLE + mod._PYNVML_AVAILABLE = True + mod.pynvml = fake_pynvml + try: + assert mod.GpuMonitor.available() is False + finally: + mod._PYNVML_AVAILABLE = orig + + +# --------------------------------------------------------------------------- +# Tests: dataclass defaults +# --------------------------------------------------------------------------- + + +class TestDataclasses: + def test_gpu_snapshot_defaults(self): + from openjarvis.telemetry.gpu_monitor import GpuSnapshot + + s = GpuSnapshot( + power_watts=300, + utilization_pct=80, + memory_used_gb=12, + temperature_c=65, + ) + assert s.device_id == 0 + + def test_gpu_sample_defaults(self): + from openjarvis.telemetry.gpu_monitor import GpuSample + + s = GpuSample() + assert s.energy_joules == 0.0 + assert s.num_snapshots == 0 + assert s.duration_seconds == 0.0 diff --git a/tests/telemetry/test_vllm_metrics.py b/tests/telemetry/test_vllm_metrics.py new file mode 100644 index 00000000..eb651e58 --- /dev/null +++ b/tests/telemetry/test_vllm_metrics.py @@ -0,0 +1,201 @@ +"""Tests for VLLMMetricsScraper — Prometheus text format parsing.""" + +from __future__ import annotations + +from unittest.mock import patch + +from openjarvis.telemetry.vllm_metrics import ( + VLLMMetrics, + VLLMMetricsScraper, + _parse_gauge, + _parse_histogram_buckets, + _percentile_from_buckets, +) + +# --------------------------------------------------------------------------- +# Sample Prometheus text fixtures +# --------------------------------------------------------------------------- + +SAMPLE_METRICS = """\ +# HELP vllm:time_to_first_token_seconds Histogram of TTFT in seconds. +# TYPE vllm:time_to_first_token_seconds histogram +vllm:time_to_first_token_seconds_bucket{le="0.01"} 5 +vllm:time_to_first_token_seconds_bucket{le="0.025"} 15 +vllm:time_to_first_token_seconds_bucket{le="0.05"} 40 +vllm:time_to_first_token_seconds_bucket{le="0.1"} 80 +vllm:time_to_first_token_seconds_bucket{le="0.25"} 95 +vllm:time_to_first_token_seconds_bucket{le="0.5"} 100 +vllm:time_to_first_token_seconds_bucket{le="+Inf"} 100 +vllm:time_to_first_token_seconds_sum 5.5 +vllm:time_to_first_token_seconds_count 100 +# HELP vllm:gpu_cache_usage_perc GPU KV-cache usage percentage. +# TYPE vllm:gpu_cache_usage_perc gauge +vllm:gpu_cache_usage_perc 0.42 +# HELP vllm:e2e_request_latency_seconds Histogram of E2E request latency. +# TYPE vllm:e2e_request_latency_seconds histogram +vllm:e2e_request_latency_seconds_bucket{le="0.1"} 10 +vllm:e2e_request_latency_seconds_bucket{le="0.5"} 50 +vllm:e2e_request_latency_seconds_bucket{le="1.0"} 80 +vllm:e2e_request_latency_seconds_bucket{le="2.5"} 95 +vllm:e2e_request_latency_seconds_bucket{le="5.0"} 100 +vllm:e2e_request_latency_seconds_bucket{le="+Inf"} 100 +vllm:e2e_request_latency_seconds_sum 75.0 +vllm:e2e_request_latency_seconds_count 100 +# HELP vllm:num_requests_waiting Number of requests waiting in queue. +# TYPE vllm:num_requests_waiting gauge +vllm:num_requests_waiting 3 +""" + + +class TestVLLMMetricsDataclass: + def test_defaults(self): + m = VLLMMetrics() + assert m.ttft_p50 == 0.0 + assert m.ttft_p95 == 0.0 + assert m.ttft_p99 == 0.0 + assert m.gpu_cache_usage_pct == 0.0 + assert m.e2e_latency_p50 == 0.0 + assert m.e2e_latency_p95 == 0.0 + assert m.queue_depth == 0.0 + + def test_custom_values(self): + m = VLLMMetrics(ttft_p50=0.05, gpu_cache_usage_pct=0.8) + assert m.ttft_p50 == 0.05 + assert m.gpu_cache_usage_pct == 0.8 + + +class TestParseHistogramBuckets: + def test_parses_ttft_buckets(self): + lines = SAMPLE_METRICS.splitlines() + buckets, sum_val, count_val = _parse_histogram_buckets( + lines, "vllm:time_to_first_token_seconds" + ) + assert len(buckets) == 7 + assert sum_val == 5.5 + assert count_val == 100 + # First finite bucket + assert buckets[0] == (0.01, 5.0) + # Last finite bucket + assert buckets[5] == (0.5, 100.0) + + def test_empty_lines(self): + buckets, s, c = _parse_histogram_buckets([], "nonexistent") + assert buckets == [] + assert s == 0.0 + assert c == 0.0 + + def test_no_matching_metric(self): + lines = SAMPLE_METRICS.splitlines() + buckets, s, c = _parse_histogram_buckets(lines, "no_such_metric") + assert buckets == [] + assert s == 0.0 + + +class TestPercentileFromBuckets: + def test_empty_buckets(self): + assert _percentile_from_buckets([], 50) == 0.0 + + def test_zero_count(self): + buckets = [(0.1, 0.0), (0.5, 0.0)] + assert _percentile_from_buckets(buckets, 50) == 0.0 + + def test_median_interpolation(self): + lines = SAMPLE_METRICS.splitlines() + buckets, _, _ = _parse_histogram_buckets( + lines, "vllm:time_to_first_token_seconds" + ) + p50 = _percentile_from_buckets(buckets, 50) + # 50th percentile: target = 50 out of 100 + # Bucket le=0.05 has 40, le=0.1 has 80 + # fraction = (50-40)/(80-40) = 0.25 + # result = 0.05 + 0.25 * (0.1-0.05) = 0.0625 + assert abs(p50 - 0.0625) < 1e-6 + + def test_p95(self): + lines = SAMPLE_METRICS.splitlines() + buckets, _, _ = _parse_histogram_buckets( + lines, "vllm:time_to_first_token_seconds" + ) + p95 = _percentile_from_buckets(buckets, 95) + # target = 95, bucket le=0.25 has 95 exactly + # le=0.1 has 80, le=0.25 has 95 + # fraction = (95-80)/(95-80) = 1.0 + # result = 0.1 + 1.0 * (0.25 - 0.1) = 0.25 + assert abs(p95 - 0.25) < 1e-6 + + +class TestParseGauge: + def test_parses_gauge(self): + lines = SAMPLE_METRICS.splitlines() + val = _parse_gauge(lines, "vllm:gpu_cache_usage_perc") + assert val == 0.42 + + def test_missing_gauge(self): + lines = SAMPLE_METRICS.splitlines() + val = _parse_gauge(lines, "nonexistent_gauge") + assert val == 0.0 + + def test_queue_depth(self): + lines = SAMPLE_METRICS.splitlines() + val = _parse_gauge(lines, "vllm:num_requests_waiting") + assert val == 3.0 + + +class TestVLLMMetricsScraper: + def test_parse_full_metrics(self): + scraper = VLLMMetricsScraper("http://localhost:8000") + metrics = scraper._parse(SAMPLE_METRICS) + + assert metrics.gpu_cache_usage_pct == 0.42 + assert metrics.queue_depth == 3.0 + assert metrics.ttft_p50 > 0 + assert metrics.ttft_p95 > 0 + assert metrics.ttft_p99 > 0 + assert metrics.e2e_latency_p50 > 0 + assert metrics.e2e_latency_p95 > 0 + + def test_scrape_connection_error(self): + scraper = VLLMMetricsScraper("http://localhost:99999") + metrics = scraper.scrape() + # Should return zeroed metrics, not raise + assert metrics == VLLMMetrics() + + def test_scrape_success_with_mock(self): + scraper = VLLMMetricsScraper("http://localhost:8000") + + class FakeResp: + status_code = 200 + text = SAMPLE_METRICS + def raise_for_status(self): + pass + + target = "openjarvis.telemetry.vllm_metrics.httpx.get" + with patch(target, return_value=FakeResp()): + metrics = scraper.scrape() + + assert metrics.gpu_cache_usage_pct == 0.42 + assert metrics.queue_depth == 3.0 + assert metrics.ttft_p50 > 0 + + def test_scrape_http_error(self): + import httpx as _httpx + + scraper = VLLMMetricsScraper("http://localhost:8000") + + def raise_status_error(*args, **kwargs): + request = _httpx.Request("GET", "http://localhost:8000/metrics") + response = _httpx.Response(500, request=request) + raise _httpx.HTTPStatusError( + "Server Error", request=request, response=response + ) + + target = "openjarvis.telemetry.vllm_metrics.httpx.get" + with patch(target, side_effect=raise_status_error): + metrics = scraper.scrape() + + assert metrics == VLLMMetrics() + + def test_empty_response(self): + scraper = VLLMMetricsScraper() + metrics = scraper._parse("") + assert metrics == VLLMMetrics() diff --git a/tests/tools/test_tool_descriptions.py b/tests/tools/test_tool_descriptions.py new file mode 100644 index 00000000..96bd8b2b --- /dev/null +++ b/tests/tools/test_tool_descriptions.py @@ -0,0 +1,182 @@ +"""Tests for the shared build_tool_descriptions() builder.""" + +from __future__ import annotations + +from openjarvis.core.types import ToolResult +from openjarvis.tools._stubs import ( + BaseTool, + ToolSpec, + build_tool_descriptions, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _CalcTool(BaseTool): + tool_id = "calculator" + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name="calculator", + description="Evaluate a mathematical expression safely.", + parameters={ + "type": "object", + "properties": { + "expression": { + "type": "string", + "description": ( + "Math expression to evaluate" + " (e.g. '2+3*4', 'sqrt(16)')" + ), + }, + }, + "required": ["expression"], + }, + category="math", + cost_estimate=0.0, + latency_estimate=0.0, + ) + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="calculator", content="0", success=True) + + +class _WebSearchTool(BaseTool): + tool_id = "web_search" + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name="web_search", + description="Search the web for real-time information.", + parameters={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query string", + }, + "max_results": { + "type": "integer", + "description": "Maximum results to return", + }, + }, + "required": ["query"], + }, + category="utility", + cost_estimate=0.001, + latency_estimate=2.0, + ) + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="web_search", content="", success=True) + + +class _NoCategoryTool(BaseTool): + tool_id = "think" + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name="think", + description="Internal reasoning scratchpad.", + parameters={ + "type": "object", + "properties": { + "thought": {"type": "string"}, + }, + }, + ) + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="think", content="", success=True) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestBuildToolDescriptions: + def test_empty_list_returns_no_tools(self): + assert build_tool_descriptions([]) == "No tools available." + + def test_single_tool_name_present(self): + result = build_tool_descriptions([_CalcTool()]) + assert "### calculator" in result + + def test_description_present(self): + result = build_tool_descriptions([_CalcTool()]) + assert "Evaluate a mathematical expression safely." in result + + def test_parameter_type_and_required(self): + result = build_tool_descriptions([_CalcTool()]) + assert "expression" in result + assert "string" in result + assert "required" in result + + def test_parameter_description(self): + result = build_tool_descriptions([_CalcTool()]) + assert "Math expression to evaluate" in result + + def test_optional_parameter_no_required_marker(self): + result = build_tool_descriptions([_WebSearchTool()]) + # max_results is optional + assert "max_results" in result + # The line for max_results should not say "required" + for line in result.splitlines(): + if "max_results" in line: + assert "required" not in line + + def test_category_shown_by_default(self): + result = build_tool_descriptions([_CalcTool()]) + assert "Category: math" in result + + def test_category_hidden_when_disabled(self): + result = build_tool_descriptions([_CalcTool()], include_category=False) + assert "Category:" not in result + + def test_empty_category_not_shown(self): + result = build_tool_descriptions([_NoCategoryTool()]) + assert "Category:" not in result + + def test_cost_hidden_by_default(self): + result = build_tool_descriptions([_WebSearchTool()]) + assert "Cost estimate:" not in result + + def test_cost_shown_when_enabled(self): + result = build_tool_descriptions([_WebSearchTool()], include_cost=True) + assert "Cost estimate:" in result + assert "Latency estimate:" in result + + def test_zero_cost_not_shown(self): + result = build_tool_descriptions([_CalcTool()], include_cost=True) + # cost_estimate is 0.0, so should not show + assert "Cost estimate:" not in result + + def test_multiple_tools(self): + result = build_tool_descriptions([_CalcTool(), _WebSearchTool()]) + assert "### calculator" in result + assert "### web_search" in result + + def test_tool_without_parameters(self): + """Tool with empty parameters dict.""" + + class _EmptyParamTool(BaseTool): + tool_id = "noop" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name="noop", description="Does nothing.") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="noop", content="", success=True) + + result = build_tool_descriptions([_EmptyParamTool()]) + assert "### noop" in result + assert "Does nothing." in result + # No Parameters section since no properties + assert "Parameters:" not in result diff --git a/uv.lock b/uv.lock index a9e52e3a..b303f140 100644 --- a/uv.lock +++ b/uv.lock @@ -2672,6 +2672,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.590.48" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/a0/f4fc18cf72f06821a9a665085435b901449986855519d5b3843532db35e9/nvidia_ml_py-13.590.48.tar.gz", hash = "sha256:8184d1be52914ac7f0991cd1c0d946c65dc88a840c754cd12c274b77b88760dd", size = 49732, upload-time = "2026-01-22T01:14:56.456Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/72/fb2af0d259a651affdce65fd6a495f0e07a685a0136baf585c5065204ee7/nvidia_ml_py-13.590.48-py3-none-any.whl", hash = "sha256:fd43d30ee9cd0b7940f5f9f9220b68d42722975e3992b6c21d14144c48760e43", size = 50680, upload-time = "2026-01-22T01:14:55.281Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.27.5" @@ -2764,6 +2773,7 @@ source = { editable = "." } dependencies = [ { name = "click" }, { name = "httpx" }, + { name = "pynvml" }, { name = "rich" }, { name = "tomli", marker = "python_full_version < '3.11'" }, ] @@ -2790,6 +2800,9 @@ docs = [ { name = "mkdocs-material" }, { name = "mkdocstrings", extra = ["python"] }, ] +gpu-metrics = [ + { name = "pynvml" }, +] inference-cloud = [ { name = "anthropic" }, { name = "openai" }, @@ -2851,6 +2864,8 @@ requires-dist = [ { name = "openhands-sdk", marker = "python_full_version >= '3.12' and extra == 'openhands'", specifier = ">=1.0" }, { name = "pdfplumber", marker = "extra == 'memory-pdf'", specifier = ">=0.10" }, { name = "pydantic", marker = "extra == 'server'", specifier = ">=2.0" }, + { name = "pynvml", specifier = ">=13.0.1" }, + { name = "pynvml", marker = "extra == 'gpu-metrics'", specifier = ">=12.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.24" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5" }, @@ -2868,7 +2883,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'orchestrator-training'", specifier = ">=4.40" }, { name = "uvicorn", marker = "extra == 'server'", specifier = ">=0.30" }, ] -provides-extras = ["dev", "inference-ollama", "inference-vllm", "inference-llamacpp", "inference-cloud", "inference-google", "inference-litellm", "tools-search", "memory-faiss", "memory-colbert", "memory-pdf", "memory-bm25", "server", "agents", "openhands", "learning", "orchestrator-training", "channel-telegram", "channel-discord", "channel-slack", "channel-webhook", "channel-email", "channel-whatsapp", "channel-signal", "channel-google-chat", "channel-irc", "channel-webchat", "channel-teams", "channel-matrix", "channel-mattermost", "channel-feishu", "channel-bluebubbles", "docs"] +provides-extras = ["dev", "inference-ollama", "inference-vllm", "inference-llamacpp", "inference-cloud", "inference-google", "inference-litellm", "tools-search", "memory-faiss", "memory-colbert", "memory-pdf", "memory-bm25", "server", "agents", "openhands", "gpu-metrics", "learning", "orchestrator-training", "channel-telegram", "channel-discord", "channel-slack", "channel-webhook", "channel-email", "channel-whatsapp", "channel-signal", "channel-google-chat", "channel-irc", "channel-webchat", "channel-teams", "channel-matrix", "channel-mattermost", "channel-feishu", "channel-bluebubbles", "docs"] [[package]] name = "opentelemetry-api" @@ -3835,6 +3850,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6f/2c/5b079febdc65e1c3fb2729bf958d18b45be7113828528e8a0b5850dd819a/pymdown_extensions-10.21-py3-none-any.whl", hash = "sha256:91b879f9f864d49794c2d9534372b10150e6141096c3908a455e45ca72ad9d3f", size = 268877, upload-time = "2026-02-15T20:44:05.464Z" }, ] +[[package]] +name = "pynvml" +version = "13.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-ml-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/57/da7dc63a79f59e082e26a66ac02d87d69ea316b35b35b7a00d82f3ce3d2f/pynvml-13.0.1.tar.gz", hash = "sha256:1245991d9db786b4d2f277ce66869bd58f38ac654e38c9397d18f243c8f6e48f", size = 35226, upload-time = "2025-09-05T20:33:25.377Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/4a/cac76c174bb439a0c46c9a4413fcbea5c6cabfb01879f7bbdb9fdfaed76c/pynvml-13.0.1-py3-none-any.whl", hash = "sha256:e2b20e0a501eeec951e2455b7ab444759cf048e0e13a57b08049fa2775266aa8", size = 28810, upload-time = "2025-09-05T20:33:24.13Z" }, +] + [[package]] name = "pypdfium2" version = "5.4.0"