mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
feat: Algolia DocSearch + learning subsystem reorganization (#43)
* chore: create learning subdirectory structure (routing, agents, intelligence) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * refactor: move routing files to learning/routing/ subdirectory Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: create LearnedRouterPolicy merging trace-driven + SFT routing Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * refactor: move agent_evolver and skill_discovery to learning/agents/ Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor: move learning/orchestrator to learning/intelligence/orchestrator Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor: delete removed learning policies, rewrite __init__.py, clean up api_routes Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: add SFT/GRPO/DSPy/GEPA config dataclasses, update LearningConfig Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: add general-purpose SFT trainer (intelligence/sft_trainer.py) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * feat: add general-purpose GRPO trainer (intelligence/grpo_trainer.py) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: add DSPy agent optimizer (agents/dspy_optimizer.py) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: add GEPA agent optimizer (agents/gepa_optimizer.py) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: add learning-dspy and learning-gepa optional dependency extras Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: update integration test to check for learned policy instead of grpo Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * fix: import SFT/GRPO trainers in intelligence/__init__.py for registry Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * chore: remove .md file changes from PR Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * chore: restore search boost frontmatter for key docs pages Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jon Saad-Falcon
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
2f553afb21
commit
05f2c02131
@@ -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.
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
---
|
||||
title: Installation
|
||||
description: Get OpenJarvis running — browser app, desktop app, CLI, or Python SDK
|
||||
search:
|
||||
boost: 3
|
||||
---
|
||||
|
||||
# Installation
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
---
|
||||
title: Quick Start
|
||||
description: Get up and running with OpenJarvis in minutes
|
||||
search:
|
||||
boost: 3
|
||||
---
|
||||
|
||||
# Quick Start
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
---
|
||||
title: OpenJarvis
|
||||
description: Personal AI, On Personal Devices
|
||||
search:
|
||||
boost: 2
|
||||
hide:
|
||||
- navigation
|
||||
---
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
})();
|
||||
@@ -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 %}
|
||||
<script>
|
||||
window.__DOCSEARCH_CONFIG__ = {
|
||||
appId: "{{ config.extra.algolia.app_id }}",
|
||||
apiKey: "{{ config.extra.algolia.api_key }}",
|
||||
indexName: "{{ config.extra.algolia.index_name }}"
|
||||
};
|
||||
</script>
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
+10
-1
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Agent learning — config optimization via DSPy, GEPA, and evolution."""
|
||||
|
||||
from __future__ import annotations
|
||||
+1
-1
@@ -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
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
+7
-7
@@ -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,
|
||||
+1
-1
@@ -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,
|
||||
+3
-3
@@ -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,
|
||||
)
|
||||
|
||||
+1
-1
@@ -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,
|
||||
+1
-1
@@ -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
|
||||
+1
-1
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Routing policies — model selection based on query characteristics."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -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"]
|
||||
+1
-1
@@ -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:
|
||||
+63
-65
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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"]
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 == []
|
||||
@@ -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?"
|
||||
@@ -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 ----
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
@@ -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 ---------------------------------------------------------------
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user