From 05f2c02131c1b0574ba4892599c606d9607e4e78 Mon Sep 17 00:00:00 2001 From: Jon Saad-Falcon <41205309+jonsaadfalcon@users.noreply.github.com> Date: Thu, 12 Mar 2026 21:34:31 -0700 Subject: [PATCH] feat: Algolia DocSearch + learning subsystem reorganization (#43) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: create learning subdirectory structure (routing, agents, intelligence) Co-Authored-By: Claude Opus 4.6 * feat: extract classify_query to routing/_utils.py Move the classify_query() function and its regex patterns into a shared utility module so multiple routing policies can import it without depending on the full trace_policy module. Co-Authored-By: Claude Opus 4.6 * refactor: move routing files to learning/routing/ subdirectory Co-Authored-By: Claude Opus 4.6 * feat: create LearnedRouterPolicy merging trace-driven + SFT routing Co-Authored-By: Claude Opus 4.6 * feat: add conditional Algolia DocSearch integration Add Algolia DocSearch as an optional search upgrade — native lunr.js search remains the default until credentials are configured. Includes CDN assets, Jinja2 conditional config injection, init script with graceful fallback, light/dark theme CSS, improved search tokenization for snake_case/dotted identifiers, and search boosts for key pages. Co-Authored-By: Claude Opus 4.6 * refactor: move agent_evolver and skill_discovery to learning/agents/ Co-Authored-By: Claude Opus 4.6 * refactor: move learning/orchestrator to learning/intelligence/orchestrator Co-Authored-By: Claude Opus 4.6 * refactor: delete removed learning policies, rewrite __init__.py, clean up api_routes Co-Authored-By: Claude Opus 4.6 * feat: add SFT/GRPO/DSPy/GEPA config dataclasses, update LearningConfig Co-Authored-By: Claude Opus 4.6 * feat: add general-purpose SFT trainer (intelligence/sft_trainer.py) Co-Authored-By: Claude Opus 4.6 * fix: update stale imports in multi_model_router example Update imports to use new learning/routing/ paths after the subdirectory reorganization. Replace BanditRouterPolicy with LearnedRouterPolicy. Co-Authored-By: Claude Opus 4.6 * feat: add general-purpose GRPO trainer (intelligence/grpo_trainer.py) Co-Authored-By: Claude Opus 4.6 * feat: add DSPy agent optimizer (agents/dspy_optimizer.py) Co-Authored-By: Claude Opus 4.6 * feat: add GEPA agent optimizer (agents/gepa_optimizer.py) Co-Authored-By: Claude Opus 4.6 * feat: add learning-dspy and learning-gepa optional dependency extras Co-Authored-By: Claude Opus 4.6 * fix: update integration test to check for learned policy instead of grpo Co-Authored-By: Claude Opus 4.6 * fix: clean up stale APIs and unused params in examples - deep_research: remove system_prompt and max_turns params not accepted by Jarvis.ask(), inline system prompt into the query instead - doc_qa: remove unused --top-k CLI arg that was never passed to the API - multi_model_router: fix select_model() call to match single-arg signature Co-Authored-By: Claude Opus 4.6 * fix: import SFT/GRPO trainers in intelligence/__init__.py for registry Co-Authored-By: Claude Opus 4.6 * chore: remove .md file changes from PR Co-Authored-By: Claude Opus 4.6 * chore: restore search boost frontmatter for key docs pages Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- docs/architecture/overview.md | 7 + docs/getting-started/installation.md | 2 + docs/getting-started/quickstart.md | 2 + docs/index.md | 2 + docs/javascripts/docsearch-init.js | 36 ++ docs/overrides/main.html | 14 + docs/stylesheets/extra.css | 87 ++++ docs/user-guide/python-sdk.md | 7 + examples/deep_research/research.py | 15 +- examples/doc_qa/doc_qa.py | 8 +- .../multi_model_router/multi_model_router.py | 11 +- mkdocs.yml | 11 +- pyproject.toml | 2 + src/openjarvis/agents/orchestrator.py | 2 +- src/openjarvis/core/config.py | 106 ++++- src/openjarvis/intelligence/router.py | 2 +- src/openjarvis/learning/__init__.py | 54 +-- src/openjarvis/learning/agent_advisor.py | 156 ------- src/openjarvis/learning/agents/__init__.py | 3 + .../learning/{ => agents}/agent_evolver.py | 2 +- .../learning/agents/dspy_optimizer.py | 225 ++++++++++ .../learning/agents/gepa_optimizer.py | 321 +++++++++++++++ .../learning/{ => agents}/skill_discovery.py | 0 src/openjarvis/learning/bandit_router.py | 174 -------- src/openjarvis/learning/grpo_policy.py | 182 --------- src/openjarvis/learning/icl_updater.py | 236 ----------- .../learning/intelligence/__init__.py | 13 + .../learning/intelligence/grpo_trainer.py | 385 ++++++++++++++++++ .../orchestrator/__init__.py | 14 +- .../orchestrator/environment.py | 2 +- .../orchestrator/grpo_trainer.py | 6 +- .../orchestrator/policy_model.py | 2 +- .../orchestrator/prompt_registry.py | 0 .../{ => intelligence}/orchestrator/reward.py | 2 +- .../orchestrator/sft_trainer.py | 2 +- .../{ => intelligence}/orchestrator/types.py | 0 .../learning/intelligence/sft_trainer.py | 149 +++++++ .../learning/learning_orchestrator.py | 2 +- src/openjarvis/learning/routing/__init__.py | 3 + src/openjarvis/learning/routing/_utils.py | 34 ++ .../{ => routing}/heuristic_policy.py | 2 +- .../{ => routing}/heuristic_reward.py | 0 .../learned_router.py} | 128 +++--- .../learning/{ => routing}/router.py | 0 src/openjarvis/learning/sft_policy.py | 101 ----- src/openjarvis/learning/training/data.py | 2 +- src/openjarvis/server/api_routes.py | 61 +-- src/openjarvis/system.py | 7 +- tests/core/test_config.py | 1 - tests/integration/test_integration.py | 12 +- tests/learning/agents/__init__.py | 0 tests/learning/agents/test_dspy_optimizer.py | 82 ++++ tests/learning/agents/test_gepa_optimizer.py | 85 ++++ tests/learning/intelligence/__init__.py | 0 .../intelligence/test_grpo_trainer.py | 88 ++++ .../learning/intelligence/test_sft_trainer.py | 66 +++ tests/learning/routing/__init__.py | 0 tests/learning/routing/test_learned_router.py | 151 +++++++ tests/learning/routing/test_utils.py | 27 ++ tests/learning/test_agent_advisor.py | 78 ---- tests/learning/test_agent_evolver.py | 2 +- tests/learning/test_bandit_router.py | 184 --------- tests/learning/test_device_selection.py | 8 +- tests/learning/test_grpo_policy.py | 158 ------- tests/learning/test_heuristic_policy.py | 4 +- tests/learning/test_heuristic_reward.py | 2 +- tests/learning/test_icl_updater.py | 136 ------- tests/learning/test_icl_updates.py | 205 ---------- tests/learning/test_learning_api.py | 31 +- tests/learning/test_router.py | 2 +- tests/learning/test_router_stubs.py | 2 +- tests/learning/test_routing_models.py | 2 +- tests/learning/test_sft_policy.py | 82 ---- tests/learning/test_skill_discovery.py | 2 +- tests/learning/test_system_learning.py | 6 +- tests/learning/test_trace_policy.py | 203 --------- .../test_environment.py | 4 +- .../test_grpo_trainer.py | 8 +- .../test_policy_model.py | 4 +- .../test_prompt_registry.py | 2 +- .../test_orchestrator_learning/test_reward.py | 4 +- .../test_sft_trainer.py | 4 +- .../test_orchestrator_learning/test_types.py | 2 +- 83 files changed, 2049 insertions(+), 2180 deletions(-) create mode 100644 docs/javascripts/docsearch-init.js create mode 100644 docs/overrides/main.html delete mode 100644 src/openjarvis/learning/agent_advisor.py create mode 100644 src/openjarvis/learning/agents/__init__.py rename src/openjarvis/learning/{ => agents}/agent_evolver.py (99%) create mode 100644 src/openjarvis/learning/agents/dspy_optimizer.py create mode 100644 src/openjarvis/learning/agents/gepa_optimizer.py rename src/openjarvis/learning/{ => agents}/skill_discovery.py (100%) delete mode 100644 src/openjarvis/learning/bandit_router.py delete mode 100644 src/openjarvis/learning/grpo_policy.py delete mode 100644 src/openjarvis/learning/icl_updater.py create mode 100644 src/openjarvis/learning/intelligence/__init__.py create mode 100644 src/openjarvis/learning/intelligence/grpo_trainer.py rename src/openjarvis/learning/{ => intelligence}/orchestrator/__init__.py (80%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/environment.py (98%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/grpo_trainer.py (98%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/policy_model.py (99%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/prompt_registry.py (100%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/reward.py (99%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/sft_trainer.py (99%) rename src/openjarvis/learning/{ => intelligence}/orchestrator/types.py (100%) create mode 100644 src/openjarvis/learning/intelligence/sft_trainer.py create mode 100644 src/openjarvis/learning/routing/__init__.py create mode 100644 src/openjarvis/learning/routing/_utils.py rename src/openjarvis/learning/{ => routing}/heuristic_policy.py (86%) rename src/openjarvis/learning/{ => routing}/heuristic_reward.py (100%) rename src/openjarvis/learning/{trace_policy.py => routing/learned_router.py} (64%) rename src/openjarvis/learning/{ => routing}/router.py (100%) delete mode 100644 src/openjarvis/learning/sft_policy.py create mode 100644 tests/learning/agents/__init__.py create mode 100644 tests/learning/agents/test_dspy_optimizer.py create mode 100644 tests/learning/agents/test_gepa_optimizer.py create mode 100644 tests/learning/intelligence/__init__.py create mode 100644 tests/learning/intelligence/test_grpo_trainer.py create mode 100644 tests/learning/intelligence/test_sft_trainer.py create mode 100644 tests/learning/routing/__init__.py create mode 100644 tests/learning/routing/test_learned_router.py create mode 100644 tests/learning/routing/test_utils.py delete mode 100644 tests/learning/test_agent_advisor.py delete mode 100644 tests/learning/test_bandit_router.py delete mode 100644 tests/learning/test_grpo_policy.py delete mode 100644 tests/learning/test_icl_updater.py delete mode 100644 tests/learning/test_icl_updates.py delete mode 100644 tests/learning/test_sft_policy.py delete mode 100644 tests/learning/test_trace_policy.py diff --git a/docs/architecture/overview.md b/docs/architecture/overview.md index b009c386..cea7c904 100644 --- a/docs/architecture/overview.md +++ b/docs/architecture/overview.md @@ -1,3 +1,10 @@ +--- +title: Architecture Overview +description: The five-primitive architecture behind OpenJarvis — Intelligence, Engine, Agents, Tools, and Learning +search: + boost: 2 +--- + # Architecture Overview OpenJarvis is a research framework for studying on-device AI systems. Its architecture is organized around **five core abstractions** -- Intelligence, Engine, Agentic Logic, Memory, and Learning -- that work together through trace-driven feedback. diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 682c2eee..9324dbe0 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -1,6 +1,8 @@ --- title: Installation description: Get OpenJarvis running — browser app, desktop app, CLI, or Python SDK +search: + boost: 3 --- # Installation diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index be7f01e3..10f7fcfc 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -1,6 +1,8 @@ --- title: Quick Start description: Get up and running with OpenJarvis in minutes +search: + boost: 3 --- # Quick Start diff --git a/docs/index.md b/docs/index.md index ca041a87..a3b0ad54 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,8 @@ --- title: OpenJarvis description: Personal AI, On Personal Devices +search: + boost: 2 hide: - navigation --- diff --git a/docs/javascripts/docsearch-init.js b/docs/javascripts/docsearch-init.js new file mode 100644 index 00000000..8db67eca --- /dev/null +++ b/docs/javascripts/docsearch-init.js @@ -0,0 +1,36 @@ +(function () { + "use strict"; + + var cfg = window.__DOCSEARCH_CONFIG__; + if (!cfg || !cfg.appId || !cfg.apiKey || !cfg.indexName) return; + + function init() { + var header = document.querySelector(".md-header__inner"); + if (!header) return; + + var container = document.createElement("div"); + container.id = "docsearch"; + + var nativeSearch = header.querySelector(".md-search"); + if (nativeSearch) { + header.insertBefore(container, nativeSearch); + nativeSearch.style.display = "none"; + } else { + header.appendChild(container); + } + + docsearch({ + appId: cfg.appId, + apiKey: cfg.apiKey, + indexName: cfg.indexName, + container: "#docsearch", + placeholder: "Search docs...", + }); + } + + if (document.readyState === "loading") { + document.addEventListener("DOMContentLoaded", init); + } else { + init(); + } +})(); diff --git a/docs/overrides/main.html b/docs/overrides/main.html new file mode 100644 index 00000000..f7796764 --- /dev/null +++ b/docs/overrides/main.html @@ -0,0 +1,14 @@ +{% extends "base.html" %} + +{% block extrahead %} + {{ super() }} + {% if config.extra.algolia and config.extra.algolia.app_id and config.extra.algolia.api_key and config.extra.algolia.index_name %} + + {% endif %} +{% endblock %} diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css index 60469369..7da9ea3f 100644 --- a/docs/stylesheets/extra.css +++ b/docs/stylesheets/extra.css @@ -410,3 +410,90 @@ font-family: var(--md-code-font-family, monospace); font-size: 13px; } + +/* ── DocSearch ───────────────────────────────────────────────────────── */ +#docsearch { + display: flex; + align-items: center; + margin-right: 0.4rem; +} + +.DocSearch-Button { + background: transparent !important; + border: 1px solid rgba(255, 255, 255, 0.2) !important; + border-radius: 6px !important; + font-family: Georgia, "Times New Roman", serif !important; + font-size: 0.82rem !important; + padding: 0 0.8rem !important; + height: 2rem !important; + color: rgba(255, 255, 255, 0.7) !important; + transition: border-color 0.2s ease !important; +} + +.DocSearch-Button:hover { + border-color: rgba(255, 255, 255, 0.5) !important; + box-shadow: none !important; +} + +.DocSearch-Button .DocSearch-Search-Icon { + width: 14px; + height: 14px; + color: rgba(255, 255, 255, 0.5); +} + +.DocSearch-Button-Keys { + display: flex !important; +} + +/* Light theme modal */ +[data-md-color-scheme="default"] { + --docsearch-primary-color: #3d5a80; + --docsearch-text-color: #2b2b2b; + --docsearch-muted-color: #666666; + --docsearch-container-background: rgba(43, 43, 43, 0.6); + --docsearch-modal-background: #f8f7f4; + --docsearch-modal-shadow: 0 4px 30px rgba(0, 0, 0, 0.15); + --docsearch-searchbox-background: #f0efeb; + --docsearch-searchbox-focus-background: #ffffff; + --docsearch-searchbox-shadow: inset 0 0 0 2px #3d5a80; + --docsearch-hit-background: #ffffff; + --docsearch-hit-shadow: 0 1px 3px rgba(0, 0, 0, 0.06); + --docsearch-hit-color: #2b2b2b; + --docsearch-highlight-color: #3d5a80; + --docsearch-logo-color: #666666; + --docsearch-footer-background: #f0efeb; + --docsearch-footer-shadow: 0 -1px 0 rgba(0, 0, 0, 0.08); + --docsearch-key-gradient: linear-gradient(-26.5deg, #e8e6e1, #f8f7f4); + --docsearch-key-shadow: inset 0 -2px 0 0 #d6d4cf, inset 0 0 1px 1px #ffffff, + 0 1px 2px 1px rgba(0, 0, 0, 0.08); +} + +/* Dark theme modal */ +[data-md-color-scheme="slate"] { + --docsearch-primary-color: #8da9c4; + --docsearch-text-color: #e8e6e1; + --docsearch-muted-color: #999999; + --docsearch-container-background: rgba(0, 0, 0, 0.7); + --docsearch-modal-background: #1a1a1a; + --docsearch-modal-shadow: 0 4px 30px rgba(0, 0, 0, 0.4); + --docsearch-searchbox-background: #222222; + --docsearch-searchbox-focus-background: #2b2b2b; + --docsearch-searchbox-shadow: inset 0 0 0 2px #8da9c4; + --docsearch-hit-background: #222222; + --docsearch-hit-shadow: 0 1px 3px rgba(0, 0, 0, 0.2); + --docsearch-hit-color: #e8e6e1; + --docsearch-highlight-color: #8da9c4; + --docsearch-logo-color: #999999; + --docsearch-footer-background: #222222; + --docsearch-footer-shadow: 0 -1px 0 rgba(255, 255, 255, 0.08); + --docsearch-key-gradient: linear-gradient(-26.5deg, #2b2b2b, #333333); + --docsearch-key-shadow: inset 0 -2px 0 0 #1a1a1a, inset 0 0 1px 1px #444444, + 0 1px 2px 1px rgba(0, 0, 0, 0.3); +} + +/* Responsive: hide placeholder on mobile */ +@media (max-width: 768px) { + .DocSearch-Button-Placeholder { + display: none !important; + } +} diff --git a/docs/user-guide/python-sdk.md b/docs/user-guide/python-sdk.md index 01dad7b6..6da7cd51 100644 --- a/docs/user-guide/python-sdk.md +++ b/docs/user-guide/python-sdk.md @@ -1,3 +1,10 @@ +--- +title: Python SDK +description: High-level Python interface for local inference, memory, and agent workflows +search: + boost: 2 +--- + # Python SDK The OpenJarvis Python SDK provides a high-level interface for interacting with local inference engines, managing memory, and running agent workflows. The primary entry point is the `Jarvis` class. diff --git a/examples/deep_research/research.py b/examples/deep_research/research.py index 4ec981f7..dd21620a 100644 --- a/examples/deep_research/research.py +++ b/examples/deep_research/research.py @@ -6,7 +6,7 @@ Usage: python examples/deep_research/research.py "climate policy" \ --model gpt-4o --engine cloud python examples/deep_research/research.py "rust vs go" \ - --output report.md --max-turns 20 + --output report.md """ from __future__ import annotations @@ -31,13 +31,6 @@ import click show_default=True, help="Engine backend (ollama, cloud, vllm, etc.).", ) -@click.option( - "--max-turns", - default=15, - show_default=True, - type=int, - help="Maximum agent loop iterations.", -) @click.option( "--output", default=None, @@ -48,7 +41,6 @@ def main( topic: str, model: str, engine_key: str, - max_turns: int, output: str | None, ) -> None: """Run a deep research session on TOPIC using an orchestrator agent. @@ -80,7 +72,7 @@ def main( ) click.echo(f"Researching: {topic}") - click.echo(f"Model: {model} | Engine: {engine_key} | Max turns: {max_turns}") + click.echo(f"Model: {model} | Engine: {engine_key}") click.echo("-" * 60) try: @@ -98,6 +90,7 @@ def main( try: prompt = ( + f"{system_prompt}\n\n" "Research the following topic in depth " f"and produce a report:\n\n{topic}" ) @@ -105,8 +98,6 @@ def main( prompt, agent="orchestrator", tools=tools, - system_prompt=system_prompt, - max_turns=max_turns, temperature=0.5, ) except Exception as exc: diff --git a/examples/doc_qa/doc_qa.py b/examples/doc_qa/doc_qa.py index d2d29ab0..c6c25481 100644 --- a/examples/doc_qa/doc_qa.py +++ b/examples/doc_qa/doc_qa.py @@ -53,12 +53,6 @@ def main() -> None: default=512, help="Chunk size for document indexing (default: 512).", ) - parser.add_argument( - "--top-k", - type=int, - default=5, - help="Number of relevant chunks to retrieve (default: 5).", - ) args = parser.parse_args() try: @@ -100,7 +94,7 @@ def main() -> None: sys.exit(1) # Step 2: Ask the question with memory context enabled - print(f"Searching for relevant context (top_k={args.top_k})...") + print("Searching for relevant context...") try: response = j.ask( args.query, diff --git a/examples/multi_model_router/multi_model_router.py b/examples/multi_model_router/multi_model_router.py index a45e75b7..ff2c3bf3 100644 --- a/examples/multi_model_router/multi_model_router.py +++ b/examples/multi_model_router/multi_model_router.py @@ -62,7 +62,10 @@ def main() -> None: try: from openjarvis import Jarvis - from openjarvis.learning.router import HeuristicRouter, build_routing_context + from openjarvis.learning.routing.router import ( + HeuristicRouter, + build_routing_context, + ) except ImportError: print( "Error: openjarvis is not installed. " @@ -108,10 +111,10 @@ def main() -> None: # Select the model using the chosen strategy if args.strategy == "bandit": - from openjarvis.learning.bandit_router import BanditRouterPolicy + from openjarvis.learning.routing.learned_router import LearnedRouterPolicy - router = BanditRouterPolicy() - selected_model = router.route(context, available_models) + router = LearnedRouterPolicy() + selected_model = router.select_model(context) else: router = HeuristicRouter(available_models) selected_model = router.select_model(context) diff --git a/mkdocs.yml b/mkdocs.yml index be1fc758..495f86aa 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -10,6 +10,7 @@ copyright: Copyright © 2026 OpenJarvis Contributors theme: name: material + custom_dir: docs/overrides language: en font: text: Georgia @@ -47,9 +48,11 @@ theme: extra_css: - stylesheets/extra.css + - https://cdn.jsdelivr.net/npm/@docsearch/css@3 plugins: - - search + - search: + separator: '[\s\-\._]+' - gen-files: scripts: - docs/gen_ref_pages.py @@ -124,8 +127,14 @@ markdown_extensions: extra_javascript: - javascripts/leaderboard.js + - https://cdn.jsdelivr.net/npm/@docsearch/js@3 + - javascripts/docsearch-init.js extra: + algolia: + app_id: "" # Fill after DocSearch approval + api_key: "" # Search-only key + index_name: "" # e.g. "openjarvis" social: - icon: fontawesome/brands/github link: https://github.com/open-jarvis/OpenJarvis diff --git a/pyproject.toml b/pyproject.toml index f6781c5a..6a50980e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,8 @@ energy-amd = ["amdsmi>=6.1"] energy-apple = ["zeus-ml[apple]"] energy-all = ["pynvml>=12.0", "amdsmi>=6.1", "zeus-ml[apple]"] orchestrator-training = ["torch>=2.0", "transformers>=4.40"] +learning-dspy = ["dspy>=2.6"] +learning-gepa = ["gepa>=0.1"] channel-telegram = ["python-telegram-bot>=21.0"] channel-discord = ["discord.py>=2.3"] channel-slack = ["slack-sdk>=3.27"] diff --git a/src/openjarvis/agents/orchestrator.py b/src/openjarvis/agents/orchestrator.py index 329b3392..b89a406d 100644 --- a/src/openjarvis/agents/orchestrator.py +++ b/src/openjarvis/agents/orchestrator.py @@ -91,7 +91,7 @@ class OrchestratorAgent(ToolUsingAgent): if self._system_prompt: sys_prompt = self._system_prompt else: - from openjarvis.learning.orchestrator.prompt_registry import ( + from openjarvis.learning.intelligence.orchestrator.prompt_registry import ( build_system_prompt, ) sys_prompt = build_system_prompt(tools=self._tools) diff --git a/src/openjarvis/core/config.py b/src/openjarvis/core/config.py index 5c3dce1b..416854ab 100644 --- a/src/openjarvis/core/config.py +++ b/src/openjarvis/core/config.py @@ -466,24 +466,111 @@ class IntelligenceConfig: class RoutingLearningConfig: """Routing sub-policy config within Learning.""" - policy: str = "heuristic" # heuristic | learned | grpo + policy: str = "heuristic" # heuristic | learned min_samples: int = 5 # Min traces before trusting learned routing +@dataclass(slots=True) +class SFTConfig: + """General-purpose SFT training config. Maps to [learning.intelligence.sft].""" + + model_name: str = "Qwen/Qwen3-1.7B" + max_seq_length: int = 4096 + num_epochs: int = 3 + batch_size: int = 8 + learning_rate: float = 2e-5 + weight_decay: float = 0.01 + warmup_ratio: float = 0.1 + max_grad_norm: float = 1.0 + gradient_checkpointing: bool = True + use_lora: bool = True + lora_rank: int = 16 + lora_alpha: int = 32 + lora_dropout: float = 0.05 + target_modules: str = "q_proj,v_proj" # comma-separated for TOML compat + use_4bit: bool = False + checkpoint_dir: str = "checkpoints/sft" + min_pairs: int = 10 + agent_filter: str = "" + + +@dataclass(slots=True) +class GRPOConfig: + """General-purpose GRPO training config. Maps to [learning.intelligence.grpo].""" + + model_name: str = "Qwen/Qwen3-1.7B" + max_seq_length: int = 4096 + max_response_length: int = 2048 + num_epochs: int = 10 + batch_size: int = 16 + learning_rate: float = 1e-6 + max_grad_norm: float = 1.0 + gradient_checkpointing: bool = True + num_samples_per_prompt: int = 8 + temperature: float = 1.0 + kl_coef: float = 0.0001 + clip_ratio: float = 0.2 + use_8bit_ref: bool = True + checkpoint_dir: str = "checkpoints/grpo" + save_every_n_epochs: int = 1 + keep_last_n: int = 3 + min_prompts: int = 10 + agent_filter: str = "" + + +@dataclass(slots=True) +class DSPyOptimizerConfig: + """DSPy agent optimizer config. Maps to [learning.agent.dspy].""" + + optimizer: str = "BootstrapFewShotWithRandomSearch" + task_lm: str = "" + teacher_lm: str = "" + max_bootstrapped_demos: int = 4 + max_labeled_demos: int = 4 + num_candidate_programs: int = 10 + max_rounds: int = 1 + optimize_system_prompt: bool = True + optimize_few_shot: bool = True + optimize_tool_descriptions: bool = True + min_traces: int = 20 + metric_threshold: float = 0.7 + agent_filter: str = "" + config_dir: str = "" + + +@dataclass(slots=True) +class GEPAOptimizerConfig: + """GEPA agent optimizer config. Maps to [learning.agent.gepa].""" + + reflection_lm: str = "" + max_metric_calls: int = 150 + population_size: int = 10 + optimize_system_prompt: bool = True + optimize_tools: bool = True + optimize_max_turns: bool = True + optimize_temperature: bool = True + min_traces: int = 20 + assessment_batch_size: int = 10 + agent_filter: str = "" + config_dir: str = "" + + @dataclass(slots=True) class IntelligenceLearningConfig: """Intelligence sub-policy config within Learning.""" - policy: str = "none" # none | sft + policy: str = "none" # none | sft | grpo + sft: SFTConfig = field(default_factory=SFTConfig) + grpo: GRPOConfig = field(default_factory=GRPOConfig) @dataclass(slots=True) class AgentLearningConfig: """Agent sub-policy config within Learning.""" - policy: str = "none" # none | agent_advisor | icl_updater - max_icl_examples: int = 20 - advisor_confidence_threshold: float = 0.7 + policy: str = "none" # none | dspy | gepa + dspy: DSPyOptimizerConfig = field(default_factory=DSPyOptimizerConfig) + gepa: GEPAOptimizerConfig = field(default_factory=GEPAOptimizerConfig) @dataclass(slots=True) @@ -501,8 +588,8 @@ class LearningConfig: """Learning system settings with per-primitive sub-policies.""" enabled: bool = False - update_interval: int = 100 # traces between automatic policy updates - auto_update: bool = False # auto-trigger updates on interval + update_interval: int = 100 + auto_update: bool = False routing: RoutingLearningConfig = field(default_factory=RoutingLearningConfig) intelligence: IntelligenceLearningConfig = field( default_factory=IntelligenceLearningConfig, @@ -512,10 +599,7 @@ class LearningConfig: # Training pipeline training_enabled: bool = False - training_schedule: str = "" # cron expression or empty for on-demand - lora_rank: int = 16 - lora_alpha: int = 32 - min_sft_pairs: int = 50 + training_schedule: str = "" min_improvement: float = 0.02 # Backward-compat properties for old flat field names diff --git a/src/openjarvis/intelligence/router.py b/src/openjarvis/intelligence/router.py index d57decfb..6ab3d94e 100644 --- a/src/openjarvis/intelligence/router.py +++ b/src/openjarvis/intelligence/router.py @@ -1,6 +1,6 @@ """Backward-compat shim — canonical location is learning.router.""" -from openjarvis.learning.router import ( # noqa: F401 +from openjarvis.learning.routing.router import ( # noqa: F401 DefaultQueryAnalyzer, HeuristicRouter, build_routing_context, diff --git a/src/openjarvis/learning/__init__.py b/src/openjarvis/learning/__init__.py index dad28f12..ec1c8315 100644 --- a/src/openjarvis/learning/__init__.py +++ b/src/openjarvis/learning/__init__.py @@ -1,4 +1,4 @@ -"""Learning primitive — router policies, reward functions, and trace-driven learning.""" +"""Learning primitive -- router policies, reward functions, learning.""" from __future__ import annotations @@ -8,13 +8,13 @@ from openjarvis.learning._stubs import ( RouterPolicy, RoutingContext, ) -from openjarvis.learning.agent_evolver import AgentConfigEvolver -from openjarvis.learning.heuristic_reward import HeuristicRewardFunction +from openjarvis.learning.agents.agent_evolver import AgentConfigEvolver from openjarvis.learning.learning_orchestrator import LearningOrchestrator from openjarvis.learning.optimize.llm_optimizer import LLMOptimizer from openjarvis.learning.optimize.optimizer import OptimizationEngine from openjarvis.learning.optimize.store import OptimizationStore -from openjarvis.learning.router import ( +from openjarvis.learning.routing.heuristic_reward import HeuristicRewardFunction +from openjarvis.learning.routing.router import ( HeuristicRouter, build_routing_context, ) @@ -23,58 +23,36 @@ from openjarvis.learning.training.lora import HAS_TORCH, LoRATrainer, LoRATraini def ensure_registered() -> None: - """Ensure all learning policies are registered in RouterPolicyRegistry. - - Imported lazily to avoid circular imports with the intelligence primitive. - """ - from openjarvis.learning.heuristic_policy import ( + """Ensure all learning policies are registered in RouterPolicyRegistry.""" + from openjarvis.learning.routing.heuristic_policy import ( ensure_registered as _reg_heuristic, ) - _reg_heuristic() - try: - from openjarvis.learning.grpo_policy import ( - ensure_registered as _reg_grpo, - ) - - _reg_grpo() - except ImportError: - pass - - try: - from openjarvis.learning.bandit_router import ( - ensure_registered as _reg_bandit, - ) - - _reg_bandit() - except ImportError: - pass - - from openjarvis.learning.trace_policy import ( - ensure_registered as _reg_trace, + from openjarvis.learning.routing.learned_router import ( + ensure_registered as _reg_learned, ) + _reg_learned() - _reg_trace() - + # Intelligence training (optional deps) try: - import openjarvis.learning.sft_policy # noqa: F401 + import openjarvis.learning.intelligence # noqa: F401 except ImportError: pass + # Orchestrator-specific training (optional deps) try: - import openjarvis.learning.agent_advisor # noqa: F401 + import openjarvis.learning.intelligence.orchestrator # noqa: F401 except ImportError: pass + # Agent optimizers (optional deps) try: - import openjarvis.learning.icl_updater # noqa: F401 + import openjarvis.learning.agents.dspy_optimizer # noqa: F401 except ImportError: pass - - # Orchestrator-native SFT & GRPO training try: - import openjarvis.learning.orchestrator # noqa: F401 + import openjarvis.learning.agents.gepa_optimizer # noqa: F401 except ImportError: pass diff --git a/src/openjarvis/learning/agent_advisor.py b/src/openjarvis/learning/agent_advisor.py deleted file mode 100644 index 661e8b4e..00000000 --- a/src/openjarvis/learning/agent_advisor.py +++ /dev/null @@ -1,156 +0,0 @@ -"""LM-guided agent restructuring — analyzes traces and suggests improvements.""" - -from __future__ import annotations - -import logging -from typing import Any, Dict, List - -from openjarvis.core.registry import LearningRegistry -from openjarvis.learning._stubs import AgentLearningPolicy - -logger = logging.getLogger(__name__) - - -@LearningRegistry.register("agent_advisor") -class AgentAdvisorPolicy(AgentLearningPolicy): - """Higher-level LM analyzes traces, suggests agent structure changes. - - Does NOT auto-apply changes — returns recommendations that can be - reviewed or applied via config. - """ - - def __init__( - self, - *, - advisor_engine: Any = None, - advisor_model: str = "", - max_traces: int = 50, - ) -> None: - self._advisor_engine = advisor_engine - self._advisor_model = advisor_model - self._max_traces = max_traces - - def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: - """Analyze traces and return agent improvement recommendations.""" - try: - traces = trace_store.list_traces() - except Exception as exc: - logger.warning("Agent advisor analysis failed: %s", exc) - return {"recommendations": [], "confidence": 0.0} - - # Collect failing or slow traces - problem_traces = [] - for trace in traces[-self._max_traces :]: - is_failing = trace.outcome != "success" - is_slow = (trace.total_latency_seconds or 0) > 5.0 - if is_failing or is_slow: - problem_traces.append(trace) - - if not problem_traces: - return { - "recommendations": [], - "confidence": 1.0, - "message": "No problematic traces found", - } - - # Analyze patterns without LM (structural analysis) - recommendations = self._analyze_patterns(problem_traces) - - # If advisor engine available, get LM-guided recommendations - if self._advisor_engine and self._advisor_model: - try: - lm_recs = self._get_lm_recommendations(problem_traces) - recommendations.extend(lm_recs) - except Exception as exc: - logger.debug("Failed to parse agent advisor recommendation: %s", exc) - - confidence = 1.0 - (len(problem_traces) / max(len(traces), 1)) - return { - "recommendations": recommendations, - "confidence": round(confidence, 2), - "analyzed_traces": len(traces), - "problem_traces": len(problem_traces), - } - - def _analyze_patterns(self, traces: List[Any]) -> List[Dict[str, Any]]: - """Structural analysis of trace patterns.""" - recs: List[Dict[str, Any]] = [] - - # Check for excessive tool calls - tool_heavy = [ - t - for t in traces - if sum(1 for s in (t.steps or []) if s.step_type.value == "tool_call") - > 5 - ] - if len(tool_heavy) > len(traces) * 0.3: - recs.append( - { - "type": "agent_structure", - "suggestion": ( - "Reduce tool call frequency" - " — many traces have >5 tool calls" - ), - "severity": "medium", - } - ) - - # Check for repeated failures on same query type - failure_classes: Dict[str, int] = {} - for t in traces: - if t.outcome != "success": - qclass = self._classify(t.query) - failure_classes[qclass] = failure_classes.get(qclass, 0) + 1 - for qclass, count in failure_classes.items(): - if count >= 3: - recs.append( - { - "type": "routing", - "suggestion": ( - f"Query class '{qclass}' has {count}" - " failures — consider different model or agent" - ), - "severity": "high", - } - ) - - return recs - - @staticmethod - def _classify(query: str) -> str: - q = query.lower() - if any(kw in q for kw in ("def ", "class ", "import ")): - return "code" - if any(kw in q for kw in ("solve", "equation")): - return "math" - return "general" - - def _get_lm_recommendations(self, traces: List[Any]) -> List[Dict[str, Any]]: - """Get LM-guided recommendations (requires advisor engine).""" - from openjarvis.core.types import Message, Role - - summaries = [] - for t in traces[:10]: - summary = ( - f"Query: {t.query[:100]}, " - f"Outcome: {t.outcome}, " - f"Latency: {t.total_latency_seconds:.1f}s, " - f"Steps: {len(t.steps or [])}" - ) - summaries.append(summary) - - prompt = ( - "Analyze these agent interaction traces and suggest improvements:\n\n" - + "\n".join(summaries) - + "\n\nProvide 1-3 specific recommendations." - ) - messages = [Message(role=Role.USER, content=prompt)] - result = self._advisor_engine.generate( - messages, - model=self._advisor_model, - ) - content = result.get("content", "") - return [{"type": "lm_guided", "suggestion": content, "severity": "info"}] - - -__all__ = ["AgentAdvisorPolicy"] diff --git a/src/openjarvis/learning/agents/__init__.py b/src/openjarvis/learning/agents/__init__.py new file mode 100644 index 00000000..c409ac2f --- /dev/null +++ b/src/openjarvis/learning/agents/__init__.py @@ -0,0 +1,3 @@ +"""Agent learning — config optimization via DSPy, GEPA, and evolution.""" + +from __future__ import annotations diff --git a/src/openjarvis/learning/agent_evolver.py b/src/openjarvis/learning/agents/agent_evolver.py similarity index 99% rename from src/openjarvis/learning/agent_evolver.py rename to src/openjarvis/learning/agents/agent_evolver.py index fc84b1b2..5515e938 100644 --- a/src/openjarvis/learning/agent_evolver.py +++ b/src/openjarvis/learning/agents/agent_evolver.py @@ -13,7 +13,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Union from openjarvis.core.types import StepType, Trace -from openjarvis.learning.trace_policy import classify_query +from openjarvis.learning.routing._utils import classify_query from openjarvis.traces.store import TraceStore diff --git a/src/openjarvis/learning/agents/dspy_optimizer.py b/src/openjarvis/learning/agents/dspy_optimizer.py new file mode 100644 index 00000000..e81106cb --- /dev/null +++ b/src/openjarvis/learning/agents/dspy_optimizer.py @@ -0,0 +1,225 @@ +"""DSPy agent optimizer -- programmatic pipeline optimization. + +Wraps an agent's reasoning pipeline as a DSPy Module and optimizes +it end-to-end using DSPy teleprompters. Outputs TOML-compatible +config updates written via AgentConfigEvolver. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List + +from openjarvis.core.config import DSPyOptimizerConfig +from openjarvis.core.registry import LearningRegistry +from openjarvis.learning._stubs import AgentLearningPolicy + +logger = logging.getLogger(__name__) + +# Optional dependency +try: + import dspy + + HAS_DSPY = True +except ImportError: + HAS_DSPY = False + dspy = None # type: ignore[assignment] + + +class DSPyAgentOptimizer: + """Optimize agent configs using DSPy teleprompters. + + Parameters + ---------- + config: + DSPyOptimizerConfig controlling optimizer type and parameters. + """ + + def __init__(self, config: DSPyOptimizerConfig) -> None: + self.config = config + + def optimize(self, trace_store: Any) -> Dict[str, Any]: + """Run DSPy optimization on traces from the store. + + 1. Extract traces and convert to DSPy Examples + 2. Build a DSPy Module mirroring the agent pipeline + 3. Run the configured teleprompter + 4. Extract optimized parameters as TOML updates + 5. Write via AgentConfigEvolver if config_dir is set + """ + # Get traces + kwargs: Dict[str, Any] = {"limit": 10_000} + if self.config.agent_filter: + kwargs["agent"] = self.config.agent_filter + traces = trace_store.list_traces(**kwargs) + + if len(traces) < self.config.min_traces: + return { + "status": "skipped", + "reason": ( + f"only {len(traces)} traces, " + f"min_traces={self.config.min_traces}" + ), + } + + if not HAS_DSPY: + return { + "status": "error", + "reason": ( + "dspy not installed " + "(pip install 'openjarvis[learning-dspy]')" + ), + } + + try: + optimized = self._run_dspy_optimization(traces) + except Exception as exc: + logger.warning("DSPy optimization failed: %s", exc) + return {"status": "error", "reason": str(exc)} + + config_updates = self._to_config_updates(optimized) + + # Write configs if config_dir is set + if self.config.config_dir: + self._write_configs(config_updates) + + return { + "status": "completed", + "traces_used": len(traces), + "config_updates": config_updates, + } + + def _run_dspy_optimization(self, traces: List[Any]) -> Dict[str, Any]: + """Run the DSPy teleprompter on converted trace examples.""" + # Convert traces to dspy.Example objects + examples = [] + for t in traces: + ex = dspy.Example( + question=t.query, + answer=t.result, + ).with_inputs("question") + examples.append(ex) + + # Define a simple metric based on trace feedback + def metric( + example: Any, prediction: Any, trace: Any = None + ) -> float: + for tr in traces: + if tr.query == example.question: + return tr.feedback if tr.feedback is not None else 0.5 + return 0.5 + + # Build a minimal DSPy program + class AgentModule(dspy.Module): + def __init__(self_inner: Any) -> None: + super().__init__() + self_inner.generate = dspy.ChainOfThought( + "question -> answer" + ) + + def forward(self_inner: Any, question: str) -> Any: + return self_inner.generate(question=question) + + program = AgentModule() + + # Select optimizer + optimizer_name = self.config.optimizer + if optimizer_name == "BootstrapFewShotWithRandomSearch": + teleprompter = dspy.BootstrapFewShotWithRandomSearch( + metric=metric, + max_bootstrapped_demos=self.config.max_bootstrapped_demos, + max_labeled_demos=self.config.max_labeled_demos, + num_candidate_programs=self.config.num_candidate_programs, + ) + elif optimizer_name == "BootstrapFewShot": + teleprompter = dspy.BootstrapFewShot( + metric=metric, + max_bootstrapped_demos=self.config.max_bootstrapped_demos, + max_labeled_demos=self.config.max_labeled_demos, + ) + else: + teleprompter = dspy.BootstrapFewShot( + metric=metric, + max_bootstrapped_demos=self.config.max_bootstrapped_demos, + ) + + # Split data for train/test + split = max(1, int(len(examples) * 0.8)) + train_set = examples[:split] + optimized_program = teleprompter.compile( + program, trainset=train_set + ) + + # Extract optimized parameters + result: Dict[str, Any] = {} + if hasattr(optimized_program, "generate") and hasattr( + optimized_program.generate, "demos" + ): + result["few_shot_examples"] = [ + {"input": d.question, "output": d.answer} + for d in optimized_program.generate.demos + if hasattr(d, "question") and hasattr(d, "answer") + ] + + return result + + def _to_config_updates(self, optimized: Dict[str, Any]) -> Dict[str, Any]: + """Convert DSPy optimization results to TOML-compatible config.""" + updates: Dict[str, Any] = {} + + if ( + self.config.optimize_system_prompt + and "system_prompt" in optimized + ): + updates["system_prompt"] = optimized["system_prompt"] + + if ( + self.config.optimize_few_shot + and "few_shot_examples" in optimized + ): + updates["few_shot_examples"] = optimized["few_shot_examples"] + + if ( + self.config.optimize_tool_descriptions + and "tool_descriptions" in optimized + ): + updates["tool_descriptions"] = optimized["tool_descriptions"] + + return updates + + def _write_configs(self, config_updates: Dict[str, Any]) -> None: + """Write updated configs via AgentConfigEvolver.""" + import pathlib + + from openjarvis.learning.agents.agent_evolver import ( + AgentConfigEvolver, + ) + + evolver = AgentConfigEvolver.__new__(AgentConfigEvolver) + evolver._config_dir = pathlib.Path(self.config.config_dir) + evolver._history_dir = evolver._config_dir / ".history" + evolver._config_dir.mkdir(parents=True, exist_ok=True) + evolver._history_dir.mkdir(parents=True, exist_ok=True) + + agent_name = self.config.agent_filter or "default" + evolver.write_config( + agent_name, + tools=config_updates.get("tools", []), + system_prompt=config_updates.get("system_prompt", ""), + ) + + +@LearningRegistry.register("dspy") +class _DSPyLearningPolicy(AgentLearningPolicy): + """Wrapper to register DSPyAgentOptimizer in the LearningRegistry.""" + + def __init__(self, **kwargs: object) -> None: + pass + + def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: + config = DSPyOptimizerConfig() + optimizer = DSPyAgentOptimizer(config) + return optimizer.optimize(trace_store) + + +__all__ = ["DSPyAgentOptimizer"] diff --git a/src/openjarvis/learning/agents/gepa_optimizer.py b/src/openjarvis/learning/agents/gepa_optimizer.py new file mode 100644 index 00000000..20dac8dd --- /dev/null +++ b/src/openjarvis/learning/agents/gepa_optimizer.py @@ -0,0 +1,321 @@ +"""GEPA agent optimizer -- Pareto-efficient evolutionary optimization. + +Uses GEPA's adapter pattern to bridge OpenJarvis traces into GEPA's +evolutionary optimization framework. Outputs TOML config updates +written via AgentConfigEvolver. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List + +from openjarvis.core.config import GEPAOptimizerConfig +from openjarvis.core.registry import LearningRegistry +from openjarvis.learning._stubs import AgentLearningPolicy + +logger = logging.getLogger(__name__) + +# Optional dependency +try: + import gepa + HAS_GEPA = True +except ImportError: + HAS_GEPA = False + gepa = None # type: ignore[assignment] + + +class OpenJarvisGEPAAdapter: + """Implements GEPA's adapter protocol for OpenJarvis agents. + + Bridges trace data into GEPA's optimization framework via + the ``assess()`` and ``make_reflective_dataset()`` methods. + """ + + def __init__( + self, + trace_store: Any, + agent_name: str, + config: GEPAOptimizerConfig, + ) -> None: + self.trace_store = trace_store + self.agent_name = agent_name + self.config = config + self._traces: List[Any] = [] + + def load_traces(self) -> None: + """Load traces from the store.""" + kwargs: Dict[str, Any] = {"limit": 10_000} + if self.agent_name: + kwargs["agent"] = self.agent_name + self._traces = self.trace_store.list_traces(**kwargs) + + def assess( + self, + batch: List[Any], + candidate: Dict[str, Any], + capture_traces: bool = False, + ) -> Dict[str, Any]: + """Score a candidate config against a batch of test cases. + + Returns a dict with at least 'scores' (list of floats) key. + """ + scores = [] + trace_data = [] + + for item in batch: + query = item if isinstance(item, str) else getattr(item, "query", str(item)) + # Find matching traces for this query + matching = [t for t in self._traces if t.query == query] + if matching: + best = max(matching, key=lambda t: t.feedback or 0.0) + scores.append(best.feedback or 0.0) + if capture_traces: + trace_data.append({ + "query": query, + "result": best.result, + "feedback": best.feedback, + "outcome": best.outcome, + "steps": [ + { + "type": str(s.step_type), + "input": s.input, + "output": s.output, + } + for s in best.steps + ], + }) + else: + scores.append(0.0) + + result: Dict[str, Any] = {"scores": scores} + if capture_traces: + result["traces"] = trace_data + return result + + def make_reflective_dataset( + self, + candidate: Dict[str, Any], + assessment_batch: List[Any], + components_to_update: List[str], + ) -> List[Dict[str, Any]]: + """Package trace diagnostics as Actionable Side Information for GEPA.""" + dataset = [] + for item in assessment_batch: + query = item if isinstance(item, str) else getattr(item, "query", str(item)) + matching = [t for t in self._traces if t.query == query] + if not matching: + continue + best = max(matching, key=lambda t: t.feedback or 0.0) + + # Build diagnostic info + tool_calls = [] + reasoning_steps = [] + errors = [] + for step in best.steps: + step_type = str(step.step_type) + if "tool_call" in step_type: + tool_calls.append(step.input.get("tool", "unknown")) + if "generate" in step_type: + reasoning_steps.append(str(step.output)[:200]) + is_error = ( + step.output + and isinstance(step.output, dict) + and "error" in step.output + ) + if is_error: + errors.append(step.output["error"]) + + dataset.append({ + "query": query, + "result": best.result, + "feedback": best.feedback, + "outcome": best.outcome, + "tool_calls": tool_calls, + "reasoning_steps": reasoning_steps, + "errors": errors, + "components_to_update": components_to_update, + }) + + return dataset + + +class GEPAAgentOptimizer: + """Optimize agent configs using GEPA evolutionary optimization. + + Parameters + ---------- + config: + GEPAOptimizerConfig controlling optimization parameters. + """ + + def __init__(self, config: GEPAOptimizerConfig) -> None: + self.config = config + + def optimize(self, trace_store: Any) -> Dict[str, Any]: + """Run GEPA optimization on traces from the store. + + 1. Load traces and build the GEPA adapter + 2. Define the search space (agent config fields) + 3. Run GEPA's evolutionary optimization + 4. Extract best candidate as TOML updates + 5. Write via AgentConfigEvolver if config_dir is set + """ + kwargs: Dict[str, Any] = {"limit": 10_000} + if self.config.agent_filter: + kwargs["agent"] = self.config.agent_filter + traces = trace_store.list_traces(**kwargs) + + if len(traces) < self.config.min_traces: + return { + "status": "skipped", + "reason": ( + f"only {len(traces)} traces, " + f"min_traces={self.config.min_traces}" + ), + } + + if not HAS_GEPA: + return { + "status": "error", + "reason": ( + "gepa not installed" + " (pip install 'openjarvis[learning-gepa]')" + ), + } + + agent_name = self.config.agent_filter or "default" + adapter = OpenJarvisGEPAAdapter(trace_store, agent_name, self.config) + adapter.load_traces() + + try: + best_candidate = self._run_gepa(adapter, traces) + except Exception as exc: + logger.warning("GEPA optimization failed: %s", exc) + return {"status": "error", "reason": str(exc)} + + config_updates = self._to_config_updates(best_candidate) + + if self.config.config_dir: + self._write_configs(agent_name, config_updates) + + return { + "status": "completed", + "traces_used": len(traces), + "config_updates": config_updates, + } + + def _run_gepa( + self, adapter: OpenJarvisGEPAAdapter, traces: List[Any], + ) -> Dict[str, Any]: + """Run the GEPA evolutionary optimization loop.""" + # Build search space from config flags + components = [] + if self.config.optimize_system_prompt: + components.append("system_prompt") + if self.config.optimize_tools: + components.append("tools") + if self.config.optimize_max_turns: + components.append("max_turns") + if self.config.optimize_temperature: + components.append("temperature") + + # Extract unique queries as test cases + queries = list({t.query for t in traces}) + + # Initialize GEPA optimizer + optimizer = gepa.GEPAOptimizer( + adapter=adapter, + max_metric_calls=self.config.max_metric_calls, + population_size=self.config.population_size, + ) + + # Define initial candidate from trace analysis + initial_candidate = self._build_initial_candidate(traces) + + # Run optimization + result = optimizer.optimize( + initial_candidate=initial_candidate, + test_cases=queries[:self.config.assessment_batch_size], + components=components, + ) + + return result.best_candidate if hasattr(result, "best_candidate") else result + + def _build_initial_candidate(self, traces: List[Any]) -> Dict[str, Any]: + """Build initial candidate config from trace analysis.""" + # Collect tool usage frequencies + tool_freq: Dict[str, int] = {} + turn_counts: List[int] = [] + + for t in traces: + n_tools = 0 + for step in t.steps: + step_type = str(step.step_type) + if "tool_call" in step_type: + n_tools += 1 + tool_name = step.input.get("tool", "") if step.input else "" + if tool_name: + tool_freq[tool_name] = tool_freq.get(tool_name, 0) + 1 + turn_counts.append(n_tools) + + ranked_tools = sorted(tool_freq, key=tool_freq.get, reverse=True) # type: ignore[arg-type] + + avg_turns = sum(turn_counts) / len(turn_counts) if turn_counts else 10 + max_turns = max(int(avg_turns * 1.5), 5) + + return { + "system_prompt": "", + "tools": ranked_tools[:10], + "max_turns": max_turns, + "temperature": 0.3, + } + + def _to_config_updates(self, candidate: Dict[str, Any]) -> Dict[str, Any]: + """Convert GEPA candidate to TOML-compatible config dict.""" + updates: Dict[str, Any] = {} + if self.config.optimize_system_prompt and "system_prompt" in candidate: + updates["system_prompt"] = candidate["system_prompt"] + if self.config.optimize_tools and "tools" in candidate: + updates["tools"] = candidate["tools"] + if self.config.optimize_max_turns and "max_turns" in candidate: + updates["max_turns"] = candidate["max_turns"] + if self.config.optimize_temperature and "temperature" in candidate: + updates["temperature"] = candidate["temperature"] + return updates + + def _write_configs(self, agent_name: str, config_updates: Dict[str, Any]) -> None: + """Write updated configs via AgentConfigEvolver.""" + import pathlib + + from openjarvis.learning.agents.agent_evolver import AgentConfigEvolver + + evolver = AgentConfigEvolver.__new__(AgentConfigEvolver) + evolver._config_dir = pathlib.Path(self.config.config_dir) + evolver._history_dir = evolver._config_dir / ".history" + evolver._config_dir.mkdir(parents=True, exist_ok=True) + evolver._history_dir.mkdir(parents=True, exist_ok=True) + + evolver.write_config( + agent_name, + tools=config_updates.get("tools", []), + max_turns=config_updates.get("max_turns", 10), + temperature=config_updates.get("temperature", 0.3), + system_prompt=config_updates.get("system_prompt", ""), + ) + + +@LearningRegistry.register("gepa") +class _GEPALearningPolicy(AgentLearningPolicy): + """Wrapper to register GEPAAgentOptimizer in the LearningRegistry.""" + + def __init__(self, **kwargs: object) -> None: + pass + + def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: + config = GEPAOptimizerConfig() + optimizer = GEPAAgentOptimizer(config) + return optimizer.optimize(trace_store) + + +__all__ = ["GEPAAgentOptimizer", "OpenJarvisGEPAAdapter"] diff --git a/src/openjarvis/learning/skill_discovery.py b/src/openjarvis/learning/agents/skill_discovery.py similarity index 100% rename from src/openjarvis/learning/skill_discovery.py rename to src/openjarvis/learning/agents/skill_discovery.py diff --git a/src/openjarvis/learning/bandit_router.py b/src/openjarvis/learning/bandit_router.py deleted file mode 100644 index 899e4b30..00000000 --- a/src/openjarvis/learning/bandit_router.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Bandit router — Thompson Sampling / UCB for query→model selection.""" - -from __future__ import annotations - -import math -import random -from collections import defaultdict -from dataclasses import dataclass -from typing import Any, Dict, List, Literal, Optional - -from openjarvis.core.registry import RouterPolicyRegistry -from openjarvis.core.types import RoutingContext - - -def _derive_query_class(context: RoutingContext) -> str: - """Derive a query class string from RoutingContext fields.""" - if context.has_code: - return "code" - if context.has_math: - return "math" - if context.query_length < 50: - return "short" - if context.query_length > 500: - return "long" - return "general" - - -@dataclass(slots=True) -class ArmStats: - """Statistics for a single arm (model).""" - successes: int = 0 # alpha for Beta distribution - failures: int = 0 # beta for Beta distribution - total_reward: float = 0.0 - pulls: int = 0 - - @property - def mean_reward(self) -> float: - return self.total_reward / self.pulls if self.pulls > 0 else 0.0 - - -class BanditRouterPolicy: - """Multi-armed bandit router using Thompson Sampling or UCB. - - Each (query_class, model) pair is an arm. Rewards come from - trace outcomes. - """ - - def __init__( - self, - *, - strategy: Literal["thompson", "ucb"] = "thompson", - exploration_factor: float = 2.0, # UCB exploration constant - min_pulls: int = 3, # minimum pulls before trusting estimates - reward_threshold: float = 0.5, # reward above this = success - ) -> None: - self._strategy = strategy - self._exploration = exploration_factor - self._min_pulls = min_pulls - self._reward_threshold = reward_threshold - # query_class -> model -> ArmStats - self._arms: Dict[str, Dict[str, ArmStats]] = defaultdict( - lambda: defaultdict(ArmStats) - ) - self._total_pulls = 0 - - def route(self, context: RoutingContext, models: List[str]) -> str: - """Select model using the configured bandit strategy.""" - if not models: - raise ValueError("No models available") - - query_class = _derive_query_class(context) - arms = self._arms[query_class] - - # Ensure all models have arms - for m in models: - if m not in arms: - arms[m] = ArmStats() - - # Check minimum pulls — explore uniformly first - under_explored = [m for m in models if arms[m].pulls < self._min_pulls] - if under_explored: - return random.choice(under_explored) - - if self._strategy == "thompson": - return self._thompson_select(models, arms) - else: - return self._ucb_select(models, arms) - - def _thompson_select(self, models: List[str], arms: Dict[str, ArmStats]) -> str: - """Thompson Sampling: sample from Beta(alpha, beta) per arm.""" - best_model = models[0] - best_sample = -1.0 - - for m in models: - stats = arms[m] - alpha = stats.successes + 1 # Prior: Beta(1,1) - beta = stats.failures + 1 - sample = random.betavariate(alpha, beta) - if sample > best_sample: - best_sample = sample - best_model = m - - return best_model - - def _ucb_select(self, models: List[str], arms: Dict[str, ArmStats]) -> str: - """UCB1: select arm with highest upper confidence bound.""" - best_model = models[0] - best_ucb = -1.0 - - total = max(self._total_pulls, 1) - for m in models: - stats = arms[m] - if stats.pulls == 0: - return m # Unexplored arm gets priority - - mean = stats.mean_reward - exploration = self._exploration * math.sqrt( - math.log(total) / stats.pulls - ) - ucb_value = mean + exploration - - if ucb_value > best_ucb: - best_ucb = ucb_value - best_model = m - - return best_model - - def update(self, query_class: str, model: str, reward: float) -> None: - """Update arm statistics with observed reward.""" - stats = self._arms[query_class][model] - stats.pulls += 1 - stats.total_reward += reward - if reward >= self._reward_threshold: - stats.successes += 1 - else: - stats.failures += 1 - self._total_pulls += 1 - - def get_stats(self, query_class: Optional[str] = None) -> Dict[str, Any]: - """Get arm statistics.""" - if query_class: - arms = self._arms.get(query_class, {}) - return { - m: { - "pulls": s.pulls, - "mean_reward": s.mean_reward, - "successes": s.successes, - "failures": s.failures, - } - for m, s in arms.items() - } - return { - qc: { - m: {"pulls": s.pulls, "mean_reward": s.mean_reward} - for m, s in arms.items() - } - for qc, arms in self._arms.items() - } - - def reset(self) -> None: - """Reset all state.""" - self._arms.clear() - self._total_pulls = 0 - - -def ensure_registered() -> None: - """Register BanditRouterPolicy if not already present.""" - if not RouterPolicyRegistry.contains("bandit"): - RouterPolicyRegistry.register_value("bandit", BanditRouterPolicy) - - -ensure_registered() - -__all__ = ["ArmStats", "BanditRouterPolicy"] diff --git a/src/openjarvis/learning/grpo_policy.py b/src/openjarvis/learning/grpo_policy.py deleted file mode 100644 index f941a8e8..00000000 --- a/src/openjarvis/learning/grpo_policy.py +++ /dev/null @@ -1,182 +0,0 @@ -"""GRPO router — Group Relative Policy Optimization for query→model routing.""" - -from __future__ import annotations - -import math -import random -from collections import defaultdict -from dataclasses import dataclass, field -from typing import Any, Dict, List - -from openjarvis.core.registry import RouterPolicyRegistry -from openjarvis.core.types import RoutingContext - - -def _derive_query_class(context: RoutingContext) -> str: - """Derive a query class string from RoutingContext fields.""" - if context.has_code: - return "code" - if context.has_math: - return "math" - if context.query_length < 50: - return "short" - if context.query_length > 500: - return "long" - return "general" - - -@dataclass(slots=True) -class GRPOSample: - """A single sample in a GRPO group.""" - query_class: str - model: str - reward: float - - -@dataclass -class GRPOState: - """Persistent state for GRPO policy weights.""" - # model -> query_class -> weight (log probability) - weights: Dict[str, Dict[str, float]] = field( - default_factory=lambda: defaultdict( - lambda: defaultdict(float) - ), - ) - # Track sample counts for min_samples threshold - sample_counts: Dict[str, int] = field(default_factory=lambda: defaultdict(int)) - total_updates: int = 0 - - -class GRPORouterPolicy: - """Group Relative Policy Optimization for routing queries to models. - - Groups samples by query_class, computes relative advantage within each - group (reward - mean_reward) / std, and updates policy weights via - softmax gradient. - - Falls back to random selection when insufficient samples exist. - """ - - def __init__( - self, - *, - learning_rate: float = 0.1, - min_samples: int = 5, - group_size: int = 4, - temperature: float = 1.0, - ) -> None: - self._lr = learning_rate - self._min_samples = min_samples - self._group_size = group_size - self._temperature = temperature - self._state = GRPOState() - self._sample_buffer: List[GRPOSample] = [] - - def route(self, context: RoutingContext, models: List[str]) -> str: - """Select the best model for the given routing context.""" - if not models: - raise ValueError("No models available for routing") - - query_class = _derive_query_class(context) - - # Check if we have enough samples - if self._state.sample_counts.get(query_class, 0) < self._min_samples: - return random.choice(models) - - # Compute softmax probabilities from weights - scores = [] - for m in models: - w = self._state.weights.get(m, {}).get(query_class, 0.0) - scores.append(w / self._temperature) - - # Softmax - max_score = max(scores) - exp_scores = [math.exp(s - max_score) for s in scores] - total = sum(exp_scores) - probs = [e / total for e in exp_scores] - - # Sample from distribution - r = random.random() - cumulative = 0.0 - for i, p in enumerate(probs): - cumulative += p - if r <= cumulative: - return models[i] - return models[-1] - - def add_sample(self, query_class: str, model: str, reward: float) -> None: - """Add a training sample to the buffer.""" - self._sample_buffer.append(GRPOSample( - query_class=query_class, model=model, reward=reward, - )) - self._state.sample_counts[query_class] = ( - self._state.sample_counts.get(query_class, 0) + 1 - ) - - def update(self) -> Dict[str, Any]: - """Run GRPO update on accumulated samples. - - Groups samples by query_class, computes relative advantages, - and updates policy weights. - - Returns stats about the update. - """ - if not self._sample_buffer: - return {"updated": False, "reason": "no samples"} - - # Group by query_class - groups: Dict[str, List[GRPOSample]] = defaultdict(list) - for sample in self._sample_buffer: - groups[sample.query_class].append(sample) - - updates_applied = 0 - for qc, samples in groups.items(): - if len(samples) < 2: - continue # Need at least 2 for relative comparison - - # Compute group statistics - rewards = [s.reward for s in samples] - mean_r = sum(rewards) / len(rewards) - var_r = sum((r - mean_r) ** 2 for r in rewards) / len(rewards) - std_r = math.sqrt(var_r) if var_r > 0 else 1.0 - - # Compute advantages and update weights - for sample in samples: - advantage = (sample.reward - mean_r) / std_r - self._state.weights[sample.model][qc] += ( - self._lr * advantage - ) - updates_applied += 1 - - self._state.total_updates += 1 - processed = len(self._sample_buffer) - self._sample_buffer.clear() - - return { - "updated": True, - "samples_processed": processed, - "groups": len(groups), - "updates_applied": updates_applied, - "total_updates": self._state.total_updates, - } - - @property - def state(self) -> GRPOState: - """Access the current policy state.""" - return self._state - - def reset(self) -> None: - """Reset all state.""" - self._state = GRPOState() - self._sample_buffer.clear() - - -def ensure_registered() -> None: - """Register GRPORouterPolicy if not already present.""" - if not RouterPolicyRegistry.contains("grpo"): - RouterPolicyRegistry.register_value("grpo", GRPORouterPolicy) - - -ensure_registered() - -__all__ = ["GRPORouterPolicy", "GRPOSample", "GRPOState"] diff --git a/src/openjarvis/learning/icl_updater.py b/src/openjarvis/learning/icl_updater.py deleted file mode 100644 index 07ff51f5..00000000 --- a/src/openjarvis/learning/icl_updater.py +++ /dev/null @@ -1,236 +0,0 @@ -"""ICL example updater + skill discovery from traces.""" - -from __future__ import annotations - -import logging -from typing import Any, Dict, List, Optional - -from openjarvis.core.registry import LearningRegistry -from openjarvis.learning._stubs import AgentLearningPolicy - -logger = logging.getLogger(__name__) - - -@LearningRegistry.register("icl_updater") -class ICLUpdaterPolicy(AgentLearningPolicy): - """Updates in-context examples and discovers skills from traces. - - Analyzes traces for successful tool call patterns, extracts - in-context learning examples, and discovers reusable multi-tool - sequences ("skills"). This updates *agent* logic (ICL examples - and tool-use strategies), not tool implementations themselves. - """ - - def __init__( - self, - *, - min_score: float = 0.7, - max_examples: int = 20, - min_skill_occurrences: int = 3, - auto_apply: bool = False, - ) -> None: - self._min_score = min_score - self._max_examples = max_examples - self._min_skill_occurrences = min_skill_occurrences - self._auto_apply = auto_apply - self._examples: List[Dict[str, Any]] = [] - self._skills: List[Dict[str, Any]] = [] - # Versioned example database for add_example / rollback - self._example_db: List[Dict[str, Any]] = [] - self._version: int = 0 - - def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: - """Analyze traces and extract ICL examples + skills.""" - try: - traces = trace_store.list_traces() - except Exception as exc: - logger.warning("ICL updater failed: %s", exc) - return {"examples": [], "skills": []} - - # Extract high-scoring traces with tool calls - new_examples: List[Dict[str, Any]] = [] - tool_sequences: List[List[str]] = [] - - for trace in traces: - # Only consider successful traces - if trace.outcome != "success": - continue - - feedback = trace.feedback if trace.feedback is not None else 0.5 - if feedback < self._min_score: - continue - - # Extract tool call steps - tool_steps = [ - s - for s in (trace.steps or []) - if s.step_type.value == "tool_call" - ] - - if tool_steps: - # Build ICL example - tool_names = [ - s.metadata.get("tool_name", "") - for s in tool_steps - ] - example = { - "query": trace.query, - "tools_used": tool_names, - "outcome": trace.outcome, - "score": feedback, - } - new_examples.append(example) - tool_sequences.append(tool_names) - - # Keep top examples by score - new_examples.sort(key=lambda x: x["score"], reverse=True) - self._examples = new_examples[: self._max_examples] - - # Discover skills: recurring multi-tool sequences - self._skills = self._discover_skills(tool_sequences) - - return { - "examples": list(self._examples), - "skills": list(self._skills), - "traces_analyzed": len(traces), - } - - def _discover_skills(self, sequences: List[List[str]]) -> List[Dict[str, Any]]: - """Find recurring tool call sequences.""" - # Count subsequences of length 2+ - seq_counts: Dict[str, int] = {} - for seq in sequences: - if len(seq) < 2: - continue - # Check all subsequences of length 2-4 - for length in range(2, min(len(seq) + 1, 5)): - for start in range(len(seq) - length + 1): - sub = tuple(seq[start : start + length]) - key = " -> ".join(sub) - seq_counts[key] = seq_counts.get(key, 0) + 1 - - # Filter by minimum occurrences - skills: List[Dict[str, Any]] = [] - for seq_key, count in seq_counts.items(): - if count >= self._min_skill_occurrences: - skills.append( - { - "sequence": seq_key, - "occurrences": count, - "tools": seq_key.split(" -> "), - } - ) - - # Sort by frequency - skills.sort(key=lambda x: x["occurrences"], reverse=True) - return skills - - # -- Versioned example database methods ------------------------------------ - - def add_example( - self, - query: str, - response: str, - outcome: float, - metadata: Optional[Dict[str, Any]] = None, - ) -> bool: - """Add an ICL example if it meets the quality threshold. - - Parameters - ---------- - query: - The user query that produced this example. - response: - The agent/model response. - outcome: - Quality score in [0, 1]. - metadata: - Optional metadata dict attached to the example. - - Returns - ------- - True if the example was accepted, False if rejected (below threshold). - """ - if outcome < self._min_score: - return False - - self._version += 1 - entry: Dict[str, Any] = { - "query": query, - "response": response, - "outcome": outcome, - "metadata": metadata or {}, - "version": self._version, - } - self._example_db.append(entry) - - # Trim to max_examples (remove oldest first) - if len(self._example_db) > self._max_examples: - self._example_db = self._example_db[-self._max_examples:] - - return True - - def rollback(self, version: int) -> None: - """Remove all examples added after the given version. - - Parameters - ---------- - version: - The version checkpoint to rollback to. All examples with - ``version > checkpoint`` are removed. - """ - self._example_db = [ - ex for ex in self._example_db if ex["version"] <= version - ] - self._version = version - - def get_examples( - self, - query_class: str = "", - top_k: int = 5, - ) -> List[Dict[str, Any]]: - """Retrieve the best examples, optionally filtered by query class. - - Parameters - ---------- - query_class: - If non-empty, only return examples whose query contains this - substring (case-insensitive). - top_k: - Maximum number of examples to return. - - Returns - ------- - Up to *top_k* examples sorted by outcome (descending). - """ - pool = self._example_db - if query_class: - lc = query_class.lower() - pool = [ex for ex in pool if lc in ex["query"].lower()] - - # Sort by outcome descending, take top_k - ranked = sorted(pool, key=lambda ex: ex["outcome"], reverse=True) - return ranked[:top_k] - - @property - def version(self) -> int: - """Current version counter.""" - return self._version - - @property - def example_db(self) -> List[Dict[str, Any]]: - """Return a copy of the versioned example database.""" - return list(self._example_db) - - # -- Original property accessors ------------------------------------------ - - @property - def examples(self) -> List[Dict[str, Any]]: - return list(self._examples) - - @property - def skills(self) -> List[Dict[str, Any]]: - return list(self._skills) - - -__all__ = ["ICLUpdaterPolicy"] diff --git a/src/openjarvis/learning/intelligence/__init__.py b/src/openjarvis/learning/intelligence/__init__.py new file mode 100644 index 00000000..a1ab2559 --- /dev/null +++ b/src/openjarvis/learning/intelligence/__init__.py @@ -0,0 +1,13 @@ +"""Intelligence learning — model fine-tuning via SFT and GRPO.""" + +from __future__ import annotations + +# Import trainers so their @LearningRegistry.register decorators execute. +try: + from openjarvis.learning.intelligence import sft_trainer as _sft # noqa: F401 +except ImportError: + pass +try: + from openjarvis.learning.intelligence import grpo_trainer as _grpo # noqa: F401 +except ImportError: + pass diff --git a/src/openjarvis/learning/intelligence/grpo_trainer.py b/src/openjarvis/learning/intelligence/grpo_trainer.py new file mode 100644 index 00000000..6ba91bc5 --- /dev/null +++ b/src/openjarvis/learning/intelligence/grpo_trainer.py @@ -0,0 +1,385 @@ +"""General-purpose GRPO trainer -- Group Relative Policy Optimization. + +Fine-tunes any local model by sampling N responses per prompt, +computing group-relative advantages, and applying a clipped policy +gradient with KL penalty vs a frozen reference model. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Protocol, runtime_checkable + +from openjarvis.core.config import GRPOConfig +from openjarvis.core.registry import LearningRegistry +from openjarvis.learning._stubs import IntelligenceLearningPolicy + +logger = logging.getLogger(__name__) + +# Optional imports +try: + import torch + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + torch = None # type: ignore[assignment] + +try: + from transformers import AutoModelForCausalLM, AutoTokenizer + + HAS_TRANSFORMERS = True +except ImportError: + HAS_TRANSFORMERS = False + + +@runtime_checkable +class RewardFn(Protocol): + """Protocol for reward functions used by GRPOTrainer.""" + + def score(self, prompt: str, response: str, ground_truth: str | None) -> float: ... + + +class DefaultRewardFn: + """Default reward function using length-normalized response quality heuristics.""" + + def score(self, prompt: str, response: str, ground_truth: str | None) -> float: + """Score a response. Higher is better, range [0, 1].""" + score = 0.5 # baseline + + # Length heuristic: prefer non-empty, not-too-long responses + if not response.strip(): + return 0.0 + resp_len = len(response) + if resp_len < 10: + score -= 0.1 + elif resp_len > 5000: + score -= 0.05 + + # Ground truth matching + if ground_truth is not None: + gt_lower = ground_truth.strip().lower() + resp_lower = response.strip().lower() + if gt_lower == resp_lower: + score += 0.4 + elif gt_lower in resp_lower: + score += 0.2 + + return max(0.0, min(1.0, score)) + + +class GRPOTrainer: + """General-purpose GRPO trainer. + + Parameters + ---------- + config: + GRPOConfig controlling model, sampling, and optimization params. + reward_fn: + Pluggable reward function. Defaults to ``DefaultRewardFn``. + """ + + def __init__( + self, + config: GRPOConfig, + reward_fn: RewardFn | None = None, + ) -> None: + self.config = config + self.reward_fn: RewardFn = reward_fn or DefaultRewardFn() + + def train(self, trace_store: Any) -> Dict[str, Any]: + """End-to-end: mine prompts from traces, then train. + + Parameters + ---------- + trace_store: + Object with ``list_traces()`` returning trace objects. + """ + prompts = self._mine_prompts(trace_store) + return self.train_on_prompts(prompts) + + def train_on_prompts( + self, + prompts: List[str], + ground_truths: List[str | None] | None = None, + ) -> Dict[str, Any]: + """Run GRPO training on a set of prompts. + + Parameters + ---------- + prompts: + List of prompt strings to train on. + ground_truths: + Optional parallel list of ground-truth answers for reward scoring. + """ + if not prompts: + return {"status": "skipped", "reason": "no training data"} + + if len(prompts) < self.config.min_prompts: + return { + "status": "skipped", + "reason": ( + f"only {len(prompts)} prompts, " + f"min_prompts={self.config.min_prompts}" + ), + } + + if not HAS_TORCH: + return {"status": "error", "reason": "torch not available"} + + if not HAS_TRANSFORMERS: + return {"status": "error", "reason": "transformers not available"} + + try: + return self._run_grpo(prompts, ground_truths) + except Exception as exc: + logger.warning("GRPO training failed: %s", exc) + return {"status": "error", "reason": str(exc)} + + def _mine_prompts(self, trace_store: Any) -> List[str]: + """Extract unique prompts from the trace store.""" + from openjarvis.learning.training.data import TrainingDataMiner + + miner = TrainingDataMiner(trace_store, min_quality=0.5) + agent_filter = self.config.agent_filter or None + pairs = miner.extract_sft_pairs(agent=agent_filter) + # Deduplicate prompts + seen: set[str] = set() + prompts: List[str] = [] + for pair in pairs: + q = pair.get("input", "") + if q and q not in seen: + seen.add(q) + prompts.append(q) + return prompts + + def _run_grpo( + self, + prompts: List[str], + ground_truths: List[str | None] | None, + ) -> Dict[str, Any]: + """Execute the GRPO training loop. + + 1. Load policy model and frozen reference model + 2. For each epoch: + a. For each prompt, sample N responses from policy + b. Score responses with reward_fn + c. Compute group-relative advantages + d. Compute clipped policy gradient + KL penalty + e. Update policy weights + """ + if ground_truths is None: + ground_truths = [None] * len(prompts) + + # Load models + tokenizer = AutoTokenizer.from_pretrained(self.config.model_name) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + policy_model = AutoModelForCausalLM.from_pretrained( + self.config.model_name, + torch_dtype=torch.bfloat16, + device_map="auto", + ) + + if self.config.gradient_checkpointing and hasattr( + policy_model, "gradient_checkpointing_enable" + ): + policy_model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + + # Frozen reference model + ref_kwargs: Dict[str, Any] = { + "torch_dtype": torch.bfloat16, + "device_map": "auto", + } + if self.config.use_8bit_ref: + try: + from transformers import BitsAndBytesConfig + + ref_kwargs["quantization_config"] = BitsAndBytesConfig( + load_in_8bit=True + ) + except ImportError: + pass + ref_model = AutoModelForCausalLM.from_pretrained( + self.config.model_name, **ref_kwargs + ) + ref_model.requires_grad_(False) + + optimizer = torch.optim.AdamW( + policy_model.parameters(), + lr=self.config.learning_rate, + ) + + total_steps = 0 + cumulative_loss = 0.0 + + policy_model.train() + + for epoch in range(self.config.num_epochs): + epoch_loss = 0.0 + epoch_steps = 0 + + for i in range(0, len(prompts), self.config.batch_size): + batch_prompts = prompts[i : i + self.config.batch_size] + batch_gts = ground_truths[i : i + self.config.batch_size] + + loss = self._grpo_step( + policy_model, + ref_model, + tokenizer, + optimizer, + batch_prompts, + batch_gts, + ) + epoch_loss += loss + epoch_steps += 1 + total_steps += 1 + + avg_epoch_loss = epoch_loss / epoch_steps if epoch_steps else 0.0 + logger.info( + "GRPO epoch %d/%d loss=%.4f", + epoch + 1, + self.config.num_epochs, + avg_epoch_loss, + ) + cumulative_loss += epoch_loss + + avg_loss = cumulative_loss / total_steps if total_steps else 0.0 + return { + "status": "completed", + "epochs": self.config.num_epochs, + "total_steps": total_steps, + "avg_loss": avg_loss, + "prompts": len(prompts), + } + + def _grpo_step( + self, + policy_model: Any, + ref_model: Any, + tokenizer: Any, + optimizer: Any, + prompts: List[str], + ground_truths: List[str | None], + ) -> float: + """Execute one GRPO gradient step on a batch of prompts.""" + all_rewards: List[List[float]] = [] + all_log_probs: List[List[Any]] = [] + all_ref_log_probs: List[List[Any]] = [] + + for prompt, gt in zip(prompts, ground_truths): + rewards = [] + log_probs = [] + ref_lps = [] + + for _ in range(self.config.num_samples_per_prompt): + # Sample response from policy + inputs = tokenizer( + prompt, + return_tensors="pt", + truncation=True, + max_length=self.config.max_seq_length, + ).to(policy_model.device) + + with torch.no_grad(): + gen_ids = policy_model.generate( + **inputs, + max_new_tokens=self.config.max_response_length, + temperature=self.config.temperature, + do_sample=True, + pad_token_id=tokenizer.pad_token_id, + ) + + response_ids = gen_ids[0, inputs["input_ids"].shape[1] :] + response = tokenizer.decode( + response_ids, skip_special_tokens=True + ) + + # Score with reward function + reward = self.reward_fn.score(prompt, response, gt) + rewards.append(reward) + + # Compute log probabilities + full_ids = gen_ids[0].unsqueeze(0) + policy_logits = policy_model(full_ids).logits + policy_lp = torch.nn.functional.log_softmax( + policy_logits, dim=-1 + ) + token_lps = torch.gather( + policy_lp[:, :-1, :], 2, full_ids[:, 1:].unsqueeze(-1) + ).squeeze(-1) + log_probs.append(token_lps.sum()) + + with torch.no_grad(): + ref_logits = ref_model( + full_ids.to(ref_model.device) + ).logits + ref_lp = torch.nn.functional.log_softmax( + ref_logits, dim=-1 + ) + ref_token_lps = torch.gather( + ref_lp[:, :-1, :], + 2, + full_ids[:, 1:] + .to(ref_model.device) + .unsqueeze(-1), + ).squeeze(-1) + ref_lps.append(ref_token_lps.sum()) + + all_rewards.append(rewards) + all_log_probs.append(log_probs) + all_ref_log_probs.append(ref_lps) + + # Compute group-relative advantages and loss + total_loss = torch.tensor( + 0.0, device=policy_model.device, requires_grad=True + ) + + for rewards, log_probs, ref_lps in zip( + all_rewards, all_log_probs, all_ref_log_probs + ): + r_tensor = torch.tensor(rewards, device=policy_model.device) + mean_r = r_tensor.mean() + std_r = r_tensor.std() + 1e-8 + advantages = (r_tensor - mean_r) / std_r + + for adv, lp, ref_lp in zip(advantages, log_probs, ref_lps): + ratio = torch.exp(lp - lp.detach()) # importance ratio + clipped = torch.clamp( + ratio, + 1 - self.config.clip_ratio, + 1 + self.config.clip_ratio, + ) + policy_loss = -torch.min(ratio * adv, clipped * adv) + kl = lp - ref_lp.to(policy_model.device) + total_loss = total_loss + policy_loss + self.config.kl_coef * kl + + optimizer.zero_grad() + total_loss.backward() + torch.nn.utils.clip_grad_norm_( + policy_model.parameters(), self.config.max_grad_norm + ) + optimizer.step() + + return total_loss.item() + + +@LearningRegistry.register("grpo") +class _GRPOLearningPolicy(IntelligenceLearningPolicy): + """Wrapper to register GRPOTrainer in the LearningRegistry.""" + + def __init__(self, **kwargs: object) -> None: + pass + + def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: + from openjarvis.core.config import GRPOConfig + + config = GRPOConfig() + trainer = GRPOTrainer(config) + return trainer.train(trace_store) + + +__all__ = ["DefaultRewardFn", "GRPOTrainer", "RewardFn"] diff --git a/src/openjarvis/learning/orchestrator/__init__.py b/src/openjarvis/learning/intelligence/orchestrator/__init__.py similarity index 80% rename from src/openjarvis/learning/orchestrator/__init__.py rename to src/openjarvis/learning/intelligence/orchestrator/__init__.py index 0cf5dbd1..2acb70d5 100644 --- a/src/openjarvis/learning/orchestrator/__init__.py +++ b/src/openjarvis/learning/intelligence/orchestrator/__init__.py @@ -14,33 +14,33 @@ Importing this module triggers registration of ``orchestrator_sft`` and ``orchestrator_grpo`` in :class:`~openjarvis.core.registry.LearningRegistry`. """ -from openjarvis.learning.orchestrator.environment import ( +from openjarvis.learning.intelligence.orchestrator.environment import ( OrchestratorEnvironment, ) -from openjarvis.learning.orchestrator.grpo_trainer import ( +from openjarvis.learning.intelligence.orchestrator.grpo_trainer import ( OrchestratorGRPOConfig, OrchestratorGRPOTrainer, ) -from openjarvis.learning.orchestrator.policy_model import ( +from openjarvis.learning.intelligence.orchestrator.policy_model import ( OrchestratorPolicyModel, ) -from openjarvis.learning.orchestrator.prompt_registry import ( +from openjarvis.learning.intelligence.orchestrator.prompt_registry import ( TOOL_DESCRIPTIONS, build_system_prompt, ) -from openjarvis.learning.orchestrator.reward import ( +from openjarvis.learning.intelligence.orchestrator.reward import ( AdaptiveRewardWeights, MultiObjectiveReward, Normalizers, RewardWeights, ) -from openjarvis.learning.orchestrator.sft_trainer import ( +from openjarvis.learning.intelligence.orchestrator.sft_trainer import ( OrchestratorSFTConfig, OrchestratorSFTDataset, OrchestratorSFTTrainer, _select_torch_device, ) -from openjarvis.learning.orchestrator.types import ( +from openjarvis.learning.intelligence.orchestrator.types import ( Episode, EpisodeState, EpisodeStep, diff --git a/src/openjarvis/learning/orchestrator/environment.py b/src/openjarvis/learning/intelligence/orchestrator/environment.py similarity index 98% rename from src/openjarvis/learning/orchestrator/environment.py rename to src/openjarvis/learning/intelligence/orchestrator/environment.py index cf87ef42..5f200586 100644 --- a/src/openjarvis/learning/orchestrator/environment.py +++ b/src/openjarvis/learning/intelligence/orchestrator/environment.py @@ -12,7 +12,7 @@ import time from typing import List, Tuple from openjarvis.core.types import ToolCall -from openjarvis.learning.orchestrator.types import ( +from openjarvis.learning.intelligence.orchestrator.types import ( EpisodeState, OrchestratorAction, OrchestratorObservation, diff --git a/src/openjarvis/learning/orchestrator/grpo_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py similarity index 98% rename from src/openjarvis/learning/orchestrator/grpo_trainer.py rename to src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py index 9219ed4a..6949fa20 100644 --- a/src/openjarvis/learning/orchestrator/grpo_trainer.py +++ b/src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py @@ -120,7 +120,7 @@ class OrchestratorGRPOTrainer: # -- Initialisation ------------------------------------------------------ def _init_model(self) -> None: - from openjarvis.learning.orchestrator.policy_model import ( + from openjarvis.learning.intelligence.orchestrator.policy_model import ( OrchestratorPolicyModel, ) @@ -232,12 +232,12 @@ class OrchestratorGRPOTrainer: all_advantages: list[float] = [] all_rewards: list[float] = [] - from openjarvis.learning.orchestrator.reward import ( + from openjarvis.learning.intelligence.orchestrator.reward import ( MultiObjectiveReward, Normalizers, RewardWeights, ) - from openjarvis.learning.orchestrator.types import ( + from openjarvis.learning.intelligence.orchestrator.types import ( Episode, ) diff --git a/src/openjarvis/learning/orchestrator/policy_model.py b/src/openjarvis/learning/intelligence/orchestrator/policy_model.py similarity index 99% rename from src/openjarvis/learning/orchestrator/policy_model.py rename to src/openjarvis/learning/intelligence/orchestrator/policy_model.py index 41e3727f..49b14379 100644 --- a/src/openjarvis/learning/orchestrator/policy_model.py +++ b/src/openjarvis/learning/intelligence/orchestrator/policy_model.py @@ -12,7 +12,7 @@ import logging import re from typing import Any, List, Optional -from openjarvis.learning.orchestrator.types import ( +from openjarvis.learning.intelligence.orchestrator.types import ( EpisodeState, OrchestratorAction, PolicyOutput, diff --git a/src/openjarvis/learning/orchestrator/prompt_registry.py b/src/openjarvis/learning/intelligence/orchestrator/prompt_registry.py similarity index 100% rename from src/openjarvis/learning/orchestrator/prompt_registry.py rename to src/openjarvis/learning/intelligence/orchestrator/prompt_registry.py diff --git a/src/openjarvis/learning/orchestrator/reward.py b/src/openjarvis/learning/intelligence/orchestrator/reward.py similarity index 99% rename from src/openjarvis/learning/orchestrator/reward.py rename to src/openjarvis/learning/intelligence/orchestrator/reward.py index 7aa1ff4c..040d3b6e 100644 --- a/src/openjarvis/learning/orchestrator/reward.py +++ b/src/openjarvis/learning/intelligence/orchestrator/reward.py @@ -10,7 +10,7 @@ from __future__ import annotations from dataclasses import dataclass from typing import Dict, List -from openjarvis.learning.orchestrator.types import Episode +from openjarvis.learning.intelligence.orchestrator.types import Episode @dataclass diff --git a/src/openjarvis/learning/orchestrator/sft_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py similarity index 99% rename from src/openjarvis/learning/orchestrator/sft_trainer.py rename to src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py index db5feeda..154fd83f 100644 --- a/src/openjarvis/learning/orchestrator/sft_trainer.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py @@ -237,7 +237,7 @@ class OrchestratorSFTTrainer: self._init_optimizer() def _init_model(self) -> None: - from openjarvis.learning.orchestrator.policy_model import ( + from openjarvis.learning.intelligence.orchestrator.policy_model import ( OrchestratorPolicyModel, ) diff --git a/src/openjarvis/learning/orchestrator/types.py b/src/openjarvis/learning/intelligence/orchestrator/types.py similarity index 100% rename from src/openjarvis/learning/orchestrator/types.py rename to src/openjarvis/learning/intelligence/orchestrator/types.py diff --git a/src/openjarvis/learning/intelligence/sft_trainer.py b/src/openjarvis/learning/intelligence/sft_trainer.py new file mode 100644 index 00000000..3cdbf596 --- /dev/null +++ b/src/openjarvis/learning/intelligence/sft_trainer.py @@ -0,0 +1,149 @@ +"""General-purpose SFT trainer -- fine-tune any local LM on trace-derived pairs. + +Delegates to :class:`LoRATrainer` from ``training/lora.py`` when ``use_lora=True``. +Supports ``train(trace_store)`` for end-to-end pipeline and ``train_on_pairs()`` +for pre-extracted data. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List + +from openjarvis.core.config import SFTConfig +from openjarvis.core.registry import LearningRegistry +from openjarvis.learning._stubs import IntelligenceLearningPolicy + +logger = logging.getLogger(__name__) + + +class SFTTrainer: + """General-purpose supervised fine-tuning trainer. + + Parameters + ---------- + config: + SFTConfig controlling model, LoRA params, and training hyperparams. + """ + + def __init__(self, config: SFTConfig) -> None: + self.config = config + + @property + def target_module_list(self) -> List[str]: + """Parse comma-separated target_modules string into a list.""" + return [m.strip() for m in self.config.target_modules.split(",") if m.strip()] + + def train(self, trace_store: Any) -> Dict[str, Any]: + """End-to-end: mine SFT pairs from traces, then train. + + Parameters + ---------- + trace_store: + Object with ``list_traces()`` returning trace objects. + + Returns + ------- + dict with at least ``status`` key. + """ + pairs = self._mine_pairs(trace_store) + return self.train_on_pairs(pairs) + + def train_on_pairs(self, pairs: List[Dict[str, Any]]) -> Dict[str, Any]: + """Train on pre-extracted SFT pairs. + + Parameters + ---------- + pairs: + List of dicts with at least ``input`` and ``output`` keys. + + Returns + ------- + dict with ``status``, ``training_samples``, and training metrics. + """ + if not pairs: + return {"status": "skipped", "reason": "no training data"} + + if len(pairs) < self.config.min_pairs: + return { + "status": "skipped", + "reason": f"only {len(pairs)} pairs, min_pairs={self.config.min_pairs}", + } + + if self.config.use_lora: + return self._train_lora(pairs) + + return self._train_full(pairs) + + def _mine_pairs(self, trace_store: Any) -> List[Dict[str, Any]]: + """Extract SFT pairs from the trace store using TrainingDataMiner.""" + from openjarvis.learning.training.data import TrainingDataMiner + + miner = TrainingDataMiner(trace_store, min_quality=0.7) + agent_filter = self.config.agent_filter or None + return miner.extract_sft_pairs(agent=agent_filter) + + def _train_lora(self, pairs: List[Dict[str, Any]]) -> Dict[str, Any]: + """Train using LoRA via the existing LoRATrainer.""" + try: + from openjarvis.learning.training.lora import ( + HAS_TORCH, + LoRATrainer, + LoRATrainingConfig, + ) + except ImportError: + return {"status": "error", "reason": "training.lora not importable"} + + if not HAS_TORCH: + return {"status": "error", "reason": "torch not available"} + + lora_config = LoRATrainingConfig( + lora_rank=self.config.lora_rank, + lora_alpha=self.config.lora_alpha, + lora_dropout=self.config.lora_dropout, + target_modules=self.target_module_list, + num_epochs=self.config.num_epochs, + batch_size=self.config.batch_size, + learning_rate=self.config.learning_rate, + weight_decay=self.config.weight_decay, + warmup_ratio=self.config.warmup_ratio, + max_grad_norm=self.config.max_grad_norm, + max_seq_length=self.config.max_seq_length, + use_4bit=self.config.use_4bit, + output_dir=self.config.checkpoint_dir, + gradient_checkpointing=self.config.gradient_checkpointing, + ) + + try: + trainer = LoRATrainer( + lora_config, model_name=self.config.model_name + ) + return trainer.train(pairs) + except Exception as exc: + logger.warning("SFT LoRA training failed: %s", exc) + return {"status": "error", "reason": str(exc)} + + def _train_full(self, pairs: List[Dict[str, Any]]) -> Dict[str, Any]: + """Full fine-tuning (no LoRA). Placeholder for future implementation.""" + return { + "status": "error", + "reason": "full fine-tuning not yet implemented; set use_lora=true", + } + + +@LearningRegistry.register("sft") +class _SFTLearningPolicy(IntelligenceLearningPolicy): + """Wrapper to register SFTTrainer in the LearningRegistry.""" + + def __init__(self, **kwargs: object) -> None: + pass + + def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: + from openjarvis.core.config import SFTConfig + + config = SFTConfig() + trainer = SFTTrainer(config) + return trainer.train(trace_store) + + +__all__ = ["SFTTrainer"] diff --git a/src/openjarvis/learning/learning_orchestrator.py b/src/openjarvis/learning/learning_orchestrator.py index 08ead5bd..6e91ff32 100644 --- a/src/openjarvis/learning/learning_orchestrator.py +++ b/src/openjarvis/learning/learning_orchestrator.py @@ -54,7 +54,7 @@ class LearningOrchestrator: lora_config: Optional[Any] = None, model_name: Optional[str] = None, ) -> None: - from openjarvis.learning.agent_evolver import AgentConfigEvolver + from openjarvis.learning.agents.agent_evolver import AgentConfigEvolver from openjarvis.learning.training.data import TrainingDataMiner self._trace_store = trace_store diff --git a/src/openjarvis/learning/routing/__init__.py b/src/openjarvis/learning/routing/__init__.py new file mode 100644 index 00000000..470b7f17 --- /dev/null +++ b/src/openjarvis/learning/routing/__init__.py @@ -0,0 +1,3 @@ +"""Routing policies — model selection based on query characteristics.""" + +from __future__ import annotations diff --git a/src/openjarvis/learning/routing/_utils.py b/src/openjarvis/learning/routing/_utils.py new file mode 100644 index 00000000..aaa70736 --- /dev/null +++ b/src/openjarvis/learning/routing/_utils.py @@ -0,0 +1,34 @@ +"""Shared query classification utilities for routing policies. + +Extracted from ``trace_policy.py`` so multiple modules can use +``classify_query()`` without depending on the full policy. +""" + +from __future__ import annotations + +import re + +_CODE_RE = re.compile( + r"```|`[^`]+`|\bdef\s|\bclass\s|\bimport\s|\bfunction\s", + re.IGNORECASE, +) +_MATH_RE = re.compile( + r"\bsolve\b|\bintegral\b|\bequation\b|\bcalculate\b|\bcompute\b", + re.IGNORECASE, +) + + +def classify_query(query: str) -> str: + """Classify a query into a broad category for routing.""" + if _CODE_RE.search(query): + return "code" + if _MATH_RE.search(query): + return "math" + if len(query) < 50: + return "short" + if len(query) > 500: + return "long" + return "general" + + +__all__ = ["classify_query"] diff --git a/src/openjarvis/learning/heuristic_policy.py b/src/openjarvis/learning/routing/heuristic_policy.py similarity index 86% rename from src/openjarvis/learning/heuristic_policy.py rename to src/openjarvis/learning/routing/heuristic_policy.py index 8614b411..fcb58471 100644 --- a/src/openjarvis/learning/heuristic_policy.py +++ b/src/openjarvis/learning/routing/heuristic_policy.py @@ -3,7 +3,7 @@ from __future__ import annotations from openjarvis.core.registry import RouterPolicyRegistry -from openjarvis.learning.router import HeuristicRouter +from openjarvis.learning.routing.router import HeuristicRouter def ensure_registered() -> None: diff --git a/src/openjarvis/learning/heuristic_reward.py b/src/openjarvis/learning/routing/heuristic_reward.py similarity index 100% rename from src/openjarvis/learning/heuristic_reward.py rename to src/openjarvis/learning/routing/heuristic_reward.py diff --git a/src/openjarvis/learning/trace_policy.py b/src/openjarvis/learning/routing/learned_router.py similarity index 64% rename from src/openjarvis/learning/trace_policy.py rename to src/openjarvis/learning/routing/learned_router.py index f7b60beb..057b54fd 100644 --- a/src/openjarvis/learning/trace_policy.py +++ b/src/openjarvis/learning/routing/learned_router.py @@ -1,60 +1,34 @@ -"""Trace-driven router policy — learns from interaction history. +"""Learned router policy — trace-driven query_class -> model mapping. -Unlike the ``HeuristicRouter`` which uses static rules, this policy -learns from accumulated traces which model/agent/tool combinations -produce the best outcomes for different query types. It maintains a -lightweight mapping of (query_class → model) that is updated -periodically from the ``TraceAnalyzer``. +Merges the runtime ``select_model()`` logic from ``TraceDrivenPolicy`` +with the batch ``update()`` logic from ``SFTRouterPolicy`` into a single +class registered as ``"learned"`` in ``RouterPolicyRegistry``. """ from __future__ import annotations -import re +import logging from typing import Any, Dict, List, Optional from openjarvis.core.registry import RouterPolicyRegistry from openjarvis.core.types import RoutingContext from openjarvis.learning._stubs import RouterPolicy -from openjarvis.traces.analyzer import TraceAnalyzer +from openjarvis.learning.routing._utils import classify_query -# Query classification for grouping traces -_CODE_RE = re.compile( - r"```|`[^`]+`|\bdef\s|\bclass\s|\bimport\s|\bfunction\s", - re.IGNORECASE, -) -_MATH_RE = re.compile( - r"\bsolve\b|\bintegral\b|\bequation\b|\bcalculate\b|\bcompute\b", - re.IGNORECASE, -) +logger = logging.getLogger(__name__) -def classify_query(query: str) -> str: - """Classify a query into a broad category for routing.""" - if _CODE_RE.search(query): - return "code" - if _MATH_RE.search(query): - return "math" - if len(query) < 50: - return "short" - if len(query) > 500: - return "long" - return "general" +class LearnedRouterPolicy(RouterPolicy): + """Trace-driven router that learns query_class -> model mappings. - -class TraceDrivenPolicy(RouterPolicy): - """Router policy that learns from historical traces. - - Maintains a mapping of ``query_class → best_model`` derived from - trace outcomes. Falls back to the provided default when no trace - data is available for a query class. - - The policy is updated by calling :meth:`update_from_traces`, which - reads the ``TraceAnalyzer`` and recomputes the mapping. + Implements ``RouterPolicy.select_model()`` for runtime routing AND + provides ``update_from_traces()`` / ``observe()`` for learning from traces, + plus ``update()`` for batch learning from a trace store. """ def __init__( self, - analyzer: Optional[TraceAnalyzer] = None, + analyzer: Optional[Any] = None, *, available_models: Optional[List[str]] = None, default_model: str = "", @@ -64,11 +38,8 @@ class TraceDrivenPolicy(RouterPolicy): self._available = available_models or [] self._default = default_model self._fallback = fallback_model - # Learned mapping: query_class → model key self._policy_map: Dict[str, str] = {} - # Track confidence: query_class → sample count self._confidence: Dict[str, int] = {} - # Minimum samples before trusting learned policy self.min_samples: int = 5 @property @@ -80,7 +51,6 @@ class TraceDrivenPolicy(RouterPolicy): """Select the best model based on learned policy or fallback.""" query_class = classify_query(context.query) - # Use learned policy if we have enough confidence if ( query_class in self._policy_map and self._confidence.get(query_class, 0) >= self.min_samples @@ -89,7 +59,6 @@ class TraceDrivenPolicy(RouterPolicy): if not self._available or model in self._available: return model - # Fallback chain avail = self._available if self._default and (not avail or self._default in avail): return self._default @@ -105,10 +74,7 @@ class TraceDrivenPolicy(RouterPolicy): since: Optional[float] = None, until: Optional[float] = None, ) -> Dict[str, Any]: - """Recompute the policy map from trace history. - - Returns a summary of what changed for logging/debugging. - """ + """Recompute the policy map from trace history via TraceAnalyzer.""" if self._analyzer is None: return {"error": "no analyzer configured"} @@ -118,7 +84,6 @@ class TraceDrivenPolicy(RouterPolicy): if not traces: return {"updated": False, "reason": "no traces"} - # Group traces by query class groups: Dict[str, list] = {} for t in traces: qclass = classify_query(t.query) @@ -128,7 +93,6 @@ class TraceDrivenPolicy(RouterPolicy): changes: Dict[str, Dict[str, str]] = {} for qclass, class_traces in groups.items(): - # Score each model for this query class model_scores: Dict[str, _ModelScore] = {} for t in class_traces: if not t.model: @@ -147,7 +111,6 @@ class TraceDrivenPolicy(RouterPolicy): if not model_scores: continue - # Pick the best model: weighted score of success_rate and feedback best_model = max( model_scores.items(), key=lambda kv: kv[1].composite_score(), @@ -176,16 +139,10 @@ class TraceDrivenPolicy(RouterPolicy): outcome: Optional[str], feedback: Optional[float], ) -> None: - """Record a single observation for online (incremental) updates. - - This is a lighter-weight alternative to :meth:`update_from_traces` - for use cases where you want to update the policy after every - interaction rather than in batch. - """ + """Record a single observation for online (incremental) updates.""" qclass = classify_query(query) current_count = self._confidence.get(qclass, 0) - # Simple exponential moving average for online update if qclass not in self._policy_map: self._policy_map[qclass] = model self._confidence[qclass] = 1 @@ -193,12 +150,55 @@ class TraceDrivenPolicy(RouterPolicy): self._confidence[qclass] = current_count + 1 - # Only switch models if the new model shows clearly better outcomes if outcome == "success" and feedback is not None and feedback > 0.7: - # Weight new evidence against existing policy if current_count < self.min_samples: self._policy_map[qclass] = model + def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: + """Batch update: analyze trace outcomes and update the policy map. + + This is the batch learning interface (from the old SFTRouterPolicy). + """ + try: + traces = trace_store.list_traces() + except Exception as exc: + logger.warning("Learned router update failed: %s", exc) + return {"updated": False, "reason": "Could not access trace store"} + + class_model_scores: Dict[str, Dict[str, List[float]]] = {} + for trace in traces: + query_class = classify_query(trace.query) + model = trace.model or "unknown" + outcome_score = 1.0 if trace.outcome == "success" else 0.0 + fb = trace.feedback if trace.feedback is not None else 0.5 + composite = 0.6 * outcome_score + 0.4 * fb + + if query_class not in class_model_scores: + class_model_scores[query_class] = {} + if model not in class_model_scores[query_class]: + class_model_scores[query_class][model] = [] + class_model_scores[query_class][model].append(composite) + + changes = {} + for qclass, model_scores in class_model_scores.items(): + best_model = None + best_score = -1.0 + for model, scores in model_scores.items(): + if len(scores) >= self.min_samples: + avg = sum(scores) / len(scores) + if avg > best_score: + best_score = avg + best_model = model + if best_model and best_model != self._policy_map.get(qclass): + self._policy_map[qclass] = best_model + changes[qclass] = best_model + + return { + "updated": bool(changes), + "changes": changes, + "policy_map": dict(self._policy_map), + } + class _ModelScore: """Accumulator for per-model scoring during policy update.""" @@ -222,17 +222,15 @@ class _ModelScore: self.feedback_sum / self.feedback_count if self.feedback_count else 0.5 ) - # 60% success rate + 40% feedback return 0.6 * sr + 0.4 * fb def ensure_registered() -> None: - """Register TraceDrivenPolicy if not already present.""" + """Register LearnedRouterPolicy if not already present.""" if not RouterPolicyRegistry.contains("learned"): - RouterPolicyRegistry.register_value("learned", TraceDrivenPolicy) + RouterPolicyRegistry.register_value("learned", LearnedRouterPolicy) ensure_registered() - -__all__ = ["TraceDrivenPolicy", "classify_query", "ensure_registered"] +__all__ = ["LearnedRouterPolicy", "ensure_registered"] diff --git a/src/openjarvis/learning/router.py b/src/openjarvis/learning/routing/router.py similarity index 100% rename from src/openjarvis/learning/router.py rename to src/openjarvis/learning/routing/router.py diff --git a/src/openjarvis/learning/sft_policy.py b/src/openjarvis/learning/sft_policy.py deleted file mode 100644 index b00b5b33..00000000 --- a/src/openjarvis/learning/sft_policy.py +++ /dev/null @@ -1,101 +0,0 @@ -"""SFT router — learns which model handles which query type best. - -Analyses historical traces and builds a ``query_class → model`` routing -table. Despite the "SFT" name, no model weights are fine-tuned — this -is a *trace-driven routing* policy. -""" - -from __future__ import annotations - -import logging -from typing import Any, Dict, List - -from openjarvis.core.registry import LearningRegistry -from openjarvis.learning._stubs import IntelligenceLearningPolicy - -logger = logging.getLogger(__name__) - - -@LearningRegistry.register("sft") -class SFTRouterPolicy(IntelligenceLearningPolicy): - """Trace-driven router that learns query_class → model mappings. - - Reads historical traces, groups by query class (code, math, short, - long, general), scores each model via a composite metric (60% - outcome + 40% feedback), and produces a routing table that maps - query classes to their best-performing model. - """ - - def __init__(self, *, min_samples: int = 5) -> None: - self._min_samples = min_samples - self._policy_map: Dict[str, str] = {} - - def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]: - """Analyze trace outcomes and update the policy map.""" - try: - traces = trace_store.list_traces() - except Exception as exc: - logger.warning("SFT policy update failed: %s", exc) - return {"updated": False, "reason": "Could not access trace store"} - - # Group traces by query class and model - class_model_scores: Dict[str, Dict[str, List[float]]] = {} - for trace in traces: - query_class = self._classify_query(trace.query) - model = trace.model or "unknown" - outcome_score = 1.0 if trace.outcome == "success" else 0.0 - feedback = trace.feedback if trace.feedback is not None else 0.5 - - composite = 0.6 * outcome_score + 0.4 * feedback - - if query_class not in class_model_scores: - class_model_scores[query_class] = {} - if model not in class_model_scores[query_class]: - class_model_scores[query_class][model] = [] - class_model_scores[query_class][model].append(composite) - - # Update policy map: best model per query class - changes = {} - for qclass, model_scores in class_model_scores.items(): - best_model = None - best_score = -1.0 - for model, scores in model_scores.items(): - if len(scores) >= self._min_samples: - avg = sum(scores) / len(scores) - if avg > best_score: - best_score = avg - best_model = model - if best_model and best_model != self._policy_map.get(qclass): - self._policy_map[qclass] = best_model - changes[qclass] = best_model - - return { - "updated": bool(changes), - "changes": changes, - "policy_map": dict(self._policy_map), - } - - @staticmethod - def _classify_query(query: str) -> str: - """Classify a query into a category.""" - q = query.lower() - if any(kw in q for kw in ("def ", "class ", "import ", "```", "function")): - return "code" - math_kws = ("solve", "integral", "equation", "derivative", "proof") - if any(kw in q for kw in math_kws): - return "math" - if len(query.split()) < 10: - return "short" - if len(query.split()) > 100: - return "long" - return "general" - - @property - def policy_map(self) -> Dict[str, str]: - return dict(self._policy_map) - - -__all__ = ["SFTRouterPolicy"] - -# Backward-compat alias -SFTPolicy = SFTRouterPolicy diff --git a/src/openjarvis/learning/training/data.py b/src/openjarvis/learning/training/data.py index 48f693a8..cfb2ad11 100644 --- a/src/openjarvis/learning/training/data.py +++ b/src/openjarvis/learning/training/data.py @@ -16,7 +16,7 @@ from collections import defaultdict from typing import Any, Dict, List from openjarvis.core.types import StepType, Trace -from openjarvis.learning.trace_policy import classify_query +from openjarvis.learning.routing._utils import classify_query class TrainingDataMiner: diff --git a/src/openjarvis/server/api_routes.py b/src/openjarvis/server/api_routes.py index ef1dd0d9..59a4fbaf 100644 --- a/src/openjarvis/server/api_routes.py +++ b/src/openjarvis/server/api_routes.py @@ -528,54 +528,9 @@ async def learning_stats(request: Request): """Return learning system statistics across all sub-policies.""" result: Dict[str, Any] = {} - # GRPO policy state - try: - from openjarvis.learning.grpo_policy import GRPORouterPolicy - policy = GRPORouterPolicy() - state = policy.state - result["grpo"] = { - "available": True, - "total_updates": state.total_updates, - "sample_counts": dict(state.sample_counts), - "weight_count": sum( - len(qc_weights) for qc_weights in state.weights.values() - ), - } - except Exception as exc: - logger.warning("Failed to load GRPO stats: %s", exc) - result["grpo"] = {"available": False} - - # Bandit router state - try: - from openjarvis.learning.bandit_router import BanditRouterPolicy - policy = BanditRouterPolicy() - stats = policy.get_stats() - result["bandit"] = { - "available": True, - "query_classes": len(stats), - "arms": stats, - } - except Exception as exc: - logger.warning("Failed to load bandit stats: %s", exc) - result["bandit"] = {"available": False} - - # ICL updater - try: - from openjarvis.learning.icl_updater import ICLUpdaterPolicy - updater = ICLUpdaterPolicy() - result["icl"] = { - "available": True, - "example_count": len(updater.examples), - "example_db_count": len(updater.example_db), - "version": updater.version, - } - except Exception as exc: - logger.warning("Failed to load ICL stats: %s", exc) - result["icl"] = {"available": False} - # Skill discovery try: - from openjarvis.learning.skill_discovery import SkillDiscovery + from openjarvis.learning.agents.skill_discovery import SkillDiscovery discovery = SkillDiscovery() result["skill_discovery"] = { "available": True, @@ -624,20 +579,6 @@ async def learning_policy(request: Request): result["agent"] = {"policy": "none"} result["metrics"] = {} - # Include GRPO weights if the routing policy is grpo - if result.get("routing", {}).get("policy") == "grpo": - try: - from openjarvis.learning.grpo_policy import GRPORouterPolicy - policy = GRPORouterPolicy() - state = policy.state - result["grpo_weights"] = { - model: dict(qc_weights) - for model, qc_weights in state.weights.items() - } - except Exception as exc: - logger.warning("Failed to load GRPO weights: %s", exc) - result["grpo_weights"] = {} - return result diff --git a/src/openjarvis/system.py b/src/openjarvis/system.py index d06a2947..25ae8e77 100644 --- a/src/openjarvis/system.py +++ b/src/openjarvis/system.py @@ -965,16 +965,17 @@ class SystemBuilder: trace_store = TraceStore(db_path=config.traces.db_path) config_dir = DEFAULT_CONFIG_DIR / "agent_configs" + sft_cfg = config.learning.intelligence.sft lora_config = LoRATrainingConfig( - lora_rank=config.learning.lora_rank, - lora_alpha=config.learning.lora_alpha, + lora_rank=sft_cfg.lora_rank, + lora_alpha=sft_cfg.lora_alpha, ) return LearningOrchestrator( trace_store=trace_store, config_dir=config_dir, min_improvement=config.learning.min_improvement, - min_sft_pairs=config.learning.min_sft_pairs, + min_sft_pairs=sft_cfg.min_pairs, lora_config=lora_config, ) except Exception as exc: diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 11b31d65..28407d26 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -260,7 +260,6 @@ class TestNestedLearningConfig: assert lc.routing.min_samples == 5 assert lc.intelligence.policy == "none" assert lc.agent.policy == "none" - assert lc.agent.max_icl_examples == 20 assert lc.metrics.accuracy_weight == 0.6 assert lc.metrics.latency_weight == 0.2 diff --git a/tests/integration/test_integration.py b/tests/integration/test_integration.py index c6594354..ced76ecc 100644 --- a/tests/integration/test_integration.py +++ b/tests/integration/test_integration.py @@ -308,7 +308,7 @@ class TestHeuristicRewardWithTelemetry: def test_reward_from_telemetry_record(self): from openjarvis.learning._stubs import RoutingContext - from openjarvis.learning.heuristic_reward import HeuristicRewardFunction + from openjarvis.learning.routing.heuristic_reward import HeuristicRewardFunction rec = TelemetryRecord( timestamp=0.0, @@ -331,14 +331,14 @@ class TestHeuristicRewardWithTelemetry: class TestRouterPolicyRegistryDiscovery: - """RouterPolicyRegistry discovers both heuristic and grpo.""" + """RouterPolicyRegistry discovers both heuristic and learned.""" def test_both_policies_registered(self): from openjarvis.learning import ensure_registered ensure_registered() assert RouterPolicyRegistry.contains("heuristic") - assert RouterPolicyRegistry.contains("grpo") + assert RouterPolicyRegistry.contains("learned") class TestTelemetryPipeline: @@ -406,8 +406,8 @@ class TestAskFlowWithRouterPolicy: def test_mocked_ask_with_registry_router(self): from openjarvis.learning._stubs import RoutingContext - from openjarvis.learning.heuristic_policy import ensure_registered - from openjarvis.learning.router import HeuristicRouter + from openjarvis.learning.routing.heuristic_policy import ensure_registered + from openjarvis.learning.routing.router import HeuristicRouter ensure_registered() router_cls = RouterPolicyRegistry.get("heuristic") @@ -429,7 +429,7 @@ class TestRewardTelemetryIntegration: import time from openjarvis.learning._stubs import RoutingContext - from openjarvis.learning.heuristic_reward import HeuristicRewardFunction + from openjarvis.learning.routing.heuristic_reward import HeuristicRewardFunction from openjarvis.telemetry.aggregator import TelemetryAggregator from openjarvis.telemetry.store import TelemetryStore diff --git a/tests/learning/agents/__init__.py b/tests/learning/agents/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/learning/agents/test_dspy_optimizer.py b/tests/learning/agents/test_dspy_optimizer.py new file mode 100644 index 00000000..e83d45f8 --- /dev/null +++ b/tests/learning/agents/test_dspy_optimizer.py @@ -0,0 +1,82 @@ +"""Tests for the DSPy agent optimizer (mocked -- no dspy dependency required).""" + +from __future__ import annotations + +import time +from unittest.mock import MagicMock, patch + + +class TestDSPyOptimizerConfig: + def test_default_config(self) -> None: + from openjarvis.core.config import DSPyOptimizerConfig + + cfg = DSPyOptimizerConfig() + assert cfg.optimizer == "BootstrapFewShotWithRandomSearch" + assert cfg.max_bootstrapped_demos == 4 + assert cfg.min_traces == 20 + + def test_optimizer_init(self) -> None: + from openjarvis.core.config import DSPyOptimizerConfig + from openjarvis.learning.agents.dspy_optimizer import DSPyAgentOptimizer + + cfg = DSPyOptimizerConfig() + optimizer = DSPyAgentOptimizer(cfg) + assert optimizer.config is cfg + + +class TestDSPyOptimizerTraceConversion: + def test_too_few_traces_skipped(self) -> None: + from openjarvis.core.config import DSPyOptimizerConfig + from openjarvis.learning.agents.dspy_optimizer import DSPyAgentOptimizer + + optimizer = DSPyAgentOptimizer(DSPyOptimizerConfig(min_traces=10)) + mock_store = MagicMock() + mock_store.list_traces.return_value = [] + + result = optimizer.optimize(mock_store) + assert result["status"] == "skipped" + + def test_optimize_returns_toml_updates(self) -> None: + from openjarvis.core.config import DSPyOptimizerConfig + from openjarvis.core.types import StepType, Trace, TraceStep + from openjarvis.learning.agents.dspy_optimizer import DSPyAgentOptimizer + + cfg = DSPyOptimizerConfig(min_traces=1) + optimizer = DSPyAgentOptimizer(cfg) + + # Create mock traces + now = time.time() + traces = [] + for i in range(5): + traces.append(Trace( + query=f"test query {i}", + agent="native_react", + model="qwen3:8b", + result=f"result {i}", + outcome="success", + feedback=0.9, + started_at=now, + ended_at=now + 1, + total_tokens=100, + total_latency_seconds=1.0, + steps=[TraceStep( + step_type=StepType.GENERATE, + timestamp=now, + duration_seconds=0.5, + )], + )) + + mock_store = MagicMock() + mock_store.list_traces.return_value = traces + + # Mock dspy so the test works without the dependency + import openjarvis.learning.agents.dspy_optimizer as mod + + with patch.object(mod, "HAS_DSPY", True): + with patch.object(optimizer, "_run_dspy_optimization", return_value={ + "system_prompt": "You are a helpful assistant.", + "few_shot_examples": [{"input": "hi", "output": "hello"}], + }): + result = optimizer.optimize(mock_store) + assert result["status"] == "completed" + assert "config_updates" in result diff --git a/tests/learning/agents/test_gepa_optimizer.py b/tests/learning/agents/test_gepa_optimizer.py new file mode 100644 index 00000000..acb59198 --- /dev/null +++ b/tests/learning/agents/test_gepa_optimizer.py @@ -0,0 +1,85 @@ +"""Tests for the GEPA agent optimizer (mocked -- no gepa dependency required).""" + +from __future__ import annotations + +import time +from unittest.mock import MagicMock, patch + + +class TestGEPAOptimizerConfig: + def test_default_config(self) -> None: + from openjarvis.core.config import GEPAOptimizerConfig + + cfg = GEPAOptimizerConfig() + assert cfg.max_metric_calls == 150 + assert cfg.population_size == 10 + assert cfg.min_traces == 20 + + def test_optimizer_init(self) -> None: + from openjarvis.core.config import GEPAOptimizerConfig + from openjarvis.learning.agents.gepa_optimizer import GEPAAgentOptimizer + + cfg = GEPAOptimizerConfig() + optimizer = GEPAAgentOptimizer(cfg) + assert optimizer.config is cfg + + +class TestGEPAOptimizerOptimize: + def test_too_few_traces_skipped(self) -> None: + from openjarvis.core.config import GEPAOptimizerConfig + from openjarvis.learning.agents.gepa_optimizer import GEPAAgentOptimizer + + optimizer = GEPAAgentOptimizer(GEPAOptimizerConfig(min_traces=10)) + mock_store = MagicMock() + mock_store.list_traces.return_value = [] + + result = optimizer.optimize(mock_store) + assert result["status"] == "skipped" + + def test_no_gepa_reports_error(self) -> None: + from openjarvis.core.config import GEPAOptimizerConfig + from openjarvis.core.types import StepType, Trace, TraceStep + from openjarvis.learning.agents.gepa_optimizer import GEPAAgentOptimizer + + cfg = GEPAOptimizerConfig(min_traces=1) + optimizer = GEPAAgentOptimizer(cfg) + + now = time.time() + traces = [Trace( + query="test", agent="native_react", model="qwen3:8b", + result="result", outcome="success", feedback=0.9, + started_at=now, ended_at=now + 1, + total_tokens=100, total_latency_seconds=1.0, + steps=[TraceStep( + step_type=StepType.GENERATE, + timestamp=now, + duration_seconds=0.5, + )], + )] + + mock_store = MagicMock() + mock_store.list_traces.return_value = traces + + # Ensure gepa is not available + with patch.dict("sys.modules", {"gepa": None}): + with patch( + "openjarvis.learning.agents.gepa_optimizer.HAS_GEPA", + False, + ): + result = optimizer.optimize(mock_store) + assert result["status"] == "error" + assert "gepa" in result["reason"].lower() + + +class TestOpenJarvisGEPAAdapter: + def test_adapter_init(self) -> None: + from openjarvis.core.config import GEPAOptimizerConfig + from openjarvis.learning.agents.gepa_optimizer import ( + OpenJarvisGEPAAdapter, + ) + + mock_store = MagicMock() + adapter = OpenJarvisGEPAAdapter( + mock_store, "native_react", GEPAOptimizerConfig(), + ) + assert adapter.agent_name == "native_react" diff --git a/tests/learning/intelligence/__init__.py b/tests/learning/intelligence/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/learning/intelligence/test_grpo_trainer.py b/tests/learning/intelligence/test_grpo_trainer.py new file mode 100644 index 00000000..19d20f42 --- /dev/null +++ b/tests/learning/intelligence/test_grpo_trainer.py @@ -0,0 +1,88 @@ +"""Tests for the general-purpose GRPO trainer.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + + +class TestGRPOConfig: + def test_default_config(self) -> None: + from openjarvis.core.config import GRPOConfig + + cfg = GRPOConfig() + assert cfg.model_name == "Qwen/Qwen3-1.7B" + assert cfg.num_samples_per_prompt == 8 + assert cfg.kl_coef == 0.0001 + assert cfg.clip_ratio == 0.2 + assert cfg.min_prompts == 10 + + +class TestDefaultRewardFn: + def test_score_returns_float(self) -> None: + from openjarvis.learning.intelligence.grpo_trainer import DefaultRewardFn + + reward = DefaultRewardFn() + score = reward.score("prompt", "response", None) + assert isinstance(score, float) + assert 0.0 <= score <= 1.0 + + def test_score_with_ground_truth(self) -> None: + from openjarvis.learning.intelligence.grpo_trainer import DefaultRewardFn + + reward = DefaultRewardFn() + # When response matches ground truth, score should be higher + score_match = reward.score("what is 2+2?", "4", "4") + score_no_match = reward.score("what is 2+2?", "5", "4") + assert score_match > score_no_match + + +class TestGRPOTrainer: + def test_init(self) -> None: + from openjarvis.core.config import GRPOConfig + from openjarvis.learning.intelligence.grpo_trainer import GRPOTrainer + + cfg = GRPOConfig() + trainer = GRPOTrainer(cfg) + assert trainer.config is cfg + + def test_train_on_prompts_empty(self) -> None: + from openjarvis.core.config import GRPOConfig + from openjarvis.learning.intelligence.grpo_trainer import GRPOTrainer + + trainer = GRPOTrainer(GRPOConfig()) + result = trainer.train_on_prompts([]) + assert result["status"] == "skipped" + + def test_train_on_prompts_too_few(self) -> None: + from openjarvis.core.config import GRPOConfig + from openjarvis.learning.intelligence.grpo_trainer import GRPOTrainer + + trainer = GRPOTrainer(GRPOConfig(min_prompts=5)) + result = trainer.train_on_prompts(["hello"]) + assert result["status"] == "skipped" + assert "min_prompts" in result.get("reason", "") + + def test_custom_reward_fn(self) -> None: + from openjarvis.core.config import GRPOConfig + from openjarvis.learning.intelligence.grpo_trainer import GRPOTrainer + + class MyReward: + def score( + self, prompt: str, response: str, ground_truth: str | None + ) -> float: + return 0.42 + + trainer = GRPOTrainer(GRPOConfig(), reward_fn=MyReward()) + assert trainer.reward_fn.score("a", "b", None) == 0.42 + + def test_train_delegates_to_miner(self) -> None: + from openjarvis.core.config import GRPOConfig + from openjarvis.learning.intelligence.grpo_trainer import GRPOTrainer + + trainer = GRPOTrainer(GRPOConfig(min_prompts=1)) + mock_store = MagicMock() + + with patch.object(trainer, "_mine_prompts", return_value=[]) as mock_mine: + result = trainer.train(mock_store) + mock_mine.assert_called_once_with(mock_store) + assert result["status"] == "skipped" diff --git a/tests/learning/intelligence/test_sft_trainer.py b/tests/learning/intelligence/test_sft_trainer.py new file mode 100644 index 00000000..59ca374d --- /dev/null +++ b/tests/learning/intelligence/test_sft_trainer.py @@ -0,0 +1,66 @@ +"""Tests for the general-purpose SFT trainer.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + + +class TestSFTTrainerConfig: + def test_default_config(self) -> None: + from openjarvis.core.config import SFTConfig + + cfg = SFTConfig() + assert cfg.model_name == "Qwen/Qwen3-1.7B" + assert cfg.use_lora is True + assert cfg.lora_rank == 16 + assert cfg.min_pairs == 10 + + def test_trainer_init(self) -> None: + from openjarvis.core.config import SFTConfig + from openjarvis.learning.intelligence.sft_trainer import SFTTrainer + + cfg = SFTConfig() + trainer = SFTTrainer(cfg) + assert trainer.config is cfg + + def test_target_modules_parsing(self) -> None: + from openjarvis.core.config import SFTConfig + from openjarvis.learning.intelligence.sft_trainer import SFTTrainer + + cfg = SFTConfig(target_modules="q_proj,v_proj,k_proj") + trainer = SFTTrainer(cfg) + assert trainer.target_module_list == ["q_proj", "v_proj", "k_proj"] + + +class TestSFTTrainerTrainOnPairs: + def test_empty_pairs_skipped(self) -> None: + from openjarvis.core.config import SFTConfig + from openjarvis.learning.intelligence.sft_trainer import SFTTrainer + + trainer = SFTTrainer(SFTConfig()) + result = trainer.train_on_pairs([]) + assert result["status"] == "skipped" + + def test_too_few_pairs_skipped(self) -> None: + from openjarvis.core.config import SFTConfig + from openjarvis.learning.intelligence.sft_trainer import SFTTrainer + + trainer = SFTTrainer(SFTConfig(min_pairs=5)) + pairs = [{"input": "hi", "output": "hello"}] + result = trainer.train_on_pairs(pairs) + assert result["status"] == "skipped" + assert "min_pairs" in result.get("reason", "") + + +class TestSFTTrainerTraceMining: + def test_train_delegates_to_miner(self) -> None: + from openjarvis.core.config import SFTConfig + from openjarvis.learning.intelligence.sft_trainer import SFTTrainer + + trainer = SFTTrainer(SFTConfig(min_pairs=1)) + mock_store = MagicMock() + + with patch.object(trainer, "_mine_pairs", return_value=[]) as mock_mine: + result = trainer.train(mock_store) + mock_mine.assert_called_once_with(mock_store) + assert result["status"] == "skipped" diff --git a/tests/learning/routing/__init__.py b/tests/learning/routing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/learning/routing/test_learned_router.py b/tests/learning/routing/test_learned_router.py new file mode 100644 index 00000000..20dc4d71 --- /dev/null +++ b/tests/learning/routing/test_learned_router.py @@ -0,0 +1,151 @@ +"""Tests for LearnedRouterPolicy (merged trace-driven + SFT routing).""" + +from __future__ import annotations + +import time +from pathlib import Path + +from openjarvis.core.types import StepType, Trace, TraceStep +from openjarvis.learning._stubs import RoutingContext +from openjarvis.learning.routing.learned_router import LearnedRouterPolicy +from openjarvis.traces.analyzer import TraceAnalyzer +from openjarvis.traces.store import TraceStore + + +def _make_trace( + query: str = "test", + model: str = "qwen3:8b", + outcome: str | None = "success", + feedback: float | None = 0.8, +) -> Trace: + now = time.time() + return Trace( + query=query, + agent="orchestrator", + model=model, + engine="ollama", + result="result", + outcome=outcome, + feedback=feedback, + started_at=now, + ended_at=now + 0.5, + total_tokens=100, + total_latency_seconds=0.5, + steps=[ + TraceStep( + step_type=StepType.GENERATE, + timestamp=now, + duration_seconds=0.5, + output={"tokens": 100}, + ), + ], + ) + + +class TestLearnedRouterPolicy: + def test_registered_as_learned(self) -> None: + from openjarvis.core.registry import RouterPolicyRegistry + from openjarvis.learning.routing.learned_router import ensure_registered + ensure_registered() + assert RouterPolicyRegistry.contains("learned") + + def test_fallback_no_traces(self) -> None: + policy = LearnedRouterPolicy(default_model="qwen3:8b") + ctx = RoutingContext(query="hello") + assert policy.select_model(ctx) == "qwen3:8b" + + def test_fallback_chain(self) -> None: + policy = LearnedRouterPolicy( + default_model="missing", + fallback_model="llama3:8b", + available_models=["llama3:8b"], + ) + ctx = RoutingContext(query="hello") + assert policy.select_model(ctx) == "llama3:8b" + + def test_update_from_traces(self, tmp_path: Path) -> None: + store = TraceStore(tmp_path / "test.db") + for _ in range(6): + store.save(_make_trace( + query="def foo(): pass", + model="codestral", + outcome="success", + feedback=0.9, + )) + for _ in range(6): + store.save(_make_trace( + query="def bar(): return 1", + model="qwen3:8b", + outcome="failure", + feedback=0.3, + )) + + analyzer = TraceAnalyzer(store) + policy = LearnedRouterPolicy( + analyzer=analyzer, + default_model="qwen3:8b", + ) + policy.min_samples = 3 + result = policy.update_from_traces() + assert result["updated"] is True + + ctx = RoutingContext(query="import os; def main(): pass") + assert policy.select_model(ctx) == "codestral" + store.close() + + def test_policy_map_readable(self, tmp_path: Path) -> None: + store = TraceStore(tmp_path / "test.db") + for _ in range(5): + store.save(_make_trace( + query="hello", model="small-model", + outcome="success", + )) + + analyzer = TraceAnalyzer(store) + policy = LearnedRouterPolicy(analyzer=analyzer, default_model="default") + policy.min_samples = 3 + policy.update_from_traces() + + pmap = policy.policy_map + assert isinstance(pmap, dict) + assert "short" in pmap + assert pmap["short"] == "small-model" + store.close() + + def test_observe_online(self) -> None: + policy = LearnedRouterPolicy(default_model="default") + policy.min_samples = 3 + policy.observe("hello", "fast-model", "success", 0.9) + assert policy.policy_map.get("short") == "fast-model" + + def test_batch_update(self) -> None: + """Test the batch update() method inherited from SFT routing logic.""" + from unittest.mock import MagicMock + + policy = LearnedRouterPolicy(default_model="default") + mock_store = MagicMock() + mock_store.list_traces.return_value = [ + _make_trace( + query="def foo(): pass", model="code-model", + outcome="success", feedback=0.9, + ), + _make_trace( + query="def bar(): pass", model="code-model", + outcome="success", feedback=0.85, + ), + _make_trace( + query="def baz(): pass", model="code-model", + outcome="success", feedback=0.88, + ), + _make_trace( + query="def qux(): pass", model="code-model", + outcome="success", feedback=0.92, + ), + _make_trace( + query="def quux(): pass", model="code-model", + outcome="success", feedback=0.87, + ), + ] + result = policy.update(mock_store) + assert isinstance(result, dict) + assert "policy_map" in result diff --git a/tests/learning/routing/test_utils.py b/tests/learning/routing/test_utils.py new file mode 100644 index 00000000..2fce7216 --- /dev/null +++ b/tests/learning/routing/test_utils.py @@ -0,0 +1,27 @@ +"""Tests for classify_query utility.""" + +from __future__ import annotations + +from openjarvis.learning.routing._utils import classify_query + + +class TestClassifyQuery: + def test_code_detection(self) -> None: + assert classify_query("def hello(): pass") == "code" + assert classify_query("```python\nprint()```") == "code" + assert classify_query("import os") == "code" + + def test_math_detection(self) -> None: + assert classify_query("solve this equation for x") == "math" + assert classify_query("compute the integral") == "math" + + def test_short(self) -> None: + assert classify_query("hello") == "short" + assert classify_query("what time is it?") == "short" + + def test_long(self) -> None: + assert classify_query("a" * 501) == "long" + + def test_general(self) -> None: + q = "Tell me about the history of artificial intelligence research" + assert classify_query(q) == "general" diff --git a/tests/learning/test_agent_advisor.py b/tests/learning/test_agent_advisor.py deleted file mode 100644 index 28b9da14..00000000 --- a/tests/learning/test_agent_advisor.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Tests for agent advisor policy.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Optional - -from openjarvis.learning.agent_advisor import AgentAdvisorPolicy - - -@dataclass -class _MockTrace: - query: str = "" - model: str = "model-a" - outcome: str = "success" - feedback: Optional[float] = 0.8 - steps: list = field(default_factory=list) - total_latency_seconds: float = 1.0 - - -class _MockTraceStore: - def __init__(self, traces): - self._traces = traces - - def list_traces(self): - return self._traces - - -class TestAgentAdvisorPolicy: - def test_no_problem_traces(self): - traces = [_MockTrace(outcome="success") for _ in range(5)] - policy = AgentAdvisorPolicy() - store = _MockTraceStore(traces) - result = policy.update(store) - assert result["recommendations"] == [] - assert result["confidence"] == 1.0 - - def test_detects_failures(self): - traces = [ - _MockTrace(query="def code()", outcome="failure"), - _MockTrace(query="def more_code()", outcome="failure"), - _MockTrace(query="import something", outcome="failure"), - ] - policy = AgentAdvisorPolicy() - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["recommendations"]) > 0 - assert result["confidence"] < 1.0 - - def test_detects_slow_traces(self): - traces = [ - _MockTrace(outcome="success", total_latency_seconds=10.0), - ] - policy = AgentAdvisorPolicy() - store = _MockTraceStore(traces) - result = policy.update(store) - assert result["problem_traces"] == 1 - - def test_empty_traces(self): - policy = AgentAdvisorPolicy() - store = _MockTraceStore([]) - result = policy.update(store) - assert result["recommendations"] == [] - - def test_is_agent_policy(self): - from openjarvis.learning._stubs import AgentLearningPolicy - assert issubclass(AgentAdvisorPolicy, AgentLearningPolicy) - assert AgentAdvisorPolicy.target == "agent" - - def test_max_traces_limit(self): - traces = [ - _MockTrace(outcome="failure") for _ in range(100) - ] - policy = AgentAdvisorPolicy(max_traces=10) - store = _MockTraceStore(traces) - result = policy.update(store) - # Should only analyze last 10 traces - assert result["problem_traces"] <= 10 diff --git a/tests/learning/test_agent_evolver.py b/tests/learning/test_agent_evolver.py index 803c47ab..696e7e72 100644 --- a/tests/learning/test_agent_evolver.py +++ b/tests/learning/test_agent_evolver.py @@ -14,7 +14,7 @@ except ModuleNotFoundError: import tomli as tomllib # type: ignore[no-redef] from openjarvis.core.types import StepType, Trace, TraceStep -from openjarvis.learning.agent_evolver import AgentConfigEvolver +from openjarvis.learning.agents.agent_evolver import AgentConfigEvolver from openjarvis.traces.store import TraceStore # --------------------------------------------------------------------------- diff --git a/tests/learning/test_bandit_router.py b/tests/learning/test_bandit_router.py deleted file mode 100644 index 71b5717b..00000000 --- a/tests/learning/test_bandit_router.py +++ /dev/null @@ -1,184 +0,0 @@ -"""Tests for BanditRouterPolicy — Thompson Sampling / UCB.""" - -from __future__ import annotations - -import pytest - -from openjarvis.core.registry import RouterPolicyRegistry -from openjarvis.core.types import RoutingContext -from openjarvis.learning.bandit_router import ( - ArmStats, - BanditRouterPolicy, - ensure_registered, -) - - -def _make_context(query: str = "hello", **kwargs) -> RoutingContext: - """Build a RoutingContext with sensible defaults.""" - return RoutingContext( - query=query, - query_length=kwargs.pop("query_length", len(query)), - has_code=kwargs.pop("has_code", False), - has_math=kwargs.pop("has_math", False), - **kwargs, - ) - - -class TestBanditRouterPolicy: - def test_route_no_models(self) -> None: - policy = BanditRouterPolicy() - ctx = _make_context() - with pytest.raises(ValueError, match="No models available"): - policy.route(ctx, []) - - def test_route_explores_uniformly(self) -> None: - """Before min_pulls, all models should get tried.""" - policy = BanditRouterPolicy(min_pulls=5) - models = ["model-a", "model-b", "model-c"] - ctx = _make_context() - - seen = set() - for _ in range(60): - selected = policy.route(ctx, models) - assert selected in models - seen.add(selected) - - # With 60 random selections from 3 models, all should appear - assert seen == set(models) - - def test_thompson_sampling(self) -> None: - """After updates, higher-reward model should be selected more often.""" - policy = BanditRouterPolicy(strategy="thompson", min_pulls=2) - models = ["good-model", "bad-model"] - - # Give good-model high rewards, bad-model low rewards - for _ in range(20): - policy.update("short", "good-model", 0.95) - policy.update("short", "bad-model", 0.05) - - ctx = _make_context("hi", query_length=5) - counts = {"good-model": 0, "bad-model": 0} - for _ in range(100): - selected = policy.route(ctx, models) - counts[selected] += 1 - - assert counts["good-model"] > 60, ( - f"good-model only selected {counts['good-model']}/100 times" - ) - - def test_ucb_selection(self) -> None: - """After updates, UCB should favor the higher-reward model.""" - policy = BanditRouterPolicy(strategy="ucb", min_pulls=2) - models = ["best-model", "worst-model"] - - # Train with clear reward difference - for _ in range(20): - policy.update("short", "best-model", 0.9) - policy.update("short", "worst-model", 0.1) - - ctx = _make_context("hi", query_length=5) - counts = {"best-model": 0, "worst-model": 0} - for _ in range(100): - selected = policy.route(ctx, models) - counts[selected] += 1 - - assert counts["best-model"] > 60, ( - f"best-model only selected {counts['best-model']}/100 times" - ) - - def test_update_increments_stats(self) -> None: - policy = BanditRouterPolicy(reward_threshold=0.5) - - policy.update("code", "model-a", 0.8) # success - policy.update("code", "model-a", 0.3) # failure - policy.update("code", "model-a", 0.9) # success - - stats = policy.get_stats("code") - assert stats["model-a"]["pulls"] == 3 - assert stats["model-a"]["successes"] == 2 - assert stats["model-a"]["failures"] == 1 - assert abs(stats["model-a"]["mean_reward"] - (0.8 + 0.3 + 0.9) / 3) < 1e-9 - - def test_get_stats(self) -> None: - policy = BanditRouterPolicy() - policy.update("code", "model-a", 0.8) - policy.update("math", "model-b", 0.6) - - # Per-class stats - code_stats = policy.get_stats("code") - assert "model-a" in code_stats - assert code_stats["model-a"]["pulls"] == 1 - - # All stats - all_stats = policy.get_stats() - assert "code" in all_stats - assert "math" in all_stats - assert "model-a" in all_stats["code"] - assert "model-b" in all_stats["math"] - - def test_get_stats_empty_class(self) -> None: - policy = BanditRouterPolicy() - stats = policy.get_stats("nonexistent") - assert stats == {} - - def test_reset(self) -> None: - policy = BanditRouterPolicy() - policy.update("code", "model-a", 0.8) - policy.update("math", "model-b", 0.6) - assert policy._total_pulls == 2 - - policy.reset() - - assert policy._total_pulls == 0 - assert policy.get_stats() == {} - - def test_registry_registration(self) -> None: - ensure_registered() - assert RouterPolicyRegistry.contains("bandit") - assert RouterPolicyRegistry.get("bandit") is BanditRouterPolicy - - def test_under_explored_arms(self) -> None: - """Models with fewer than min_pulls should be explored first.""" - policy = BanditRouterPolicy(strategy="thompson", min_pulls=5) - models = ["explored", "unexplored"] - - # Give "explored" enough pulls - for _ in range(10): - policy.update("short", "explored", 0.9) - - ctx = _make_context("hi", query_length=5) - # "unexplored" has 0 pulls < min_pulls=5, so it should always be chosen - for _ in range(10): - selected = policy.route(ctx, models) - assert selected == "unexplored" - - def test_arm_stats_mean_reward_zero_pulls(self) -> None: - stats = ArmStats() - assert stats.mean_reward == 0.0 - - def test_arm_stats_mean_reward(self) -> None: - stats = ArmStats(total_reward=3.0, pulls=6) - assert stats.mean_reward == 0.5 - - def test_different_query_classes_independent(self) -> None: - """Arms for different query classes should be independent.""" - policy = BanditRouterPolicy(min_pulls=1) - - # model-a is good for code, model-b is good for math - for _ in range(20): - policy.update("code", "model-a", 0.95) - policy.update("code", "model-b", 0.1) - policy.update("math", "model-b", 0.95) - policy.update("math", "model-a", 0.1) - - code_stats = policy.get_stats("code") - math_stats = policy.get_stats("math") - - assert ( - code_stats["model-a"]["mean_reward"] - > code_stats["model-b"]["mean_reward"] - ) - assert ( - math_stats["model-b"]["mean_reward"] - > math_stats["model-a"]["mean_reward"] - ) diff --git a/tests/learning/test_device_selection.py b/tests/learning/test_device_selection.py index 25c17d4c..94000df5 100644 --- a/tests/learning/test_device_selection.py +++ b/tests/learning/test_device_selection.py @@ -13,7 +13,7 @@ class TestSelectTorchDevice: def test_no_torch_returns_none(self): """Without torch, _select_torch_device returns None.""" - from openjarvis.learning.orchestrator.sft_trainer import ( + from openjarvis.learning.intelligence.orchestrator.sft_trainer import ( _select_torch_device, ) @@ -64,10 +64,10 @@ class TestSelectTorchDevice: def test_function_exists_in_both_trainers(self): """_select_torch_device is defined in both trainers.""" - from openjarvis.learning.orchestrator.grpo_trainer import ( + from openjarvis.learning.intelligence.orchestrator.grpo_trainer import ( _select_torch_device as grpo_fn, ) - from openjarvis.learning.orchestrator.sft_trainer import ( + from openjarvis.learning.intelligence.orchestrator.sft_trainer import ( _select_torch_device as sft_fn, ) @@ -76,7 +76,7 @@ class TestSelectTorchDevice: def test_exported_from_orchestrator_init(self): """_select_torch_device is exported from orchestrator package.""" - from openjarvis.learning.orchestrator import ( + from openjarvis.learning.intelligence.orchestrator import ( _select_torch_device, ) diff --git a/tests/learning/test_grpo_policy.py b/tests/learning/test_grpo_policy.py deleted file mode 100644 index ad7238fd..00000000 --- a/tests/learning/test_grpo_policy.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Tests for GRPORouterPolicy — Group Relative Policy Optimization.""" - -from __future__ import annotations - -import pytest - -from openjarvis.core.registry import RouterPolicyRegistry -from openjarvis.core.types import RoutingContext -from openjarvis.learning.grpo_policy import ( - GRPORouterPolicy, - GRPOSample, - GRPOState, - ensure_registered, -) - - -def _make_context(query: str = "hello", **kwargs) -> RoutingContext: - """Build a RoutingContext with sensible defaults.""" - return RoutingContext( - query=query, - query_length=kwargs.pop("query_length", len(query)), - has_code=kwargs.pop("has_code", False), - has_math=kwargs.pop("has_math", False), - **kwargs, - ) - - -class TestGRPORouterPolicy: - def test_route_with_no_models(self) -> None: - policy = GRPORouterPolicy() - ctx = _make_context() - with pytest.raises(ValueError, match="No models available"): - policy.route(ctx, []) - - def test_route_random_before_min_samples(self) -> None: - policy = GRPORouterPolicy(min_samples=5) - models = ["model-a", "model-b", "model-c"] - ctx = _make_context() - # Before any samples, should still return a model from the list - for _ in range(20): - result = policy.route(ctx, models) - assert result in models - - def test_add_sample_increments_count(self) -> None: - policy = GRPORouterPolicy() - assert policy.state.sample_counts.get("code", 0) == 0 - policy.add_sample("code", "model-a", 0.9) - assert policy.state.sample_counts["code"] == 1 - policy.add_sample("code", "model-b", 0.7) - assert policy.state.sample_counts["code"] == 2 - - def test_update_no_samples(self) -> None: - policy = GRPORouterPolicy() - result = policy.update() - assert result["updated"] is False - assert result["reason"] == "no samples" - - def test_update_with_samples(self) -> None: - policy = GRPORouterPolicy() - policy.add_sample("code", "model-a", 0.9) - policy.add_sample("code", "model-b", 0.3) - policy.add_sample("math", "model-a", 0.5) - policy.add_sample("math", "model-b", 0.8) - - result = policy.update() - assert result["updated"] is True - assert result["samples_processed"] == 4 - assert result["groups"] == 2 - assert result["updates_applied"] == 4 - assert result["total_updates"] == 1 - - def test_update_shifts_weights(self) -> None: - policy = GRPORouterPolicy(learning_rate=0.5) - # Add many high-reward samples for model-a on "code" - for _ in range(10): - policy.add_sample("code", "model-a", 0.95) - policy.add_sample("code", "model-b", 0.1) - - policy.update() - - # model-a should have higher weight for "code" than model-b - weight_a = policy.state.weights.get("model-a", {}).get("code", 0.0) - weight_b = policy.state.weights.get("model-b", {}).get("code", 0.0) - assert weight_a > weight_b - - def test_route_after_training(self) -> None: - policy = GRPORouterPolicy(learning_rate=1.0, min_samples=5, temperature=0.1) - models = ["model-a", "model-b"] - - # Train: model-a is much better for short queries - for _ in range(20): - policy.add_sample("short", "model-a", 0.95) - policy.add_sample("short", "model-b", 0.1) - policy.update() - - # Route 100 times — model-a should be selected significantly more often - ctx = _make_context("hi", query_length=5) - counts = {"model-a": 0, "model-b": 0} - for _ in range(100): - selected = policy.route(ctx, models) - counts[selected] += 1 - - # model-a should win the majority of the time - assert counts["model-a"] > 70, ( - f"model-a only selected {counts['model-a']}/100 times" - ) - - def test_reset(self) -> None: - policy = GRPORouterPolicy() - policy.add_sample("code", "model-a", 0.9) - policy.add_sample("code", "model-b", 0.3) - policy.update() - - assert policy.state.total_updates == 1 - assert len(policy.state.weights) > 0 - - policy.reset() - - assert policy.state.total_updates == 0 - assert policy.state.sample_counts.get("code", 0) == 0 - # After reset, weights should be empty (fresh defaultdict) - assert len(policy.state.weights) == 0 - - def test_registry_registration(self) -> None: - ensure_registered() - assert RouterPolicyRegistry.contains("grpo") - assert RouterPolicyRegistry.get("grpo") is GRPORouterPolicy - - def test_update_single_sample_group_skipped(self) -> None: - """Groups with only 1 sample are skipped (need >= 2 for comparison).""" - policy = GRPORouterPolicy() - policy.add_sample("code", "model-a", 0.9) - result = policy.update() - # update() still runs but the single-sample group is skipped - assert result["updated"] is True - assert result["updates_applied"] == 0 - - def test_multiple_updates_accumulate(self) -> None: - policy = GRPORouterPolicy(learning_rate=0.1) - for batch in range(3): - policy.add_sample("general", "model-a", 0.8) - policy.add_sample("general", "model-b", 0.2) - result = policy.update() - assert result["total_updates"] == batch + 1 - - assert policy.state.total_updates == 3 - - def test_grpo_sample_dataclass(self) -> None: - sample = GRPOSample(query_class="code", model="m1", reward=0.75) - assert sample.query_class == "code" - assert sample.model == "m1" - assert sample.reward == 0.75 - - def test_grpo_state_dataclass(self) -> None: - state = GRPOState() - assert state.total_updates == 0 - assert state.sample_counts["any"] == 0 - assert state.weights["any"]["class"] == 0.0 diff --git a/tests/learning/test_heuristic_policy.py b/tests/learning/test_heuristic_policy.py index 6392db03..b8e9170c 100644 --- a/tests/learning/test_heuristic_policy.py +++ b/tests/learning/test_heuristic_policy.py @@ -3,8 +3,8 @@ from __future__ import annotations from openjarvis.core.registry import RouterPolicyRegistry -from openjarvis.learning.heuristic_policy import ensure_registered -from openjarvis.learning.router import HeuristicRouter +from openjarvis.learning.routing.heuristic_policy import ensure_registered +from openjarvis.learning.routing.router import HeuristicRouter class TestHeuristicPolicy: diff --git a/tests/learning/test_heuristic_reward.py b/tests/learning/test_heuristic_reward.py index b8c5e34c..e4534eac 100644 --- a/tests/learning/test_heuristic_reward.py +++ b/tests/learning/test_heuristic_reward.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest from openjarvis.learning._stubs import RoutingContext -from openjarvis.learning.heuristic_reward import HeuristicRewardFunction +from openjarvis.learning.routing.heuristic_reward import HeuristicRewardFunction class TestHeuristicRewardFunction: diff --git a/tests/learning/test_icl_updater.py b/tests/learning/test_icl_updater.py deleted file mode 100644 index 0082f54c..00000000 --- a/tests/learning/test_icl_updater.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Tests for ICL updater policy — example extraction + skill discovery.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Optional - -from openjarvis.core.types import StepType, TraceStep -from openjarvis.learning.icl_updater import ICLUpdaterPolicy - - -@dataclass -class _MockTrace: - query: str = "" - model: str = "model-a" - outcome: str = "success" - feedback: Optional[float] = 0.8 - steps: list = field(default_factory=list) - total_latency_seconds: float = 1.0 - - -def _make_tool_step(tool_name: str) -> TraceStep: - return TraceStep( - step_type=StepType.TOOL_CALL, - timestamp=0.0, - input={"tool": tool_name}, - output={"result": "ok"}, - metadata={"tool_name": tool_name}, - ) - - -class _MockTraceStore: - def __init__(self, traces): - self._traces = traces - - def list_traces(self): - return self._traces - - -class TestICLUpdaterPolicy: - def test_empty_traces(self): - policy = ICLUpdaterPolicy() - store = _MockTraceStore([]) - result = policy.update(store) - assert result["examples"] == [] - assert result["skills"] == [] - - def test_extracts_examples(self): - traces = [ - _MockTrace( - query="What is 2+2?", - outcome="success", - feedback=0.9, - steps=[_make_tool_step("calculator")], - ), - ] - policy = ICLUpdaterPolicy(min_score=0.5) - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["examples"]) == 1 - assert result["examples"][0]["query"] == "What is 2+2?" - - def test_filters_low_score(self): - traces = [ - _MockTrace( - query="test", - outcome="success", - feedback=0.3, - steps=[_make_tool_step("calculator")], - ), - ] - policy = ICLUpdaterPolicy(min_score=0.7) - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["examples"]) == 0 - - def test_filters_failures(self): - traces = [ - _MockTrace( - outcome="failure", - feedback=0.9, - steps=[_make_tool_step("calculator")], - ), - ] - policy = ICLUpdaterPolicy() - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["examples"]) == 0 - - def test_discovers_skills(self): - # Create multiple traces with same tool sequence - traces = [ - _MockTrace( - query=f"Query {i}", - outcome="success", - feedback=0.9, - steps=[_make_tool_step("retrieval"), _make_tool_step("llm")], - ) - for i in range(5) - ] - policy = ICLUpdaterPolicy(min_score=0.5, min_skill_occurrences=3) - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["skills"]) > 0 - skill = result["skills"][0] - assert "retrieval" in skill["sequence"] - assert "llm" in skill["sequence"] - assert skill["occurrences"] >= 3 - - def test_max_examples(self): - traces = [ - _MockTrace( - query=f"Query {i}", - outcome="success", - feedback=0.9, - steps=[_make_tool_step("calculator")], - ) - for i in range(30) - ] - policy = ICLUpdaterPolicy(min_score=0.5, max_examples=5) - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["examples"]) <= 5 - - def test_is_agent_policy(self): - from openjarvis.learning._stubs import AgentLearningPolicy - assert issubclass(ICLUpdaterPolicy, AgentLearningPolicy) - assert ICLUpdaterPolicy.target == "agent" - - def test_examples_property(self): - policy = ICLUpdaterPolicy() - assert policy.examples == [] - - def test_skills_property(self): - policy = ICLUpdaterPolicy() - assert policy.skills == [] diff --git a/tests/learning/test_icl_updates.py b/tests/learning/test_icl_updates.py deleted file mode 100644 index d8339c2c..00000000 --- a/tests/learning/test_icl_updates.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Tests for the new ICLUpdaterPolicy versioned example database features.""" - -from __future__ import annotations - -from openjarvis.learning.icl_updater import ICLUpdaterPolicy - - -class TestICLUpdaterAddExample: - def test_add_example_quality_gate(self): - """Examples below the quality threshold should be rejected.""" - policy = ICLUpdaterPolicy(min_score=0.7) - accepted = policy.add_example( - query="low quality", - response="bad answer", - outcome=0.3, - ) - assert accepted is False - assert len(policy.example_db) == 0 - assert policy.version == 0 - - def test_add_example_stores(self): - """Examples above the threshold should be stored.""" - policy = ICLUpdaterPolicy(min_score=0.5) - accepted = policy.add_example( - query="What is 2+2?", - response="4", - outcome=0.9, - metadata={"source": "test"}, - ) - assert accepted is True - assert len(policy.example_db) == 1 - ex = policy.example_db[0] - assert ex["query"] == "What is 2+2?" - assert ex["response"] == "4" - assert ex["outcome"] == 0.9 - assert ex["metadata"] == {"source": "test"} - assert ex["version"] == 1 - - def test_add_example_at_threshold(self): - """An example exactly at the threshold should be accepted.""" - policy = ICLUpdaterPolicy(min_score=0.7) - accepted = policy.add_example( - query="borderline", - response="ok", - outcome=0.7, - ) - assert accepted is True - assert len(policy.example_db) == 1 - - -class TestICLUpdaterRollback: - def test_rollback(self): - """Rollback should remove examples added after the given version.""" - policy = ICLUpdaterPolicy(min_score=0.5) - policy.add_example("q1", "r1", 0.8) - policy.add_example("q2", "r2", 0.9) - v2 = policy.version - assert v2 == 2 - - policy.add_example("q3", "r3", 0.85) - policy.add_example("q4", "r4", 0.95) - assert len(policy.example_db) == 4 - assert policy.version == 4 - - policy.rollback(v2) - assert policy.version == 2 - assert len(policy.example_db) == 2 - queries = [ex["query"] for ex in policy.example_db] - assert "q1" in queries - assert "q2" in queries - assert "q3" not in queries - assert "q4" not in queries - - def test_rollback_to_zero(self): - """Rollback to version 0 should remove all examples.""" - policy = ICLUpdaterPolicy(min_score=0.5) - policy.add_example("q1", "r1", 0.8) - policy.add_example("q2", "r2", 0.9) - policy.rollback(0) - assert policy.version == 0 - assert len(policy.example_db) == 0 - - -class TestICLUpdaterGetExamples: - def test_get_examples(self): - """Should return examples sorted by outcome.""" - policy = ICLUpdaterPolicy(min_score=0.5) - policy.add_example("math q1", "r1", 0.7) - policy.add_example("math q2", "r2", 0.95) - policy.add_example("math q3", "r3", 0.8) - - results = policy.get_examples(top_k=2) - assert len(results) == 2 - assert results[0]["outcome"] >= results[1]["outcome"] - assert results[0]["outcome"] == 0.95 - - def test_get_examples_with_query_class(self): - """Filtering by query_class should narrow results.""" - policy = ICLUpdaterPolicy(min_score=0.5) - policy.add_example("math problem", "42", 0.9) - policy.add_example("code review", "looks good", 0.85) - policy.add_example("math equation", "x=5", 0.8) - - results = policy.get_examples(query_class="math", top_k=10) - assert len(results) == 2 - assert all("math" in ex["query"].lower() for ex in results) - - def test_get_examples_empty(self): - """No examples should return empty list.""" - policy = ICLUpdaterPolicy(min_score=0.5) - results = policy.get_examples(top_k=5) - assert results == [] - - -class TestICLUpdaterVersioning: - def test_version_increments(self): - """Each successful add should increment the version.""" - policy = ICLUpdaterPolicy(min_score=0.5) - assert policy.version == 0 - - policy.add_example("q1", "r1", 0.8) - assert policy.version == 1 - - policy.add_example("q2", "r2", 0.9) - assert policy.version == 2 - - # Rejected add should NOT increment version - policy.add_example("q3", "r3", 0.1) - assert policy.version == 2 - - def test_max_examples_limit(self): - """Adding beyond max_examples should trim the oldest entries.""" - policy = ICLUpdaterPolicy(min_score=0.5, max_examples=3) - - policy.add_example("q1", "r1", 0.8) - policy.add_example("q2", "r2", 0.85) - policy.add_example("q3", "r3", 0.9) - assert len(policy.example_db) == 3 - - # Adding a 4th should trim q1 (the oldest) - policy.add_example("q4", "r4", 0.95) - assert len(policy.example_db) == 3 - queries = [ex["query"] for ex in policy.example_db] - assert "q1" not in queries - assert "q4" in queries - assert "q2" in queries - assert "q3" in queries - - -class TestICLUpdaterAutoApply: - def test_auto_apply_default_false(self): - """auto_apply should default to False.""" - policy = ICLUpdaterPolicy() - assert policy._auto_apply is False - - def test_auto_apply_configurable(self): - """auto_apply can be set via constructor.""" - policy = ICLUpdaterPolicy(auto_apply=True) - assert policy._auto_apply is True - - -class TestICLUpdaterBackwardCompat: - def test_existing_update_still_works(self): - """The original update() method must still function.""" - from dataclasses import dataclass, field - from typing import Optional - - from openjarvis.core.types import StepType, TraceStep - - @dataclass - class _MockTrace: - query: str = "" - model: str = "model-a" - outcome: str = "success" - feedback: Optional[float] = 0.8 - steps: list = field(default_factory=list) - total_latency_seconds: float = 1.0 - - class _MockTraceStore: - def __init__(self, traces): - self._traces = traces - - def list_traces(self): - return self._traces - - step = TraceStep( - step_type=StepType.TOOL_CALL, - timestamp=0.0, - input={"tool": "calculator"}, - output={"result": "ok"}, - metadata={"tool_name": "calculator"}, - ) - traces = [ - _MockTrace( - query="What is 2+2?", - outcome="success", - feedback=0.9, - steps=[step], - ), - ] - policy = ICLUpdaterPolicy(min_score=0.5) - store = _MockTraceStore(traces) - result = policy.update(store) - assert len(result["examples"]) == 1 - assert result["examples"][0]["query"] == "What is 2+2?" diff --git a/tests/learning/test_learning_api.py b/tests/learning/test_learning_api.py index 6724af9b..2d2df3f1 100644 --- a/tests/learning/test_learning_api.py +++ b/tests/learning/test_learning_api.py @@ -35,36 +35,11 @@ def test_learning_stats_returns_200(): def test_learning_stats_has_all_sections(): - """Response must contain grpo, bandit, icl, and skill_discovery sections.""" + """Response must contain skill_discovery section.""" client = _client() data = client.get("/v1/learning/stats").json() - for section in ("grpo", "bandit", "icl", "skill_discovery"): - assert section in data, f"Missing section: {section}" - assert "available" in data[section], ( - f"Section '{section}' missing 'available' key" - ) - - -def test_learning_stats_grpo_fields(): - """When GRPO is available, verify expected stat fields are present.""" - client = _client() - data = client.get("/v1/learning/stats").json() - grpo = data["grpo"] - if grpo["available"]: - assert "total_updates" in grpo - assert "sample_counts" in grpo - assert "weight_count" in grpo - - -def test_learning_stats_icl_fields(): - """When ICL is available, verify expected stat fields are present.""" - client = _client() - data = client.get("/v1/learning/stats").json() - icl = data["icl"] - if icl["available"]: - assert "example_count" in icl - assert "example_db_count" in icl - assert "version" in icl + assert "skill_discovery" in data + assert "available" in data["skill_discovery"] # ---- /v1/learning/policy tests ---- diff --git a/tests/learning/test_router.py b/tests/learning/test_router.py index 5919c75a..5e242153 100644 --- a/tests/learning/test_router.py +++ b/tests/learning/test_router.py @@ -5,7 +5,7 @@ from __future__ import annotations from openjarvis.core.registry import ModelRegistry from openjarvis.core.types import ModelSpec from openjarvis.learning._stubs import RoutingContext -from openjarvis.learning.router import ( +from openjarvis.learning.routing.router import ( HeuristicRouter, build_routing_context, ) diff --git a/tests/learning/test_router_stubs.py b/tests/learning/test_router_stubs.py index 1820b9c0..91ef44d6 100644 --- a/tests/learning/test_router_stubs.py +++ b/tests/learning/test_router_stubs.py @@ -6,7 +6,7 @@ import pytest from openjarvis.core.types import RoutingContext from openjarvis.learning._stubs import QueryAnalyzer, RouterPolicy -from openjarvis.learning.router import DefaultQueryAnalyzer +from openjarvis.learning.routing.router import DefaultQueryAnalyzer class _DummyRouter(RouterPolicy): diff --git a/tests/learning/test_routing_models.py b/tests/learning/test_routing_models.py index cc707f69..309004d0 100644 --- a/tests/learning/test_routing_models.py +++ b/tests/learning/test_routing_models.py @@ -6,7 +6,7 @@ import pytest from openjarvis.intelligence.model_catalog import register_builtin_models from openjarvis.learning._stubs import RoutingContext -from openjarvis.learning.router import ( +from openjarvis.learning.routing.router import ( HeuristicRouter, build_routing_context, ) diff --git a/tests/learning/test_sft_policy.py b/tests/learning/test_sft_policy.py deleted file mode 100644 index ea81cec0..00000000 --- a/tests/learning/test_sft_policy.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Tests for SFT policy — learning from traces.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Optional - -from openjarvis.learning.sft_policy import SFTRouterPolicy - - -@dataclass -class _MockTrace: - query: str = "" - model: str = "model-a" - outcome: str = "success" - feedback: Optional[float] = 0.8 - steps: list = field(default_factory=list) - total_latency_seconds: float = 1.0 - - -class _MockTraceStore: - def __init__(self, traces): - self._traces = traces - - def list_traces(self): - return self._traces - - -class TestSFTRouterPolicy: - def test_empty_traces(self): - policy = SFTRouterPolicy() - store = _MockTraceStore([]) - result = policy.update(store) - assert result["updated"] is False - - def test_learns_from_traces(self): - traces = [ - _MockTrace( - query=f"def foo{i}(): pass", - model="code-model", outcome="success", - feedback=0.9, - ) - for i in range(6) - ] - policy = SFTRouterPolicy(min_samples=5) - store = _MockTraceStore(traces) - result = policy.update(store) - assert result["updated"] is True - assert "code" in result["policy_map"] - assert result["policy_map"]["code"] == "code-model" - - def test_min_samples_threshold(self): - traces = [ - _MockTrace(query="def foo(): pass", model="code-model", outcome="success") - for _ in range(3) - ] - policy = SFTRouterPolicy(min_samples=5) - store = _MockTraceStore(traces) - result = policy.update(store) - assert result["updated"] is False - - def test_classify_code(self): - assert SFTRouterPolicy._classify_query("def hello(): pass") == "code" - - def test_classify_math(self): - assert SFTRouterPolicy._classify_query("solve the integral") == "math" - - def test_classify_short(self): - assert SFTRouterPolicy._classify_query("hello world") == "short" - - def test_classify_general(self): - query = "tell me about " + " ".join(["something"] * 20) - assert SFTRouterPolicy._classify_query(query) == "general" - - def test_policy_map_property(self): - policy = SFTRouterPolicy() - assert policy.policy_map == {} - - def test_is_intelligence_policy(self): - from openjarvis.learning._stubs import IntelligenceLearningPolicy - assert issubclass(SFTRouterPolicy, IntelligenceLearningPolicy) - assert SFTRouterPolicy.target == "intelligence" diff --git a/tests/learning/test_skill_discovery.py b/tests/learning/test_skill_discovery.py index c168fc4a..7873fc4c 100644 --- a/tests/learning/test_skill_discovery.py +++ b/tests/learning/test_skill_discovery.py @@ -5,7 +5,7 @@ from __future__ import annotations from dataclasses import dataclass, field from typing import List -from openjarvis.learning.skill_discovery import SkillDiscovery +from openjarvis.learning.agents.skill_discovery import SkillDiscovery # --------------------------------------------------------------------------- # Helpers diff --git a/tests/learning/test_system_learning.py b/tests/learning/test_system_learning.py index 3f71856a..9588b098 100644 --- a/tests/learning/test_system_learning.py +++ b/tests/learning/test_system_learning.py @@ -30,9 +30,9 @@ class TestSystemLearningIntegration: config = LearningConfig() assert config.training_enabled is False assert config.training_schedule == "" - assert config.lora_rank == 16 - assert config.lora_alpha == 32 - assert config.min_sft_pairs == 50 + assert config.intelligence.sft.lora_rank == 16 + assert config.intelligence.sft.lora_alpha == 32 + assert config.intelligence.sft.min_pairs == 10 assert config.min_improvement == 0.02 def test_training_components_exported(self): diff --git a/tests/learning/test_trace_policy.py b/tests/learning/test_trace_policy.py deleted file mode 100644 index 50fe4587..00000000 --- a/tests/learning/test_trace_policy.py +++ /dev/null @@ -1,203 +0,0 @@ -"""Tests for the trace-driven router policy.""" - -from __future__ import annotations - -import time -from pathlib import Path - -from openjarvis.core.types import StepType, Trace, TraceStep -from openjarvis.learning._stubs import RoutingContext -from openjarvis.learning.trace_policy import TraceDrivenPolicy, classify_query -from openjarvis.traces.analyzer import TraceAnalyzer -from openjarvis.traces.store import TraceStore - - -def _make_trace( - query: str = "test", - model: str = "qwen3:8b", - outcome: str | None = "success", - feedback: float | None = 0.8, -) -> Trace: - now = time.time() - return Trace( - query=query, - agent="orchestrator", - model=model, - engine="ollama", - result="result", - outcome=outcome, - feedback=feedback, - started_at=now, - ended_at=now + 0.5, - total_tokens=100, - total_latency_seconds=0.5, - steps=[ - TraceStep( - step_type=StepType.GENERATE, - timestamp=now, - duration_seconds=0.5, - output={"tokens": 100}, - ), - ], - ) - - -class TestClassifyQuery: - def test_code(self) -> None: - assert classify_query("def hello(): pass") == "code" - assert classify_query("```python\nprint()```") == "code" - - def test_math(self) -> None: - assert classify_query("solve this equation for x") == "math" - assert classify_query("compute the integral") == "math" - - def test_short(self) -> None: - assert classify_query("hello") == "short" - assert classify_query("what time is it?") == "short" - - def test_long(self) -> None: - assert classify_query("a" * 501) == "long" - - def test_general(self) -> None: - q = "Tell me about the history of artificial intelligence research" - assert classify_query(q) == "general" - - -class TestTraceDrivenPolicy: - def test_fallback_no_traces(self) -> None: - policy = TraceDrivenPolicy(default_model="qwen3:8b") - ctx = RoutingContext(query="hello") - assert policy.select_model(ctx) == "qwen3:8b" - - def test_fallback_chain(self) -> None: - policy = TraceDrivenPolicy( - default_model="missing", - fallback_model="llama3:8b", - available_models=["llama3:8b"], - ) - ctx = RoutingContext(query="hello") - assert policy.select_model(ctx) == "llama3:8b" - - def test_fallback_first_available(self) -> None: - policy = TraceDrivenPolicy(available_models=["modelA", "modelB"]) - ctx = RoutingContext(query="hello") - assert policy.select_model(ctx) == "modelA" - - def test_update_from_traces(self, tmp_path: Path) -> None: - store = TraceStore(tmp_path / "test.db") - # Create traces: code queries succeed more with codestral - for _ in range(6): - store.save(_make_trace( - query="def foo(): pass", - model="codestral", - outcome="success", - feedback=0.9, - )) - for _ in range(6): - store.save(_make_trace( - query="def bar(): return 1", - model="qwen3:8b", - outcome="failure", - feedback=0.3, - )) - - analyzer = TraceAnalyzer(store) - policy = TraceDrivenPolicy( - analyzer=analyzer, - default_model="qwen3:8b", - ) - policy.min_samples = 3 - - result = policy.update_from_traces() - assert result["updated"] is True - - # Policy should now route code to codestral - ctx = RoutingContext(query="import os; def main(): pass") - assert policy.select_model(ctx) == "codestral" - store.close() - - def test_policy_map_readable(self, tmp_path: Path) -> None: - store = TraceStore(tmp_path / "test.db") - for _ in range(5): - store.save(_make_trace( - query="hello", model="small-model", - outcome="success", - )) - - analyzer = TraceAnalyzer(store) - policy = TraceDrivenPolicy(analyzer=analyzer, default_model="default") - policy.min_samples = 3 - policy.update_from_traces() - - pmap = policy.policy_map - assert isinstance(pmap, dict) - assert "short" in pmap - assert pmap["short"] == "small-model" - store.close() - - def test_respects_min_samples(self, tmp_path: Path) -> None: - store = TraceStore(tmp_path / "test.db") - # Only 2 traces — below threshold - store.save(_make_trace(query="hello", model="small", outcome="success")) - store.save(_make_trace(query="hi", model="small", outcome="success")) - - analyzer = TraceAnalyzer(store) - policy = TraceDrivenPolicy( - analyzer=analyzer, - default_model="default", - ) - policy.min_samples = 5 - policy.update_from_traces() - - ctx = RoutingContext(query="hey") - # Should fallback since not enough confidence - assert policy.select_model(ctx) == "default" - store.close() - - def test_respects_available_models(self, tmp_path: Path) -> None: - store = TraceStore(tmp_path / "test.db") - for _ in range(10): - store.save(_make_trace( - query="hello", model="unavailable-model", - outcome="success", - )) - - analyzer = TraceAnalyzer(store) - policy = TraceDrivenPolicy( - analyzer=analyzer, - available_models=["qwen3:8b", "llama3:8b"], - default_model="qwen3:8b", - ) - policy.min_samples = 3 - policy.update_from_traces() - - ctx = RoutingContext(query="hey") - # Should fallback since learned model not in available_models - assert policy.select_model(ctx) == "qwen3:8b" - store.close() - - def test_observe_online(self) -> None: - policy = TraceDrivenPolicy(default_model="default") - policy.min_samples = 3 - - # First observation creates the entry - policy.observe("hello", "fast-model", "success", 0.9) - assert policy.policy_map.get("short") == "fast-model" - - # Not enough samples yet for high confidence - ctx = RoutingContext(query="hi") - # Confidence is 1 < min_samples=3, so fallback - assert policy.select_model(ctx) == "default" - - def test_update_no_analyzer(self) -> None: - policy = TraceDrivenPolicy() - result = policy.update_from_traces() - assert result["error"] == "no analyzer configured" - - def test_update_empty_store(self, tmp_path: Path) -> None: - store = TraceStore(tmp_path / "test.db") - analyzer = TraceAnalyzer(store) - policy = TraceDrivenPolicy(analyzer=analyzer) - result = policy.update_from_traces() - assert result["updated"] is False - store.close() diff --git a/tests/test_orchestrator_learning/test_environment.py b/tests/test_orchestrator_learning/test_environment.py index d2e580c9..eee9b532 100644 --- a/tests/test_orchestrator_learning/test_environment.py +++ b/tests/test_orchestrator_learning/test_environment.py @@ -5,10 +5,10 @@ from __future__ import annotations import pytest from openjarvis.core.types import ToolResult -from openjarvis.learning.orchestrator.environment import ( +from openjarvis.learning.intelligence.orchestrator.environment import ( OrchestratorEnvironment, ) -from openjarvis.learning.orchestrator.types import OrchestratorAction +from openjarvis.learning.intelligence.orchestrator.types import OrchestratorAction from openjarvis.tools._stubs import BaseTool, ToolSpec # -- Mock tool --------------------------------------------------------------- diff --git a/tests/test_orchestrator_learning/test_grpo_trainer.py b/tests/test_orchestrator_learning/test_grpo_trainer.py index cb4a2db0..066e5c24 100644 --- a/tests/test_orchestrator_learning/test_grpo_trainer.py +++ b/tests/test_orchestrator_learning/test_grpo_trainer.py @@ -2,10 +2,10 @@ from __future__ import annotations -from openjarvis.learning.orchestrator.grpo_trainer import ( +from openjarvis.learning.intelligence.orchestrator.grpo_trainer import ( OrchestratorGRPOConfig, ) -from openjarvis.learning.orchestrator.types import Episode +from openjarvis.learning.intelligence.orchestrator.types import Episode class TestOrchestratorGRPOConfig: @@ -74,7 +74,7 @@ class TestGroupAdvantageNormalization: class TestRewardIntegration: def test_episode_reward(self): - from openjarvis.learning.orchestrator.reward import ( + from openjarvis.learning.intelligence.orchestrator.reward import ( MultiObjectiveReward, Normalizers, RewardWeights, @@ -97,6 +97,6 @@ class TestRewardIntegration: class TestGRPORegistration: def test_registered_in_learning_registry(self): - import openjarvis.learning.orchestrator.grpo_trainer # noqa: F401 + import openjarvis.learning.intelligence.orchestrator.grpo_trainer # noqa: F401 from openjarvis.core.registry import LearningRegistry assert LearningRegistry.contains("orchestrator_grpo") diff --git a/tests/test_orchestrator_learning/test_policy_model.py b/tests/test_orchestrator_learning/test_policy_model.py index 4bdbc741..59890251 100644 --- a/tests/test_orchestrator_learning/test_policy_model.py +++ b/tests/test_orchestrator_learning/test_policy_model.py @@ -4,10 +4,10 @@ from __future__ import annotations import pytest -from openjarvis.learning.orchestrator.policy_model import ( +from openjarvis.learning.intelligence.orchestrator.policy_model import ( OrchestratorPolicyModel, ) -from openjarvis.learning.orchestrator.types import ( +from openjarvis.learning.intelligence.orchestrator.types import ( EpisodeState, OrchestratorAction, OrchestratorObservation, diff --git a/tests/test_orchestrator_learning/test_prompt_registry.py b/tests/test_orchestrator_learning/test_prompt_registry.py index 41af766b..30028ebc 100644 --- a/tests/test_orchestrator_learning/test_prompt_registry.py +++ b/tests/test_orchestrator_learning/test_prompt_registry.py @@ -2,7 +2,7 @@ from __future__ import annotations -from openjarvis.learning.orchestrator.prompt_registry import ( +from openjarvis.learning.intelligence.orchestrator.prompt_registry import ( TOOL_DESCRIPTIONS, build_system_prompt, ) diff --git a/tests/test_orchestrator_learning/test_reward.py b/tests/test_orchestrator_learning/test_reward.py index a4c03115..80cfda8d 100644 --- a/tests/test_orchestrator_learning/test_reward.py +++ b/tests/test_orchestrator_learning/test_reward.py @@ -4,13 +4,13 @@ from __future__ import annotations import pytest -from openjarvis.learning.orchestrator.reward import ( +from openjarvis.learning.intelligence.orchestrator.reward import ( AdaptiveRewardWeights, MultiObjectiveReward, Normalizers, RewardWeights, ) -from openjarvis.learning.orchestrator.types import ( +from openjarvis.learning.intelligence.orchestrator.types import ( Episode, ) diff --git a/tests/test_orchestrator_learning/test_sft_trainer.py b/tests/test_orchestrator_learning/test_sft_trainer.py index 0e5a2361..b9768eab 100644 --- a/tests/test_orchestrator_learning/test_sft_trainer.py +++ b/tests/test_orchestrator_learning/test_sft_trainer.py @@ -4,7 +4,7 @@ from __future__ import annotations from unittest.mock import MagicMock -from openjarvis.learning.orchestrator.sft_trainer import ( +from openjarvis.learning.intelligence.orchestrator.sft_trainer import ( OrchestratorSFTConfig, OrchestratorSFTDataset, ) @@ -126,6 +126,6 @@ class TestOrchestratorSFTDataset: class TestSFTRegistration: def test_registered_in_learning_registry(self): # Import to trigger registration - import openjarvis.learning.orchestrator.sft_trainer # noqa: F401 + import openjarvis.learning.intelligence.orchestrator.sft_trainer # noqa: F401 from openjarvis.core.registry import LearningRegistry assert LearningRegistry.contains("orchestrator_sft") diff --git a/tests/test_orchestrator_learning/test_types.py b/tests/test_orchestrator_learning/test_types.py index f1079626..63930f60 100644 --- a/tests/test_orchestrator_learning/test_types.py +++ b/tests/test_orchestrator_learning/test_types.py @@ -4,7 +4,7 @@ from __future__ import annotations import pytest -from openjarvis.learning.orchestrator.types import ( +from openjarvis.learning.intelligence.orchestrator.types import ( Episode, EpisodeState, OrchestratorAction,