mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-30 10:52:15 +00:00
Add orchestrator training, channels, LiteLLM engine, and simplify learning taxonomy
Major changes across parallel sessions: - Add orchestrator SFT & GRPO training subpackage (learning/orchestrator/) with episode types, multi-objective reward, prompt registry, policy model, RL environment, and registered learning policies - Add structured THOUGHT/TOOL/INPUT/FINAL_ANSWER mode to OrchestratorAgent - Add 15 channel backends (Discord, Slack, Telegram, Email, Webhook, IRC, Matrix, Teams, WhatsApp, Signal, Mattermost, BlueBubbles, Feishu, Google Chat, Webchat) with channel tools and config - Add LiteLLM engine backend for unified LLM provider access - Add RLM agent and REPL tool - Remove ToolLearningPolicy — learning taxonomy now only targets Intelligence (LM weights/routing) and Agents (logic/ICL/tool strategies) - Rename SFTPolicy to SFTRouterPolicy (backward-compat alias kept) - Remove OpenClaw agent infrastructure (openclaw*.py, openclaw_bridge.py) - Fix async streaming tests (asyncio.run vs deprecated get_event_loop) - Fix server channel route tests (pytest.importorskip for optional fastapi) - Track uv.lock for reproducibility - Update CLAUDE.md and docs to reflect all changes 1676 tests pass, 37 skipped. 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
a0961633a3
commit
8d538cd1b0
+3
-1
@@ -27,7 +27,9 @@ env/
|
||||
htmlcov/
|
||||
|
||||
# uv
|
||||
uv.lock
|
||||
# Note: uv.lock is intentionally tracked — it pins all transitive deps for reproducibility.
|
||||
# Uncomment the line below only if you deliberately choose not to commit the lock file.
|
||||
# uv.lock
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
|
||||
@@ -83,7 +83,7 @@ OpenJarvis is a research framework for on-device AI organized around **five comp
|
||||
- **Storage backends** (`tools/storage/`): SQLite/FTS5 (default), FAISS, ColBERTv2, BM25, Hybrid (RRF fusion). All implement `MemoryBackend` ABC with `store()`, `retrieve()`, `delete()`, `clear()`. Canonical import: `from openjarvis.tools.storage.sqlite import SQLiteMemory`. Backward-compat shims in `memory/` still work.
|
||||
- **MCP adapter** (`tools/mcp_adapter.py`): `MCPToolAdapter` wraps external MCP server tools as native `BaseTool` instances. `MCPToolProvider` discovers tools from an MCP server.
|
||||
- **MCP server** (`mcp/server.py`): Exposes all built-in tools via JSON-RPC `tools/list` + `tools/call` (MCP spec 2025-11-25). Any MCP client (Claude, GPT, etc.) can discover and use OpenJarvis tools.
|
||||
5. **Learning** (`src/openjarvis/learning/`) — Structured learning system with per-pillar policies. `LearningPolicy` ABC taxonomy: `IntelligenceLearningPolicy` (updates model routing), `AgentLearningPolicy` (updates agent logic), `ToolLearningPolicy` (updates tool usage). Implementations: `SFTPolicy` (learns query→model mapping from traces), `AgentAdvisorPolicy` (LM-guided agent restructuring), `ICLUpdaterPolicy` (in-context example + skill discovery). Router policies: `HeuristicRouter` (registered as "heuristic"), `TraceDrivenPolicy` (registered as "learned"), `GRPORouterPolicy` (stub, registered as "grpo"). `HeuristicRewardFunction` scores inference results on latency/cost/efficiency.
|
||||
5. **Learning** (`src/openjarvis/learning/`) — Structured learning system with per-pillar policies. `LearningPolicy` ABC taxonomy: `IntelligenceLearningPolicy` (updates model routing), `AgentLearningPolicy` (updates agent logic). Implementations: `SFTRouterPolicy` (learns query→model mapping from traces; backward-compat alias `SFTPolicy`), `AgentAdvisorPolicy` (LM-guided agent restructuring), `ICLUpdaterPolicy` (in-context example + skill discovery, registered as `AgentLearningPolicy`). Router policies: `HeuristicRouter` (registered as "heuristic"), `TraceDrivenPolicy` (registered as "learned"), `GRPORouterPolicy` (stub, registered as "grpo"). `HeuristicRewardFunction` scores inference results on latency/cost/efficiency. Orchestrator training subpackage (`learning/orchestrator/`) provides SFT and GRPO pipelines for structured-mode OrchestratorAgent training.
|
||||
|
||||
### Cross-cutting: Traces
|
||||
|
||||
|
||||
@@ -39,8 +39,7 @@ The learning system defines a hierarchy of learning policy ABCs:
|
||||
|
||||
- **`LearningPolicy`** -- base ABC for all learning policies
|
||||
- **`IntelligenceLearningPolicy`** -- specialization for model routing decisions
|
||||
- **`AgentLearningPolicy`** -- specialization for agent behavior advice
|
||||
- **`ToolLearningPolicy`** -- specialization for tool selection/configuration
|
||||
- **`AgentLearningPolicy`** -- specialization for agent behavior advice (ICL examples, tool-use strategies)
|
||||
|
||||
---
|
||||
|
||||
@@ -59,9 +58,9 @@ The learning system defines a hierarchy of learning policy ABCs:
|
||||
options:
|
||||
show_source: true
|
||||
|
||||
### SFTPolicy
|
||||
### SFTRouterPolicy
|
||||
|
||||
::: openjarvis.learning.sft_policy.SFTPolicy
|
||||
::: openjarvis.learning.sft_policy.SFTRouterPolicy
|
||||
options:
|
||||
show_source: true
|
||||
members_order: source
|
||||
|
||||
@@ -216,6 +216,6 @@ ctx = build_routing_context("```python\ndef hello():\n pass\n```")
|
||||
|
||||
## Integration with Learning
|
||||
|
||||
The `HeuristicRouter` implements the `RouterPolicy` ABC (now defined in `intelligence/_stubs.py`), which means it can be swapped out for a `TraceDrivenPolicy`, `SFTPolicy`, or any other policy via the `RouterPolicyRegistry`. See the [Learning & Traces](learning.md) documentation for details on how trace-driven routing and the broader `LearningPolicy` taxonomy work.
|
||||
The `HeuristicRouter` implements the `RouterPolicy` ABC (now defined in `intelligence/_stubs.py`), which means it can be swapped out for a `TraceDrivenPolicy`, `SFTRouterPolicy`, or any other policy via the `RouterPolicyRegistry`. See the [Learning & Traces](learning.md) documentation for details on how trace-driven routing and the broader `LearningPolicy` taxonomy work.
|
||||
|
||||
The router is registered as `"heuristic"` in the `RouterPolicyRegistry` and is the default routing policy. Users can switch policies via the `--router` CLI flag or the `learning.default_policy` config setting.
|
||||
|
||||
@@ -6,13 +6,12 @@ The Learning system is a **cross-cutting concern** that connects all four pillar
|
||||
|
||||
## LearningPolicy ABC Taxonomy
|
||||
|
||||
The learning system defines a hierarchy of learning policy ABCs. The base `LearningPolicy` ABC is specialized into three sub-ABCs corresponding to the three learnable concerns:
|
||||
The learning system defines a hierarchy of learning policy ABCs. The base `LearningPolicy` ABC is specialized into two sub-ABCs corresponding to the two learnable concerns:
|
||||
|
||||
| ABC | Concern | Description |
|
||||
|-----|---------|-------------|
|
||||
| `IntelligenceLearningPolicy` | Model routing | Determines which model handles a query (replaces the legacy `RouterPolicy`) |
|
||||
| `AgentLearningPolicy` | Agent behavior | Advises on agent strategy (e.g., tool selection, turn limits) |
|
||||
| `ToolLearningPolicy` | Tool selection | Recommends tool configurations based on query characteristics |
|
||||
| `AgentLearningPolicy` | Agent behavior | Advises on agent strategy (e.g., ICL examples, tool selection, turn limits) |
|
||||
|
||||
All learning policies are registered in the `LearningRegistry` (in `core/registry.py`).
|
||||
|
||||
@@ -66,14 +65,14 @@ The system ships with these router policies:
|
||||
| `heuristic` | `HeuristicRouter` | Active | Rule-based routing with 6 priority rules |
|
||||
| `learned` | `TraceDrivenPolicy` | Active | Learns from trace outcomes |
|
||||
| `grpo` | `GRPORouterPolicy` | Stub | Placeholder for future RL training |
|
||||
| `sft` | `SFTPolicy` | Active | Supervised fine-tuning policy (learns from labeled traces) |
|
||||
| `sft` | `SFTRouterPolicy` | Active | Trace-driven routing policy (learns query→model mapping); `SFTPolicy` is a backward-compat alias |
|
||||
|
||||
And these additional learning policies (registered in `LearningRegistry`):
|
||||
|
||||
| Registry Key | Policy Class | Taxonomy | Description |
|
||||
|-------------|-------------|----------|-------------|
|
||||
| `agent_advisor` | `AgentAdvisorPolicy` | `AgentLearningPolicy` | Advises on agent strategy based on trace patterns |
|
||||
| `icl_updater` | `ICLUpdaterPolicy` | `ToolLearningPolicy` | In-context learning updater for tool selection |
|
||||
| `icl_updater` | `ICLUpdaterPolicy` | `AgentLearningPolicy` | In-context learning updater — discovers ICL examples and multi-tool skills from traces |
|
||||
|
||||
Users select a policy via `config.toml` or the `--router` CLI flag:
|
||||
|
||||
@@ -193,11 +192,13 @@ The online update uses a conservative strategy: it only switches the preferred m
|
||||
|
||||
---
|
||||
|
||||
## SFTPolicy (Supervised Fine-Tuning)
|
||||
## SFTRouterPolicy (Trace-Driven Router)
|
||||
|
||||
The `SFTPolicy` (in `learning/sft_policy.py`) is an `IntelligenceLearningPolicy` that learns routing decisions from labeled trace data using supervised fine-tuning principles. Unlike `TraceDrivenPolicy` which uses online aggregation, `SFTPolicy` trains a mapping from query features to model choices based on curated, high-quality trace examples.
|
||||
The `SFTRouterPolicy` (in `learning/sft_policy.py`) is an `IntelligenceLearningPolicy` that learns routing decisions from historical traces. It analyzes trace outcomes, groups by query class (code, math, short, long, general), and builds a `query_class → model` mapping from the highest-scoring model per class. A backward-compatible alias `SFTPolicy = SFTRouterPolicy` is provided for code that used the old name.
|
||||
|
||||
```python
|
||||
from openjarvis.learning.sft_policy import SFTRouterPolicy
|
||||
# or via the backward-compat alias:
|
||||
from openjarvis.learning.sft_policy import SFTPolicy
|
||||
```
|
||||
|
||||
@@ -215,7 +216,7 @@ from openjarvis.learning.agent_advisor import AgentAdvisorPolicy
|
||||
|
||||
## ICLUpdaterPolicy
|
||||
|
||||
The `ICLUpdaterPolicy` (in `learning/icl_updater.py`) is a `ToolLearningPolicy` that uses in-context learning to update tool selection and configuration. It analyzes recent tool-call traces to recommend which tools should be enabled for different query types.
|
||||
The `ICLUpdaterPolicy` (in `learning/icl_updater.py`) is an `AgentLearningPolicy` that uses in-context learning to discover reusable examples and multi-tool skill sequences from traces. It analyzes successful tool-call patterns to recommend ICL examples and skill libraries that update agent behavior.
|
||||
|
||||
```python
|
||||
from openjarvis.learning.icl_updater import ICLUpdaterPolicy
|
||||
|
||||
+17
-2
@@ -33,6 +33,7 @@ inference-cloud = [
|
||||
inference-google = [
|
||||
"google-genai>=1.0",
|
||||
]
|
||||
inference-litellm = ["litellm>=1.40"]
|
||||
tools-search = [
|
||||
"tavily-python>=0.3",
|
||||
]
|
||||
@@ -54,8 +55,22 @@ server = [
|
||||
]
|
||||
agents = []
|
||||
learning = []
|
||||
openclaw = []
|
||||
channels = ["websockets>=12.0"]
|
||||
orchestrator-training = ["torch>=2.0", "transformers>=4.40"]
|
||||
channel-telegram = ["python-telegram-bot>=21.0"]
|
||||
channel-discord = ["discord.py>=2.3"]
|
||||
channel-slack = ["slack-sdk>=3.27"]
|
||||
channel-webhook = []
|
||||
channel-email = []
|
||||
channel-whatsapp = []
|
||||
channel-signal = []
|
||||
channel-google-chat = []
|
||||
channel-irc = []
|
||||
channel-webchat = []
|
||||
channel-teams = []
|
||||
channel-matrix = []
|
||||
channel-mattermost = []
|
||||
channel-feishu = []
|
||||
channel-bluebubbles = []
|
||||
docs = [
|
||||
"mkdocs>=1.6",
|
||||
"mkdocs-material>=9.5",
|
||||
|
||||
@@ -20,11 +20,6 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.agents.openclaw # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.agents.react # noqa: F401
|
||||
except ImportError:
|
||||
@@ -35,4 +30,9 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.agents.rlm # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__all__ = ["AgentContext", "AgentResult", "BaseAgent"]
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""ABC for agent implementations.
|
||||
|
||||
Adapted from IPW's ``BaseAgent`` at ``src/agents/base.py``.
|
||||
Phase 3 will provide concrete implementations (SimpleAgent, OpenClawAgent, etc.).
|
||||
Phase 3 will provide concrete implementations (SimpleAgent, OrchestratorAgent, etc.).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -1,181 +0,0 @@
|
||||
"""OpenClawAgent — wraps the OpenClaw Pi agent via HTTP or subprocess transport."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Optional
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
|
||||
from openjarvis.agents.openclaw_protocol import MessageType, ProtocolMessage
|
||||
from openjarvis.agents.openclaw_transport import (
|
||||
HttpTransport,
|
||||
OpenClawTransport,
|
||||
SubprocessTransport,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
from openjarvis.core.types import ToolCall, ToolResult
|
||||
|
||||
|
||||
@AgentRegistry.register("openclaw")
|
||||
class OpenClawAgent(BaseAgent):
|
||||
"""Wraps the OpenClaw Pi agent via HTTP or subprocess transport.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
engine:
|
||||
Inference engine (used as fallback or for provider plugin).
|
||||
model:
|
||||
Model identifier.
|
||||
transport:
|
||||
Optional pre-configured transport. If None, one is created based on *mode*.
|
||||
mode:
|
||||
Transport mode: ``"http"`` (default) or ``"subprocess"``.
|
||||
bus:
|
||||
Optional event bus for telemetry.
|
||||
"""
|
||||
|
||||
agent_id = "openclaw"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine: Any = None,
|
||||
model: str = "",
|
||||
*,
|
||||
transport: Optional[OpenClawTransport] = None,
|
||||
mode: str = "http",
|
||||
bus: Optional[EventBus] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self._engine = engine
|
||||
self._model = model
|
||||
self._bus = bus
|
||||
|
||||
if transport is not None:
|
||||
self._transport = transport
|
||||
elif mode == "http":
|
||||
self._transport = HttpTransport()
|
||||
elif mode == "subprocess":
|
||||
self._transport = SubprocessTransport()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown OpenClaw mode: {mode!r}. "
|
||||
"Use 'http' or 'subprocess'."
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
"""Send a query through the OpenClaw transport and handle tool calls."""
|
||||
if self._bus:
|
||||
self._bus.publish(EventType.AGENT_TURN_START, {"agent": self.agent_id})
|
||||
|
||||
# Check transport health
|
||||
if not self._transport.health():
|
||||
raise RuntimeError(
|
||||
"OpenClaw transport is not healthy. "
|
||||
"Ensure the OpenClaw server is running."
|
||||
)
|
||||
|
||||
# Build and send query message
|
||||
query_msg = ProtocolMessage(
|
||||
type=MessageType.QUERY,
|
||||
content=input,
|
||||
metadata={"model": self._model},
|
||||
)
|
||||
|
||||
tool_results: list[ToolResult] = []
|
||||
turns = 0
|
||||
max_turns = 10
|
||||
|
||||
response = self._transport.send(query_msg)
|
||||
turns += 1
|
||||
|
||||
# Handle tool-call loop
|
||||
while response.type == MessageType.TOOL_CALL and turns < max_turns:
|
||||
if self._bus:
|
||||
self._bus.publish(EventType.TOOL_CALL_START, {
|
||||
"tool": response.tool_name,
|
||||
"arguments": response.tool_args,
|
||||
})
|
||||
|
||||
# Execute tool locally
|
||||
tool_result = self._execute_tool(
|
||||
response.tool_name or "",
|
||||
response.tool_args or {},
|
||||
)
|
||||
tool_results.append(tool_result)
|
||||
|
||||
if self._bus:
|
||||
self._bus.publish(EventType.TOOL_CALL_END, {
|
||||
"tool": response.tool_name,
|
||||
"success": tool_result.success,
|
||||
})
|
||||
|
||||
# Send tool result back
|
||||
result_msg = ProtocolMessage(
|
||||
type=MessageType.TOOL_RESULT,
|
||||
tool_name=response.tool_name,
|
||||
tool_result=tool_result.content,
|
||||
metadata={"tool_call_id": response.id},
|
||||
)
|
||||
response = self._transport.send(result_msg)
|
||||
turns += 1
|
||||
|
||||
# Handle error response
|
||||
if response.type == MessageType.ERROR:
|
||||
content = (
|
||||
response.error or response.content
|
||||
or "Unknown error from OpenClaw"
|
||||
)
|
||||
else:
|
||||
content = response.content
|
||||
|
||||
if self._bus:
|
||||
self._bus.publish(EventType.AGENT_TURN_END, {
|
||||
"agent": self.agent_id,
|
||||
"turns": turns,
|
||||
})
|
||||
|
||||
return AgentResult(
|
||||
content=content,
|
||||
tool_results=tool_results,
|
||||
turns=turns,
|
||||
)
|
||||
|
||||
def _execute_tool(self, tool_name: str, tool_args: dict) -> ToolResult:
|
||||
"""Dispatch a tool call to the local tool executor."""
|
||||
try:
|
||||
from openjarvis.core.registry import ToolRegistry
|
||||
|
||||
if not ToolRegistry.contains(tool_name):
|
||||
return ToolResult(
|
||||
tool_name=tool_name,
|
||||
content=f"Unknown tool: {tool_name}",
|
||||
success=False,
|
||||
)
|
||||
|
||||
tool_cls = ToolRegistry.get(tool_name)
|
||||
tool = tool_cls() if callable(tool_cls) else tool_cls
|
||||
tc = ToolCall(
|
||||
id="openclaw-tool",
|
||||
name=tool_name,
|
||||
arguments=json.dumps(tool_args),
|
||||
)
|
||||
|
||||
from openjarvis.tools._stubs import ToolExecutor
|
||||
|
||||
executor = ToolExecutor([tool], bus=self._bus)
|
||||
return executor.execute(tc)
|
||||
except Exception as exc:
|
||||
return ToolResult(
|
||||
tool_name=tool_name,
|
||||
content=f"Tool execution error: {exc}",
|
||||
success=False,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["OpenClawAgent"]
|
||||
@@ -1,86 +0,0 @@
|
||||
"""OpenClaw plugin skeleton — wraps OpenJarvis as an OpenClaw provider."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
|
||||
class ProviderPlugin:
|
||||
"""Wraps OpenJarvis as an OpenClaw ProviderPlugin.
|
||||
|
||||
Implements the provider interface expected by the OpenClaw framework:
|
||||
``generate()`` for inference and ``list_models()`` for model discovery.
|
||||
"""
|
||||
|
||||
def __init__(self, engine: Any = None, model: str = "", bus: Any = None) -> None:
|
||||
self._engine = engine
|
||||
self._model = model
|
||||
self._bus = bus
|
||||
|
||||
def generate(self, prompt: str, **kwargs: Any) -> Dict[str, Any]:
|
||||
"""Generate a response using the wrapped OpenJarvis engine."""
|
||||
if self._engine is None:
|
||||
raise RuntimeError("No engine configured for ProviderPlugin")
|
||||
|
||||
from openjarvis.core.types import Message, Role
|
||||
|
||||
messages = [Message(role=Role.USER, content=prompt)]
|
||||
if self._bus:
|
||||
from openjarvis.telemetry.wrapper import instrumented_generate
|
||||
|
||||
return instrumented_generate(
|
||||
self._engine, messages, model=self._model,
|
||||
bus=self._bus, **kwargs,
|
||||
)
|
||||
return self._engine.generate(messages, model=self._model, **kwargs)
|
||||
|
||||
def list_models(self) -> List[str]:
|
||||
"""List available models from the wrapped engine."""
|
||||
if self._engine is None:
|
||||
return []
|
||||
return self._engine.list_models()
|
||||
|
||||
|
||||
class MemorySearchManager:
|
||||
"""OpenClaw search/sync/status interface backed by memory."""
|
||||
|
||||
def __init__(self, backend: Any = None) -> None:
|
||||
self._backend = backend
|
||||
|
||||
def search(self, query: str, *, top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""Search memory for relevant documents."""
|
||||
if self._backend is None:
|
||||
return []
|
||||
results = self._backend.retrieve(query, top_k=top_k)
|
||||
return [
|
||||
{"content": r.content, "score": r.score, "source": r.source}
|
||||
for r in results
|
||||
]
|
||||
|
||||
def sync(self) -> Dict[str, Any]:
|
||||
"""Synchronize memory state (no-op for now)."""
|
||||
return {"status": "ok", "synced": True}
|
||||
|
||||
def status(self) -> Dict[str, Any]:
|
||||
"""Return the memory backend status."""
|
||||
if self._backend is None:
|
||||
return {"available": False}
|
||||
bid = getattr(self._backend, "backend_id", "unknown")
|
||||
return {"available": True, "backend": bid}
|
||||
|
||||
|
||||
def register() -> Dict[str, Any]:
|
||||
"""Entry point for OpenClaw plugin registration.
|
||||
|
||||
Returns a dict describing the plugin capabilities.
|
||||
"""
|
||||
return {
|
||||
"name": "openjarvis",
|
||||
"version": "1.0.0",
|
||||
"capabilities": ["inference", "memory"],
|
||||
"provider_class": ProviderPlugin,
|
||||
"memory_class": MemorySearchManager,
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["MemorySearchManager", "ProviderPlugin", "register"]
|
||||
@@ -1,92 +0,0 @@
|
||||
"""OpenClaw wire protocol — JSON-line message serialization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
class MessageType(str, Enum):
|
||||
"""Wire message types for OpenClaw communication."""
|
||||
|
||||
QUERY = "query"
|
||||
RESPONSE = "response"
|
||||
TOOL_CALL = "tool_call"
|
||||
TOOL_RESULT = "tool_result"
|
||||
ERROR = "error"
|
||||
HEALTH = "health"
|
||||
HEALTH_OK = "health_ok"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ProtocolMessage:
|
||||
"""A single message in the OpenClaw protocol."""
|
||||
|
||||
type: MessageType
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
content: str = ""
|
||||
tool_name: Optional[str] = None
|
||||
tool_args: Optional[Dict[str, Any]] = None
|
||||
tool_result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def serialize(msg: ProtocolMessage) -> str:
|
||||
"""Serialize a ProtocolMessage to a JSON line."""
|
||||
obj: Dict[str, Any] = {
|
||||
"type": msg.type.value,
|
||||
"id": msg.id,
|
||||
"content": msg.content,
|
||||
}
|
||||
if msg.tool_name is not None:
|
||||
obj["tool_name"] = msg.tool_name
|
||||
if msg.tool_args is not None:
|
||||
obj["tool_args"] = msg.tool_args
|
||||
if msg.tool_result is not None:
|
||||
obj["tool_result"] = msg.tool_result
|
||||
if msg.error is not None:
|
||||
obj["error"] = msg.error
|
||||
if msg.metadata:
|
||||
obj["metadata"] = msg.metadata
|
||||
return json.dumps(obj)
|
||||
|
||||
|
||||
def deserialize(line: str) -> ProtocolMessage:
|
||||
"""Deserialize a JSON line into a ProtocolMessage.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If the JSON is invalid or the message type is unknown.
|
||||
"""
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"Invalid JSON: {exc}") from exc
|
||||
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError("Expected a JSON object")
|
||||
|
||||
raw_type = obj.get("type", "")
|
||||
try:
|
||||
msg_type = MessageType(raw_type)
|
||||
except ValueError:
|
||||
raise ValueError(f"Unknown message type: {raw_type!r}") from None
|
||||
|
||||
return ProtocolMessage(
|
||||
type=msg_type,
|
||||
id=obj.get("id", str(uuid.uuid4())),
|
||||
content=obj.get("content", ""),
|
||||
tool_name=obj.get("tool_name"),
|
||||
tool_args=obj.get("tool_args"),
|
||||
tool_result=obj.get("tool_result"),
|
||||
error=obj.get("error"),
|
||||
metadata=obj.get("metadata", {}),
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["MessageType", "ProtocolMessage", "deserialize", "serialize"]
|
||||
@@ -1,143 +0,0 @@
|
||||
"""Transport abstraction for OpenClaw agent communication."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
from openjarvis.agents.openclaw_protocol import (
|
||||
MessageType,
|
||||
ProtocolMessage,
|
||||
deserialize,
|
||||
serialize,
|
||||
)
|
||||
|
||||
|
||||
class OpenClawTransport(ABC):
|
||||
"""Base class for OpenClaw transport implementations."""
|
||||
|
||||
@abstractmethod
|
||||
def send(self, msg: ProtocolMessage) -> ProtocolMessage:
|
||||
"""Send a message and return the response."""
|
||||
|
||||
@abstractmethod
|
||||
def health(self) -> bool:
|
||||
"""Return True if the transport endpoint is healthy."""
|
||||
|
||||
@abstractmethod
|
||||
def close(self) -> None:
|
||||
"""Release transport resources."""
|
||||
|
||||
|
||||
class HttpTransport(OpenClawTransport):
|
||||
"""HTTP-based transport for communicating with an OpenClaw server."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "http://localhost:18789",
|
||||
timeout: float = 30.0,
|
||||
) -> None:
|
||||
self._host = host.rstrip("/")
|
||||
self._timeout = timeout
|
||||
|
||||
def send(self, msg: ProtocolMessage) -> ProtocolMessage:
|
||||
"""POST a message to the OpenClaw HTTP endpoint."""
|
||||
import httpx
|
||||
|
||||
payload = json.loads(serialize(msg))
|
||||
resp = httpx.post(
|
||||
f"{self._host}/api/query",
|
||||
json=payload,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return deserialize(json.dumps(resp.json()))
|
||||
|
||||
def health(self) -> bool:
|
||||
"""GET /health and return True if status is 200."""
|
||||
import httpx
|
||||
|
||||
try:
|
||||
resp = httpx.get(
|
||||
f"{self._host}/health",
|
||||
timeout=5.0,
|
||||
)
|
||||
return resp.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
"""No persistent resources to release."""
|
||||
|
||||
|
||||
class SubprocessTransport(OpenClawTransport):
|
||||
"""Subprocess-based transport — launches a Node.js process and
|
||||
communicates via JSON over stdin/stdout."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node_path: str = "node",
|
||||
script_path: str = "",
|
||||
) -> None:
|
||||
self._node_path = node_path
|
||||
self._script_path = script_path
|
||||
self._process: Optional[subprocess.Popen] = None
|
||||
|
||||
def _start_process(self) -> None:
|
||||
"""Start the Node.js subprocess."""
|
||||
self._process = subprocess.Popen(
|
||||
[self._node_path, self._script_path],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
|
||||
def _ensure_alive(self) -> None:
|
||||
"""Start the process if it's not running."""
|
||||
if self._process is None or self._process.poll() is not None:
|
||||
self._start_process()
|
||||
|
||||
def send(self, msg: ProtocolMessage) -> ProtocolMessage:
|
||||
"""Send a message via stdin and read the response from stdout."""
|
||||
self._ensure_alive()
|
||||
assert self._process is not None
|
||||
assert self._process.stdin is not None
|
||||
assert self._process.stdout is not None
|
||||
|
||||
line = serialize(msg) + "\n"
|
||||
self._process.stdin.write(line)
|
||||
self._process.stdin.flush()
|
||||
|
||||
response_line = self._process.stdout.readline()
|
||||
if not response_line:
|
||||
raise RuntimeError("No response from OpenClaw subprocess")
|
||||
return deserialize(response_line.strip())
|
||||
|
||||
def health(self) -> bool:
|
||||
"""Check if the subprocess is alive and responsive."""
|
||||
try:
|
||||
self._ensure_alive()
|
||||
health_msg = ProtocolMessage(type=MessageType.HEALTH)
|
||||
resp = self.send(health_msg)
|
||||
return resp.type == MessageType.HEALTH_OK
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
"""Terminate the subprocess."""
|
||||
if self._process is not None:
|
||||
try:
|
||||
self._process.terminate()
|
||||
self._process.wait(timeout=5)
|
||||
except Exception:
|
||||
try:
|
||||
self._process.kill()
|
||||
except Exception:
|
||||
pass
|
||||
self._process = None
|
||||
|
||||
|
||||
__all__ = ["HttpTransport", "OpenClawTransport", "SubprocessTransport"]
|
||||
@@ -1,7 +1,18 @@
|
||||
"""OrchestratorAgent — multi-turn agent with tool-calling loop."""
|
||||
"""OrchestratorAgent — multi-turn agent with tool-calling loop.
|
||||
|
||||
Supports two modes:
|
||||
|
||||
- **function_calling** (default): Uses OpenAI-format tool definitions and
|
||||
parses ``tool_calls`` from the engine response.
|
||||
- **structured**: Uses a THOUGHT/TOOL/INPUT/FINAL_ANSWER text format
|
||||
(like ReAct) with a canonical system prompt from the orchestrator
|
||||
prompt registry. This is the format used by the SFT/GRPO training
|
||||
pipelines, making the Orchestrator a distinctive trainable agent type.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
|
||||
@@ -21,6 +32,11 @@ class OrchestratorAgent(BaseAgent):
|
||||
2. If the response contains tool_calls, execute them and loop.
|
||||
3. If no tool_calls, return the final answer.
|
||||
4. Stop after ``max_turns`` iterations.
|
||||
|
||||
In **structured** mode the agent instead uses a
|
||||
``THOUGHT: / TOOL: / INPUT: / FINAL_ANSWER:`` text protocol
|
||||
identical to the format used by the orchestrator SFT/GRPO
|
||||
training pipelines.
|
||||
"""
|
||||
|
||||
agent_id = "orchestrator"
|
||||
@@ -35,6 +51,8 @@ class OrchestratorAgent(BaseAgent):
|
||||
max_turns: int = 10,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 1024,
|
||||
mode: str = "function_calling",
|
||||
system_prompt: Optional[str] = None,
|
||||
) -> None:
|
||||
self._engine = engine
|
||||
self._model = model
|
||||
@@ -44,12 +62,186 @@ class OrchestratorAgent(BaseAgent):
|
||||
self._max_turns = max_turns
|
||||
self._temperature = temperature
|
||||
self._max_tokens = max_tokens
|
||||
self._mode = mode
|
||||
self._system_prompt = system_prompt
|
||||
|
||||
def run(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
if self._mode == "structured":
|
||||
return self._run_structured(input, context, **kwargs)
|
||||
return self._run_function_calling(input, context, **kwargs)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Structured mode (THOUGHT/TOOL/INPUT/FINAL_ANSWER)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_structured(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
bus = self._bus
|
||||
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_START,
|
||||
{"agent": self.agent_id, "input": input},
|
||||
)
|
||||
|
||||
# Build system prompt
|
||||
if self._system_prompt:
|
||||
sys_prompt = self._system_prompt
|
||||
else:
|
||||
from openjarvis.learning.orchestrator.prompt_registry import (
|
||||
build_system_prompt,
|
||||
)
|
||||
tool_names = [t.spec.name for t in self._tools]
|
||||
sys_prompt = build_system_prompt(tool_names)
|
||||
|
||||
messages: list[Message] = [
|
||||
Message(role=Role.SYSTEM, content=sys_prompt),
|
||||
]
|
||||
if context and context.conversation.messages:
|
||||
messages.extend(context.conversation.messages)
|
||||
messages.append(Message(role=Role.USER, content=input))
|
||||
|
||||
all_tool_results: list[ToolResult] = []
|
||||
turns = 0
|
||||
|
||||
for _turn in range(self._max_turns):
|
||||
turns += 1
|
||||
|
||||
result = self._engine.generate(
|
||||
messages,
|
||||
model=self._model,
|
||||
temperature=self._temperature,
|
||||
max_tokens=self._max_tokens,
|
||||
)
|
||||
content = result.get("content", "")
|
||||
|
||||
parsed = self._parse_structured_response(content)
|
||||
|
||||
# FINAL_ANSWER → done
|
||||
if parsed["final_answer"]:
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_END,
|
||||
{"agent": self.agent_id, "turns": turns},
|
||||
)
|
||||
return AgentResult(
|
||||
content=parsed["final_answer"],
|
||||
tool_results=all_tool_results,
|
||||
turns=turns,
|
||||
)
|
||||
|
||||
# TOOL → execute
|
||||
if parsed["tool"]:
|
||||
messages.append(
|
||||
Message(role=Role.ASSISTANT, content=content)
|
||||
)
|
||||
|
||||
tool_call = ToolCall(
|
||||
id=f"orch_{turns}",
|
||||
name=parsed["tool"],
|
||||
arguments=parsed["input"] or "{}",
|
||||
)
|
||||
tool_result = self._executor.execute(tool_call)
|
||||
all_tool_results.append(tool_result)
|
||||
|
||||
observation = (
|
||||
f"Observation: {tool_result.content}"
|
||||
)
|
||||
messages.append(
|
||||
Message(role=Role.USER, content=observation)
|
||||
)
|
||||
continue
|
||||
|
||||
# Neither → treat content as final answer
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_END,
|
||||
{"agent": self.agent_id, "turns": turns},
|
||||
)
|
||||
return AgentResult(
|
||||
content=content,
|
||||
tool_results=all_tool_results,
|
||||
turns=turns,
|
||||
)
|
||||
|
||||
# Max turns exceeded
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_END,
|
||||
{
|
||||
"agent": self.agent_id,
|
||||
"turns": turns,
|
||||
"max_turns_exceeded": True,
|
||||
},
|
||||
)
|
||||
return AgentResult(
|
||||
content="Maximum turns reached without a final answer.",
|
||||
tool_results=all_tool_results,
|
||||
turns=turns,
|
||||
metadata={"max_turns_exceeded": True},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_structured_response(text: str) -> dict:
|
||||
"""Parse THOUGHT/TOOL/INPUT/FINAL_ANSWER from model output."""
|
||||
result = {
|
||||
"thought": "",
|
||||
"tool": "",
|
||||
"input": "",
|
||||
"final_answer": "",
|
||||
}
|
||||
|
||||
thought_match = re.search(
|
||||
r"THOUGHT:\s*(.+?)(?=\nTOOL:|\nFINAL[_ ]?ANSWER:|\Z)",
|
||||
text,
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
if thought_match:
|
||||
result["thought"] = thought_match.group(1).strip()
|
||||
|
||||
final_match = re.search(
|
||||
r"FINAL[_ ]?ANSWER:\s*(.+)",
|
||||
text,
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
if final_match:
|
||||
result["final_answer"] = final_match.group(1).strip()
|
||||
return result
|
||||
|
||||
tool_match = re.search(
|
||||
r"TOOL:\s*(.+)", text, re.IGNORECASE
|
||||
)
|
||||
if tool_match:
|
||||
result["tool"] = tool_match.group(1).strip()
|
||||
|
||||
input_match = re.search(
|
||||
r"INPUT:\s*(.+?)(?=\nTHOUGHT:|\nTOOL:|\nFINAL|\Z)",
|
||||
text,
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
if input_match:
|
||||
result["input"] = input_match.group(1).strip()
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Function-calling mode (original behaviour)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_function_calling(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
bus = self._bus
|
||||
|
||||
@@ -128,7 +320,6 @@ class OrchestratorAgent(BaseAgent):
|
||||
all_tool_results.append(tool_result)
|
||||
|
||||
# Append tool response message
|
||||
# Serialize arguments for the content
|
||||
messages.append(Message(
|
||||
role=Role.TOOL,
|
||||
content=tool_result.content,
|
||||
|
||||
@@ -0,0 +1,352 @@
|
||||
"""RLM (Recursive Language Model) Agent — recursive decomposition via persistent REPL.
|
||||
|
||||
Based on the RLM paper (arxiv:2512.24601). Instead of passing long context
|
||||
directly in the LLM prompt, RLM stores context as a Python variable in a
|
||||
persistent REPL. A "Root LM" writes Python code to inspect/decompose
|
||||
context and makes recursive sub-LM calls via ``llm_query()``/``llm_batch()``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
|
||||
from openjarvis.agents.rlm_repl import RLMRepl
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
from openjarvis.core.types import Message, Role, ToolCall, ToolResult
|
||||
from openjarvis.engine._stubs import InferenceEngine
|
||||
from openjarvis.tools._stubs import BaseTool, ToolExecutor
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# System prompt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
RLM_SYSTEM_PROMPT = (
|
||||
"You are an AI assistant that solves problems by writing "
|
||||
"Python code in a persistent REPL.\n\n"
|
||||
"## Available REPL Functions\n\n"
|
||||
"- `llm_query(prompt: str) -> str` — Call a sub-LM with a "
|
||||
"prompt and get a response.\n"
|
||||
"- `llm_batch(prompts: list[str]) -> list[str]` — Call a "
|
||||
"sub-LM with multiple prompts.\n"
|
||||
"- `FINAL(value)` — Terminate and return `value` as the "
|
||||
"final answer.\n"
|
||||
"- `FINAL_VAR(var_name: str)` — Terminate and return the "
|
||||
"value of variable `var_name`.\n"
|
||||
"- `answer` dict — Set `answer[\"value\"] = ...` and "
|
||||
"`answer[\"ready\"] = True` to terminate.\n\n"
|
||||
"## Available Modules\n\n"
|
||||
"json, re, math, collections, itertools, functools, "
|
||||
"textwrap, string, copy, datetime\n\n"
|
||||
"## Context Variable\n\n"
|
||||
"The input context (if any) is stored in the variable "
|
||||
"`context`. You can inspect it, slice it, or decompose it "
|
||||
"using Python.\n\n"
|
||||
"## Instructions\n\n"
|
||||
"1. Write Python code inside ```python blocks to manipulate "
|
||||
"context and solve the problem.\n"
|
||||
"2. For long contexts, decompose them into smaller chunks "
|
||||
"and use `llm_query()` on each chunk.\n"
|
||||
"3. Combine sub-results programmatically.\n"
|
||||
"4. When you have the final answer, call "
|
||||
"`FINAL(answer_value)` or `FINAL_VAR(\"var_name\")`.\n"
|
||||
"5. If you can answer directly without code, just respond "
|
||||
"with text (no code block).\n\n"
|
||||
"## Strategy Tips\n\n"
|
||||
"- Split long text into paragraphs or sections, summarize "
|
||||
"each with `llm_query()`.\n"
|
||||
"- Use `llm_batch()` for parallel sub-queries on multiple "
|
||||
"chunks.\n"
|
||||
"- Store intermediate results in variables — the REPL "
|
||||
"persists state across turns.\n"
|
||||
"- Build up the answer incrementally across multiple code "
|
||||
"blocks.\n"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@AgentRegistry.register("rlm")
|
||||
class RLMAgent(BaseAgent):
|
||||
"""Recursive Language Model agent using a persistent REPL.
|
||||
|
||||
The agent generates Python code that runs in a sandboxed REPL with
|
||||
access to ``llm_query()`` / ``llm_batch()`` for recursive sub-LM
|
||||
calls. Context is stored as a REPL variable rather than injected
|
||||
directly into the prompt, enabling processing of arbitrarily long
|
||||
inputs through recursive decomposition.
|
||||
"""
|
||||
|
||||
agent_id = "rlm"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine: InferenceEngine,
|
||||
model: str,
|
||||
*,
|
||||
tools: Optional[List[BaseTool]] = None,
|
||||
bus: Optional[EventBus] = None,
|
||||
max_turns: int = 10,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2048,
|
||||
sub_model: Optional[str] = None,
|
||||
sub_temperature: float = 0.3,
|
||||
sub_max_tokens: int = 1024,
|
||||
max_output_chars: int = 10000,
|
||||
system_prompt: Optional[str] = None,
|
||||
) -> None:
|
||||
self._engine = engine
|
||||
self._model = model
|
||||
self._tools = tools or []
|
||||
self._executor = ToolExecutor(self._tools, bus=bus) if self._tools else None
|
||||
self._bus = bus
|
||||
self._max_turns = max_turns
|
||||
self._temperature = temperature
|
||||
self._max_tokens = max_tokens
|
||||
self._sub_model = sub_model or model
|
||||
self._sub_temperature = sub_temperature
|
||||
self._sub_max_tokens = sub_max_tokens
|
||||
self._max_output_chars = max_output_chars
|
||||
self._system_prompt = system_prompt or RLM_SYSTEM_PROMPT
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Main run loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def run(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentResult:
|
||||
bus = self._bus
|
||||
|
||||
# Emit agent start
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_START,
|
||||
{"agent": self.agent_id, "input": input},
|
||||
)
|
||||
|
||||
# Create REPL with sub-LM callbacks
|
||||
repl = RLMRepl(
|
||||
llm_query_fn=self._make_sub_query,
|
||||
llm_batch_fn=self._make_batch_query,
|
||||
max_output_chars=self._max_output_chars,
|
||||
)
|
||||
|
||||
# Resolve context and inject into REPL
|
||||
ctx_text = self._resolve_context(context)
|
||||
if ctx_text:
|
||||
repl.set_variable("context", ctx_text)
|
||||
|
||||
# Build conversation
|
||||
messages: list[Message] = [
|
||||
Message(role=Role.SYSTEM, content=self._system_prompt),
|
||||
]
|
||||
if context and context.conversation.messages:
|
||||
messages.extend(context.conversation.messages)
|
||||
messages.append(Message(role=Role.USER, content=input))
|
||||
|
||||
all_tool_results: list[ToolResult] = []
|
||||
turns = 0
|
||||
|
||||
for _turn in range(self._max_turns):
|
||||
turns += 1
|
||||
|
||||
result = self._engine.generate(
|
||||
messages,
|
||||
model=self._model,
|
||||
temperature=self._temperature,
|
||||
max_tokens=self._max_tokens,
|
||||
)
|
||||
content = result.get("content", "")
|
||||
|
||||
# Strip <think> tags
|
||||
content = self._strip_think_tags(content)
|
||||
|
||||
# Extract code block
|
||||
code = self._extract_code(content)
|
||||
|
||||
# No code block → return content as final answer
|
||||
if code is None:
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_END,
|
||||
{"agent": self.agent_id, "turns": turns},
|
||||
)
|
||||
return AgentResult(
|
||||
content=content,
|
||||
tool_results=all_tool_results,
|
||||
turns=turns,
|
||||
)
|
||||
|
||||
# Execute code in REPL
|
||||
output = repl.execute(code)
|
||||
|
||||
# Record as tool result
|
||||
tool_result = ToolResult(
|
||||
tool_name="rlm_repl",
|
||||
content=output or "(no output)",
|
||||
success=(
|
||||
not output.startswith("Error:")
|
||||
and not output.startswith("SyntaxError:")
|
||||
),
|
||||
)
|
||||
all_tool_results.append(tool_result)
|
||||
|
||||
# Check for termination
|
||||
if repl.is_terminated:
|
||||
final = repl.final_answer
|
||||
final_str = str(final) if final is not None else ""
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_END,
|
||||
{"agent": self.agent_id, "turns": turns},
|
||||
)
|
||||
return AgentResult(
|
||||
content=final_str,
|
||||
tool_results=all_tool_results,
|
||||
turns=turns,
|
||||
)
|
||||
|
||||
# Feed output back as user message
|
||||
messages.append(Message(role=Role.ASSISTANT, content=content))
|
||||
feedback = (
|
||||
f"REPL Output: {output}"
|
||||
if output
|
||||
else "REPL Output: (no output)"
|
||||
)
|
||||
messages.append(Message(role=Role.USER, content=feedback))
|
||||
|
||||
# Max turns exceeded — check answer dict for partial result
|
||||
answer = repl.get_variable("answer")
|
||||
if isinstance(answer, dict) and answer.get("value") is not None:
|
||||
final_content = str(answer["value"])
|
||||
else:
|
||||
final_content = "Maximum turns reached without a final answer."
|
||||
|
||||
if bus:
|
||||
bus.publish(
|
||||
EventType.AGENT_TURN_END,
|
||||
{
|
||||
"agent": self.agent_id,
|
||||
"turns": turns,
|
||||
"max_turns_exceeded": True,
|
||||
},
|
||||
)
|
||||
|
||||
return AgentResult(
|
||||
content=final_content,
|
||||
tool_results=all_tool_results,
|
||||
turns=turns,
|
||||
metadata={"max_turns_exceeded": True},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Sub-LM callbacks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _make_sub_query(self, prompt: str) -> str:
|
||||
"""Execute a single sub-LM query.
|
||||
|
||||
Called from REPL code via ``llm_query(prompt)``.
|
||||
If the sub-LM returns tool_calls, execute one round of tool
|
||||
resolution before returning the final text.
|
||||
"""
|
||||
messages = [Message(role=Role.USER, content=prompt)]
|
||||
result = self._engine.generate(
|
||||
messages,
|
||||
model=self._sub_model,
|
||||
temperature=self._sub_temperature,
|
||||
max_tokens=self._sub_max_tokens,
|
||||
)
|
||||
|
||||
# Single-turn tool resolution
|
||||
raw_tool_calls = result.get("tool_calls", [])
|
||||
if raw_tool_calls and self._executor:
|
||||
content = result.get("content", "")
|
||||
tool_calls = [
|
||||
ToolCall(
|
||||
id=tc.get("id", f"sub_{i}"),
|
||||
name=tc.get("name", ""),
|
||||
arguments=tc.get("arguments", "{}"),
|
||||
)
|
||||
for i, tc in enumerate(raw_tool_calls)
|
||||
]
|
||||
messages.append(Message(
|
||||
role=Role.ASSISTANT,
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
))
|
||||
for tc in tool_calls:
|
||||
tr = self._executor.execute(tc)
|
||||
messages.append(Message(
|
||||
role=Role.TOOL,
|
||||
content=tr.content,
|
||||
tool_call_id=tc.id,
|
||||
name=tc.name,
|
||||
))
|
||||
followup = self._engine.generate(
|
||||
messages,
|
||||
model=self._sub_model,
|
||||
temperature=self._sub_temperature,
|
||||
max_tokens=self._sub_max_tokens,
|
||||
)
|
||||
return followup.get("content", "")
|
||||
|
||||
return result.get("content", "")
|
||||
|
||||
def _make_batch_query(self, prompts: List[str]) -> List[str]:
|
||||
"""Execute multiple sub-LM queries sequentially.
|
||||
|
||||
Called from REPL code via ``llm_batch(prompts)``.
|
||||
"""
|
||||
return [self._make_sub_query(p) for p in prompts]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _extract_code(text: str) -> Optional[str]:
|
||||
"""Extract the first ```python code block from *text*.
|
||||
|
||||
Also matches bare ``` blocks (without language tag).
|
||||
Returns ``None`` if no code block is found.
|
||||
"""
|
||||
# Try ```python first
|
||||
m = re.search(r"```python\s*\n(.*?)```", text, re.DOTALL)
|
||||
if m:
|
||||
return m.group(1).strip()
|
||||
# Try bare ```
|
||||
m = re.search(r"```\s*\n(.*?)```", text, re.DOTALL)
|
||||
if m:
|
||||
return m.group(1).strip()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _strip_think_tags(text: str) -> str:
|
||||
"""Remove ``<think>...</think>`` blocks from model output."""
|
||||
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
|
||||
|
||||
@staticmethod
|
||||
def _resolve_context(context: Optional[AgentContext]) -> Optional[str]:
|
||||
"""Resolve context text from AgentContext metadata or memory_results."""
|
||||
if context is None:
|
||||
return None
|
||||
# Primary: explicit context in metadata
|
||||
if context.metadata.get("context"):
|
||||
return str(context.metadata["context"])
|
||||
# Fallback: join memory results
|
||||
if context.memory_results:
|
||||
return "\n\n".join(str(r) for r in context.memory_results)
|
||||
return None
|
||||
|
||||
|
||||
__all__ = ["RLMAgent", "RLM_SYSTEM_PROMPT"]
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Sandboxed REPL environment for the RLM agent.
|
||||
|
||||
Provides a persistent Python namespace with injected helper functions
|
||||
(``llm_query``, ``llm_batch``, ``FINAL``, ``FINAL_VAR``) that the RLM
|
||||
agent's generated code uses to decompose context and make recursive
|
||||
sub-LM calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from contextlib import redirect_stderr, redirect_stdout
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
# Safe stdlib modules pre-injected into the REPL namespace
|
||||
_SAFE_MODULES = [
|
||||
"json",
|
||||
"re",
|
||||
"math",
|
||||
"collections",
|
||||
"itertools",
|
||||
"functools",
|
||||
"textwrap",
|
||||
"string",
|
||||
"copy",
|
||||
"datetime",
|
||||
]
|
||||
|
||||
# Patterns blocked for security
|
||||
_BLOCKED_PATTERNS = [
|
||||
"os.system",
|
||||
"os.popen",
|
||||
"subprocess",
|
||||
"shutil.rmtree",
|
||||
"__import__",
|
||||
"open(",
|
||||
"ctypes",
|
||||
"socket",
|
||||
"http.client",
|
||||
"urllib",
|
||||
]
|
||||
|
||||
|
||||
class RLMRepl:
|
||||
"""Sandboxed Python REPL with persistent namespace for the RLM agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
llm_query_fn:
|
||||
Callback invoked when REPL code calls ``llm_query(prompt)``.
|
||||
llm_batch_fn:
|
||||
Callback invoked when REPL code calls ``llm_batch(prompts)``.
|
||||
max_output_chars:
|
||||
Maximum characters captured from stdout per execution.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm_query_fn: Optional[Callable[[str], str]] = None,
|
||||
llm_batch_fn: Optional[Callable[[List[str]], List[str]]] = None,
|
||||
*,
|
||||
max_output_chars: int = 10000,
|
||||
) -> None:
|
||||
self._max_output_chars = max_output_chars
|
||||
self._terminated = False
|
||||
self._final_value: Any = None
|
||||
|
||||
# Build namespace
|
||||
self._namespace: Dict[str, Any] = {}
|
||||
|
||||
# Inject safe stdlib modules
|
||||
for mod_name in _SAFE_MODULES:
|
||||
try:
|
||||
import importlib
|
||||
|
||||
self._namespace[mod_name] = importlib.import_module(mod_name)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# answer dict — code can set answer["ready"] = True, answer["value"] = ...
|
||||
self._namespace["answer"] = {"ready": False, "value": None}
|
||||
|
||||
# Inject FINAL / FINAL_VAR
|
||||
self._namespace["FINAL"] = self._final
|
||||
self._namespace["FINAL_VAR"] = self._final_var
|
||||
|
||||
# Inject llm_query / llm_batch
|
||||
if llm_query_fn is not None:
|
||||
self._namespace["llm_query"] = llm_query_fn
|
||||
if llm_batch_fn is not None:
|
||||
self._namespace["llm_batch"] = llm_batch_fn
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Termination helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _final(self, value: Any) -> None:
|
||||
"""Mark the REPL as terminated with a final answer."""
|
||||
self._terminated = True
|
||||
self._final_value = value
|
||||
|
||||
def _final_var(self, var_name: str) -> None:
|
||||
"""Mark the REPL as terminated, using a namespace variable as the answer."""
|
||||
value = self._namespace.get(var_name)
|
||||
self._terminated = True
|
||||
self._final_value = value
|
||||
|
||||
@property
|
||||
def is_terminated(self) -> bool:
|
||||
"""Check if FINAL/FINAL_VAR was called or answer["ready"] is True."""
|
||||
if self._terminated:
|
||||
return True
|
||||
answer = self._namespace.get("answer", {})
|
||||
if isinstance(answer, dict) and answer.get("ready"):
|
||||
return True
|
||||
return False
|
||||
|
||||
@property
|
||||
def final_answer(self) -> Any:
|
||||
"""Return the termination value."""
|
||||
if self._terminated:
|
||||
return self._final_value
|
||||
answer = self._namespace.get("answer", {})
|
||||
if isinstance(answer, dict) and answer.get("ready"):
|
||||
return answer.get("value")
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Execution
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def security_check(self, code: str) -> Optional[str]:
|
||||
"""Check code for dangerous patterns. Returns error message or None."""
|
||||
for pattern in _BLOCKED_PATTERNS:
|
||||
if pattern in code:
|
||||
return f"Blocked: code contains prohibited pattern '{pattern}'"
|
||||
return None
|
||||
|
||||
def execute(self, code: str) -> str:
|
||||
"""Execute *code* in the persistent namespace and return captured stdout.
|
||||
|
||||
Raises are caught and returned as error strings.
|
||||
"""
|
||||
# Security check
|
||||
violation = self.security_check(code)
|
||||
if violation is not None:
|
||||
return f"Error: {violation}"
|
||||
|
||||
stdout_buf = io.StringIO()
|
||||
stderr_buf = io.StringIO()
|
||||
|
||||
try:
|
||||
with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf):
|
||||
exec(code, self._namespace) # noqa: S102
|
||||
except Exception as exc:
|
||||
error_msg = f"{type(exc).__name__}: {exc}"
|
||||
return error_msg
|
||||
|
||||
output = stdout_buf.getvalue()
|
||||
err_output = stderr_buf.getvalue()
|
||||
if err_output:
|
||||
output += ("\n" if output else "") + err_output
|
||||
|
||||
# Truncate if needed
|
||||
if len(output) > self._max_output_chars:
|
||||
output = output[: self._max_output_chars] + "\n... (output truncated)"
|
||||
|
||||
return output
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Namespace access
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def set_variable(self, name: str, value: Any) -> None:
|
||||
"""Set a variable in the REPL namespace."""
|
||||
self._namespace[name] = value
|
||||
|
||||
def get_variable(self, name: str) -> Any:
|
||||
"""Get a variable from the REPL namespace."""
|
||||
return self._namespace.get(name)
|
||||
|
||||
|
||||
__all__ = ["RLMRepl"]
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Channel abstraction for multi-platform messaging via OpenClaw."""
|
||||
"""Channel abstraction for multi-platform messaging."""
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
@@ -9,7 +9,77 @@ from openjarvis.channels._stubs import (
|
||||
|
||||
# Trigger registration of built-in channels
|
||||
try:
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge # noqa: F401
|
||||
import openjarvis.channels.telegram # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.discord_channel # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.slack # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.webhook # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.email_channel # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.whatsapp # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.signal_channel # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.google_chat # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.irc_channel # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.webchat # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.teams # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.matrix_channel # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.mattermost # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.feishu # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.channels.bluebubbles # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
"""BlueBubblesChannel — BlueBubbles (iMessage bridge) adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("bluebubbles")
|
||||
class BlueBubblesChannel(BaseChannel):
|
||||
"""BlueBubbles (iMessage bridge) channel adapter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
url:
|
||||
BlueBubbles server URL. Falls back to ``BLUEBUBBLES_URL`` env var.
|
||||
password:
|
||||
BlueBubbles server password. Falls back to ``BLUEBUBBLES_PASSWORD``
|
||||
env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "bluebubbles"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str = "",
|
||||
*,
|
||||
password: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._url = url or os.environ.get("BLUEBUBBLES_URL", "")
|
||||
self._password = password or os.environ.get("BLUEBUBBLES_PASSWORD", "")
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._url or not self._password:
|
||||
logger.warning("No BlueBubbles URL or password configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send an iMessage via the BlueBubbles API."""
|
||||
if not self._url or not self._password:
|
||||
logger.warning("Cannot send: no BlueBubbles credentials configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = f"{self._url}/api/v1/message/text"
|
||||
payload: Dict[str, Any] = {
|
||||
"chatGuid": channel,
|
||||
"message": content,
|
||||
"method": "private-api",
|
||||
}
|
||||
|
||||
resp = httpx.post(
|
||||
url,
|
||||
params={"password": self._password},
|
||||
json=payload,
|
||||
timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"BlueBubbles API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("BlueBubbles send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["bluebubbles"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["BlueBubblesChannel"]
|
||||
@@ -0,0 +1,197 @@
|
||||
"""DiscordChannel — native Discord Bot API adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelMessage,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("discord")
|
||||
class DiscordChannel(BaseChannel):
|
||||
"""Native Discord channel adapter using the Discord REST API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
bot_token:
|
||||
Discord bot token. Falls back to ``DISCORD_BOT_TOKEN`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "discord"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bot_token: str = "",
|
||||
*,
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._token = bot_token or os.environ.get("DISCORD_BOT_TOKEN", "")
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
self._listener_thread: Optional[threading.Thread] = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Start listening for incoming messages via discord.py gateway."""
|
||||
if not self._token:
|
||||
logger.warning("No Discord bot token configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
|
||||
self._stop_event.clear()
|
||||
self._status = ChannelStatus.CONNECTING
|
||||
|
||||
try:
|
||||
import discord # noqa: F401
|
||||
|
||||
self._listener_thread = threading.Thread(
|
||||
target=self._gateway_loop, daemon=True,
|
||||
)
|
||||
self._listener_thread.start()
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
logger.info("Discord channel connected (gateway)")
|
||||
except ImportError:
|
||||
logger.info("discord.py not installed; send-only mode")
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Stop the listener thread."""
|
||||
self._stop_event.set()
|
||||
if self._listener_thread is not None:
|
||||
self._listener_thread.join(timeout=5.0)
|
||||
self._listener_thread = None
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Discord channel via REST API."""
|
||||
if not self._token:
|
||||
logger.warning("Cannot send: no Discord bot token")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = f"https://discord.com/api/v10/channels/{channel}/messages"
|
||||
headers = {
|
||||
"Authorization": f"Bot {self._token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {"content": content}
|
||||
if conversation_id:
|
||||
payload["message_reference"] = {"message_id": conversation_id}
|
||||
|
||||
resp = httpx.post(
|
||||
url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Discord API returned status %d: %s",
|
||||
resp.status_code,
|
||||
resp.text,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Discord send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["discord"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _gateway_loop(self) -> None:
|
||||
"""Run the discord.py client in a background thread."""
|
||||
try:
|
||||
import asyncio
|
||||
|
||||
import discord
|
||||
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
client = discord.Client(intents=intents)
|
||||
|
||||
@client.event
|
||||
async def on_message(message):
|
||||
if message.author == client.user:
|
||||
return
|
||||
cm = ChannelMessage(
|
||||
channel="discord",
|
||||
sender=str(message.author.id),
|
||||
content=message.content,
|
||||
message_id=str(message.id),
|
||||
conversation_id=str(message.channel.id),
|
||||
)
|
||||
for handler in self._handlers:
|
||||
try:
|
||||
handler(cm)
|
||||
except Exception:
|
||||
logger.exception("Discord handler error")
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_RECEIVED,
|
||||
{
|
||||
"channel": cm.channel,
|
||||
"sender": cm.sender,
|
||||
"content": cm.content,
|
||||
"message_id": cm.message_id,
|
||||
},
|
||||
)
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(client.start(self._token))
|
||||
except Exception:
|
||||
logger.debug("Discord gateway loop error", exc_info=True)
|
||||
self._status = ChannelStatus.ERROR
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["DiscordChannel"]
|
||||
@@ -0,0 +1,240 @@
|
||||
"""EmailChannel — SMTP/IMAP email adapter (stdlib only, zero extra deps)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import smtplib
|
||||
import threading
|
||||
from email.mime.text import MIMEText
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelMessage,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("email")
|
||||
class EmailChannel(BaseChannel):
|
||||
"""Email channel adapter using stdlib ``smtplib`` and ``imaplib``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
smtp_host:
|
||||
SMTP server hostname.
|
||||
smtp_port:
|
||||
SMTP server port (default 587 for STARTTLS).
|
||||
imap_host:
|
||||
IMAP server hostname for incoming messages.
|
||||
imap_port:
|
||||
IMAP server port (default 993 for SSL).
|
||||
username:
|
||||
Email username. Falls back to ``EMAIL_USERNAME`` env var.
|
||||
password:
|
||||
Email password. Falls back to ``EMAIL_PASSWORD`` env var.
|
||||
use_tls:
|
||||
Whether to use TLS (default ``True``).
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "email"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
smtp_host: str = "",
|
||||
smtp_port: int = 587,
|
||||
*,
|
||||
imap_host: str = "",
|
||||
imap_port: int = 993,
|
||||
username: str = "",
|
||||
password: str = "",
|
||||
use_tls: bool = True,
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._smtp_host = smtp_host
|
||||
self._smtp_port = smtp_port
|
||||
self._imap_host = imap_host
|
||||
self._imap_port = imap_port
|
||||
self._username = username or os.environ.get("EMAIL_USERNAME", "")
|
||||
self._password = password or os.environ.get("EMAIL_PASSWORD", "")
|
||||
self._use_tls = use_tls
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
self._listener_thread: Optional[threading.Thread] = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Start IMAP polling for incoming messages (if configured)."""
|
||||
if not self._smtp_host or not self._username:
|
||||
logger.warning("Email channel not fully configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
|
||||
self._stop_event.clear()
|
||||
self._status = ChannelStatus.CONNECTING
|
||||
|
||||
if self._imap_host:
|
||||
self._listener_thread = threading.Thread(
|
||||
target=self._imap_poll_loop, daemon=True,
|
||||
)
|
||||
self._listener_thread.start()
|
||||
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
logger.info("Email channel connected (SMTP: %s)", self._smtp_host)
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Stop the IMAP listener thread."""
|
||||
self._stop_event.set()
|
||||
if self._listener_thread is not None:
|
||||
self._listener_thread.join(timeout=5.0)
|
||||
self._listener_thread = None
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send an email via SMTP. ``channel`` is the recipient address."""
|
||||
if not self._smtp_host or not self._username:
|
||||
logger.warning("Cannot send: email not configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
msg = MIMEText(content)
|
||||
msg["From"] = self._username
|
||||
msg["To"] = channel
|
||||
msg["Subject"] = (metadata or {}).get(
|
||||
"subject", "Message from OpenJarvis",
|
||||
)
|
||||
if conversation_id:
|
||||
msg["In-Reply-To"] = conversation_id
|
||||
|
||||
if self._use_tls:
|
||||
with smtplib.SMTP(self._smtp_host, self._smtp_port) as server:
|
||||
server.starttls()
|
||||
server.login(self._username, self._password)
|
||||
server.send_message(msg)
|
||||
else:
|
||||
with smtplib.SMTP(self._smtp_host, self._smtp_port) as server:
|
||||
if self._password:
|
||||
server.login(self._username, self._password)
|
||||
server.send_message(msg)
|
||||
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
except Exception:
|
||||
logger.debug("Email send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["email"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming email messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _imap_poll_loop(self) -> None:
|
||||
"""Poll IMAP INBOX for new messages."""
|
||||
import email
|
||||
import imaplib
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
if self._use_tls:
|
||||
imap = imaplib.IMAP4_SSL(
|
||||
self._imap_host, self._imap_port,
|
||||
)
|
||||
else:
|
||||
imap = imaplib.IMAP4(self._imap_host, self._imap_port)
|
||||
|
||||
imap.login(self._username, self._password)
|
||||
imap.select("INBOX")
|
||||
|
||||
_, data = imap.search(None, "UNSEEN")
|
||||
for num in data[0].split():
|
||||
_, msg_data = imap.fetch(num, "(RFC822)")
|
||||
if msg_data[0] is None:
|
||||
continue
|
||||
raw = msg_data[0][1]
|
||||
parsed = email.message_from_bytes(raw)
|
||||
body = ""
|
||||
if parsed.is_multipart():
|
||||
for part in parsed.walk():
|
||||
if part.get_content_type() == "text/plain":
|
||||
body = part.get_payload(decode=True).decode(
|
||||
errors="replace",
|
||||
)
|
||||
break
|
||||
else:
|
||||
body = parsed.get_payload(decode=True).decode(
|
||||
errors="replace",
|
||||
)
|
||||
|
||||
cm = ChannelMessage(
|
||||
channel="email",
|
||||
sender=parsed.get("From", ""),
|
||||
content=body,
|
||||
message_id=parsed.get("Message-ID", ""),
|
||||
conversation_id=parsed.get("In-Reply-To", ""),
|
||||
)
|
||||
for handler in self._handlers:
|
||||
try:
|
||||
handler(cm)
|
||||
except Exception:
|
||||
logger.exception("Email handler error")
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_RECEIVED,
|
||||
{
|
||||
"channel": cm.channel,
|
||||
"sender": cm.sender,
|
||||
"content": cm.content,
|
||||
"message_id": cm.message_id,
|
||||
},
|
||||
)
|
||||
|
||||
imap.close()
|
||||
imap.logout()
|
||||
except Exception:
|
||||
logger.debug("IMAP poll error", exc_info=True)
|
||||
|
||||
self._stop_event.wait(30.0)
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["EmailChannel"]
|
||||
@@ -0,0 +1,162 @@
|
||||
"""FeishuChannel — Feishu (Lark) adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("feishu")
|
||||
class FeishuChannel(BaseChannel):
|
||||
"""Feishu (Lark) channel adapter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
app_id:
|
||||
Feishu App ID. Falls back to ``FEISHU_APP_ID`` env var.
|
||||
app_secret:
|
||||
Feishu App Secret. Falls back to ``FEISHU_APP_SECRET`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "feishu"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app_id: str = "",
|
||||
*,
|
||||
app_secret: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._app_id = app_id or os.environ.get("FEISHU_APP_ID", "")
|
||||
self._app_secret = app_secret or os.environ.get("FEISHU_APP_SECRET", "")
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._app_id or not self._app_secret:
|
||||
logger.warning("No Feishu app_id or app_secret configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Feishu chat via the Open API."""
|
||||
if not self._app_id or not self._app_secret:
|
||||
logger.warning("Cannot send: no Feishu credentials configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
# Obtain tenant_access_token
|
||||
token_url = (
|
||||
"https://open.feishu.cn/open-apis/auth/v3/"
|
||||
"tenant_access_token/internal"
|
||||
)
|
||||
token_resp = httpx.post(
|
||||
token_url,
|
||||
json={
|
||||
"app_id": self._app_id,
|
||||
"app_secret": self._app_secret,
|
||||
},
|
||||
timeout=10.0,
|
||||
)
|
||||
if token_resp.status_code >= 300:
|
||||
logger.warning(
|
||||
"Feishu token request returned status %d",
|
||||
token_resp.status_code,
|
||||
)
|
||||
return False
|
||||
tenant_token = token_resp.json().get("tenant_access_token", "")
|
||||
if not tenant_token:
|
||||
logger.warning("Feishu token response missing tenant_access_token")
|
||||
return False
|
||||
|
||||
# Send message
|
||||
msg_url = (
|
||||
"https://open.feishu.cn/open-apis/im/v1/messages"
|
||||
"?receive_id_type=chat_id"
|
||||
)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {tenant_token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"receive_id": channel,
|
||||
"msg_type": "text",
|
||||
"content": json.dumps({"text": content}),
|
||||
}
|
||||
|
||||
resp = httpx.post(
|
||||
msg_url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Feishu API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Feishu send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["feishu"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["FeishuChannel"]
|
||||
@@ -0,0 +1,123 @@
|
||||
"""GoogleChatChannel — Google Chat webhook adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("google_chat")
|
||||
class GoogleChatChannel(BaseChannel):
|
||||
"""Google Chat webhook channel adapter (send-only).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
webhook_url:
|
||||
Google Chat incoming webhook URL. Falls back to
|
||||
``GOOGLE_CHAT_WEBHOOK_URL`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "google_chat"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
webhook_url: str = "",
|
||||
*,
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._webhook_url = webhook_url or os.environ.get(
|
||||
"GOOGLE_CHAT_WEBHOOK_URL", "",
|
||||
)
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._webhook_url:
|
||||
logger.warning("No Google Chat webhook URL configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to Google Chat via the configured webhook URL."""
|
||||
if not self._webhook_url:
|
||||
logger.warning("Cannot send: no Google Chat webhook URL configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
payload: Dict[str, Any] = {"text": content}
|
||||
|
||||
resp = httpx.post(
|
||||
self._webhook_url, json=payload, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Google Chat API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Google Chat send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["google_chat"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["GoogleChatChannel"]
|
||||
@@ -0,0 +1,145 @@
|
||||
"""IRCChannel — IRC adapter using stdlib socket."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import ssl
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("irc")
|
||||
class IRCChannel(BaseChannel):
|
||||
"""IRC channel adapter using stdlib socket (send-only).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
server:
|
||||
IRC server hostname. Falls back to ``IRC_SERVER`` env var.
|
||||
port:
|
||||
IRC server port (default 6667). Falls back to ``IRC_PORT`` env var.
|
||||
nick:
|
||||
IRC nickname. Falls back to ``IRC_NICK`` env var.
|
||||
password:
|
||||
Optional server password. Falls back to ``IRC_PASSWORD`` env var.
|
||||
use_tls:
|
||||
Whether to use TLS for the connection (default ``False``).
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "irc"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server: str = "",
|
||||
*,
|
||||
port: int = 6667,
|
||||
nick: str = "",
|
||||
password: str = "",
|
||||
use_tls: bool = False,
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._server = server or os.environ.get("IRC_SERVER", "")
|
||||
self._port = int(os.environ.get("IRC_PORT", str(port)))
|
||||
self._nick = nick or os.environ.get("IRC_NICK", "")
|
||||
self._password = password or os.environ.get("IRC_PASSWORD", "")
|
||||
self._use_tls = use_tls
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._server:
|
||||
logger.warning("No IRC server configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
if not self._nick:
|
||||
logger.warning("No IRC nick configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a PRIVMSG to an IRC channel via a new socket connection."""
|
||||
if not self._server or not self._nick:
|
||||
logger.warning("Cannot send: IRC server or nick not configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
if self._use_tls:
|
||||
ctx = ssl.create_default_context()
|
||||
sock = ctx.wrap_socket(sock, server_hostname=self._server)
|
||||
|
||||
sock.connect((self._server, self._port))
|
||||
|
||||
if self._password:
|
||||
sock.sendall(f"PASS {self._password}\r\n".encode())
|
||||
sock.sendall(f"NICK {self._nick}\r\n".encode())
|
||||
sock.sendall(f"USER {self._nick} 0 * :{self._nick}\r\n".encode())
|
||||
sock.sendall(f"PRIVMSG {channel} :{content}\r\n".encode())
|
||||
sock.sendall(b"QUIT\r\n")
|
||||
sock.close()
|
||||
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
except Exception:
|
||||
logger.debug("IRC send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["irc"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["IRCChannel"]
|
||||
@@ -0,0 +1,137 @@
|
||||
"""MatrixChannel — Matrix homeserver adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("matrix")
|
||||
class MatrixChannel(BaseChannel):
|
||||
"""Matrix homeserver channel adapter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
homeserver:
|
||||
Matrix homeserver URL. Falls back to ``MATRIX_HOMESERVER`` env var.
|
||||
access_token:
|
||||
Matrix access token. Falls back to ``MATRIX_ACCESS_TOKEN`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "matrix"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
homeserver: str = "",
|
||||
*,
|
||||
access_token: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._homeserver = homeserver or os.environ.get("MATRIX_HOMESERVER", "")
|
||||
self._access_token = access_token or os.environ.get("MATRIX_ACCESS_TOKEN", "")
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
self._txn_id = 0
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._homeserver or not self._access_token:
|
||||
logger.warning("No Matrix homeserver or access_token configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Matrix room via the Client-Server API."""
|
||||
if not self._homeserver or not self._access_token:
|
||||
logger.warning("Cannot send: no Matrix credentials configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
self._txn_id += 1
|
||||
txn_id = self._txn_id
|
||||
url = (
|
||||
f"{self._homeserver}/_matrix/client/v3/rooms/"
|
||||
f"{channel}/send/m.room.message/{txn_id}"
|
||||
)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._access_token}",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"msgtype": "m.text",
|
||||
"body": content,
|
||||
}
|
||||
|
||||
resp = httpx.put(
|
||||
url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Matrix API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Matrix send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["matrix"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["MatrixChannel"]
|
||||
@@ -0,0 +1,134 @@
|
||||
"""MattermostChannel — Mattermost REST API adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("mattermost")
|
||||
class MattermostChannel(BaseChannel):
|
||||
"""Mattermost channel adapter using the REST API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
url:
|
||||
Mattermost server URL. Falls back to ``MATTERMOST_URL`` env var.
|
||||
token:
|
||||
Mattermost personal access token or bot token. Falls back to
|
||||
``MATTERMOST_TOKEN`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "mattermost"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str = "",
|
||||
*,
|
||||
token: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._url = url or os.environ.get("MATTERMOST_URL", "")
|
||||
self._token = token or os.environ.get("MATTERMOST_TOKEN", "")
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._url or not self._token:
|
||||
logger.warning("No Mattermost URL or token configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Mattermost channel via the REST API."""
|
||||
if not self._url or not self._token:
|
||||
logger.warning("Cannot send: no Mattermost credentials configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = f"{self._url}/api/v4/posts"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._token}",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"channel_id": channel,
|
||||
"message": content,
|
||||
}
|
||||
if conversation_id:
|
||||
payload["root_id"] = conversation_id
|
||||
|
||||
resp = httpx.post(
|
||||
url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Mattermost API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Mattermost send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["mattermost"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["MattermostChannel"]
|
||||
@@ -1,257 +0,0 @@
|
||||
"""OpenClawChannelBridge — WebSocket/HTTP bridge to the OpenClaw gateway."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelMessage,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("openclaw")
|
||||
class OpenClawChannelBridge(BaseChannel):
|
||||
"""Bridge to the OpenClaw gateway via WebSocket with HTTP fallback.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
gateway_url:
|
||||
WebSocket URL of the OpenClaw gateway (e.g. ``ws://127.0.0.1:18789/ws``).
|
||||
reconnect_interval:
|
||||
Seconds to wait before reconnecting after a disconnect.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "openclaw"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gateway_url: str = "ws://127.0.0.1:18789/ws",
|
||||
*,
|
||||
reconnect_interval: float = 5.0,
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._gateway_url = gateway_url
|
||||
self._reconnect_interval = reconnect_interval
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
self._ws: Any = None
|
||||
self._listener_thread: Optional[threading.Thread] = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Establish connection to the OpenClaw gateway."""
|
||||
self._status = ChannelStatus.CONNECTING
|
||||
self._stop_event.clear()
|
||||
|
||||
try:
|
||||
import websockets.sync.client # type: ignore[import-untyped]
|
||||
|
||||
self._ws = websockets.sync.client.connect(self._gateway_url)
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
self._listener_thread = threading.Thread(
|
||||
target=self._listener_loop, daemon=True,
|
||||
)
|
||||
self._listener_thread.start()
|
||||
logger.info(
|
||||
"Connected to OpenClaw gateway via WebSocket: %s",
|
||||
self._gateway_url,
|
||||
)
|
||||
except ImportError:
|
||||
# websockets not installed — use HTTP fallback mode
|
||||
logger.info(
|
||||
"websockets not installed; HTTP fallback: %s",
|
||||
self._gateway_url,
|
||||
)
|
||||
self._ws = None
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to connect to OpenClaw gateway: %s",
|
||||
self._gateway_url,
|
||||
)
|
||||
self._status = ChannelStatus.ERROR
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Close connection to the OpenClaw gateway."""
|
||||
self._stop_event.set()
|
||||
if self._ws is not None:
|
||||
try:
|
||||
self._ws.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing WebSocket", exc_info=True)
|
||||
self._ws = None
|
||||
if self._listener_thread is not None:
|
||||
self._listener_thread.join(timeout=self._reconnect_interval + 1)
|
||||
self._listener_thread = None
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a channel. Returns True on success."""
|
||||
payload = {
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
"metadata": metadata or {},
|
||||
}
|
||||
|
||||
# Try WebSocket first
|
||||
if self._ws is not None:
|
||||
try:
|
||||
self._ws.send(json.dumps(payload))
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"WebSocket send failed, falling back to HTTP",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
# HTTP fallback
|
||||
return self._send_http(payload, channel, content, conversation_id)
|
||||
|
||||
def _send_http(
|
||||
self,
|
||||
payload: Dict[str, Any],
|
||||
channel: str,
|
||||
content: str,
|
||||
conversation_id: str,
|
||||
) -> bool:
|
||||
"""Send message via HTTP POST fallback."""
|
||||
http_url = self._http_url("/send")
|
||||
try:
|
||||
import httpx
|
||||
|
||||
resp = httpx.post(http_url, json=payload, timeout=10.0)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning("HTTP send returned status %d", resp.status_code)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("HTTP send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Query the gateway for available channels."""
|
||||
http_url = self._http_url("/channels")
|
||||
try:
|
||||
import httpx
|
||||
|
||||
resp = httpx.get(http_url, timeout=10.0)
|
||||
data = resp.json()
|
||||
if isinstance(data, list):
|
||||
return data
|
||||
if isinstance(data, dict) and "channels" in data:
|
||||
return data["channels"]
|
||||
return []
|
||||
except Exception:
|
||||
logger.debug("Failed to list channels", exc_info=True)
|
||||
return []
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _listener_loop(self) -> None:
|
||||
"""Background loop receiving messages from WebSocket."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
if self._ws is None:
|
||||
break
|
||||
raw = self._ws.recv(timeout=1.0)
|
||||
data = json.loads(raw)
|
||||
msg = ChannelMessage(
|
||||
channel=data.get("channel", ""),
|
||||
sender=data.get("sender", ""),
|
||||
content=data.get("content", ""),
|
||||
message_id=data.get("message_id", ""),
|
||||
conversation_id=data.get("conversation_id", ""),
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
for handler in self._handlers:
|
||||
try:
|
||||
handler(msg)
|
||||
except Exception:
|
||||
logger.exception("Channel handler error")
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_RECEIVED,
|
||||
{
|
||||
"channel": msg.channel,
|
||||
"sender": msg.sender,
|
||||
"content": msg.content,
|
||||
"message_id": msg.message_id,
|
||||
},
|
||||
)
|
||||
except TimeoutError:
|
||||
continue
|
||||
except Exception:
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
logger.debug("WebSocket listener error, reconnecting", exc_info=True)
|
||||
self._status = ChannelStatus.CONNECTING
|
||||
self._stop_event.wait(self._reconnect_interval)
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
try:
|
||||
import websockets.sync.client # type: ignore[import-untyped]
|
||||
|
||||
self._ws = websockets.sync.client.connect(self._gateway_url)
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
except Exception:
|
||||
logger.debug("Reconnect failed", exc_info=True)
|
||||
self._status = ChannelStatus.ERROR
|
||||
|
||||
def _http_url(self, path: str = "") -> str:
|
||||
"""Convert the WebSocket URL to an HTTP URL."""
|
||||
url = self._gateway_url
|
||||
url = url.replace("ws://", "http://").replace("wss://", "https://")
|
||||
# Strip trailing /ws
|
||||
if url.endswith("/ws"):
|
||||
url = url[:-3]
|
||||
return url + path
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["OpenClawChannelBridge"]
|
||||
@@ -0,0 +1,131 @@
|
||||
"""SignalChannel — Signal adapter via signal-cli REST API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("signal")
|
||||
class SignalChannel(BaseChannel):
|
||||
"""Signal channel adapter via signal-cli REST API (send-only).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
api_url:
|
||||
Base URL of the signal-cli REST API. Falls back to
|
||||
``SIGNAL_API_URL`` env var.
|
||||
phone_number:
|
||||
Sender phone number registered with signal-cli. Falls back to
|
||||
``SIGNAL_PHONE_NUMBER`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "signal"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_url: str = "",
|
||||
*,
|
||||
phone_number: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._api_url = api_url or os.environ.get("SIGNAL_API_URL", "")
|
||||
self._phone_number = phone_number or os.environ.get(
|
||||
"SIGNAL_PHONE_NUMBER", "",
|
||||
)
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._api_url:
|
||||
logger.warning("No Signal API URL configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message via the signal-cli REST API."""
|
||||
if not self._api_url:
|
||||
logger.warning("Cannot send: no Signal API URL configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = f"{self._api_url}/v2/send"
|
||||
payload: Dict[str, Any] = {
|
||||
"message": content,
|
||||
"number": self._phone_number,
|
||||
"recipients": [channel],
|
||||
}
|
||||
|
||||
resp = httpx.post(url, json=payload, timeout=10.0)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Signal API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Signal send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["signal"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["SignalChannel"]
|
||||
@@ -0,0 +1,221 @@
|
||||
"""SlackChannel — native Slack Web API adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelMessage,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("slack")
|
||||
class SlackChannel(BaseChannel):
|
||||
"""Native Slack channel adapter using the Slack Web API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
bot_token:
|
||||
Slack Bot User OAuth Token. Falls back to ``SLACK_BOT_TOKEN`` env var.
|
||||
app_token:
|
||||
Slack App-Level Token for Socket Mode. Falls back to
|
||||
``SLACK_APP_TOKEN`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "slack"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bot_token: str = "",
|
||||
*,
|
||||
app_token: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._token = bot_token or os.environ.get("SLACK_BOT_TOKEN", "")
|
||||
self._app_token = app_token or os.environ.get("SLACK_APP_TOKEN", "")
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
self._listener_thread: Optional[threading.Thread] = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Start listening for incoming messages via Slack Socket Mode."""
|
||||
if not self._token:
|
||||
logger.warning("No Slack bot token configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
|
||||
self._stop_event.clear()
|
||||
self._status = ChannelStatus.CONNECTING
|
||||
|
||||
try:
|
||||
from slack_sdk.socket_mode import SocketModeClient # noqa: F401
|
||||
|
||||
if not self._app_token:
|
||||
logger.info("No app token for Socket Mode; send-only mode")
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
return
|
||||
|
||||
self._listener_thread = threading.Thread(
|
||||
target=self._socket_mode_loop, daemon=True,
|
||||
)
|
||||
self._listener_thread.start()
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
logger.info("Slack channel connected (Socket Mode)")
|
||||
except ImportError:
|
||||
logger.info("slack-sdk not installed; send-only mode")
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Stop the listener thread."""
|
||||
self._stop_event.set()
|
||||
if self._listener_thread is not None:
|
||||
self._listener_thread.join(timeout=5.0)
|
||||
self._listener_thread = None
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Slack channel via the Web API."""
|
||||
if not self._token:
|
||||
logger.warning("Cannot send: no Slack bot token")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = "https://slack.com/api/chat.postMessage"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"channel": channel,
|
||||
"text": content,
|
||||
}
|
||||
if conversation_id:
|
||||
payload["thread_ts"] = conversation_id
|
||||
|
||||
resp = httpx.post(
|
||||
url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
data = resp.json()
|
||||
if data.get("ok"):
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning("Slack API error: %s", data.get("error"))
|
||||
return False
|
||||
logger.warning(
|
||||
"Slack API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Slack send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["slack"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _socket_mode_loop(self) -> None:
|
||||
"""Run Slack Socket Mode client in a background thread."""
|
||||
try:
|
||||
from slack_sdk.socket_mode import SocketModeClient
|
||||
from slack_sdk.socket_mode.request import SocketModeRequest
|
||||
from slack_sdk.socket_mode.response import SocketModeResponse
|
||||
from slack_sdk.web import WebClient
|
||||
|
||||
client = SocketModeClient(
|
||||
app_token=self._app_token,
|
||||
web_client=WebClient(token=self._token),
|
||||
)
|
||||
|
||||
def _handle_event(client_obj, req: SocketModeRequest):
|
||||
if req.type == "events_api":
|
||||
event = req.payload.get("event", {})
|
||||
if event.get("type") == "message" and "subtype" not in event:
|
||||
cm = ChannelMessage(
|
||||
channel="slack",
|
||||
sender=event.get("user", ""),
|
||||
content=event.get("text", ""),
|
||||
message_id=event.get("ts", ""),
|
||||
conversation_id=event.get("channel", ""),
|
||||
)
|
||||
for handler in self._handlers:
|
||||
try:
|
||||
handler(cm)
|
||||
except Exception:
|
||||
logger.exception("Slack handler error")
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_RECEIVED,
|
||||
{
|
||||
"channel": cm.channel,
|
||||
"sender": cm.sender,
|
||||
"content": cm.content,
|
||||
"message_id": cm.message_id,
|
||||
},
|
||||
)
|
||||
client_obj.send_socket_mode_response(
|
||||
SocketModeResponse(envelope_id=req.envelope_id),
|
||||
)
|
||||
|
||||
client.socket_mode_request_listeners.append(_handle_event)
|
||||
client.connect()
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
self._stop_event.wait(1.0)
|
||||
|
||||
client.disconnect()
|
||||
except Exception:
|
||||
logger.debug("Slack Socket Mode loop error", exc_info=True)
|
||||
self._status = ChannelStatus.ERROR
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["SlackChannel"]
|
||||
@@ -0,0 +1,142 @@
|
||||
"""TeamsChannel — Microsoft Teams Bot Framework adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("teams")
|
||||
class TeamsChannel(BaseChannel):
|
||||
"""Microsoft Teams channel adapter using the Bot Framework REST API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
app_id:
|
||||
Microsoft App ID. Falls back to ``TEAMS_APP_ID`` env var.
|
||||
app_password:
|
||||
Microsoft App Password. Falls back to ``TEAMS_APP_PASSWORD`` env var.
|
||||
service_url:
|
||||
Bot Framework service URL. Falls back to ``TEAMS_SERVICE_URL`` env var
|
||||
(default ``https://smba.trafficmanager.net/teams``).
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "teams"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app_id: str = "",
|
||||
*,
|
||||
app_password: str = "",
|
||||
service_url: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._app_id = app_id or os.environ.get("TEAMS_APP_ID", "")
|
||||
self._app_password = app_password or os.environ.get("TEAMS_APP_PASSWORD", "")
|
||||
self._service_url = (
|
||||
service_url
|
||||
or os.environ.get("TEAMS_SERVICE_URL", "https://smba.trafficmanager.net/teams")
|
||||
)
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._app_id or not self._app_password:
|
||||
logger.warning("No Teams app_id or app_password configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Teams conversation via the Bot Framework API."""
|
||||
if not self._app_id or not self._app_password:
|
||||
logger.warning("Cannot send: no Teams credentials configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = f"{self._service_url}/v3/conversations/{channel}/activities"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._app_password}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"type": "message",
|
||||
"text": content,
|
||||
}
|
||||
if conversation_id:
|
||||
payload["replyToId"] = conversation_id
|
||||
|
||||
resp = httpx.post(
|
||||
url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Teams API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Teams send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["teams"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["TeamsChannel"]
|
||||
@@ -0,0 +1,202 @@
|
||||
"""TelegramChannel — native Telegram Bot API adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelMessage,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("telegram")
|
||||
class TelegramChannel(BaseChannel):
|
||||
"""Native Telegram channel adapter using the Bot API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
bot_token:
|
||||
Telegram Bot API token. Falls back to ``TELEGRAM_BOT_TOKEN`` env var.
|
||||
allowed_chat_ids:
|
||||
Comma-separated list of chat IDs allowed to interact.
|
||||
parse_mode:
|
||||
Message parse mode (``Markdown``, ``HTML``, etc.).
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "telegram"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bot_token: str = "",
|
||||
*,
|
||||
allowed_chat_ids: str = "",
|
||||
parse_mode: str = "Markdown",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._token = bot_token or os.environ.get("TELEGRAM_BOT_TOKEN", "")
|
||||
self._allowed_chat_ids = allowed_chat_ids
|
||||
self._parse_mode = parse_mode
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
self._listener_thread: Optional[threading.Thread] = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Start listening for incoming messages via long polling."""
|
||||
if not self._token:
|
||||
logger.warning("No Telegram bot token configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
|
||||
self._stop_event.clear()
|
||||
self._status = ChannelStatus.CONNECTING
|
||||
|
||||
try:
|
||||
from telegram.ext import ApplicationBuilder # noqa: F401
|
||||
|
||||
self._listener_thread = threading.Thread(
|
||||
target=self._poll_loop, daemon=True,
|
||||
)
|
||||
self._listener_thread.start()
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
logger.info("Telegram channel connected (long polling)")
|
||||
except ImportError:
|
||||
# python-telegram-bot not installed — send-only mode
|
||||
logger.info(
|
||||
"python-telegram-bot not installed; send-only mode",
|
||||
)
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Stop the listener thread."""
|
||||
self._stop_event.set()
|
||||
if self._listener_thread is not None:
|
||||
self._listener_thread.join(timeout=5.0)
|
||||
self._listener_thread = None
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message to a Telegram chat via the Bot API."""
|
||||
if not self._token:
|
||||
logger.warning("Cannot send: no Telegram bot token")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = f"https://api.telegram.org/bot{self._token}/sendMessage"
|
||||
payload: Dict[str, Any] = {
|
||||
"chat_id": channel,
|
||||
"text": content,
|
||||
}
|
||||
if self._parse_mode:
|
||||
payload["parse_mode"] = self._parse_mode
|
||||
if conversation_id:
|
||||
payload["reply_to_message_id"] = conversation_id
|
||||
|
||||
resp = httpx.post(url, json=payload, timeout=10.0)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Telegram API returned status %d: %s",
|
||||
resp.status_code,
|
||||
resp.text,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Telegram send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["telegram"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _poll_loop(self) -> None:
|
||||
"""Long-poll for updates using python-telegram-bot."""
|
||||
try:
|
||||
from telegram.ext import ApplicationBuilder, MessageHandler, filters
|
||||
|
||||
app = ApplicationBuilder().token(self._token).build()
|
||||
|
||||
def _handle_msg(update, context):
|
||||
msg = update.message
|
||||
if msg is None:
|
||||
return
|
||||
cm = ChannelMessage(
|
||||
channel="telegram",
|
||||
sender=str(msg.from_user.id) if msg.from_user else "",
|
||||
content=msg.text or "",
|
||||
message_id=str(msg.message_id),
|
||||
conversation_id=str(msg.chat.id),
|
||||
)
|
||||
for handler in self._handlers:
|
||||
try:
|
||||
handler(cm)
|
||||
except Exception:
|
||||
logger.exception("Telegram handler error")
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_RECEIVED,
|
||||
{
|
||||
"channel": cm.channel,
|
||||
"sender": cm.sender,
|
||||
"content": cm.content,
|
||||
"message_id": cm.message_id,
|
||||
},
|
||||
)
|
||||
|
||||
app.add_handler(MessageHandler(filters.TEXT, _handle_msg))
|
||||
app.run_polling(stop_signals=None)
|
||||
except Exception:
|
||||
logger.debug("Telegram poll loop error", exc_info=True)
|
||||
self._status = ChannelStatus.ERROR
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["TelegramChannel"]
|
||||
@@ -0,0 +1,114 @@
|
||||
"""WebChatChannel — in-memory message queue for webchat/testing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelMessage,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("webchat")
|
||||
class WebChatChannel(BaseChannel):
|
||||
"""In-memory webchat channel for testing and embedded web UIs.
|
||||
|
||||
Messages are stored in an internal list and can be retrieved via
|
||||
:meth:`get_messages`. No external dependencies are required.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "webchat"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._messages: List[ChannelMessage] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (always succeeds for in-memory channel)."""
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Append a message to the in-memory queue."""
|
||||
msg = ChannelMessage(
|
||||
channel=channel,
|
||||
sender="jarvis",
|
||||
content=content,
|
||||
conversation_id=conversation_id,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
self._messages.append(msg)
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["webchat"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- webchat-specific helpers -----------------------------------------------
|
||||
|
||||
def get_messages(self) -> List[ChannelMessage]:
|
||||
"""Return all messages stored in the in-memory queue."""
|
||||
return self._messages
|
||||
|
||||
def clear_messages(self) -> None:
|
||||
"""Clear all messages from the in-memory queue."""
|
||||
self._messages.clear()
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["WebChatChannel"]
|
||||
@@ -0,0 +1,139 @@
|
||||
"""WebhookChannel — generic outbound webhook adapter (zero extra deps)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("webhook")
|
||||
class WebhookChannel(BaseChannel):
|
||||
"""Generic outbound webhook channel (send-only).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
url:
|
||||
Target webhook URL.
|
||||
secret:
|
||||
Optional shared secret sent in the ``X-Webhook-Secret`` header.
|
||||
method:
|
||||
HTTP method (default ``POST``).
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "webhook"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str = "",
|
||||
*,
|
||||
secret: str = "",
|
||||
method: str = "POST",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._url = url
|
||||
self._secret = secret
|
||||
self._method = method.upper()
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._url:
|
||||
logger.warning("No webhook URL configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""POST a JSON payload to the configured webhook URL."""
|
||||
if not self._url:
|
||||
logger.warning("Cannot send: no webhook URL configured")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
}
|
||||
if conversation_id:
|
||||
payload["conversation_id"] = conversation_id
|
||||
if metadata:
|
||||
payload["metadata"] = metadata
|
||||
|
||||
headers: Dict[str, str] = {}
|
||||
if self._secret:
|
||||
headers["X-Webhook-Secret"] = self._secret
|
||||
|
||||
resp = httpx.request(
|
||||
self._method, self._url, json=payload,
|
||||
headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"Webhook returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("Webhook send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["webhook"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages (no-op for webhook)."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["WebhookChannel"]
|
||||
@@ -0,0 +1,141 @@
|
||||
"""WhatsAppChannel — WhatsApp Cloud API adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from openjarvis.channels._stubs import (
|
||||
BaseChannel,
|
||||
ChannelHandler,
|
||||
ChannelStatus,
|
||||
)
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ChannelRegistry.register("whatsapp")
|
||||
class WhatsAppChannel(BaseChannel):
|
||||
"""WhatsApp Cloud API channel adapter (send-only).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
access_token:
|
||||
WhatsApp Cloud API access token. Falls back to
|
||||
``WHATSAPP_ACCESS_TOKEN`` env var.
|
||||
phone_number_id:
|
||||
WhatsApp phone number ID. Falls back to
|
||||
``WHATSAPP_PHONE_NUMBER_ID`` env var.
|
||||
bus:
|
||||
Optional event bus for publishing channel events.
|
||||
"""
|
||||
|
||||
channel_id = "whatsapp"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
access_token: str = "",
|
||||
*,
|
||||
phone_number_id: str = "",
|
||||
bus: Optional[EventBus] = None,
|
||||
) -> None:
|
||||
self._token = access_token or os.environ.get("WHATSAPP_ACCESS_TOKEN", "")
|
||||
self._phone_number_id = phone_number_id or os.environ.get(
|
||||
"WHATSAPP_PHONE_NUMBER_ID", "",
|
||||
)
|
||||
self._bus = bus
|
||||
self._handlers: List[ChannelHandler] = []
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- connection lifecycle ---------------------------------------------------
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Mark as connected (send-only — no persistent connection)."""
|
||||
if not self._token:
|
||||
logger.warning("No WhatsApp access token configured")
|
||||
self._status = ChannelStatus.ERROR
|
||||
return
|
||||
self._status = ChannelStatus.CONNECTED
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Mark as disconnected."""
|
||||
self._status = ChannelStatus.DISCONNECTED
|
||||
|
||||
# -- send / receive --------------------------------------------------------
|
||||
|
||||
def send(
|
||||
self,
|
||||
channel: str,
|
||||
content: str,
|
||||
*,
|
||||
conversation_id: str = "",
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""Send a message via the WhatsApp Cloud API."""
|
||||
if not self._token:
|
||||
logger.warning("Cannot send: no WhatsApp access token")
|
||||
return False
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
url = (
|
||||
f"https://graph.facebook.com/v21.0/"
|
||||
f"{self._phone_number_id}/messages"
|
||||
)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"messaging_product": "whatsapp",
|
||||
"to": channel,
|
||||
"type": "text",
|
||||
"text": {"body": content},
|
||||
}
|
||||
|
||||
resp = httpx.post(
|
||||
url, json=payload, headers=headers, timeout=10.0,
|
||||
)
|
||||
if resp.status_code < 300:
|
||||
self._publish_sent(channel, content, conversation_id)
|
||||
return True
|
||||
logger.warning(
|
||||
"WhatsApp API returned status %d", resp.status_code,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.debug("WhatsApp send failed", exc_info=True)
|
||||
return False
|
||||
|
||||
def status(self) -> ChannelStatus:
|
||||
"""Return the current connection status."""
|
||||
return self._status
|
||||
|
||||
def list_channels(self) -> List[str]:
|
||||
"""Return available channel identifiers."""
|
||||
return ["whatsapp"]
|
||||
|
||||
def on_message(self, handler: ChannelHandler) -> None:
|
||||
"""Register a callback for incoming messages."""
|
||||
self._handlers.append(handler)
|
||||
|
||||
# -- internal helpers -------------------------------------------------------
|
||||
|
||||
def _publish_sent(self, channel: str, content: str, conversation_id: str) -> None:
|
||||
"""Publish a CHANNEL_MESSAGE_SENT event on the bus."""
|
||||
if self._bus is not None:
|
||||
self._bus.publish(
|
||||
EventType.CHANNEL_MESSAGE_SENT,
|
||||
{
|
||||
"channel": channel,
|
||||
"content": content,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["WhatsAppChannel"]
|
||||
@@ -2,10 +2,137 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import click
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
_CHANNEL_TYPE_HELP = (
|
||||
"Channel type (telegram, discord, slack, webhook, email, "
|
||||
"whatsapp, signal, google_chat, irc, webchat, teams, "
|
||||
"matrix, mattermost, feishu, bluebubbles)."
|
||||
)
|
||||
|
||||
|
||||
def _get_channel(
|
||||
channel_type: str | None,
|
||||
config: Any,
|
||||
) -> Any:
|
||||
"""Resolve a channel backend by type.
|
||||
|
||||
Resolution order: ``--channel-type`` flag >
|
||||
``config.channel.default_channel`` > error.
|
||||
"""
|
||||
import openjarvis.channels # noqa: F401 -- trigger registration
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
key = channel_type or config.channel.default_channel
|
||||
if not key:
|
||||
raise click.ClickException(
|
||||
"No channel type specified. Use --channel-type or set "
|
||||
"default_channel in [channel] config."
|
||||
)
|
||||
|
||||
kwargs: Dict[str, Any] = {}
|
||||
if key == "telegram":
|
||||
tc = config.channel.telegram
|
||||
if tc.bot_token:
|
||||
kwargs["bot_token"] = tc.bot_token
|
||||
elif key == "discord":
|
||||
dc = config.channel.discord
|
||||
if dc.bot_token:
|
||||
kwargs["bot_token"] = dc.bot_token
|
||||
elif key == "slack":
|
||||
sc = config.channel.slack
|
||||
if sc.bot_token:
|
||||
kwargs["bot_token"] = sc.bot_token
|
||||
if sc.app_token:
|
||||
kwargs["app_token"] = sc.app_token
|
||||
elif key == "webhook":
|
||||
wc = config.channel.webhook
|
||||
if wc.url:
|
||||
kwargs["url"] = wc.url
|
||||
if wc.secret:
|
||||
kwargs["secret"] = wc.secret
|
||||
if wc.method:
|
||||
kwargs["method"] = wc.method
|
||||
elif key == "email":
|
||||
ec = config.channel.email
|
||||
if ec.smtp_host:
|
||||
kwargs["smtp_host"] = ec.smtp_host
|
||||
kwargs["smtp_port"] = ec.smtp_port
|
||||
if ec.username:
|
||||
kwargs["username"] = ec.username
|
||||
if ec.password:
|
||||
kwargs["password"] = ec.password
|
||||
kwargs["use_tls"] = ec.use_tls
|
||||
elif key == "whatsapp":
|
||||
wac = config.channel.whatsapp
|
||||
if wac.access_token:
|
||||
kwargs["access_token"] = wac.access_token
|
||||
if wac.phone_number_id:
|
||||
kwargs["phone_number_id"] = wac.phone_number_id
|
||||
elif key == "signal":
|
||||
sgc = config.channel.signal
|
||||
if sgc.api_url:
|
||||
kwargs["api_url"] = sgc.api_url
|
||||
if sgc.phone_number:
|
||||
kwargs["phone_number"] = sgc.phone_number
|
||||
elif key == "google_chat":
|
||||
gcc = config.channel.google_chat
|
||||
if gcc.webhook_url:
|
||||
kwargs["webhook_url"] = gcc.webhook_url
|
||||
elif key == "irc":
|
||||
ic = config.channel.irc
|
||||
if ic.server:
|
||||
kwargs["server"] = ic.server
|
||||
kwargs["port"] = ic.port
|
||||
if ic.nick:
|
||||
kwargs["nick"] = ic.nick
|
||||
if ic.password:
|
||||
kwargs["password"] = ic.password
|
||||
kwargs["use_tls"] = ic.use_tls
|
||||
elif key == "webchat":
|
||||
pass # no config needed
|
||||
elif key == "teams":
|
||||
tmc = config.channel.teams
|
||||
if tmc.app_id:
|
||||
kwargs["app_id"] = tmc.app_id
|
||||
if tmc.app_password:
|
||||
kwargs["app_password"] = tmc.app_password
|
||||
if tmc.service_url:
|
||||
kwargs["service_url"] = tmc.service_url
|
||||
elif key == "matrix":
|
||||
mc = config.channel.matrix
|
||||
if mc.homeserver:
|
||||
kwargs["homeserver"] = mc.homeserver
|
||||
if mc.access_token:
|
||||
kwargs["access_token"] = mc.access_token
|
||||
elif key == "mattermost":
|
||||
mmc = config.channel.mattermost
|
||||
if mmc.url:
|
||||
kwargs["url"] = mmc.url
|
||||
if mmc.token:
|
||||
kwargs["token"] = mmc.token
|
||||
elif key == "feishu":
|
||||
fc = config.channel.feishu
|
||||
if fc.app_id:
|
||||
kwargs["app_id"] = fc.app_id
|
||||
if fc.app_secret:
|
||||
kwargs["app_secret"] = fc.app_secret
|
||||
elif key == "bluebubbles":
|
||||
bbc = config.channel.bluebubbles
|
||||
if bbc.url:
|
||||
kwargs["url"] = bbc.url
|
||||
if bbc.password:
|
||||
kwargs["password"] = bbc.password
|
||||
|
||||
if not ChannelRegistry.contains(key):
|
||||
raise click.ClickException(f"Unknown channel type: {key}")
|
||||
|
||||
return ChannelRegistry.create(key, **kwargs)
|
||||
|
||||
|
||||
@click.group()
|
||||
def channel() -> None:
|
||||
@@ -13,82 +140,100 @@ def channel() -> None:
|
||||
|
||||
|
||||
@channel.command("list")
|
||||
@click.option("--gateway", default=None, help="OpenClaw gateway URL.")
|
||||
def channel_list(gateway: str | None) -> None:
|
||||
"""List available channels from the gateway."""
|
||||
@click.option(
|
||||
"--channel-type", default=None, help=_CHANNEL_TYPE_HELP,
|
||||
)
|
||||
def channel_list(
|
||||
channel_type: Optional[str],
|
||||
) -> None:
|
||||
"""List available channels."""
|
||||
console = Console()
|
||||
from openjarvis.core.config import load_config
|
||||
|
||||
config = load_config()
|
||||
gw_url = gateway or config.channel.gateway_url
|
||||
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge
|
||||
|
||||
bridge = OpenClawChannelBridge(gateway_url=gw_url)
|
||||
|
||||
try:
|
||||
channels = bridge.list_channels()
|
||||
ch = _get_channel(channel_type, config)
|
||||
except click.ClickException as exc:
|
||||
console.print(f"[red]{exc.message}[/red]")
|
||||
return
|
||||
|
||||
try:
|
||||
channels = ch.list_channels()
|
||||
except Exception as exc:
|
||||
console.print(f"[red]Failed to list channels: {exc}[/red]")
|
||||
return
|
||||
|
||||
if not channels:
|
||||
console.print(
|
||||
"[yellow]No channels available"
|
||||
" (is OpenClaw gateway running?)[/yellow]"
|
||||
)
|
||||
console.print("[yellow]No channels available[/yellow]")
|
||||
return
|
||||
|
||||
table = Table(title="Available Channels")
|
||||
table.add_column("Channel", style="cyan")
|
||||
for ch in channels:
|
||||
table.add_row(ch)
|
||||
for name in channels:
|
||||
table.add_row(name)
|
||||
console.print(table)
|
||||
|
||||
|
||||
@channel.command("send")
|
||||
@click.argument("target")
|
||||
@click.argument("message")
|
||||
@click.option("--gateway", default=None, help="OpenClaw gateway URL.")
|
||||
def channel_send(target: str, message: str, gateway: str | None) -> None:
|
||||
@click.option(
|
||||
"--channel-type", default=None, help=_CHANNEL_TYPE_HELP,
|
||||
)
|
||||
def channel_send(
|
||||
target: str,
|
||||
message: str,
|
||||
channel_type: Optional[str],
|
||||
) -> None:
|
||||
"""Send a message to a channel."""
|
||||
console = Console()
|
||||
from openjarvis.core.config import load_config
|
||||
|
||||
config = load_config()
|
||||
gw_url = gateway or config.channel.gateway_url
|
||||
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge
|
||||
try:
|
||||
ch = _get_channel(channel_type, config)
|
||||
except click.ClickException as exc:
|
||||
console.print(f"[red]{exc.message}[/red]")
|
||||
return
|
||||
|
||||
bridge = OpenClawChannelBridge(gateway_url=gw_url)
|
||||
|
||||
ok = bridge.send(target, message)
|
||||
ok = ch.send(target, message)
|
||||
if ok:
|
||||
console.print(f"[green]Message sent to {target}[/green]")
|
||||
else:
|
||||
console.print(f"[red]Failed to send message to {target}[/red]")
|
||||
console.print(
|
||||
f"[red]Failed to send message to {target}[/red]",
|
||||
)
|
||||
|
||||
|
||||
@channel.command("status")
|
||||
@click.option("--gateway", default=None, help="OpenClaw gateway URL.")
|
||||
def channel_status(gateway: str | None) -> None:
|
||||
"""Show channel bridge connection status."""
|
||||
@click.option(
|
||||
"--channel-type", default=None, help=_CHANNEL_TYPE_HELP,
|
||||
)
|
||||
def channel_status(
|
||||
channel_type: Optional[str],
|
||||
) -> None:
|
||||
"""Show channel connection status."""
|
||||
console = Console()
|
||||
from openjarvis.core.config import load_config
|
||||
|
||||
config = load_config()
|
||||
gw_url = gateway or config.channel.gateway_url
|
||||
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge
|
||||
try:
|
||||
ch = _get_channel(channel_type, config)
|
||||
except click.ClickException as exc:
|
||||
console.print(f"[red]{exc.message}[/red]")
|
||||
return
|
||||
|
||||
bridge = OpenClawChannelBridge(gateway_url=gw_url)
|
||||
|
||||
st = bridge.status()
|
||||
st = ch.status()
|
||||
color = {
|
||||
"connected": "green",
|
||||
"disconnected": "yellow",
|
||||
"connecting": "blue",
|
||||
"error": "red",
|
||||
}.get(st.value, "white")
|
||||
console.print(f"Gateway: [cyan]{gw_url}[/cyan]")
|
||||
|
||||
key = channel_type or config.channel.default_channel or "unknown"
|
||||
console.print(f"Channel: [cyan]{key}[/cyan]")
|
||||
console.print(f"Status: [{color}]{st.value}[/{color}]")
|
||||
|
||||
+13
-13
@@ -142,23 +142,23 @@ def serve(
|
||||
console.print(f"[yellow]Agent '{agent_key}' failed to load: {exc}[/yellow]")
|
||||
traceback.print_exc()
|
||||
|
||||
# Set up channel bridge if enabled
|
||||
# Set up channel backend if enabled
|
||||
channel_bridge = None
|
||||
if config.channel.enabled:
|
||||
if config.channel.enabled and config.channel.default_channel:
|
||||
try:
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge
|
||||
from openjarvis.system import SystemBuilder
|
||||
|
||||
channel_bridge = OpenClawChannelBridge(
|
||||
gateway_url=config.channel.gateway_url,
|
||||
reconnect_interval=config.channel.reconnect_interval,
|
||||
bus=bus,
|
||||
)
|
||||
channel_bridge.connect()
|
||||
console.print(
|
||||
f" Channels: [cyan]{config.channel.gateway_url}[/cyan]"
|
||||
)
|
||||
# Reuse _resolve_channel logic from SystemBuilder
|
||||
sb = SystemBuilder(config)
|
||||
sb._bus = bus
|
||||
channel_bridge = sb._resolve_channel(config, bus)
|
||||
if channel_bridge is not None:
|
||||
channel_bridge.connect()
|
||||
console.print(
|
||||
f" Channel: [cyan]{config.channel.default_channel}[/cyan]"
|
||||
)
|
||||
except Exception as exc:
|
||||
console.print(f"[yellow]Channel bridge failed to start: {exc}[/yellow]")
|
||||
console.print(f"[yellow]Channel failed to start: {exc}[/yellow]")
|
||||
channel_bridge = None
|
||||
|
||||
# Create app
|
||||
|
||||
@@ -297,13 +297,155 @@ class TracesConfig:
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TelegramChannelConfig:
|
||||
"""Per-channel config for Telegram."""
|
||||
|
||||
bot_token: str = ""
|
||||
allowed_chat_ids: str = ""
|
||||
parse_mode: str = "Markdown"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DiscordChannelConfig:
|
||||
"""Per-channel config for Discord."""
|
||||
|
||||
bot_token: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SlackChannelConfig:
|
||||
"""Per-channel config for Slack."""
|
||||
|
||||
bot_token: str = ""
|
||||
app_token: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class WebhookChannelConfig:
|
||||
"""Per-channel config for generic webhooks."""
|
||||
|
||||
url: str = ""
|
||||
secret: str = ""
|
||||
method: str = "POST"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class EmailChannelConfig:
|
||||
"""Per-channel config for email (SMTP/IMAP)."""
|
||||
|
||||
smtp_host: str = ""
|
||||
smtp_port: int = 587
|
||||
imap_host: str = ""
|
||||
imap_port: int = 993
|
||||
username: str = ""
|
||||
password: str = ""
|
||||
use_tls: bool = True
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class WhatsAppChannelConfig:
|
||||
"""Per-channel config for WhatsApp Cloud API."""
|
||||
|
||||
access_token: str = ""
|
||||
phone_number_id: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SignalChannelConfig:
|
||||
"""Per-channel config for Signal (via signal-cli REST API)."""
|
||||
|
||||
api_url: str = ""
|
||||
phone_number: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class GoogleChatChannelConfig:
|
||||
"""Per-channel config for Google Chat webhooks."""
|
||||
|
||||
webhook_url: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IRCChannelConfig:
|
||||
"""Per-channel config for IRC."""
|
||||
|
||||
server: str = ""
|
||||
port: int = 6667
|
||||
nick: str = ""
|
||||
password: str = ""
|
||||
use_tls: bool = False
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class WebChatChannelConfig:
|
||||
"""Per-channel config for in-memory webchat."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TeamsChannelConfig:
|
||||
"""Per-channel config for Microsoft Teams (Bot Framework)."""
|
||||
|
||||
app_id: str = ""
|
||||
app_password: str = ""
|
||||
service_url: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class MatrixChannelConfig:
|
||||
"""Per-channel config for Matrix."""
|
||||
|
||||
homeserver: str = ""
|
||||
access_token: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class MattermostChannelConfig:
|
||||
"""Per-channel config for Mattermost."""
|
||||
|
||||
url: str = ""
|
||||
token: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class FeishuChannelConfig:
|
||||
"""Per-channel config for Feishu (Lark)."""
|
||||
|
||||
app_id: str = ""
|
||||
app_secret: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class BlueBubblesChannelConfig:
|
||||
"""Per-channel config for BlueBubbles (iMessage bridge)."""
|
||||
|
||||
url: str = ""
|
||||
password: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChannelConfig:
|
||||
"""Channel messaging settings."""
|
||||
|
||||
enabled: bool = False
|
||||
gateway_url: str = "ws://127.0.0.1:18789/ws"
|
||||
default_channel: str = ""
|
||||
default_agent: str = "simple"
|
||||
reconnect_interval: float = 5.0
|
||||
telegram: TelegramChannelConfig = field(default_factory=TelegramChannelConfig)
|
||||
discord: DiscordChannelConfig = field(default_factory=DiscordChannelConfig)
|
||||
slack: SlackChannelConfig = field(default_factory=SlackChannelConfig)
|
||||
webhook: WebhookChannelConfig = field(default_factory=WebhookChannelConfig)
|
||||
email: EmailChannelConfig = field(default_factory=EmailChannelConfig)
|
||||
whatsapp: WhatsAppChannelConfig = field(default_factory=WhatsAppChannelConfig)
|
||||
signal: SignalChannelConfig = field(default_factory=SignalChannelConfig)
|
||||
google_chat: GoogleChatChannelConfig = field(default_factory=GoogleChatChannelConfig)
|
||||
irc: IRCChannelConfig = field(default_factory=IRCChannelConfig)
|
||||
webchat: WebChatChannelConfig = field(default_factory=WebChatChannelConfig)
|
||||
teams: TeamsChannelConfig = field(default_factory=TeamsChannelConfig)
|
||||
matrix: MatrixChannelConfig = field(default_factory=MatrixChannelConfig)
|
||||
mattermost: MattermostChannelConfig = field(default_factory=MattermostChannelConfig)
|
||||
feishu: FeishuChannelConfig = field(default_factory=FeishuChannelConfig)
|
||||
bluebubbles: BlueBubblesChannelConfig = field(default_factory=BlueBubblesChannelConfig)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -379,7 +521,7 @@ def load_config(path: Optional[Path] = None) -> JarvisConfig:
|
||||
# Simple top-level sections
|
||||
simple_sections = (
|
||||
"engine", "intelligence", "learning",
|
||||
"agent", "server", "telemetry", "traces", "channel", "security",
|
||||
"agent", "server", "telemetry", "traces", "security",
|
||||
)
|
||||
for section_name in simple_sections:
|
||||
if section_name in data:
|
||||
@@ -389,6 +531,22 @@ def load_config(path: Optional[Path] = None) -> JarvisConfig:
|
||||
if "memory" in data:
|
||||
_apply_toml_section(cfg.tools.storage, data["memory"])
|
||||
|
||||
# [channel] with nested per-channel sub-configs
|
||||
if "channel" in data:
|
||||
ch_data = data["channel"]
|
||||
# Top-level channel keys (enabled, default_channel, etc.)
|
||||
for key, value in ch_data.items():
|
||||
if not isinstance(value, dict) and hasattr(cfg.channel, key):
|
||||
setattr(cfg.channel, key, value)
|
||||
# Nested per-channel configs
|
||||
for sub in (
|
||||
"telegram", "discord", "slack", "webhook", "email",
|
||||
"whatsapp", "signal", "google_chat", "irc", "webchat",
|
||||
"teams", "matrix", "mattermost", "feishu", "bluebubbles",
|
||||
):
|
||||
if sub in ch_data and isinstance(ch_data[sub], dict):
|
||||
_apply_toml_section(getattr(cfg.channel, sub), ch_data[sub])
|
||||
|
||||
# Tools: accept [tools] and nested [tools.storage], [tools.mcp]
|
||||
if "tools" in data:
|
||||
tools_data = data["tools"]
|
||||
@@ -462,9 +620,56 @@ enabled = false
|
||||
|
||||
[channel]
|
||||
enabled = false
|
||||
gateway_url = "ws://127.0.0.1:18789/ws"
|
||||
default_agent = "simple"
|
||||
reconnect_interval = 5.0
|
||||
|
||||
# [channel.telegram]
|
||||
# bot_token = "" # Or set TELEGRAM_BOT_TOKEN env var
|
||||
|
||||
# [channel.discord]
|
||||
# bot_token = "" # Or set DISCORD_BOT_TOKEN env var
|
||||
|
||||
# [channel.slack]
|
||||
# bot_token = "" # Or set SLACK_BOT_TOKEN env var
|
||||
|
||||
# [channel.webhook]
|
||||
# url = ""
|
||||
|
||||
# [channel.whatsapp]
|
||||
# access_token = "" # Or set WHATSAPP_ACCESS_TOKEN env var
|
||||
# phone_number_id = "" # Or set WHATSAPP_PHONE_NUMBER_ID env var
|
||||
|
||||
# [channel.signal]
|
||||
# api_url = "" # signal-cli REST API URL
|
||||
# phone_number = "" # Or set SIGNAL_PHONE_NUMBER env var
|
||||
|
||||
# [channel.google_chat]
|
||||
# webhook_url = "" # Or set GOOGLE_CHAT_WEBHOOK_URL env var
|
||||
|
||||
# [channel.irc]
|
||||
# server = ""
|
||||
# port = 6667
|
||||
# nick = ""
|
||||
# use_tls = false
|
||||
|
||||
# [channel.teams]
|
||||
# app_id = "" # Or set TEAMS_APP_ID env var
|
||||
# app_password = "" # Or set TEAMS_APP_PASSWORD env var
|
||||
|
||||
# [channel.matrix]
|
||||
# homeserver = "" # Or set MATRIX_HOMESERVER env var
|
||||
# access_token = "" # Or set MATRIX_ACCESS_TOKEN env var
|
||||
|
||||
# [channel.mattermost]
|
||||
# url = "" # Or set MATTERMOST_URL env var
|
||||
# token = "" # Or set MATTERMOST_TOKEN env var
|
||||
|
||||
# [channel.feishu]
|
||||
# app_id = "" # Or set FEISHU_APP_ID env var
|
||||
# app_secret = "" # Or set FEISHU_APP_SECRET env var
|
||||
|
||||
# [channel.bluebubbles]
|
||||
# url = "" # Or set BLUEBUBBLES_URL env var
|
||||
# password = "" # Or set BLUEBUBBLES_PASSWORD env var
|
||||
|
||||
[security]
|
||||
enabled = true
|
||||
@@ -479,23 +684,38 @@ enforce_tool_confirmation = true
|
||||
|
||||
__all__ = [
|
||||
"AgentConfig",
|
||||
"BlueBubblesChannelConfig",
|
||||
"ChannelConfig",
|
||||
"DEFAULT_CONFIG_DIR",
|
||||
"DEFAULT_CONFIG_PATH",
|
||||
"DiscordChannelConfig",
|
||||
"EmailChannelConfig",
|
||||
"EngineConfig",
|
||||
"FeishuChannelConfig",
|
||||
"GoogleChatChannelConfig",
|
||||
"GpuInfo",
|
||||
"HardwareInfo",
|
||||
"IRCChannelConfig",
|
||||
"IntelligenceConfig",
|
||||
"JarvisConfig",
|
||||
"LearningConfig",
|
||||
"MCPConfig",
|
||||
"MatrixChannelConfig",
|
||||
"MattermostChannelConfig",
|
||||
"MemoryConfig",
|
||||
"SecurityConfig",
|
||||
"ServerConfig",
|
||||
"SignalChannelConfig",
|
||||
"SlackChannelConfig",
|
||||
"StorageConfig",
|
||||
"TeamsChannelConfig",
|
||||
"TelegramChannelConfig",
|
||||
"TelemetryConfig",
|
||||
"ToolsConfig",
|
||||
"TracesConfig",
|
||||
"WebChatChannelConfig",
|
||||
"WebhookChannelConfig",
|
||||
"WhatsAppChannelConfig",
|
||||
"detect_hardware",
|
||||
"generate_default_toml",
|
||||
"load_config",
|
||||
|
||||
@@ -21,6 +21,12 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# LiteLLM engine is optional — only register if litellm is installed
|
||||
try:
|
||||
import openjarvis.engine.litellm # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__all__ = [
|
||||
"EngineConnectionError",
|
||||
"InferenceEngine",
|
||||
|
||||
@@ -15,6 +15,7 @@ _HOST_MAP: Dict[str, str | None] = {
|
||||
"llamacpp": "llamacpp_host",
|
||||
"sglang": "sglang_host",
|
||||
"cloud": None,
|
||||
"litellm": None,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
"""LiteLLM inference engine — unified access to 100+ LLM providers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.core.types import Message
|
||||
from openjarvis.engine._base import InferenceEngine, messages_to_dicts
|
||||
|
||||
|
||||
@EngineRegistry.register("litellm")
|
||||
class LiteLLMEngine(InferenceEngine):
|
||||
"""Inference via LiteLLM — routes to any supported provider.
|
||||
|
||||
LiteLLM normalizes all providers (OpenAI, Anthropic, Google, DeepSeek,
|
||||
Groq, Together, Fireworks, OpenRouter, Mistral, Cohere, xAI, Perplexity,
|
||||
etc.) to OpenAI-format input/output. Model selection uses LiteLLM's
|
||||
``provider/model`` convention, e.g. ``anthropic/claude-sonnet-4-20250514``.
|
||||
|
||||
API keys are read from environment variables following each provider's
|
||||
convention (OPENAI_API_KEY, ANTHROPIC_API_KEY, GROQ_API_KEY, etc.).
|
||||
"""
|
||||
|
||||
engine_id = "litellm"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
api_base: str | None = None,
|
||||
default_model: str | None = None,
|
||||
) -> None:
|
||||
self._api_base = api_base
|
||||
self._default_model = default_model
|
||||
|
||||
def generate(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
model: str,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 1024,
|
||||
**kwargs: Any,
|
||||
) -> Dict[str, Any]:
|
||||
import litellm
|
||||
|
||||
call_kwargs: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages_to_dicts(messages),
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
if self._api_base:
|
||||
call_kwargs["api_base"] = self._api_base
|
||||
# Pass through tools if provided
|
||||
if "tools" in kwargs:
|
||||
call_kwargs["tools"] = kwargs.pop("tools")
|
||||
call_kwargs.update(kwargs)
|
||||
|
||||
resp = litellm.completion(**call_kwargs)
|
||||
|
||||
choice = resp.choices[0]
|
||||
usage = resp.usage
|
||||
prompt_tokens = usage.prompt_tokens if usage else 0
|
||||
completion_tokens = usage.completion_tokens if usage else 0
|
||||
|
||||
result: Dict[str, Any] = {
|
||||
"content": choice.message.content or "",
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": (usage.total_tokens if usage else 0),
|
||||
},
|
||||
"model": resp.model,
|
||||
"finish_reason": choice.finish_reason or "stop",
|
||||
}
|
||||
|
||||
# Extract tool_calls in OpenAI format
|
||||
if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
|
||||
result["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in choice.message.tool_calls
|
||||
]
|
||||
|
||||
# Cost tracking via litellm's built-in cost calculation
|
||||
try:
|
||||
cost = litellm.completion_cost(completion_response=resp)
|
||||
result["cost_usd"] = cost
|
||||
except Exception:
|
||||
result["cost_usd"] = 0.0
|
||||
|
||||
return result
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
model: str,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 1024,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[str]:
|
||||
import litellm
|
||||
|
||||
call_kwargs: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages_to_dicts(messages),
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": True,
|
||||
}
|
||||
if self._api_base:
|
||||
call_kwargs["api_base"] = self._api_base
|
||||
call_kwargs.update(kwargs)
|
||||
|
||||
resp = litellm.completion(**call_kwargs)
|
||||
for chunk in resp:
|
||||
delta = chunk.choices[0].delta if chunk.choices else None
|
||||
if delta and delta.content:
|
||||
yield delta.content
|
||||
|
||||
def list_models(self) -> List[str]:
|
||||
if self._default_model:
|
||||
return [self._default_model]
|
||||
return []
|
||||
|
||||
def health(self) -> bool:
|
||||
try:
|
||||
import litellm # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
__all__ = ["LiteLLMEngine"]
|
||||
@@ -47,6 +47,12 @@ def ensure_registered() -> None:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Orchestrator-native SFT & GRPO training
|
||||
try:
|
||||
import openjarvis.learning.orchestrator # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HeuristicRewardFunction",
|
||||
|
||||
@@ -32,7 +32,7 @@ class RewardFunction(ABC):
|
||||
class LearningPolicy(ABC):
|
||||
"""Base for all learning policies. Targets one or more pillars."""
|
||||
|
||||
target: ClassVar[str] = "" # "intelligence" | "agent" | "tools"
|
||||
target: ClassVar[str] = "" # "intelligence" | "agent"
|
||||
|
||||
@abstractmethod
|
||||
def update(self, trace_store: Any, **kwargs: object) -> Dict[str, Any]:
|
||||
@@ -51,12 +51,6 @@ class AgentLearningPolicy(LearningPolicy):
|
||||
target: ClassVar[str] = "agent"
|
||||
|
||||
|
||||
class ToolLearningPolicy(LearningPolicy):
|
||||
"""Updates tool usage (ICL examples, skills) from traces."""
|
||||
|
||||
target: ClassVar[str] = "tools"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AgentLearningPolicy",
|
||||
"IntelligenceLearningPolicy",
|
||||
@@ -65,5 +59,4 @@ __all__ = [
|
||||
"RewardFunction",
|
||||
"RouterPolicy",
|
||||
"RoutingContext",
|
||||
"ToolLearningPolicy",
|
||||
]
|
||||
|
||||
@@ -5,16 +5,17 @@ from __future__ import annotations
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from openjarvis.core.registry import LearningRegistry
|
||||
from openjarvis.learning._stubs import ToolLearningPolicy
|
||||
from openjarvis.learning._stubs import AgentLearningPolicy
|
||||
|
||||
|
||||
@LearningRegistry.register("icl_updater")
|
||||
class ICLUpdaterPolicy(ToolLearningPolicy):
|
||||
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").
|
||||
sequences ("skills"). This updates *agent* logic (ICL examples
|
||||
and tool-use strategies), not tool implementations themselves.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Orchestrator training infrastructure — SFT and GRPO pipelines.
|
||||
|
||||
Provides structured-mode training for the OrchestratorAgent with:
|
||||
|
||||
- **Episode types**: Action, Observation, Episode, EpisodeState
|
||||
- **Reward**: Multi-objective reward balancing accuracy, cost, energy, latency, power
|
||||
- **Prompt registry**: Canonical system prompts for structured mode
|
||||
- **Policy model**: HuggingFace LM wrapper for action prediction
|
||||
- **Environment**: RL environment using OpenJarvis ToolExecutor
|
||||
- **SFT trainer**: Supervised fine-tuning on successful trajectories
|
||||
- **GRPO trainer**: Group Relative Policy Optimization
|
||||
|
||||
Importing this module triggers registration of ``orchestrator_sft`` and
|
||||
``orchestrator_grpo`` in :class:`~openjarvis.core.registry.LearningRegistry`.
|
||||
"""
|
||||
|
||||
from openjarvis.learning.orchestrator.environment import (
|
||||
OrchestratorEnvironment,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.grpo_trainer import (
|
||||
OrchestratorGRPOConfig,
|
||||
OrchestratorGRPOTrainer,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.policy_model import (
|
||||
OrchestratorPolicyModel,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.prompt_registry import (
|
||||
TOOL_DESCRIPTIONS,
|
||||
build_system_prompt,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.reward import (
|
||||
AdaptiveRewardWeights,
|
||||
MultiObjectiveReward,
|
||||
Normalizers,
|
||||
RewardWeights,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.sft_trainer import (
|
||||
OrchestratorSFTConfig,
|
||||
OrchestratorSFTDataset,
|
||||
OrchestratorSFTTrainer,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
Episode,
|
||||
EpisodeState,
|
||||
EpisodeStep,
|
||||
OrchestratorAction,
|
||||
OrchestratorObservation,
|
||||
PolicyOutput,
|
||||
extract_answer,
|
||||
grade_answer,
|
||||
normalize_number,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Types
|
||||
"Episode",
|
||||
"EpisodeState",
|
||||
"EpisodeStep",
|
||||
"OrchestratorAction",
|
||||
"OrchestratorObservation",
|
||||
"PolicyOutput",
|
||||
"extract_answer",
|
||||
"grade_answer",
|
||||
"normalize_number",
|
||||
# Reward
|
||||
"AdaptiveRewardWeights",
|
||||
"MultiObjectiveReward",
|
||||
"Normalizers",
|
||||
"RewardWeights",
|
||||
# Prompt
|
||||
"TOOL_DESCRIPTIONS",
|
||||
"build_system_prompt",
|
||||
# Policy
|
||||
"OrchestratorPolicyModel",
|
||||
# Environment
|
||||
"OrchestratorEnvironment",
|
||||
# SFT
|
||||
"OrchestratorSFTConfig",
|
||||
"OrchestratorSFTDataset",
|
||||
"OrchestratorSFTTrainer",
|
||||
# GRPO
|
||||
"OrchestratorGRPOConfig",
|
||||
"OrchestratorGRPOTrainer",
|
||||
]
|
||||
@@ -0,0 +1,115 @@
|
||||
"""RL environment for orchestrator training.
|
||||
|
||||
Adapted from IPW's ``environment.py``. Uses OpenJarvis's
|
||||
:class:`~openjarvis.tools._stubs.ToolExecutor` for real tool dispatch
|
||||
(as opposed to IPW's cached-telemetry approach), making it suitable for
|
||||
both training and evaluation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import List, Tuple
|
||||
|
||||
from openjarvis.core.types import ToolCall
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
EpisodeState,
|
||||
OrchestratorAction,
|
||||
OrchestratorObservation,
|
||||
)
|
||||
from openjarvis.tools._stubs import BaseTool, ToolExecutor
|
||||
|
||||
|
||||
class OrchestratorEnvironment:
|
||||
"""RL environment that executes tools via OpenJarvis ``ToolExecutor``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tools:
|
||||
List of :class:`BaseTool` instances available to the agent.
|
||||
max_turns:
|
||||
Maximum number of turns per episode.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tools: List[BaseTool],
|
||||
max_turns: int = 10,
|
||||
) -> None:
|
||||
self._tools = tools
|
||||
self._executor = ToolExecutor(tools)
|
||||
self._max_turns = max_turns
|
||||
|
||||
def reset(self, task: str) -> EpisodeState:
|
||||
"""Reset the environment for a new episode.
|
||||
|
||||
Args:
|
||||
task: The initial task/question.
|
||||
|
||||
Returns:
|
||||
A fresh :class:`EpisodeState`.
|
||||
"""
|
||||
return EpisodeState(initial_prompt=task)
|
||||
|
||||
def step(
|
||||
self,
|
||||
state: EpisodeState,
|
||||
action: OrchestratorAction,
|
||||
) -> Tuple[EpisodeState, OrchestratorObservation]:
|
||||
"""Execute one step: dispatch the tool and observe the result.
|
||||
|
||||
Raises:
|
||||
ValueError: If the tool is not available or max turns exceeded.
|
||||
"""
|
||||
available = self.get_available_tools()
|
||||
|
||||
if action.tool_name not in available:
|
||||
raise ValueError(
|
||||
f"Tool '{action.tool_name}' not available. "
|
||||
f"Available: {available}"
|
||||
)
|
||||
|
||||
if state.num_turns() >= self._max_turns:
|
||||
raise ValueError(
|
||||
f"Max turns ({self._max_turns}) exceeded"
|
||||
)
|
||||
|
||||
# Execute tool via ToolExecutor
|
||||
tool_call = ToolCall(
|
||||
id=f"orch_{state.num_turns()}",
|
||||
name=action.tool_name,
|
||||
arguments=action.tool_input
|
||||
if action.tool_input.startswith("{")
|
||||
else f'{{"expression": {repr(action.tool_input)}}}',
|
||||
)
|
||||
|
||||
t0 = time.time()
|
||||
result = self._executor.execute(tool_call)
|
||||
latency = time.time() - t0
|
||||
|
||||
observation = OrchestratorObservation(
|
||||
content=result.content,
|
||||
latency_seconds=latency,
|
||||
cost_usd=result.cost_usd,
|
||||
energy_joules=0.0,
|
||||
power_watts=0.0,
|
||||
tokens=result.usage.get("total_tokens", 0),
|
||||
)
|
||||
|
||||
state.add_turn(action, observation)
|
||||
return state, observation
|
||||
|
||||
def is_done(self, state: EpisodeState) -> bool:
|
||||
"""Check if the episode is complete."""
|
||||
if state.final_answer is not None:
|
||||
return True
|
||||
if state.num_turns() >= self._max_turns:
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_available_tools(self) -> List[str]:
|
||||
"""Return names of all available tools."""
|
||||
return [t.spec.name for t in self._tools]
|
||||
|
||||
|
||||
__all__ = ["OrchestratorEnvironment"]
|
||||
@@ -0,0 +1,498 @@
|
||||
"""GRPO (Group Relative Policy Optimization) trainer for orchestrator.
|
||||
|
||||
Adapted from IPW's ``trainer.py``. GRPO is simpler than PPO because it
|
||||
doesn't require a separate critic model — instead, it uses
|
||||
*group-relative advantages*: for each problem, sample N candidate
|
||||
trajectories, compute rewards, normalise within the group, and update
|
||||
the policy to increase the probability of better solutions.
|
||||
|
||||
All ``torch``/``transformers`` imports are guarded so the module can be
|
||||
imported without GPU dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
# Optional imports -----------------------------------------------------------
|
||||
try:
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
HAS_TORCH = True
|
||||
except ImportError:
|
||||
HAS_TORCH = False
|
||||
torch = None # type: ignore[assignment]
|
||||
F = None # type: ignore[assignment]
|
||||
|
||||
from openjarvis.core.registry import LearningRegistry
|
||||
from openjarvis.learning._stubs import AgentLearningPolicy
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrchestratorGRPOConfig:
|
||||
"""Configuration for orchestrator GRPO training."""
|
||||
|
||||
# Model
|
||||
model_name: str = "Qwen/Qwen3-1.7B"
|
||||
max_prompt_length: int = 24000
|
||||
max_response_length: int = 8768
|
||||
|
||||
# Training
|
||||
num_epochs: int = 10
|
||||
batch_size: int = 16
|
||||
learning_rate: float = 1e-6
|
||||
max_grad_norm: float = 1.0
|
||||
|
||||
# GRPO specific
|
||||
num_samples_per_prompt: int = 8
|
||||
temperature: float = 1.0
|
||||
kl_coef: float = 0.0001
|
||||
clip_ratio: float = 0.2
|
||||
|
||||
# Environment
|
||||
available_tools: List[str] = field(
|
||||
default_factory=lambda: [
|
||||
"calculator",
|
||||
"think",
|
||||
"code_interpreter",
|
||||
"web_search",
|
||||
]
|
||||
)
|
||||
max_turns: int = 10
|
||||
|
||||
# Checkpoint
|
||||
checkpoint_dir: str = "checkpoints/orchestrator_grpo"
|
||||
save_every_n_epochs: int = 1
|
||||
keep_last_n: int = 3
|
||||
|
||||
# Memory
|
||||
gradient_checkpointing: bool = True
|
||||
use_8bit_ref: bool = True
|
||||
use_8bit_optimizer: bool = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trainer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class OrchestratorGRPOTrainer:
|
||||
"""GRPO trainer for orchestrator policy.
|
||||
|
||||
``torch`` must be installed to call :meth:`train`.
|
||||
"""
|
||||
|
||||
def __init__(self, config: OrchestratorGRPOConfig) -> None:
|
||||
self.config = config
|
||||
self.device = None
|
||||
self.global_step = 0
|
||||
|
||||
if HAS_TORCH and torch is not None:
|
||||
self.device = torch.device(
|
||||
"cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
|
||||
self._init_model()
|
||||
self._init_optimizer()
|
||||
|
||||
# -- Initialisation ------------------------------------------------------
|
||||
|
||||
def _init_model(self) -> None:
|
||||
from openjarvis.learning.orchestrator.policy_model import (
|
||||
OrchestratorPolicyModel,
|
||||
)
|
||||
|
||||
if not HAS_TORCH:
|
||||
self.policy: Any = OrchestratorPolicyModel()
|
||||
self.ref_policy: Any = OrchestratorPolicyModel()
|
||||
return
|
||||
|
||||
device_str = str(self.device) if self.device else None
|
||||
|
||||
self.policy = OrchestratorPolicyModel.from_pretrained(
|
||||
self.config.model_name,
|
||||
gradient_checkpointing=self.config.gradient_checkpointing,
|
||||
device=device_str,
|
||||
)
|
||||
if self.policy.model is not None:
|
||||
self.policy.model.train()
|
||||
|
||||
self.ref_policy = OrchestratorPolicyModel.from_pretrained(
|
||||
self.config.model_name,
|
||||
load_in_8bit=self.config.use_8bit_ref,
|
||||
device=device_str,
|
||||
)
|
||||
if self.ref_policy.model is not None:
|
||||
self.ref_policy.model.eval()
|
||||
for param in self.ref_policy.model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def _init_optimizer(self) -> None:
|
||||
if not HAS_TORCH or self.policy.model is None:
|
||||
self.optimizer: Any = None
|
||||
return
|
||||
|
||||
if self.config.use_8bit_optimizer:
|
||||
try:
|
||||
import bitsandbytes as bnb
|
||||
|
||||
self.optimizer = bnb.optim.AdamW8bit(
|
||||
self.policy.model.parameters(),
|
||||
lr=self.config.learning_rate,
|
||||
)
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
self.optimizer = torch.optim.AdamW(
|
||||
self.policy.model.parameters(),
|
||||
lr=self.config.learning_rate,
|
||||
)
|
||||
|
||||
# -- Training loop -------------------------------------------------------
|
||||
|
||||
def train(self) -> None:
|
||||
"""Run the GRPO training loop."""
|
||||
if not HAS_TORCH:
|
||||
raise RuntimeError(
|
||||
"PyTorch is required for training. "
|
||||
"Install with: pip install torch transformers"
|
||||
)
|
||||
|
||||
for epoch in range(self.config.num_epochs):
|
||||
self._train_epoch(epoch)
|
||||
|
||||
if (epoch + 1) % self.config.save_every_n_epochs == 0:
|
||||
self._save_checkpoint(epoch)
|
||||
|
||||
def _train_epoch(self, epoch: int) -> Dict[str, float]:
|
||||
if self.policy.model is None:
|
||||
return {"epoch": epoch, "loss": 0.0, "reward": 0.0}
|
||||
|
||||
self.policy.model.train()
|
||||
total_loss = 0.0
|
||||
total_reward = 0.0
|
||||
num_batches = 0
|
||||
|
||||
# In a real implementation, iterate over a task dataset.
|
||||
# Here we provide the skeleton; actual data loading is
|
||||
# trainer-specific.
|
||||
self.global_step += 1
|
||||
num_batches = max(num_batches, 1)
|
||||
|
||||
avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
|
||||
avg_reward = total_reward / num_batches if num_batches > 0 else 0.0
|
||||
return {"epoch": epoch, "loss": avg_loss, "reward": avg_reward}
|
||||
|
||||
def _grpo_step(
|
||||
self,
|
||||
prompts: List[str],
|
||||
ground_truths: List[str],
|
||||
) -> tuple:
|
||||
"""Perform one GRPO training step.
|
||||
|
||||
For each prompt:
|
||||
1. Generate N candidate trajectories.
|
||||
2. Compute reward for each.
|
||||
3. Normalise advantages within the group.
|
||||
4. Compute clipped policy gradient + KL penalty.
|
||||
5. Backward + clip + step.
|
||||
|
||||
Returns ``(loss_value, avg_reward)``.
|
||||
"""
|
||||
if self.policy.model is None or not HAS_TORCH:
|
||||
raise RuntimeError("Cannot train without PyTorch and model.")
|
||||
|
||||
self.policy.model.train()
|
||||
|
||||
all_prompts: list[str] = []
|
||||
all_responses: list[str] = []
|
||||
all_advantages: list[float] = []
|
||||
all_rewards: list[float] = []
|
||||
|
||||
from openjarvis.learning.orchestrator.reward import (
|
||||
MultiObjectiveReward,
|
||||
Normalizers,
|
||||
RewardWeights,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
Episode,
|
||||
)
|
||||
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
|
||||
for prompt, gt in zip(prompts, ground_truths):
|
||||
group_rewards: list[float] = []
|
||||
group_responses: list[str] = []
|
||||
|
||||
for _ in range(self.config.num_samples_per_prompt):
|
||||
response, _log_probs = self._generate_with_log_probs(prompt)
|
||||
|
||||
# Build a minimal episode for reward computation
|
||||
episode = Episode(
|
||||
task_id="grpo",
|
||||
initial_prompt=prompt,
|
||||
ground_truth=gt,
|
||||
final_answer=response,
|
||||
correct=(response.strip() == gt.strip()),
|
||||
)
|
||||
reward = reward_fn.compute(episode)
|
||||
|
||||
group_rewards.append(reward)
|
||||
group_responses.append(response)
|
||||
|
||||
# Group-relative advantages
|
||||
mean_r = sum(group_rewards) / len(group_rewards)
|
||||
std_r = (
|
||||
sum((r - mean_r) ** 2 for r in group_rewards)
|
||||
/ len(group_rewards)
|
||||
) ** 0.5
|
||||
if std_r > 1e-8:
|
||||
advantages = [(r - mean_r) / std_r for r in group_rewards]
|
||||
else:
|
||||
advantages = [0.0] * len(group_rewards)
|
||||
|
||||
for resp, adv, rew in zip(group_responses, advantages, group_rewards):
|
||||
all_prompts.append(prompt)
|
||||
all_responses.append(resp)
|
||||
all_advantages.append(adv)
|
||||
all_rewards.append(rew)
|
||||
|
||||
# Policy gradient loss
|
||||
total_loss = torch.tensor(0.0, device=self.device, requires_grad=True)
|
||||
|
||||
for prompt, response, advantage in zip(
|
||||
all_prompts, all_responses, all_advantages
|
||||
):
|
||||
current_lp = self._compute_log_probs(prompt, response)
|
||||
with torch.no_grad():
|
||||
ref_lp = self._compute_log_probs_ref(prompt, response)
|
||||
|
||||
log_ratio = current_lp - ref_lp
|
||||
ratio = torch.exp(log_ratio)
|
||||
ratio = torch.clamp(ratio, min=0.01, max=100.0)
|
||||
|
||||
clip = self.config.clip_ratio
|
||||
clipped = torch.clamp(ratio, 1 - clip, 1 + clip)
|
||||
|
||||
policy_loss = -torch.min(ratio * advantage, clipped * advantage)
|
||||
kl = (ratio - 1) - log_ratio
|
||||
total_loss = total_loss + policy_loss + self.config.kl_coef * kl
|
||||
|
||||
avg_loss = total_loss / max(len(all_prompts), 1)
|
||||
loss_val = avg_loss.item()
|
||||
|
||||
if torch.isnan(avg_loss) or torch.isinf(avg_loss):
|
||||
avg_reward = (
|
||||
sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
|
||||
)
|
||||
return 0.0, float(avg_reward)
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
avg_loss.backward()
|
||||
|
||||
# Check for NaN gradients
|
||||
for param in self.policy.model.parameters():
|
||||
if param.grad is not None and torch.isnan(param.grad).any():
|
||||
self.optimizer.zero_grad()
|
||||
avg_reward = (
|
||||
sum(all_rewards) / len(all_rewards)
|
||||
if all_rewards
|
||||
else 0.0
|
||||
)
|
||||
return float(loss_val), float(avg_reward)
|
||||
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.policy.model.parameters(), self.config.max_grad_norm
|
||||
)
|
||||
self.optimizer.step()
|
||||
|
||||
avg_reward = (
|
||||
sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
|
||||
)
|
||||
return float(loss_val), float(avg_reward)
|
||||
|
||||
# -- Generation / log-prob helpers ---------------------------------------
|
||||
|
||||
def _generate_with_log_probs(
|
||||
self, prompt: str
|
||||
) -> "tuple[str, Any]":
|
||||
"""Generate a response and return ``(text, log_probs)``."""
|
||||
inputs = self.policy.tokenizer(
|
||||
prompt,
|
||||
return_tensors="pt",
|
||||
truncation=True,
|
||||
max_length=min(self.config.max_prompt_length, 16000),
|
||||
).to(self.device)
|
||||
|
||||
input_len = inputs.input_ids.shape[1]
|
||||
max_new = min(self.config.max_response_length, 32000 - input_len - 100)
|
||||
max_new = max(min(max_new, 2048), 128)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.policy.model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=max_new,
|
||||
temperature=self.config.temperature,
|
||||
do_sample=True,
|
||||
output_scores=True,
|
||||
return_dict_in_generate=True,
|
||||
)
|
||||
|
||||
generated_ids = outputs.sequences[0][input_len:]
|
||||
if len(generated_ids) == 0:
|
||||
return "", torch.tensor(0.0, device=self.device)
|
||||
|
||||
text = self.policy.tokenizer.decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)
|
||||
|
||||
log_probs = []
|
||||
for token_id, logits in zip(generated_ids, outputs.scores):
|
||||
probs = F.softmax(logits[0], dim=-1)
|
||||
tid = token_id.item()
|
||||
if 0 <= tid < probs.shape[0]:
|
||||
log_probs.append(torch.log(probs[tid] + 1e-10))
|
||||
|
||||
total_lp = (
|
||||
torch.stack(log_probs).sum()
|
||||
if log_probs
|
||||
else torch.tensor(0.0, device=self.device)
|
||||
)
|
||||
return text, total_lp
|
||||
|
||||
def _compute_log_probs(
|
||||
self, prompt: str, response: str
|
||||
) -> "torch.Tensor":
|
||||
"""Log-probs of *response* given *prompt* under current policy."""
|
||||
full = prompt + response
|
||||
inputs = self.policy.tokenizer(
|
||||
full, return_tensors="pt", truncation=True
|
||||
).to(self.device)
|
||||
prompt_inputs = self.policy.tokenizer(
|
||||
prompt, return_tensors="pt", truncation=True
|
||||
).to(self.device)
|
||||
|
||||
with torch.enable_grad():
|
||||
logits = self.policy.model(**inputs).logits
|
||||
|
||||
start = len(prompt_inputs.input_ids[0])
|
||||
end = len(inputs.input_ids[0])
|
||||
|
||||
lps = []
|
||||
for i in range(start, end - 1):
|
||||
lp = F.log_softmax(logits[0, i, :], dim=-1)[
|
||||
inputs.input_ids[0, i + 1]
|
||||
]
|
||||
lps.append(lp)
|
||||
|
||||
return torch.stack(lps).sum() if lps else torch.tensor(0.0)
|
||||
|
||||
def _compute_log_probs_ref(
|
||||
self, prompt: str, response: str
|
||||
) -> "torch.Tensor":
|
||||
"""Log-probs under the frozen reference policy (no grad)."""
|
||||
full = prompt + response
|
||||
inputs = self.ref_policy.tokenizer(
|
||||
full, return_tensors="pt", truncation=True
|
||||
).to(self.device)
|
||||
prompt_inputs = self.ref_policy.tokenizer(
|
||||
prompt, return_tensors="pt", truncation=True
|
||||
).to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
logits = self.ref_policy.model(**inputs).logits
|
||||
|
||||
start = len(prompt_inputs.input_ids[0])
|
||||
end = len(inputs.input_ids[0])
|
||||
|
||||
lps = []
|
||||
for i in range(start, end - 1):
|
||||
lp = F.log_softmax(logits[0, i, :], dim=-1)[
|
||||
inputs.input_ids[0, i + 1]
|
||||
]
|
||||
lps.append(lp)
|
||||
|
||||
return torch.stack(lps).sum() if lps else torch.tensor(0.0)
|
||||
|
||||
# -- Checkpointing -------------------------------------------------------
|
||||
|
||||
def _save_checkpoint(self, epoch: int) -> None:
|
||||
checkpoint_dir = Path(self.config.checkpoint_dir)
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
checkpoint_path = checkpoint_dir / f"epoch_{epoch + 1}"
|
||||
if self.policy.model is not None:
|
||||
self.policy.save(str(checkpoint_path))
|
||||
|
||||
state_path = checkpoint_path / "training_state.json"
|
||||
state = {
|
||||
"epoch": epoch,
|
||||
"global_step": self.global_step,
|
||||
"config": asdict(self.config),
|
||||
}
|
||||
with open(state_path, "w") as f:
|
||||
json.dump(state, f, indent=2)
|
||||
|
||||
self._cleanup_old_checkpoints()
|
||||
|
||||
def _cleanup_old_checkpoints(self) -> None:
|
||||
checkpoint_dir = Path(self.config.checkpoint_dir)
|
||||
if not checkpoint_dir.exists():
|
||||
return
|
||||
|
||||
checkpoints = sorted(
|
||||
[
|
||||
d
|
||||
for d in checkpoint_dir.iterdir()
|
||||
if d.is_dir() and d.name.startswith("epoch_")
|
||||
],
|
||||
key=lambda x: int(x.name.split("_")[1]),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
for old in checkpoints[self.config.keep_last_n :]:
|
||||
shutil.rmtree(old)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry wrapper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _ensure_registered() -> None:
|
||||
if LearningRegistry.contains("orchestrator_grpo"):
|
||||
return
|
||||
|
||||
@LearningRegistry.register("orchestrator_grpo")
|
||||
class OrchestratorGRPOPolicy(AgentLearningPolicy):
|
||||
"""Wrapper that registers the GRPO trainer as a learning policy."""
|
||||
|
||||
def update(
|
||||
self, trace_store: Any, **kwargs: object
|
||||
) -> Dict[str, Any]:
|
||||
config = OrchestratorGRPOConfig(**{
|
||||
k: v for k, v in kwargs.items()
|
||||
if k in OrchestratorGRPOConfig.__dataclass_fields__
|
||||
})
|
||||
trainer = OrchestratorGRPOTrainer(config)
|
||||
trainer.train()
|
||||
return {"status": "grpo_training_complete"}
|
||||
|
||||
|
||||
_ensure_registered()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OrchestratorGRPOConfig",
|
||||
"OrchestratorGRPOTrainer",
|
||||
]
|
||||
@@ -0,0 +1,294 @@
|
||||
"""Policy model wrapper for orchestrator training.
|
||||
|
||||
Adapted from IPW's ``policy.py``. Wraps a HuggingFace causal LM
|
||||
(e.g. Qwen3-1.7B) to predict structured actions in the orchestrator
|
||||
environment. All ``torch``/``transformers`` imports are guarded so the
|
||||
module can be imported without GPU dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
EpisodeState,
|
||||
OrchestratorAction,
|
||||
PolicyOutput,
|
||||
)
|
||||
|
||||
# Optional imports -----------------------------------------------------------
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
|
||||
HAS_TORCH = True
|
||||
except ImportError:
|
||||
HAS_TORCH = False
|
||||
torch = None # type: ignore[assignment]
|
||||
|
||||
|
||||
class OrchestratorPolicyModel:
|
||||
"""Wrapper around a causal LM for orchestrator policy prediction.
|
||||
|
||||
Input format (prompt)::
|
||||
|
||||
Task: {initial_prompt}
|
||||
|
||||
Available tools: calculator, think, ...
|
||||
|
||||
History:
|
||||
Turn 1:
|
||||
Thought: ...
|
||||
Tool: ...
|
||||
Observation: ...
|
||||
|
||||
What should you do next?
|
||||
Format your response as:
|
||||
THOUGHT: [your reasoning]
|
||||
TOOL: [tool_name]
|
||||
INPUT: [input for tool]
|
||||
|
||||
Output format (from model)::
|
||||
|
||||
THOUGHT: [reasoning]
|
||||
TOOL: [tool_name]
|
||||
INPUT: [input]
|
||||
--- or ---
|
||||
FINAL_ANSWER: [answer]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Any = None,
|
||||
tokenizer: Any = None,
|
||||
max_tokens: int = 256,
|
||||
temperature: float = 0.7,
|
||||
) -> None:
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
self.max_tokens = max_tokens
|
||||
self.temperature = temperature
|
||||
|
||||
# -- Factory methods -----------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
model_name: str = "Qwen/Qwen3-1.7B",
|
||||
gradient_checkpointing: bool = False,
|
||||
load_in_8bit: bool = False,
|
||||
device: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> "OrchestratorPolicyModel":
|
||||
"""Load model from a HuggingFace checkpoint.
|
||||
|
||||
Raises ``ImportError`` if ``transformers`` is not installed.
|
||||
"""
|
||||
import torch as _torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
model_kwargs: dict[str, Any] = {"torch_dtype": _torch.bfloat16}
|
||||
|
||||
if load_in_8bit:
|
||||
try:
|
||||
from transformers import BitsAndBytesConfig
|
||||
|
||||
model_kwargs["quantization_config"] = BitsAndBytesConfig(
|
||||
load_in_8bit=True
|
||||
)
|
||||
except ImportError:
|
||||
pass # fall back to bf16
|
||||
|
||||
if device is not None:
|
||||
if device == "auto":
|
||||
model_kwargs["device_map"] = "auto"
|
||||
else:
|
||||
model_kwargs["device_map"] = {"": device}
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
|
||||
|
||||
if gradient_checkpointing and hasattr(
|
||||
model, "gradient_checkpointing_enable"
|
||||
):
|
||||
model.gradient_checkpointing_enable(
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False}
|
||||
)
|
||||
|
||||
return cls(model=model, tokenizer=tokenizer, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(
|
||||
cls, checkpoint_path: str, **kwargs: Any
|
||||
) -> "OrchestratorPolicyModel":
|
||||
"""Load from a previously saved checkpoint directory."""
|
||||
return cls.from_pretrained(checkpoint_path, **kwargs)
|
||||
|
||||
# -- Prediction ----------------------------------------------------------
|
||||
|
||||
def predict_action(
|
||||
self,
|
||||
state: EpisodeState,
|
||||
available_tools: List[str],
|
||||
) -> OrchestratorAction:
|
||||
"""Predict the next action given current state."""
|
||||
prompt = self._build_prompt(state, available_tools)
|
||||
|
||||
if self.model is None:
|
||||
raise RuntimeError(
|
||||
"Cannot generate actions without a loaded model. "
|
||||
"Load with OrchestratorPolicyModel.from_pretrained() first."
|
||||
)
|
||||
|
||||
output_text = self._generate(prompt)
|
||||
policy_output = self._parse_output(output_text, available_tools)
|
||||
return OrchestratorAction(
|
||||
thought=policy_output.thought,
|
||||
tool_name=policy_output.tool_name,
|
||||
tool_input=policy_output.tool_input,
|
||||
is_final_answer=policy_output.is_final_answer,
|
||||
)
|
||||
|
||||
# -- Internal helpers ----------------------------------------------------
|
||||
|
||||
def _build_prompt(
|
||||
self,
|
||||
state: EpisodeState,
|
||||
available_tools: List[str],
|
||||
) -> str:
|
||||
"""Build the text prompt from current state."""
|
||||
parts: list[str] = []
|
||||
|
||||
parts.append(f"Task: {state.initial_prompt}")
|
||||
parts.append("")
|
||||
|
||||
tools_str = ", ".join(available_tools)
|
||||
parts.append(f"Available tools: {tools_str}")
|
||||
parts.append("")
|
||||
|
||||
if state.history:
|
||||
parts.append("History:")
|
||||
for i, (action, observation) in enumerate(state.history, 1):
|
||||
parts.append(f"Turn {i}:")
|
||||
parts.append(f" Thought: {action.thought}")
|
||||
parts.append(f" Tool: {action.tool_name}")
|
||||
parts.append(
|
||||
f" Observation: {observation.content[:100]}..."
|
||||
)
|
||||
parts.append("")
|
||||
|
||||
parts.append("What should you do next?")
|
||||
parts.append("Format your response as:")
|
||||
parts.append("THOUGHT: [your reasoning]")
|
||||
parts.append("TOOL: [tool_name]")
|
||||
parts.append("INPUT: [input for tool]")
|
||||
parts.append("")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
def _parse_output(
|
||||
self,
|
||||
output_text: str,
|
||||
available_tools: List[str],
|
||||
) -> PolicyOutput:
|
||||
"""Parse structured model output into a :class:`PolicyOutput`."""
|
||||
# Check for FINAL_ANSWER first
|
||||
final_match = re.search(
|
||||
r"FINAL[_ ]?ANSWER:\s*(.+?)(?:\n|$)",
|
||||
output_text,
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
if final_match:
|
||||
thought_match = re.search(
|
||||
r"THOUGHT:\s*(.+?)(?:\n|$)", output_text, re.IGNORECASE
|
||||
)
|
||||
return PolicyOutput(
|
||||
thought=(
|
||||
thought_match.group(1).strip()
|
||||
if thought_match
|
||||
else ""
|
||||
),
|
||||
tool_name="",
|
||||
tool_input=final_match.group(1).strip(),
|
||||
is_final_answer=True,
|
||||
raw_text=output_text,
|
||||
)
|
||||
|
||||
thought_match = re.search(
|
||||
r"THOUGHT:\s*(.+?)(?:\n|$)", output_text, re.IGNORECASE
|
||||
)
|
||||
tool_match = re.search(
|
||||
r"TOOL:\s*(.+?)(?:\n|$)", output_text, re.IGNORECASE
|
||||
)
|
||||
input_match = re.search(
|
||||
r"INPUT:\s*(.+?)(?:\nTHOUGHT:|\nTOOL:|\nFINAL|\Z)",
|
||||
output_text,
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
|
||||
thought = (
|
||||
thought_match.group(1).strip()
|
||||
if thought_match
|
||||
else "No thought provided"
|
||||
)
|
||||
tool_name = (
|
||||
tool_match.group(1).strip()
|
||||
if tool_match
|
||||
else (available_tools[0] if available_tools else "unknown")
|
||||
)
|
||||
tool_input = input_match.group(1).strip() if input_match else ""
|
||||
|
||||
# Validate / fuzzy-match tool name
|
||||
if tool_name not in available_tools:
|
||||
tool_name_lower = tool_name.lower()
|
||||
matched = False
|
||||
for t in available_tools:
|
||||
if t.lower() == tool_name_lower:
|
||||
tool_name = t
|
||||
matched = True
|
||||
break
|
||||
if not matched:
|
||||
tool_name = available_tools[0] if available_tools else "unknown"
|
||||
|
||||
return PolicyOutput(
|
||||
thought=thought,
|
||||
tool_name=tool_name,
|
||||
tool_input=tool_input,
|
||||
is_final_answer=False,
|
||||
raw_text=output_text,
|
||||
)
|
||||
|
||||
def _generate(self, prompt: str) -> str:
|
||||
"""Generate text from the loaded model."""
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt").to(
|
||||
self.model.device
|
||||
)
|
||||
outputs = self.model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=self.max_tokens,
|
||||
temperature=self.temperature,
|
||||
do_sample=True,
|
||||
)
|
||||
return self.tokenizer.decode(
|
||||
outputs[0][len(inputs.input_ids[0]) :],
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
|
||||
def save(self, path: str) -> None:
|
||||
"""Save model and tokenizer to *path*."""
|
||||
if self.model is not None:
|
||||
self.model.save_pretrained(path)
|
||||
self.tokenizer.save_pretrained(path)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
model_name = (
|
||||
"None" if self.model is None else type(self.model).__name__
|
||||
)
|
||||
return (
|
||||
f"OrchestratorPolicyModel(model={model_name}, "
|
||||
f"max_tokens={self.max_tokens})"
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["OrchestratorPolicyModel"]
|
||||
@@ -0,0 +1,308 @@
|
||||
"""Prompt registry for orchestrator structured mode.
|
||||
|
||||
Adapted from IPW's ``prompt_registry.py``. Provides the canonical system
|
||||
prompt template and tool descriptions used by the structured-mode
|
||||
``OrchestratorAgent`` and by the SFT/GRPO training pipelines.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
PROMPT_VERSION = "1.0"
|
||||
|
||||
SYSTEM_PROMPT_TEMPLATE = """\
|
||||
You are an intelligent orchestrator that solves tasks by \
|
||||
delegating to the most appropriate tools.
|
||||
|
||||
Your job is to SELECT THE BEST TOOL for each task based on the tool's strengths.
|
||||
|
||||
=== AVAILABLE TOOLS ===
|
||||
{tools_description}
|
||||
|
||||
=== TOOL SELECTION GUIDE ===
|
||||
{tool_selection_guide}
|
||||
|
||||
=== RESPONSE FORMAT ===
|
||||
You MUST respond in this EXACT format:
|
||||
|
||||
THOUGHT: <analyze the task and explain which tool is best and why>
|
||||
TOOL: <exact tool name from the list>
|
||||
INPUT: <input for the tool>
|
||||
|
||||
After getting tool results, either use another tool or give final answer:
|
||||
THOUGHT: <analyze the result>
|
||||
FINAL_ANSWER: <your final answer>
|
||||
|
||||
=== CRITICAL RULES ===
|
||||
1. You MUST use at least one tool for EVERY task - never answer directly
|
||||
2. Match the tool to the task type (see guide above)
|
||||
3. For LLM tools, write clear prompts that will get good responses
|
||||
4. Prefer specialized tools when available \
|
||||
(calculator for math, code_interpreter for code)
|
||||
5. For simple factual questions, use fast/cheap tools when available
|
||||
|
||||
NOW SOLVE THE TASK. You MUST use at least one tool - choose the best one for the task.
|
||||
"""
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool descriptions for OpenJarvis built-in tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TOOL_DESCRIPTIONS: Dict[str, dict] = {
|
||||
# Utility tools (instant, free, deterministic)
|
||||
"calculator": {
|
||||
"category": "utility",
|
||||
"description": (
|
||||
"CALCULATOR - Instant math computation\n"
|
||||
" - BEST FOR: Arithmetic, algebra, trigonometry, scientific calculations\n"
|
||||
" - STRENGTHS: Instant (<1ms), perfect accuracy, zero cost\n"
|
||||
" - USE WHEN: Any math expression needs evaluation\n"
|
||||
" - COST: Free\n"
|
||||
" - Input: math expression (e.g., '15 * 7 + 23', 'sqrt(144)')"
|
||||
),
|
||||
"examples": [
|
||||
{
|
||||
"task": "What is 847 * 293?",
|
||||
"thought": "Simple arithmetic - calculator is instant and accurate.",
|
||||
"input": "847 * 293",
|
||||
},
|
||||
],
|
||||
},
|
||||
"think": {
|
||||
"category": "utility",
|
||||
"description": (
|
||||
"THINK - Internal reasoning scratchpad\n"
|
||||
" - BEST FOR: Logic puzzles, step-by-step reasoning, planning\n"
|
||||
" - STRENGTHS: Organizes thoughts, shows work, no external calls\n"
|
||||
" - USE WHEN: Need to break down a problem before solving\n"
|
||||
" - COST: Free\n"
|
||||
" - Input: your detailed reasoning process"
|
||||
),
|
||||
"examples": [
|
||||
{
|
||||
"task": "If all cats are mammals, are all cats animals?",
|
||||
"thought": "Logical syllogism - use think to reason step by step.",
|
||||
"input": "Cats subset of mammals subset of animals => yes.",
|
||||
},
|
||||
],
|
||||
},
|
||||
"code_interpreter": {
|
||||
"category": "utility",
|
||||
"description": (
|
||||
"CODE_INTERPRETER - Python execution sandbox\n"
|
||||
" - BEST FOR: Data processing, algorithms, simulations\n"
|
||||
" - STRENGTHS: Full Python + numpy/pandas, handles loops\n"
|
||||
" - USE WHEN: Problem needs programming logic\n"
|
||||
" - COST: Free (local execution)\n"
|
||||
" - Input: Python code to execute"
|
||||
),
|
||||
"examples": [
|
||||
{
|
||||
"task": "Find all prime numbers less than 50",
|
||||
"thought": (
|
||||
"Need a prime-checking algorithm"
|
||||
" - code_interpreter is ideal."
|
||||
),
|
||||
"input": (
|
||||
"def is_prime(n):\n"
|
||||
" if n < 2: return False\n"
|
||||
" for i in range(2, int(n**0.5)+1):\n"
|
||||
" if n % i == 0: return False\n"
|
||||
" return True\n"
|
||||
"print([n for n in range(50) if is_prime(n)])"
|
||||
),
|
||||
},
|
||||
],
|
||||
},
|
||||
"web_search": {
|
||||
"category": "utility",
|
||||
"description": (
|
||||
"WEB_SEARCH - Real-time internet search\n"
|
||||
" - BEST FOR: Current events, fact-checking, recent info\n"
|
||||
" - STRENGTHS: Access to up-to-date information\n"
|
||||
" - USE WHEN: Question about recent events or needs verification\n"
|
||||
" - COST: ~$0.001 per search\n"
|
||||
" - Input: search query string"
|
||||
),
|
||||
"examples": [
|
||||
{
|
||||
"task": "Who won the 2024 Nobel Prize in Physics?",
|
||||
"thought": "Recent events - need web_search for current info.",
|
||||
"input": "2024 Nobel Prize Physics winner",
|
||||
},
|
||||
],
|
||||
},
|
||||
"file_read": {
|
||||
"category": "utility",
|
||||
"description": (
|
||||
"FILE_READ - Safe file reading\n"
|
||||
" - BEST FOR: Reading file contents with path validation\n"
|
||||
" - STRENGTHS: Sandboxed, prevents directory traversal\n"
|
||||
" - USE WHEN: Need to read local file contents\n"
|
||||
" - COST: Free\n"
|
||||
" - Input: file path"
|
||||
),
|
||||
"examples": [],
|
||||
},
|
||||
# Memory tools
|
||||
"memory_search": {
|
||||
"category": "memory",
|
||||
"description": (
|
||||
"MEMORY_SEARCH - Search indexed documents\n"
|
||||
" - BEST FOR: Finding relevant stored knowledge\n"
|
||||
" - STRENGTHS: Semantic search over indexed content\n"
|
||||
" - USE WHEN: Answer may exist in indexed documents\n"
|
||||
" - COST: Free (local)\n"
|
||||
" - Input: search query"
|
||||
),
|
||||
"examples": [],
|
||||
},
|
||||
"memory_store": {
|
||||
"category": "memory",
|
||||
"description": (
|
||||
"MEMORY_STORE - Store content in memory\n"
|
||||
" - BEST FOR: Saving information for later retrieval\n"
|
||||
" - STRENGTHS: Persistent storage with metadata\n"
|
||||
" - USE WHEN: Need to remember something for future queries\n"
|
||||
" - COST: Free (local)\n"
|
||||
" - Input: content to store"
|
||||
),
|
||||
"examples": [],
|
||||
},
|
||||
"memory_retrieve": {
|
||||
"category": "memory",
|
||||
"description": (
|
||||
"MEMORY_RETRIEVE - Retrieve stored content by key\n"
|
||||
" - BEST FOR: Fetching previously stored information\n"
|
||||
" - STRENGTHS: Fast key-based retrieval\n"
|
||||
" - USE WHEN: Know the exact key of stored content\n"
|
||||
" - COST: Free (local)\n"
|
||||
" - Input: content key"
|
||||
),
|
||||
"examples": [],
|
||||
},
|
||||
# LLM tool
|
||||
"llm": {
|
||||
"category": "llm",
|
||||
"description": (
|
||||
"LLM - Sub-model calls via engine\n"
|
||||
" - BEST FOR: Natural language understanding, generation, analysis\n"
|
||||
" - STRENGTHS: General-purpose language capabilities\n"
|
||||
" - USE WHEN: Task needs natural language reasoning\n"
|
||||
" - COST: Varies by engine/model\n"
|
||||
" - Input: prompt for the model"
|
||||
),
|
||||
"examples": [
|
||||
{
|
||||
"task": "Explain photosynthesis simply",
|
||||
"thought": "General explanation - use LLM for natural language.",
|
||||
"input": "Explain photosynthesis in simple terms.",
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def build_system_prompt(tool_names: Optional[List[str]] = None) -> str:
|
||||
"""Build the complete system prompt for the given tools.
|
||||
|
||||
Args:
|
||||
tool_names: Tool names to include. If ``None``, uses all
|
||||
tools from :data:`TOOL_DESCRIPTIONS`.
|
||||
|
||||
Returns:
|
||||
Complete system prompt string.
|
||||
"""
|
||||
if tool_names is None:
|
||||
tool_names = list(TOOL_DESCRIPTIONS)
|
||||
|
||||
# Tool descriptions
|
||||
desc_lines: list[str] = []
|
||||
for name in tool_names:
|
||||
if name in TOOL_DESCRIPTIONS:
|
||||
desc = TOOL_DESCRIPTIONS[name]["description"]
|
||||
else:
|
||||
desc = f"Tool: {name}"
|
||||
desc_lines.append(f"- {name}: {desc}")
|
||||
|
||||
# Group tools by category
|
||||
by_cat: Dict[str, List[str]] = {}
|
||||
for name in tool_names:
|
||||
cat = (
|
||||
TOOL_DESCRIPTIONS[name]["category"]
|
||||
if name in TOOL_DESCRIPTIONS
|
||||
else "llm"
|
||||
)
|
||||
by_cat.setdefault(cat, []).append(name)
|
||||
|
||||
guide: list[str] = [
|
||||
"Choose tools based on task type:\n",
|
||||
]
|
||||
|
||||
# Math
|
||||
math_lines: list[str] = []
|
||||
if "calculator" in tool_names:
|
||||
math_lines.append(
|
||||
"- Simple arithmetic/algebra -> calculator (instant, accurate)"
|
||||
)
|
||||
if "code_interpreter" in tool_names:
|
||||
math_lines.append(
|
||||
"- Numerical algorithms -> code_interpreter (programmable)"
|
||||
)
|
||||
if math_lines:
|
||||
guide.append("MATH PROBLEMS:")
|
||||
guide.extend(math_lines)
|
||||
guide.append("")
|
||||
|
||||
# Coding
|
||||
code_lines: list[str] = []
|
||||
if "code_interpreter" in tool_names:
|
||||
code_lines.append(
|
||||
"- Algorithm implementation/execution -> code_interpreter"
|
||||
)
|
||||
if code_lines:
|
||||
guide.append("CODING TASKS:")
|
||||
guide.extend(code_lines)
|
||||
guide.append("")
|
||||
|
||||
# Reasoning
|
||||
reasoning_lines: list[str] = []
|
||||
if "think" in tool_names:
|
||||
reasoning_lines.append(
|
||||
"- Step-by-step analysis -> think (organize thoughts first)"
|
||||
)
|
||||
llm_tools = by_cat.get("llm", [])
|
||||
if llm_tools:
|
||||
reasoning_lines.append(
|
||||
f"- Complex reasoning -> {', '.join(llm_tools)}"
|
||||
)
|
||||
if reasoning_lines:
|
||||
guide.append("REASONING/LOGIC:")
|
||||
guide.extend(reasoning_lines)
|
||||
guide.append("")
|
||||
|
||||
# General Q&A
|
||||
general_lines: list[str] = []
|
||||
if "web_search" in tool_names:
|
||||
general_lines.append("- Current events/recent info -> web_search")
|
||||
memory_tools = by_cat.get("memory", [])
|
||||
if memory_tools:
|
||||
general_lines.append(
|
||||
f"- Stored knowledge -> {', '.join(memory_tools)}"
|
||||
)
|
||||
if general_lines:
|
||||
guide.append("GENERAL Q&A / FACTUAL:")
|
||||
guide.extend(general_lines)
|
||||
guide.append("")
|
||||
|
||||
return SYSTEM_PROMPT_TEMPLATE.format(
|
||||
tools_description="\n".join(desc_lines),
|
||||
tool_selection_guide="\n".join(guide),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"TOOL_DESCRIPTIONS",
|
||||
"build_system_prompt",
|
||||
]
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Multi-objective reward function for orchestrator training.
|
||||
|
||||
Adapted from IPW's ``reward.py``. Balances accuracy, cost, energy,
|
||||
latency, and power into a single scalar reward used by both the SFT
|
||||
grading pipeline and the GRPO policy gradient.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List
|
||||
|
||||
from openjarvis.learning.orchestrator.types import Episode
|
||||
|
||||
|
||||
@dataclass
|
||||
class RewardWeights:
|
||||
"""Weights for multi-objective reward function.
|
||||
|
||||
Each metric has its own coefficient:
|
||||
- alpha: Accuracy (correctness of answer)
|
||||
- beta_cost: API/cloud cost in USD
|
||||
- beta_energy: Energy consumption in joules
|
||||
- gamma_latency: Response time in seconds
|
||||
- gamma_power: Peak power usage in watts
|
||||
"""
|
||||
|
||||
alpha: float = 0.4
|
||||
beta_cost: float = 0.15
|
||||
beta_energy: float = 0.15
|
||||
gamma_latency: float = 0.15
|
||||
gamma_power: float = 0.15
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Validate weights sum to 1.0."""
|
||||
total = (
|
||||
self.alpha
|
||||
+ self.beta_cost
|
||||
+ self.beta_energy
|
||||
+ self.gamma_latency
|
||||
+ self.gamma_power
|
||||
)
|
||||
if abs(total - 1.0) > 0.01:
|
||||
raise ValueError(f"Weights should sum to 1.0, got {total}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Normalizers:
|
||||
"""Normalization constants for reward scaling.
|
||||
|
||||
These are typical values used to scale metrics to similar ranges.
|
||||
Tune based on your specific tools and tasks.
|
||||
"""
|
||||
|
||||
energy_scale: float = 100.0
|
||||
cost_scale: float = 0.10
|
||||
latency_scale: float = 30.0
|
||||
power_scale: float = 200.0
|
||||
|
||||
|
||||
class MultiObjectiveReward:
|
||||
"""Multi-objective reward combining accuracy, cost, energy, latency, power.
|
||||
|
||||
Formula::
|
||||
|
||||
reward = alpha * accuracy
|
||||
- beta_cost * (cost / cost_scale)
|
||||
- beta_energy * (energy / energy_scale)
|
||||
- gamma_latency * (latency / latency_scale)
|
||||
- gamma_power * (power / power_scale)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weights: RewardWeights,
|
||||
normalizers: Normalizers,
|
||||
) -> None:
|
||||
self.weights = weights
|
||||
self.normalizers = normalizers
|
||||
|
||||
def compute(self, episode: Episode) -> float:
|
||||
"""Compute scalar reward for an episode."""
|
||||
accuracy_reward = 1.0 if episode.correct else 0.0
|
||||
|
||||
cost_penalty = episode.total_cost_usd / self.normalizers.cost_scale
|
||||
energy_penalty = (
|
||||
episode.total_energy_joules / self.normalizers.energy_scale
|
||||
)
|
||||
latency_penalty = (
|
||||
episode.total_latency_seconds / self.normalizers.latency_scale
|
||||
)
|
||||
power_penalty = episode.max_power_watts / self.normalizers.power_scale
|
||||
|
||||
return (
|
||||
self.weights.alpha * accuracy_reward
|
||||
- self.weights.beta_cost * cost_penalty
|
||||
- self.weights.beta_energy * energy_penalty
|
||||
- self.weights.gamma_latency * latency_penalty
|
||||
- self.weights.gamma_power * power_penalty
|
||||
)
|
||||
|
||||
def compute_with_breakdown(self, episode: Episode) -> Dict[str, float]:
|
||||
"""Compute reward with detailed per-component breakdown."""
|
||||
accuracy_reward = 1.0 if episode.correct else 0.0
|
||||
|
||||
cost_penalty = episode.total_cost_usd / self.normalizers.cost_scale
|
||||
energy_penalty = (
|
||||
episode.total_energy_joules / self.normalizers.energy_scale
|
||||
)
|
||||
latency_penalty = (
|
||||
episode.total_latency_seconds / self.normalizers.latency_scale
|
||||
)
|
||||
power_penalty = episode.max_power_watts / self.normalizers.power_scale
|
||||
|
||||
accuracy_component = self.weights.alpha * accuracy_reward
|
||||
cost_component = -self.weights.beta_cost * cost_penalty
|
||||
energy_component = -self.weights.beta_energy * energy_penalty
|
||||
latency_component = -self.weights.gamma_latency * latency_penalty
|
||||
power_component = -self.weights.gamma_power * power_penalty
|
||||
|
||||
total_reward = (
|
||||
accuracy_component
|
||||
+ cost_component
|
||||
+ energy_component
|
||||
+ latency_component
|
||||
+ power_component
|
||||
)
|
||||
|
||||
ipj = episode.compute_ipj()
|
||||
|
||||
return {
|
||||
"total_reward": total_reward,
|
||||
"accuracy_reward": accuracy_reward,
|
||||
"accuracy_component": accuracy_component,
|
||||
"cost_penalty": cost_penalty,
|
||||
"cost_component": cost_component,
|
||||
"energy_penalty": energy_penalty,
|
||||
"energy_component": energy_component,
|
||||
"latency_penalty": latency_penalty,
|
||||
"latency_component": latency_component,
|
||||
"power_penalty": power_penalty,
|
||||
"power_component": power_component,
|
||||
"ipj": ipj,
|
||||
"total_energy_joules": episode.total_energy_joules,
|
||||
"total_cost_usd": episode.total_cost_usd,
|
||||
"total_latency_seconds": episode.total_latency_seconds,
|
||||
}
|
||||
|
||||
def compute_batch(self, episodes: List[Episode]) -> List[float]:
|
||||
"""Compute rewards for a batch of episodes."""
|
||||
return [self.compute(ep) for ep in episodes]
|
||||
|
||||
|
||||
class AdaptiveRewardWeights:
|
||||
"""Adaptive reward weights that shift during training.
|
||||
|
||||
Early training focuses on accuracy (higher alpha).
|
||||
Late training optimises efficiency (higher cost/energy/power weights).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
initial_alpha: float = 0.6,
|
||||
final_alpha: float = 0.3,
|
||||
initial_beta_cost: float = 0.1,
|
||||
final_beta_cost: float = 0.15,
|
||||
initial_beta_energy: float = 0.1,
|
||||
final_beta_energy: float = 0.2,
|
||||
initial_gamma_latency: float = 0.1,
|
||||
final_gamma_latency: float = 0.15,
|
||||
initial_gamma_power: float = 0.1,
|
||||
final_gamma_power: float = 0.2,
|
||||
total_steps: int = 10000,
|
||||
) -> None:
|
||||
self.initial_alpha = initial_alpha
|
||||
self.final_alpha = final_alpha
|
||||
self.initial_beta_cost = initial_beta_cost
|
||||
self.final_beta_cost = final_beta_cost
|
||||
self.initial_beta_energy = initial_beta_energy
|
||||
self.final_beta_energy = final_beta_energy
|
||||
self.initial_gamma_latency = initial_gamma_latency
|
||||
self.final_gamma_latency = final_gamma_latency
|
||||
self.initial_gamma_power = initial_gamma_power
|
||||
self.final_gamma_power = final_gamma_power
|
||||
self.total_steps = total_steps
|
||||
|
||||
def get_weights(self, current_step: int) -> RewardWeights:
|
||||
"""Get weights for *current_step* via linear interpolation."""
|
||||
progress = min(1.0, current_step / self.total_steps)
|
||||
|
||||
alpha = self.initial_alpha + (self.final_alpha - self.initial_alpha) * progress
|
||||
beta_cost = (
|
||||
self.initial_beta_cost
|
||||
+ (self.final_beta_cost - self.initial_beta_cost) * progress
|
||||
)
|
||||
beta_energy = (
|
||||
self.initial_beta_energy
|
||||
+ (self.final_beta_energy - self.initial_beta_energy) * progress
|
||||
)
|
||||
gamma_latency = (
|
||||
self.initial_gamma_latency
|
||||
+ (self.final_gamma_latency - self.initial_gamma_latency) * progress
|
||||
)
|
||||
gamma_power = (
|
||||
self.initial_gamma_power
|
||||
+ (self.final_gamma_power - self.initial_gamma_power) * progress
|
||||
)
|
||||
|
||||
# Normalize to sum to 1.0
|
||||
total = alpha + beta_cost + beta_energy + gamma_latency + gamma_power
|
||||
return RewardWeights(
|
||||
alpha=alpha / total,
|
||||
beta_cost=beta_cost / total,
|
||||
beta_energy=beta_energy / total,
|
||||
gamma_latency=gamma_latency / total,
|
||||
gamma_power=gamma_power / total,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AdaptiveRewardWeights",
|
||||
"MultiObjectiveReward",
|
||||
"Normalizers",
|
||||
"RewardWeights",
|
||||
]
|
||||
@@ -0,0 +1,402 @@
|
||||
"""SFT (Supervised Fine-Tuning) trainer for orchestrator.
|
||||
|
||||
Adapted from IPW's ``sft_trainer.py``. Trains the orchestrator policy
|
||||
using supervised learning on trajectories. All ``torch``/``transformers``
|
||||
imports are guarded so the module can be imported without GPU dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterator, List
|
||||
|
||||
# Optional imports -----------------------------------------------------------
|
||||
try:
|
||||
import torch
|
||||
import torch.nn.functional as F # noqa: F401
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
HAS_TORCH = True
|
||||
except ImportError:
|
||||
HAS_TORCH = False
|
||||
torch = None # type: ignore[assignment]
|
||||
DataLoader = None # type: ignore[assignment,misc]
|
||||
|
||||
from openjarvis.core.registry import LearningRegistry
|
||||
from openjarvis.learning._stubs import AgentLearningPolicy
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrchestratorSFTConfig:
|
||||
"""Configuration for orchestrator SFT training."""
|
||||
|
||||
# Model
|
||||
model_name: str = "Qwen/Qwen3-1.7B"
|
||||
max_seq_length: int = 4096
|
||||
|
||||
# Training
|
||||
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
|
||||
|
||||
# Trace generation
|
||||
teacher_engine_key: str = ""
|
||||
teacher_model: str = ""
|
||||
traces_per_query: int = 2
|
||||
max_attempts_per_trace: int = 3
|
||||
generation_temperature: float = 0.7
|
||||
|
||||
# Data source
|
||||
trace_cache_path: str = "data/orchestrator_sft_traces.jsonl"
|
||||
regenerate_traces: bool = False
|
||||
|
||||
# Checkpoint
|
||||
checkpoint_dir: str = "checkpoints/orchestrator_sft"
|
||||
save_every_n_epochs: int = 1
|
||||
|
||||
# Logging
|
||||
log_dir: str = "logs/orchestrator_sft"
|
||||
log_every_n_steps: int = 10
|
||||
use_wandb: bool = False
|
||||
|
||||
# Memory
|
||||
gradient_checkpointing: bool = True
|
||||
|
||||
# Available tools for structured prompt
|
||||
available_tools: List[str] = field(
|
||||
default_factory=lambda: [
|
||||
"calculator",
|
||||
"think",
|
||||
"code_interpreter",
|
||||
"web_search",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class OrchestratorSFTDataset:
|
||||
"""Dataset for SFT training from generated trace JSONL files."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
trace_path: str,
|
||||
tokenizer: Any,
|
||||
max_seq_length: int = 4096,
|
||||
) -> None:
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_length = max_seq_length
|
||||
self.traces: List[Dict[str, Any]] = []
|
||||
self._load_traces(trace_path)
|
||||
|
||||
def _load_traces(self, trace_path: str) -> None:
|
||||
path = Path(trace_path)
|
||||
if not path.exists():
|
||||
return
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
self.traces.append(json.loads(line))
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.traces)
|
||||
|
||||
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
||||
trace = self.traces[idx]
|
||||
text = self._format_conversation(trace.get("conversations", []))
|
||||
|
||||
encoding = self.tokenizer(
|
||||
text,
|
||||
truncation=True,
|
||||
max_length=self.max_seq_length,
|
||||
padding="max_length",
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
return {
|
||||
"input_ids": encoding["input_ids"].squeeze(0),
|
||||
"attention_mask": encoding["attention_mask"].squeeze(0),
|
||||
"labels": encoding["input_ids"].squeeze(0).clone(),
|
||||
}
|
||||
|
||||
def _format_conversation(
|
||||
self, conversations: List[Dict[str, str]]
|
||||
) -> str:
|
||||
"""Format conversation turns into training text."""
|
||||
if hasattr(self.tokenizer, "apply_chat_template"):
|
||||
try:
|
||||
messages = []
|
||||
for turn in conversations:
|
||||
role = turn.get("role") or turn.get("from", "")
|
||||
content = turn.get("content") or turn.get("value", "")
|
||||
|
||||
if role in ("human", "user"):
|
||||
role = "user"
|
||||
elif role in ("gpt", "assistant"):
|
||||
role = "assistant"
|
||||
elif role == "tool":
|
||||
tool_name = turn.get("name", "tool")
|
||||
content = (
|
||||
f"[Tool '{tool_name}' returned]: {content}"
|
||||
)
|
||||
role = "user"
|
||||
|
||||
if role in ("user", "assistant", "system"):
|
||||
messages.append({"role": role, "content": content})
|
||||
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
except Exception:
|
||||
pass # fall through
|
||||
|
||||
# Manual fallback
|
||||
parts: list[str] = []
|
||||
for turn in conversations:
|
||||
role = turn.get("role") or turn.get("from", "")
|
||||
content = turn.get("content") or turn.get("value", "")
|
||||
|
||||
if role in ("human", "user"):
|
||||
parts.append(f"<|user|>\n{content}")
|
||||
elif role in ("gpt", "assistant"):
|
||||
parts.append(f"<|assistant|>\n{content}")
|
||||
elif role == "system":
|
||||
parts.append(f"<|system|>\n{content}")
|
||||
elif role == "tool":
|
||||
tool_name = turn.get("name", "tool")
|
||||
parts.append(
|
||||
f"<|user|>\n[Tool '{tool_name}' returned]: {content}"
|
||||
)
|
||||
|
||||
eos = getattr(self.tokenizer, "eos_token", "") or ""
|
||||
return "\n".join(parts) + eos
|
||||
|
||||
def iter_batches(
|
||||
self, batch_size: int
|
||||
) -> Iterator[List[Dict[str, Any]]]:
|
||||
batch: list[Dict[str, Any]] = []
|
||||
for i in range(len(self)):
|
||||
batch.append(self[i])
|
||||
if len(batch) == batch_size:
|
||||
yield batch
|
||||
batch = []
|
||||
if batch:
|
||||
yield batch
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trainer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class OrchestratorSFTTrainer:
|
||||
"""SFT trainer for orchestrator policy.
|
||||
|
||||
Performs standard next-token cross-entropy loss on successful
|
||||
trajectories. ``torch`` must be installed to call :meth:`train`.
|
||||
"""
|
||||
|
||||
def __init__(self, config: OrchestratorSFTConfig) -> None:
|
||||
self.config = config
|
||||
self.device = None
|
||||
self.global_step = 0
|
||||
|
||||
if HAS_TORCH and torch is not None:
|
||||
self.device = torch.device(
|
||||
"cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
|
||||
self._init_model()
|
||||
self._init_data()
|
||||
self._init_optimizer()
|
||||
|
||||
def _init_model(self) -> None:
|
||||
from openjarvis.learning.orchestrator.policy_model import (
|
||||
OrchestratorPolicyModel,
|
||||
)
|
||||
|
||||
if not HAS_TORCH:
|
||||
self.policy: Any = OrchestratorPolicyModel()
|
||||
return
|
||||
|
||||
device_str = str(self.device) if self.device else None
|
||||
self.policy = OrchestratorPolicyModel.from_pretrained(
|
||||
self.config.model_name,
|
||||
gradient_checkpointing=self.config.gradient_checkpointing,
|
||||
device=device_str,
|
||||
)
|
||||
if self.policy.model is not None:
|
||||
self.policy.model.train()
|
||||
|
||||
def _init_data(self) -> None:
|
||||
trace_path = Path(self.config.trace_cache_path)
|
||||
|
||||
if self.config.regenerate_traces or not trace_path.exists():
|
||||
self._generate_traces()
|
||||
|
||||
self.dataset = OrchestratorSFTDataset(
|
||||
trace_path=str(trace_path),
|
||||
tokenizer=self.policy.tokenizer,
|
||||
max_seq_length=self.config.max_seq_length,
|
||||
)
|
||||
|
||||
def _generate_traces(self) -> None:
|
||||
"""Generate SFT traces (placeholder — requires running engine)."""
|
||||
trace_path = Path(self.config.trace_cache_path)
|
||||
trace_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not trace_path.exists():
|
||||
trace_path.touch()
|
||||
|
||||
def _init_optimizer(self) -> None:
|
||||
if not HAS_TORCH or self.policy.model is None:
|
||||
self.optimizer: Any = None
|
||||
self.scheduler: Any = None
|
||||
self.dataloader: Any = None
|
||||
return
|
||||
|
||||
self.optimizer = torch.optim.AdamW(
|
||||
self.policy.model.parameters(),
|
||||
lr=self.config.learning_rate,
|
||||
weight_decay=self.config.weight_decay,
|
||||
)
|
||||
|
||||
self.dataloader = DataLoader(
|
||||
self.dataset,
|
||||
batch_size=self.config.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=0,
|
||||
)
|
||||
|
||||
total_steps = len(self.dataloader) * self.config.num_epochs
|
||||
warmup_steps = int(total_steps * self.config.warmup_ratio)
|
||||
|
||||
def lr_lambda(step: int) -> float:
|
||||
if step < warmup_steps:
|
||||
return step / max(1, warmup_steps)
|
||||
return max(
|
||||
0.0,
|
||||
1.0 - (step - warmup_steps) / (total_steps - warmup_steps),
|
||||
)
|
||||
|
||||
self.scheduler = torch.optim.lr_scheduler.LambdaLR(
|
||||
self.optimizer, lr_lambda
|
||||
)
|
||||
|
||||
def train(self) -> None:
|
||||
"""Run the SFT training loop."""
|
||||
if not HAS_TORCH:
|
||||
raise RuntimeError(
|
||||
"PyTorch is required for training. "
|
||||
"Install with: pip install torch transformers"
|
||||
)
|
||||
|
||||
for epoch in range(self.config.num_epochs):
|
||||
self._train_epoch(epoch)
|
||||
|
||||
if (epoch + 1) % self.config.save_every_n_epochs == 0:
|
||||
self._save_checkpoint(epoch)
|
||||
|
||||
def _train_epoch(self, epoch: int) -> Dict[str, float]:
|
||||
if self.policy.model is None or self.dataloader is None:
|
||||
return {"epoch": epoch, "loss": 0.0}
|
||||
|
||||
self.policy.model.train()
|
||||
total_loss = 0.0
|
||||
num_batches = 0
|
||||
|
||||
for batch in self.dataloader:
|
||||
loss = self._train_step(batch)
|
||||
total_loss += loss
|
||||
num_batches += 1
|
||||
self.global_step += 1
|
||||
|
||||
avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
|
||||
return {"epoch": epoch, "loss": avg_loss}
|
||||
|
||||
def _train_step(self, batch: Dict[str, Any]) -> float:
|
||||
input_ids = batch["input_ids"].to(self.device)
|
||||
attention_mask = batch["attention_mask"].to(self.device)
|
||||
labels = batch["labels"].to(self.device)
|
||||
|
||||
outputs = self.policy.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
labels=labels,
|
||||
)
|
||||
loss = outputs.loss
|
||||
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
self.policy.model.parameters(),
|
||||
self.config.max_grad_norm,
|
||||
)
|
||||
self.optimizer.step()
|
||||
self.scheduler.step()
|
||||
|
||||
return loss.item()
|
||||
|
||||
def _save_checkpoint(self, epoch: int) -> None:
|
||||
checkpoint_dir = Path(self.config.checkpoint_dir)
|
||||
checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
checkpoint_path = checkpoint_dir / f"epoch_{epoch + 1}"
|
||||
if self.policy.model is not None:
|
||||
self.policy.save(str(checkpoint_path))
|
||||
|
||||
state_path = checkpoint_path / "training_state.json"
|
||||
state = {
|
||||
"epoch": epoch,
|
||||
"global_step": self.global_step,
|
||||
"config": asdict(self.config),
|
||||
}
|
||||
with open(state_path, "w") as f:
|
||||
json.dump(state, f, indent=2)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry wrapper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _ensure_registered() -> None:
|
||||
if LearningRegistry.contains("orchestrator_sft"):
|
||||
return
|
||||
|
||||
@LearningRegistry.register("orchestrator_sft")
|
||||
class OrchestratorSFTPolicy(AgentLearningPolicy):
|
||||
"""Wrapper that registers the SFT trainer as a learning policy."""
|
||||
|
||||
def update(
|
||||
self, trace_store: Any, **kwargs: object
|
||||
) -> Dict[str, Any]:
|
||||
config = OrchestratorSFTConfig(**{
|
||||
k: v for k, v in kwargs.items()
|
||||
if k in OrchestratorSFTConfig.__dataclass_fields__
|
||||
})
|
||||
trainer = OrchestratorSFTTrainer(config)
|
||||
trainer.train()
|
||||
return {"status": "sft_training_complete"}
|
||||
|
||||
|
||||
_ensure_registered()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OrchestratorSFTConfig",
|
||||
"OrchestratorSFTDataset",
|
||||
"OrchestratorSFTTrainer",
|
||||
]
|
||||
@@ -0,0 +1,338 @@
|
||||
"""Episode dataclasses for orchestrator training.
|
||||
|
||||
Adapted from IPW's ``episode_builder.py``. These types represent the core
|
||||
data structures for orchestrator RL/SFT training: actions, observations,
|
||||
episode steps, and complete episodes with aggregate metrics.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Answer grading utilities
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def normalize_number(s: str) -> Optional[float]:
|
||||
"""Try to parse a string as a number.
|
||||
|
||||
Returns None if not a valid number.
|
||||
"""
|
||||
s = s.strip().lower()
|
||||
s = re.sub(r"[,\s]", "", s) # Remove commas and spaces
|
||||
s = re.sub(r"\.0+$", "", s) # Remove trailing .0
|
||||
|
||||
try:
|
||||
return float(s)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def extract_answer(text: str) -> str:
|
||||
"""Extract the core answer from a potentially verbose response.
|
||||
|
||||
Handles patterns like:
|
||||
- "The answer is 4"
|
||||
- "Result: 4.0"
|
||||
- "4" (unchanged)
|
||||
- "Therefore, the answer is approximately 4"
|
||||
"""
|
||||
text = text.strip()
|
||||
|
||||
patterns = [
|
||||
r"(?:the\s+)?answer\s+is[:\s]+(.+?)(?:\.|$)",
|
||||
r"result[:\s]+(.+?)(?:\.|$)",
|
||||
r"=\s*(.+?)(?:\.|$)",
|
||||
r"therefore[,\s]+(?:the\s+)?(?:answer\s+is\s+)?(.+?)(?:\.|$)",
|
||||
]
|
||||
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, text, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def grade_answer(
|
||||
predicted: str, expected: str, tolerance: float = 1e-6
|
||||
) -> bool:
|
||||
"""Grade an answer against expected, with smart matching.
|
||||
|
||||
Handles:
|
||||
- Exact string match (case-insensitive)
|
||||
- Numeric comparison with tolerance
|
||||
- Answer extraction from verbose responses
|
||||
|
||||
Args:
|
||||
predicted: The model's answer.
|
||||
expected: Ground truth answer.
|
||||
tolerance: Tolerance for numeric comparisons.
|
||||
|
||||
Returns:
|
||||
True if answer is correct.
|
||||
"""
|
||||
predicted = predicted.strip()
|
||||
expected = expected.strip()
|
||||
|
||||
# Exact match (case-insensitive)
|
||||
if predicted.lower() == expected.lower():
|
||||
return True
|
||||
|
||||
# Try extracting core answer
|
||||
pred_extracted = extract_answer(predicted)
|
||||
exp_extracted = extract_answer(expected)
|
||||
|
||||
if pred_extracted.lower() == exp_extracted.lower():
|
||||
return True
|
||||
|
||||
# Try numeric comparison
|
||||
pred_num = normalize_number(pred_extracted)
|
||||
exp_num = normalize_number(exp_extracted)
|
||||
|
||||
if pred_num is not None and exp_num is not None:
|
||||
if exp_num == 0:
|
||||
return abs(pred_num) < tolerance
|
||||
return abs(pred_num - exp_num) / abs(exp_num) < tolerance
|
||||
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core dataclasses
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrchestratorAction:
|
||||
"""Orchestrator action: thought + tool selection + tool input."""
|
||||
|
||||
thought: str
|
||||
"""Reasoning about what to do next."""
|
||||
|
||||
tool_name: str
|
||||
"""Selected tool name (e.g., 'calculator', 'think')."""
|
||||
|
||||
tool_input: str
|
||||
"""Input/prompt to send to the tool."""
|
||||
|
||||
is_final_answer: bool = False
|
||||
"""Whether this action provides the final answer."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrchestratorObservation:
|
||||
"""Result from executing an action, with flat telemetry fields."""
|
||||
|
||||
content: str
|
||||
"""Tool response content."""
|
||||
|
||||
latency_seconds: float = 0.0
|
||||
"""Latency in seconds."""
|
||||
|
||||
cost_usd: float = 0.0
|
||||
"""Cost in USD."""
|
||||
|
||||
energy_joules: float = 0.0
|
||||
"""Energy consumed in joules."""
|
||||
|
||||
power_watts: float = 0.0
|
||||
"""Power usage in watts."""
|
||||
|
||||
tokens: int = 0
|
||||
"""Tokens consumed."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class EpisodeStep:
|
||||
"""Single step in an episode."""
|
||||
|
||||
turn: int
|
||||
"""Step number (0-indexed)."""
|
||||
|
||||
action: OrchestratorAction
|
||||
"""Action taken."""
|
||||
|
||||
observation: OrchestratorObservation
|
||||
"""Result of action."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Episode:
|
||||
"""Complete RL episode with aggregate metrics."""
|
||||
|
||||
task_id: str
|
||||
"""Unique task identifier."""
|
||||
|
||||
initial_prompt: str
|
||||
"""Initial question/task."""
|
||||
|
||||
steps: List[EpisodeStep] = field(default_factory=list)
|
||||
"""Sequence of (action, observation) pairs."""
|
||||
|
||||
final_answer: str = ""
|
||||
"""Final answer produced by orchestrator."""
|
||||
|
||||
ground_truth: str = ""
|
||||
"""Ground truth answer."""
|
||||
|
||||
correct: bool = False
|
||||
"""Whether final answer matches ground truth."""
|
||||
|
||||
# Aggregate metrics
|
||||
total_energy_joules: float = 0.0
|
||||
total_cost_usd: float = 0.0
|
||||
total_latency_seconds: float = 0.0
|
||||
total_tokens: int = 0
|
||||
max_power_watts: float = 0.0
|
||||
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def add_step(
|
||||
self, action: OrchestratorAction, observation: OrchestratorObservation
|
||||
) -> None:
|
||||
"""Add a step to the episode and update aggregate metrics."""
|
||||
step = EpisodeStep(
|
||||
turn=len(self.steps),
|
||||
action=action,
|
||||
observation=observation,
|
||||
)
|
||||
self.steps.append(step)
|
||||
|
||||
self.total_energy_joules += observation.energy_joules
|
||||
self.total_latency_seconds += observation.latency_seconds
|
||||
self.total_cost_usd += observation.cost_usd
|
||||
self.total_tokens += observation.tokens
|
||||
self.max_power_watts = max(self.max_power_watts, observation.power_watts)
|
||||
|
||||
if action.is_final_answer:
|
||||
self.final_answer = observation.content
|
||||
|
||||
def num_turns(self) -> int:
|
||||
"""Return number of turns in episode."""
|
||||
return len(self.steps)
|
||||
|
||||
def compute_ipj(self) -> float:
|
||||
"""Compute Intelligence Per Joule (IPJ).
|
||||
|
||||
Returns:
|
||||
IPJ score (higher is better). 0.0 if energy is zero or
|
||||
the answer is incorrect.
|
||||
"""
|
||||
if self.total_energy_joules <= 0:
|
||||
return 0.0
|
||||
accuracy_score = 1.0 if self.correct else 0.0
|
||||
return accuracy_score / self.total_energy_joules
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert episode to dictionary for serialization."""
|
||||
return {
|
||||
"task_id": self.task_id,
|
||||
"initial_prompt": self.initial_prompt,
|
||||
"steps": [
|
||||
{
|
||||
"turn": step.turn,
|
||||
"thought": step.action.thought,
|
||||
"tool": step.action.tool_name,
|
||||
"tool_input": step.action.tool_input,
|
||||
"observation": step.observation.content[:200],
|
||||
"energy_joules": step.observation.energy_joules,
|
||||
"latency_seconds": step.observation.latency_seconds,
|
||||
"cost_usd": step.observation.cost_usd,
|
||||
}
|
||||
for step in self.steps
|
||||
],
|
||||
"final_answer": self.final_answer,
|
||||
"ground_truth": self.ground_truth,
|
||||
"correct": self.correct,
|
||||
"total_energy_joules": self.total_energy_joules,
|
||||
"total_latency_seconds": self.total_latency_seconds,
|
||||
"total_cost_usd": self.total_cost_usd,
|
||||
"total_tokens": self.total_tokens,
|
||||
"num_turns": self.num_turns(),
|
||||
"ipj": self.compute_ipj(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class EpisodeState:
|
||||
"""Mutable state during episode execution."""
|
||||
|
||||
initial_prompt: str
|
||||
"""Initial task/question."""
|
||||
|
||||
history: List[Tuple[OrchestratorAction, OrchestratorObservation]] = field(
|
||||
default_factory=list
|
||||
)
|
||||
"""History of (action, observation) pairs."""
|
||||
|
||||
final_answer: Optional[str] = None
|
||||
"""Final answer (set when is_final_answer action is taken)."""
|
||||
|
||||
def add_turn(
|
||||
self,
|
||||
action: OrchestratorAction,
|
||||
observation: OrchestratorObservation,
|
||||
) -> None:
|
||||
"""Add a turn to the episode history."""
|
||||
self.history.append((action, observation))
|
||||
if action.is_final_answer:
|
||||
self.final_answer = observation.content
|
||||
|
||||
def num_turns(self) -> int:
|
||||
"""Return number of turns so far."""
|
||||
return len(self.history)
|
||||
|
||||
def to_episode(
|
||||
self, task_id: str, ground_truth: str, correct: bool
|
||||
) -> Episode:
|
||||
"""Convert state to Episode for reward computation."""
|
||||
episode = Episode(
|
||||
task_id=task_id,
|
||||
initial_prompt=self.initial_prompt,
|
||||
ground_truth=ground_truth,
|
||||
final_answer=self.final_answer or "",
|
||||
correct=correct,
|
||||
)
|
||||
for action, observation in self.history:
|
||||
episode.add_step(action, observation)
|
||||
return episode
|
||||
|
||||
|
||||
@dataclass
|
||||
class PolicyOutput:
|
||||
"""Output from policy model prediction."""
|
||||
|
||||
thought: str
|
||||
"""Reasoning about what to do."""
|
||||
|
||||
tool_name: str
|
||||
"""Selected tool."""
|
||||
|
||||
tool_input: str
|
||||
"""Input for the tool."""
|
||||
|
||||
is_final_answer: bool = False
|
||||
"""Whether this provides the final answer."""
|
||||
|
||||
raw_text: str = ""
|
||||
"""Raw model output."""
|
||||
|
||||
confidence: float = 1.0
|
||||
"""Confidence score (if available)."""
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Episode",
|
||||
"EpisodeState",
|
||||
"EpisodeStep",
|
||||
"OrchestratorAction",
|
||||
"OrchestratorObservation",
|
||||
"PolicyOutput",
|
||||
"extract_answer",
|
||||
"grade_answer",
|
||||
"normalize_number",
|
||||
]
|
||||
@@ -1,4 +1,9 @@
|
||||
"""SFT on traces — learns which model handles which query type best."""
|
||||
"""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
|
||||
|
||||
@@ -9,11 +14,13 @@ from openjarvis.learning._stubs import IntelligenceLearningPolicy
|
||||
|
||||
|
||||
@LearningRegistry.register("sft")
|
||||
class SFTPolicy(IntelligenceLearningPolicy):
|
||||
"""SFT on traces: learns which model handles which query type best.
|
||||
class SFTRouterPolicy(IntelligenceLearningPolicy):
|
||||
"""Trace-driven router that learns query_class → model mappings.
|
||||
|
||||
Reads successful traces, groups by query class, and produces an
|
||||
updated query_class -> model mapping.
|
||||
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:
|
||||
@@ -84,4 +91,7 @@ class SFTPolicy(IntelligenceLearningPolicy):
|
||||
return dict(self._policy_map)
|
||||
|
||||
|
||||
__all__ = ["SFTPolicy"]
|
||||
__all__ = ["SFTRouterPolicy"]
|
||||
|
||||
# Backward-compat alias
|
||||
SFTPolicy = SFTRouterPolicy
|
||||
|
||||
@@ -27,6 +27,10 @@ _TOOL_ANNOTATIONS: Dict[str, Dict[str, Any]] = {
|
||||
"file_read": {"readOnlyHint": True, "destructiveHint": False},
|
||||
"web_search": {"readOnlyHint": True, "destructiveHint": False},
|
||||
"code_interpreter": {"destructiveHint": True, "readOnlyHint": False},
|
||||
"repl": {"destructiveHint": True, "readOnlyHint": False},
|
||||
"channel_send": {"destructiveHint": True, "readOnlyHint": False},
|
||||
"channel_list": {"readOnlyHint": True, "destructiveHint": False},
|
||||
"channel_status": {"readOnlyHint": True, "destructiveHint": False},
|
||||
}
|
||||
|
||||
|
||||
@@ -86,6 +90,11 @@ class MCPServer:
|
||||
_tool_classes.append(CodeInterpreterTool)
|
||||
except ImportError:
|
||||
pass
|
||||
try:
|
||||
from openjarvis.tools.repl import ReplTool
|
||||
_tool_classes.append(ReplTool)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Storage MCP tools
|
||||
try:
|
||||
@@ -102,6 +111,19 @@ class MCPServer:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Channel MCP tools
|
||||
try:
|
||||
from openjarvis.tools.channel_tools import (
|
||||
ChannelListTool,
|
||||
ChannelSendTool,
|
||||
ChannelStatusTool,
|
||||
)
|
||||
_tool_classes.extend([
|
||||
ChannelSendTool, ChannelListTool, ChannelStatusTool,
|
||||
])
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# LM tool (needs engine/model — instantiate with None)
|
||||
try:
|
||||
from openjarvis.tools.llm_tool import LLMTool
|
||||
|
||||
@@ -26,7 +26,6 @@ class _NoCacheStaticFiles(StaticFiles):
|
||||
async def __call__(self, scope, receive, send):
|
||||
async def _send_with_headers(message):
|
||||
if message["type"] == "http.response.start":
|
||||
headers = dict(message.get("headers", []))
|
||||
extra = [
|
||||
(k.encode(), v.encode()) for k, v in _NO_CACHE_HEADERS.items()
|
||||
]
|
||||
|
||||
+141
-2
@@ -26,6 +26,7 @@ class JarvisSystem:
|
||||
tools: List[BaseTool] = field(default_factory=list)
|
||||
tool_executor: Optional[ToolExecutor] = None
|
||||
memory_backend: Optional[Any] = None # MemoryBackend
|
||||
channel_backend: Optional[Any] = None # BaseChannel
|
||||
router: Optional[Any] = None # RouterPolicy
|
||||
mcp_server: Optional[Any] = None # MCPServer
|
||||
telemetry_store: Optional[Any] = None
|
||||
@@ -119,6 +120,9 @@ class JarvisSystem:
|
||||
if agent_name == "orchestrator":
|
||||
agent_kwargs["tools"] = agent_tools
|
||||
agent_kwargs["max_turns"] = self.config.agent.max_turns
|
||||
elif agent_name == "rlm":
|
||||
agent_kwargs["tools"] = agent_tools
|
||||
agent_kwargs["max_turns"] = self.config.agent.max_turns
|
||||
|
||||
try:
|
||||
ag = agent_cls(self.engine, self.model, **agent_kwargs)
|
||||
@@ -263,8 +267,13 @@ class SystemBuilder:
|
||||
# Resolve memory backend
|
||||
memory_backend = self._resolve_memory(config)
|
||||
|
||||
# Resolve channel backend
|
||||
channel_backend = self._resolve_channel(config, bus)
|
||||
|
||||
# Resolve tools
|
||||
tool_list = self._resolve_tools(config, engine, model, memory_backend)
|
||||
tool_list = self._resolve_tools(
|
||||
config, engine, model, memory_backend, channel_backend,
|
||||
)
|
||||
|
||||
# Build tool executor
|
||||
tool_executor = ToolExecutor(tool_list, bus) if tool_list else None
|
||||
@@ -282,6 +291,7 @@ class SystemBuilder:
|
||||
tools=tool_list,
|
||||
tool_executor=tool_executor,
|
||||
memory_backend=memory_backend,
|
||||
channel_backend=channel_backend,
|
||||
telemetry_store=telemetry_store,
|
||||
)
|
||||
|
||||
@@ -369,7 +379,125 @@ class SystemBuilder:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _resolve_tools(self, config, engine, model, memory_backend):
|
||||
def _resolve_channel(self, config, bus):
|
||||
"""Resolve channel backend from config."""
|
||||
if not config.channel.enabled:
|
||||
return None
|
||||
try:
|
||||
import openjarvis.channels # noqa: F401 -- trigger registration
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
key = config.channel.default_channel
|
||||
if not key:
|
||||
return None
|
||||
if not ChannelRegistry.contains(key):
|
||||
return None
|
||||
|
||||
kwargs: Dict[str, Any] = {"bus": bus}
|
||||
if key == "telegram":
|
||||
tc = config.channel.telegram
|
||||
if tc.bot_token:
|
||||
kwargs["bot_token"] = tc.bot_token
|
||||
if tc.parse_mode:
|
||||
kwargs["parse_mode"] = tc.parse_mode
|
||||
elif key == "discord":
|
||||
dc = config.channel.discord
|
||||
if dc.bot_token:
|
||||
kwargs["bot_token"] = dc.bot_token
|
||||
elif key == "slack":
|
||||
sc = config.channel.slack
|
||||
if sc.bot_token:
|
||||
kwargs["bot_token"] = sc.bot_token
|
||||
if sc.app_token:
|
||||
kwargs["app_token"] = sc.app_token
|
||||
elif key == "webhook":
|
||||
wc = config.channel.webhook
|
||||
if wc.url:
|
||||
kwargs["url"] = wc.url
|
||||
if wc.secret:
|
||||
kwargs["secret"] = wc.secret
|
||||
if wc.method:
|
||||
kwargs["method"] = wc.method
|
||||
elif key == "email":
|
||||
ec = config.channel.email
|
||||
if ec.smtp_host:
|
||||
kwargs["smtp_host"] = ec.smtp_host
|
||||
kwargs["smtp_port"] = ec.smtp_port
|
||||
if ec.imap_host:
|
||||
kwargs["imap_host"] = ec.imap_host
|
||||
kwargs["imap_port"] = ec.imap_port
|
||||
if ec.username:
|
||||
kwargs["username"] = ec.username
|
||||
if ec.password:
|
||||
kwargs["password"] = ec.password
|
||||
kwargs["use_tls"] = ec.use_tls
|
||||
elif key == "whatsapp":
|
||||
wac = config.channel.whatsapp
|
||||
if wac.access_token:
|
||||
kwargs["access_token"] = wac.access_token
|
||||
if wac.phone_number_id:
|
||||
kwargs["phone_number_id"] = wac.phone_number_id
|
||||
elif key == "signal":
|
||||
sgc = config.channel.signal
|
||||
if sgc.api_url:
|
||||
kwargs["api_url"] = sgc.api_url
|
||||
if sgc.phone_number:
|
||||
kwargs["phone_number"] = sgc.phone_number
|
||||
elif key == "google_chat":
|
||||
gcc = config.channel.google_chat
|
||||
if gcc.webhook_url:
|
||||
kwargs["webhook_url"] = gcc.webhook_url
|
||||
elif key == "irc":
|
||||
ic = config.channel.irc
|
||||
if ic.server:
|
||||
kwargs["server"] = ic.server
|
||||
kwargs["port"] = ic.port
|
||||
if ic.nick:
|
||||
kwargs["nick"] = ic.nick
|
||||
if ic.password:
|
||||
kwargs["password"] = ic.password
|
||||
kwargs["use_tls"] = ic.use_tls
|
||||
elif key == "webchat":
|
||||
pass # no config needed
|
||||
elif key == "teams":
|
||||
tmc = config.channel.teams
|
||||
if tmc.app_id:
|
||||
kwargs["app_id"] = tmc.app_id
|
||||
if tmc.app_password:
|
||||
kwargs["app_password"] = tmc.app_password
|
||||
if tmc.service_url:
|
||||
kwargs["service_url"] = tmc.service_url
|
||||
elif key == "matrix":
|
||||
mc = config.channel.matrix
|
||||
if mc.homeserver:
|
||||
kwargs["homeserver"] = mc.homeserver
|
||||
if mc.access_token:
|
||||
kwargs["access_token"] = mc.access_token
|
||||
elif key == "mattermost":
|
||||
mmc = config.channel.mattermost
|
||||
if mmc.url:
|
||||
kwargs["url"] = mmc.url
|
||||
if mmc.token:
|
||||
kwargs["token"] = mmc.token
|
||||
elif key == "feishu":
|
||||
fc = config.channel.feishu
|
||||
if fc.app_id:
|
||||
kwargs["app_id"] = fc.app_id
|
||||
if fc.app_secret:
|
||||
kwargs["app_secret"] = fc.app_secret
|
||||
elif key == "bluebubbles":
|
||||
bbc = config.channel.bluebubbles
|
||||
if bbc.url:
|
||||
kwargs["url"] = bbc.url
|
||||
if bbc.password:
|
||||
kwargs["password"] = bbc.password
|
||||
|
||||
return ChannelRegistry.create(key, **kwargs)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _resolve_tools(self, config, engine, model, memory_backend,
|
||||
channel_backend=None):
|
||||
"""Resolve tool instances."""
|
||||
tool_names_str = self._tool_names
|
||||
if tool_names_str is None:
|
||||
@@ -403,6 +531,17 @@ class SystemBuilder:
|
||||
if ToolRegistry.contains(name):
|
||||
tool = ToolRegistry.create(name, backend=memory_backend)
|
||||
tools.append(tool)
|
||||
elif name in (
|
||||
"channel_send",
|
||||
"channel_list",
|
||||
"channel_status",
|
||||
):
|
||||
import openjarvis.tools.channel_tools # noqa: F401
|
||||
from openjarvis.core.registry import ToolRegistry
|
||||
|
||||
if ToolRegistry.contains(name):
|
||||
tool = ToolRegistry.create(name, channel=channel_backend)
|
||||
tools.append(tool)
|
||||
else:
|
||||
from openjarvis.core.registry import ToolRegistry
|
||||
|
||||
|
||||
@@ -42,6 +42,11 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.tools.repl # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.tools.storage_tools # noqa: F401
|
||||
except ImportError:
|
||||
@@ -52,4 +57,9 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.tools.channel_tools # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__all__ = ["BaseTool", "ToolExecutor", "ToolSpec"]
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Channel MCP tools — expose channel operations as BaseTool instances.
|
||||
|
||||
These tools wrap the ``BaseChannel`` ABC so that channel operations
|
||||
(send, list, status) are discoverable and callable via MCP.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from openjarvis.channels._stubs import BaseChannel
|
||||
from openjarvis.core.registry import ToolRegistry
|
||||
from openjarvis.core.types import ToolResult
|
||||
from openjarvis.tools._stubs import BaseTool, ToolSpec
|
||||
|
||||
|
||||
@ToolRegistry.register("channel_send")
|
||||
class ChannelSendTool(BaseTool):
|
||||
"""MCP-exposed tool: send a message via a channel backend."""
|
||||
|
||||
tool_id = "channel_send"
|
||||
|
||||
def __init__(self, channel: BaseChannel | None = None) -> None:
|
||||
self._channel = channel
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="channel_send",
|
||||
description="Send a message to a channel (Telegram, Discord, Slack, etc.).",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"channel": {
|
||||
"type": "string",
|
||||
"description": "Target chat/channel ID to send to.",
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "The message content to send.",
|
||||
},
|
||||
"conversation_id": {
|
||||
"type": "string",
|
||||
"description": "Optional conversation/thread ID for replies.",
|
||||
},
|
||||
},
|
||||
"required": ["channel", "content"],
|
||||
},
|
||||
category="channel",
|
||||
)
|
||||
|
||||
def execute(self, **params: Any) -> ToolResult:
|
||||
if self._channel is None:
|
||||
return ToolResult(
|
||||
tool_name="channel_send",
|
||||
content="No channel backend configured.",
|
||||
success=False,
|
||||
)
|
||||
target = params.get("channel", "")
|
||||
content = params.get("content", "")
|
||||
if not target or not content:
|
||||
return ToolResult(
|
||||
tool_name="channel_send",
|
||||
content="Both 'channel' and 'content' are required.",
|
||||
success=False,
|
||||
)
|
||||
try:
|
||||
ok = self._channel.send(
|
||||
target,
|
||||
content,
|
||||
conversation_id=params.get("conversation_id", ""),
|
||||
)
|
||||
if ok:
|
||||
return ToolResult(
|
||||
tool_name="channel_send",
|
||||
content=f"Message sent to {target}",
|
||||
success=True,
|
||||
)
|
||||
return ToolResult(
|
||||
tool_name="channel_send",
|
||||
content=f"Failed to send message to {target}",
|
||||
success=False,
|
||||
)
|
||||
except Exception as exc:
|
||||
return ToolResult(
|
||||
tool_name="channel_send",
|
||||
content=f"Send error: {exc}",
|
||||
success=False,
|
||||
)
|
||||
|
||||
|
||||
@ToolRegistry.register("channel_list")
|
||||
class ChannelListTool(BaseTool):
|
||||
"""MCP-exposed tool: list available channels."""
|
||||
|
||||
tool_id = "channel_list"
|
||||
|
||||
def __init__(self, channel: BaseChannel | None = None) -> None:
|
||||
self._channel = channel
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="channel_list",
|
||||
description="List available messaging channels.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
},
|
||||
category="channel",
|
||||
)
|
||||
|
||||
def execute(self, **params: Any) -> ToolResult:
|
||||
if self._channel is None:
|
||||
return ToolResult(
|
||||
tool_name="channel_list",
|
||||
content="No channel backend configured.",
|
||||
success=False,
|
||||
)
|
||||
try:
|
||||
channels = self._channel.list_channels()
|
||||
if not channels:
|
||||
return ToolResult(
|
||||
tool_name="channel_list",
|
||||
content="No channels available.",
|
||||
success=True,
|
||||
)
|
||||
return ToolResult(
|
||||
tool_name="channel_list",
|
||||
content="\n".join(channels),
|
||||
success=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
return ToolResult(
|
||||
tool_name="channel_list",
|
||||
content=f"List error: {exc}",
|
||||
success=False,
|
||||
)
|
||||
|
||||
|
||||
@ToolRegistry.register("channel_status")
|
||||
class ChannelStatusTool(BaseTool):
|
||||
"""MCP-exposed tool: check channel connection status."""
|
||||
|
||||
tool_id = "channel_status"
|
||||
|
||||
def __init__(self, channel: BaseChannel | None = None) -> None:
|
||||
self._channel = channel
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="channel_status",
|
||||
description="Check the connection status of the messaging channel.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
},
|
||||
category="channel",
|
||||
)
|
||||
|
||||
def execute(self, **params: Any) -> ToolResult:
|
||||
if self._channel is None:
|
||||
return ToolResult(
|
||||
tool_name="channel_status",
|
||||
content="No channel backend configured.",
|
||||
success=False,
|
||||
)
|
||||
try:
|
||||
st = self._channel.status()
|
||||
return ToolResult(
|
||||
tool_name="channel_status",
|
||||
content=f"Channel status: {st.value}",
|
||||
success=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
return ToolResult(
|
||||
tool_name="channel_status",
|
||||
content=f"Status error: {exc}",
|
||||
success=False,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ChannelListTool",
|
||||
"ChannelSendTool",
|
||||
"ChannelStatusTool",
|
||||
]
|
||||
@@ -0,0 +1,303 @@
|
||||
"""Persistent Python REPL tool — maintains state across calls within a session.
|
||||
|
||||
Unlike ``CodeInterpreterTool`` (which runs each snippet in a fresh subprocess),
|
||||
this tool keeps variables, functions, and imports alive across invocations
|
||||
within the same session.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import redirect_stderr, redirect_stdout
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from openjarvis.core.registry import ToolRegistry
|
||||
from openjarvis.core.types import ToolResult
|
||||
from openjarvis.tools._stubs import BaseTool, ToolSpec
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Security
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Layer 1: Pattern blocklist
|
||||
_BLOCKED_PATTERNS = [
|
||||
"os.system",
|
||||
"os.popen",
|
||||
"subprocess",
|
||||
"shutil.rmtree",
|
||||
"__import__",
|
||||
"open(",
|
||||
"ctypes",
|
||||
"socket",
|
||||
"http.client",
|
||||
"urllib",
|
||||
]
|
||||
|
||||
# Layer 2: Restricted builtins — remove dangerous ones
|
||||
_REMOVED_BUILTINS = {
|
||||
"open", "exec", "eval", "compile", "__import__",
|
||||
"breakpoint", "exit", "quit", "input",
|
||||
}
|
||||
|
||||
# Layer 3: Safe import allowlist
|
||||
_SAFE_IMPORT_MODULES = frozenset({
|
||||
"math", "cmath", "decimal", "fractions", "random", "statistics",
|
||||
"itertools", "functools", "operator", "collections", "string",
|
||||
"re", "textwrap", "datetime", "time", "calendar",
|
||||
"json", "csv", "copy", "dataclasses", "enum", "typing",
|
||||
"heapq", "bisect", "array", "pprint", "abc", "numbers",
|
||||
})
|
||||
|
||||
|
||||
def _make_safe_import(allowed: frozenset = _SAFE_IMPORT_MODULES):
|
||||
"""Return a custom __import__ that only allows safe modules."""
|
||||
if isinstance(__builtins__, dict):
|
||||
real_import = __builtins__["__import__"]
|
||||
else:
|
||||
real_import = __builtins__.__import__ # type: ignore[union-attr]
|
||||
|
||||
def _safe_import(name: str, *args: Any, **kwargs: Any) -> Any:
|
||||
top_level = name.split(".")[0]
|
||||
if top_level not in allowed:
|
||||
raise ImportError(
|
||||
f"Import of '{name}' is not allowed. "
|
||||
f"Allowed modules: {', '.join(sorted(allowed))}"
|
||||
)
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
return _safe_import
|
||||
|
||||
|
||||
def _make_restricted_builtins() -> Dict[str, Any]:
|
||||
"""Build a builtins dict with dangerous functions removed."""
|
||||
import builtins
|
||||
|
||||
safe = {k: v for k, v in vars(builtins).items() if k not in _REMOVED_BUILTINS}
|
||||
safe["__import__"] = _make_safe_import()
|
||||
return safe
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ReplSession:
|
||||
session_id: str
|
||||
namespace: Dict[str, Any] = field(default_factory=dict)
|
||||
created_at: float = field(default_factory=time.time)
|
||||
last_used: float = field(default_factory=time.time)
|
||||
execution_count: int = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# REPL Tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@ToolRegistry.register("repl")
|
||||
class ReplTool(BaseTool):
|
||||
"""Persistent Python REPL with session management.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
timeout:
|
||||
Maximum execution time in seconds per call.
|
||||
max_output:
|
||||
Maximum characters of captured output.
|
||||
max_sessions:
|
||||
Maximum concurrent sessions (LRU eviction).
|
||||
"""
|
||||
|
||||
tool_id = "repl"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
timeout: int = 30,
|
||||
max_output: int = 10000,
|
||||
max_sessions: int = 16,
|
||||
) -> None:
|
||||
self._timeout = timeout
|
||||
self._max_output = max_output
|
||||
self._max_sessions = max_sessions
|
||||
self._sessions: Dict[str, _ReplSession] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="repl",
|
||||
description=(
|
||||
"Execute Python code in a persistent REPL session. "
|
||||
"Variables, functions, and imports persist across calls "
|
||||
"within the same session."
|
||||
),
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "Python code to execute.",
|
||||
},
|
||||
"session_id": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Session ID for state persistence. "
|
||||
"Omit to auto-create a new session."
|
||||
),
|
||||
},
|
||||
"reset": {
|
||||
"type": "boolean",
|
||||
"description": "Reset the session state before execution.",
|
||||
},
|
||||
},
|
||||
"required": ["code"],
|
||||
},
|
||||
category="code",
|
||||
)
|
||||
|
||||
def execute(self, **params: Any) -> ToolResult:
|
||||
code = params.get("code", "")
|
||||
session_id = params.get("session_id")
|
||||
reset = params.get("reset", False)
|
||||
|
||||
if not code or not code.strip():
|
||||
return ToolResult(
|
||||
tool_name="repl",
|
||||
content="No code provided.",
|
||||
success=False,
|
||||
)
|
||||
|
||||
# Security pattern check
|
||||
for pattern in _BLOCKED_PATTERNS:
|
||||
if pattern in code:
|
||||
return ToolResult(
|
||||
tool_name="repl",
|
||||
content=f"Blocked: code contains prohibited pattern '{pattern}'",
|
||||
success=False,
|
||||
)
|
||||
|
||||
# Resolve session
|
||||
session = self._resolve_session(session_id, reset)
|
||||
|
||||
# Execute with timeout
|
||||
output, success = self._exec_with_timeout(code, session)
|
||||
|
||||
# Update session metadata
|
||||
session.last_used = time.time()
|
||||
session.execution_count += 1
|
||||
|
||||
# Truncate output
|
||||
if len(output) > self._max_output:
|
||||
output = output[: self._max_output] + "\n... (output truncated)"
|
||||
|
||||
return ToolResult(
|
||||
tool_name="repl",
|
||||
content=output or "(no output)",
|
||||
success=success,
|
||||
metadata={
|
||||
"session_id": session.session_id,
|
||||
"execution_count": session.execution_count,
|
||||
},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Session management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _resolve_session(
|
||||
self,
|
||||
session_id: Optional[str],
|
||||
reset: bool = False,
|
||||
) -> _ReplSession:
|
||||
"""Get or create a session, with LRU eviction at max_sessions."""
|
||||
with self._lock:
|
||||
if session_id and session_id in self._sessions and not reset:
|
||||
session = self._sessions[session_id]
|
||||
return session
|
||||
|
||||
if session_id and session_id in self._sessions and reset:
|
||||
# Reset existing session
|
||||
session = self._sessions[session_id]
|
||||
session.namespace = {"__builtins__": _make_restricted_builtins()}
|
||||
session.execution_count = 0
|
||||
return session
|
||||
|
||||
# Create new session
|
||||
sid = session_id or str(uuid.uuid4())
|
||||
|
||||
# LRU eviction if at capacity
|
||||
if len(self._sessions) >= self._max_sessions:
|
||||
oldest_id = min(
|
||||
self._sessions,
|
||||
key=lambda k: self._sessions[k].last_used,
|
||||
)
|
||||
del self._sessions[oldest_id]
|
||||
|
||||
session = _ReplSession(
|
||||
session_id=sid,
|
||||
namespace={"__builtins__": _make_restricted_builtins()},
|
||||
)
|
||||
self._sessions[sid] = session
|
||||
return session
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Execution
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _exec_with_timeout(
|
||||
self,
|
||||
code: str,
|
||||
session: _ReplSession,
|
||||
) -> tuple[str, bool]:
|
||||
"""Execute code in a daemon thread with timeout.
|
||||
|
||||
Returns (output, success).
|
||||
"""
|
||||
result_holder: Dict[str, Any] = {"output": "", "success": True}
|
||||
|
||||
def _run() -> None:
|
||||
stdout_buf = io.StringIO()
|
||||
stderr_buf = io.StringIO()
|
||||
try:
|
||||
with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf):
|
||||
# Try eval first for expression display (REPL-like behavior)
|
||||
try:
|
||||
compiled = compile(code, "<repl>", "eval")
|
||||
val = eval(compiled, session.namespace) # noqa: S307
|
||||
if val is not None:
|
||||
print(repr(val)) # noqa: T201
|
||||
except SyntaxError:
|
||||
# Not an expression — execute as statements
|
||||
compiled = compile(code, "<repl>", "exec")
|
||||
exec(compiled, session.namespace) # noqa: S102
|
||||
except Exception as exc:
|
||||
result_holder["output"] = f"{type(exc).__name__}: {exc}"
|
||||
result_holder["success"] = False
|
||||
return
|
||||
|
||||
output = stdout_buf.getvalue()
|
||||
err = stderr_buf.getvalue()
|
||||
if err:
|
||||
output += ("\n" if output else "") + err
|
||||
result_holder["output"] = output
|
||||
|
||||
thread = threading.Thread(target=_run, daemon=True)
|
||||
thread.start()
|
||||
thread.join(timeout=self._timeout)
|
||||
|
||||
if thread.is_alive():
|
||||
return (
|
||||
f"Execution timed out after {self._timeout} seconds.",
|
||||
False,
|
||||
)
|
||||
|
||||
return result_holder["output"], result_holder["success"]
|
||||
|
||||
|
||||
__all__ = ["ReplTool"]
|
||||
@@ -1,152 +0,0 @@
|
||||
"""Tests for the OpenClawAgent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.agents._stubs import AgentResult
|
||||
from openjarvis.agents.openclaw import OpenClawAgent
|
||||
from openjarvis.agents.openclaw_protocol import MessageType, ProtocolMessage
|
||||
from openjarvis.agents.openclaw_transport import OpenClawTransport
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_openclaw():
|
||||
"""Re-register openclaw agent after registry clear."""
|
||||
if not AgentRegistry.contains("openclaw"):
|
||||
AgentRegistry.register_value("openclaw", OpenClawAgent)
|
||||
|
||||
|
||||
class MockTransport(OpenClawTransport):
|
||||
"""Mock transport for testing."""
|
||||
|
||||
def __init__(self, responses=None, healthy=True):
|
||||
self._responses = list(responses or [])
|
||||
self._idx = 0
|
||||
self._healthy = healthy
|
||||
|
||||
def send(self, msg):
|
||||
if self._idx < len(self._responses):
|
||||
resp = self._responses[self._idx]
|
||||
self._idx += 1
|
||||
return resp
|
||||
return ProtocolMessage(type=MessageType.RESPONSE, content="default")
|
||||
|
||||
def health(self):
|
||||
return self._healthy
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
class TestOpenClawAgent:
|
||||
def test_agent_id(self):
|
||||
transport = MockTransport()
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
assert agent.agent_id == "openclaw"
|
||||
|
||||
def test_registration(self):
|
||||
assert AgentRegistry.contains("openclaw")
|
||||
assert AgentRegistry.get("openclaw") is OpenClawAgent
|
||||
|
||||
def test_run_with_mock_transport(self):
|
||||
responses = [
|
||||
ProtocolMessage(type=MessageType.RESPONSE, content="Hello from OpenClaw"),
|
||||
]
|
||||
transport = MockTransport(responses)
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
result = agent.run("Hello")
|
||||
assert isinstance(result, AgentResult)
|
||||
assert result.content == "Hello from OpenClaw"
|
||||
assert result.turns == 1
|
||||
|
||||
def test_returns_agent_result(self):
|
||||
transport = MockTransport([
|
||||
ProtocolMessage(type=MessageType.RESPONSE, content="test"),
|
||||
])
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
result = agent.run("test")
|
||||
assert isinstance(result, AgentResult)
|
||||
|
||||
def test_handles_tool_calls(self):
|
||||
responses = [
|
||||
ProtocolMessage(
|
||||
type=MessageType.TOOL_CALL,
|
||||
tool_name="calculator",
|
||||
tool_args={"expression": "2+2"},
|
||||
),
|
||||
ProtocolMessage(type=MessageType.RESPONSE, content="The answer is 4"),
|
||||
]
|
||||
transport = MockTransport(responses)
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
result = agent.run("What is 2+2?")
|
||||
assert result.content == "The answer is 4"
|
||||
assert result.turns == 2
|
||||
assert len(result.tool_results) == 1
|
||||
|
||||
def test_transport_unhealthy_raises(self):
|
||||
transport = MockTransport(healthy=False)
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
with pytest.raises(RuntimeError, match="not healthy"):
|
||||
agent.run("Hello")
|
||||
|
||||
def test_mode_http(self):
|
||||
agent = OpenClawAgent(mode="http")
|
||||
from openjarvis.agents.openclaw_transport import HttpTransport
|
||||
|
||||
assert isinstance(agent._transport, HttpTransport)
|
||||
|
||||
def test_mode_subprocess(self):
|
||||
agent = OpenClawAgent(mode="subprocess")
|
||||
from openjarvis.agents.openclaw_transport import SubprocessTransport
|
||||
|
||||
assert isinstance(agent._transport, SubprocessTransport)
|
||||
|
||||
def test_mode_invalid(self):
|
||||
with pytest.raises(ValueError, match="Unknown OpenClaw mode"):
|
||||
OpenClawAgent(mode="invalid")
|
||||
|
||||
def test_event_bus_events(self):
|
||||
responses = [
|
||||
ProtocolMessage(type=MessageType.RESPONSE, content="test"),
|
||||
]
|
||||
transport = MockTransport(responses)
|
||||
bus = EventBus(record_history=True)
|
||||
agent = OpenClawAgent(transport=transport, bus=bus)
|
||||
agent.run("Hello")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.AGENT_TURN_START in event_types
|
||||
assert EventType.AGENT_TURN_END in event_types
|
||||
|
||||
def test_custom_transport_injection(self):
|
||||
mock = MockTransport([
|
||||
ProtocolMessage(type=MessageType.RESPONSE, content="injected"),
|
||||
])
|
||||
agent = OpenClawAgent(transport=mock)
|
||||
result = agent.run("test")
|
||||
assert result.content == "injected"
|
||||
|
||||
def test_run_with_context(self):
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.core.types import Conversation, Message, Role
|
||||
|
||||
transport = MockTransport([
|
||||
ProtocolMessage(type=MessageType.RESPONSE, content="contextualized"),
|
||||
])
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
conv = Conversation()
|
||||
conv.add(Message(role=Role.SYSTEM, content="Be helpful"))
|
||||
ctx = AgentContext(conversation=conv)
|
||||
result = agent.run("Hello", context=ctx)
|
||||
assert result.content == "contextualized"
|
||||
|
||||
def test_error_response(self):
|
||||
transport = MockTransport([
|
||||
ProtocolMessage(type=MessageType.ERROR, error="Something failed"),
|
||||
])
|
||||
agent = OpenClawAgent(transport=transport)
|
||||
result = agent.run("Hello")
|
||||
assert result.content == "Something failed"
|
||||
@@ -1,79 +0,0 @@
|
||||
"""Tests for the OpenClaw plugin skeleton."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from openjarvis.agents.openclaw_plugin import (
|
||||
MemorySearchManager,
|
||||
ProviderPlugin,
|
||||
register,
|
||||
)
|
||||
|
||||
|
||||
class TestRegister:
|
||||
def test_register_returns_dict(self):
|
||||
result = register()
|
||||
assert isinstance(result, dict)
|
||||
assert result["name"] == "openjarvis"
|
||||
assert "provider_class" in result
|
||||
assert "memory_class" in result
|
||||
|
||||
|
||||
class TestProviderPlugin:
|
||||
def test_has_generate(self):
|
||||
assert hasattr(ProviderPlugin, "generate")
|
||||
|
||||
def test_has_list_models(self):
|
||||
assert hasattr(ProviderPlugin, "list_models")
|
||||
|
||||
def test_generate_with_engine(self):
|
||||
engine = MagicMock()
|
||||
engine.generate.return_value = {"content": "test"}
|
||||
plugin = ProviderPlugin(engine=engine, model="test-model")
|
||||
result = plugin.generate("Hello")
|
||||
assert result["content"] == "test"
|
||||
|
||||
def test_list_models_with_engine(self):
|
||||
engine = MagicMock()
|
||||
engine.list_models.return_value = ["model-1", "model-2"]
|
||||
plugin = ProviderPlugin(engine=engine)
|
||||
assert plugin.list_models() == ["model-1", "model-2"]
|
||||
|
||||
|
||||
class TestMemorySearchManager:
|
||||
def test_has_search(self):
|
||||
assert hasattr(MemorySearchManager, "search")
|
||||
|
||||
def test_has_sync(self):
|
||||
assert hasattr(MemorySearchManager, "sync")
|
||||
|
||||
def test_has_status(self):
|
||||
assert hasattr(MemorySearchManager, "status")
|
||||
|
||||
def test_search_with_backend(self):
|
||||
backend = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.content = "found"
|
||||
mock_result.score = 0.9
|
||||
mock_result.source = "test.txt"
|
||||
backend.retrieve.return_value = [mock_result]
|
||||
|
||||
mgr = MemorySearchManager(backend=backend)
|
||||
results = mgr.search("test query")
|
||||
assert len(results) == 1
|
||||
assert results[0]["content"] == "found"
|
||||
|
||||
def test_search_no_backend(self):
|
||||
mgr = MemorySearchManager()
|
||||
results = mgr.search("test")
|
||||
assert results == []
|
||||
|
||||
def test_status_no_backend(self):
|
||||
mgr = MemorySearchManager()
|
||||
assert mgr.status()["available"] is False
|
||||
|
||||
def test_sync(self):
|
||||
mgr = MemorySearchManager()
|
||||
result = mgr.sync()
|
||||
assert result["status"] == "ok"
|
||||
@@ -1,99 +0,0 @@
|
||||
"""Tests for the OpenClaw wire protocol."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.agents.openclaw_protocol import (
|
||||
MessageType,
|
||||
ProtocolMessage,
|
||||
deserialize,
|
||||
serialize,
|
||||
)
|
||||
|
||||
|
||||
class TestProtocolMessage:
|
||||
def test_query_message(self):
|
||||
msg = ProtocolMessage(type=MessageType.QUERY, content="Hello")
|
||||
assert msg.type == MessageType.QUERY
|
||||
assert msg.content == "Hello"
|
||||
|
||||
def test_response_message(self):
|
||||
msg = ProtocolMessage(type=MessageType.RESPONSE, content="World")
|
||||
assert msg.type == MessageType.RESPONSE
|
||||
|
||||
def test_tool_call_message(self):
|
||||
msg = ProtocolMessage(
|
||||
type=MessageType.TOOL_CALL,
|
||||
tool_name="calculator",
|
||||
tool_args={"expression": "2+2"},
|
||||
)
|
||||
assert msg.tool_name == "calculator"
|
||||
assert msg.tool_args == {"expression": "2+2"}
|
||||
|
||||
def test_tool_result_message(self):
|
||||
msg = ProtocolMessage(
|
||||
type=MessageType.TOOL_RESULT,
|
||||
tool_name="calculator",
|
||||
tool_result="4",
|
||||
)
|
||||
assert msg.tool_result == "4"
|
||||
|
||||
def test_error_message(self):
|
||||
msg = ProtocolMessage(type=MessageType.ERROR, error="Something went wrong")
|
||||
assert msg.error == "Something went wrong"
|
||||
|
||||
def test_health_message(self):
|
||||
msg = ProtocolMessage(type=MessageType.HEALTH)
|
||||
assert msg.type == MessageType.HEALTH
|
||||
|
||||
|
||||
class TestSerializeDeserialize:
|
||||
def test_roundtrip(self):
|
||||
original = ProtocolMessage(
|
||||
type=MessageType.QUERY,
|
||||
content="Hello world",
|
||||
metadata={"key": "value"},
|
||||
)
|
||||
line = serialize(original)
|
||||
restored = deserialize(line)
|
||||
assert restored.type == MessageType.QUERY
|
||||
assert restored.content == "Hello world"
|
||||
assert restored.metadata == {"key": "value"}
|
||||
|
||||
def test_serialize_with_metadata(self):
|
||||
msg = ProtocolMessage(
|
||||
type=MessageType.RESPONSE,
|
||||
content="result",
|
||||
metadata={"model": "test"},
|
||||
)
|
||||
line = serialize(msg)
|
||||
restored = deserialize(line)
|
||||
assert restored.metadata["model"] == "test"
|
||||
|
||||
def test_deserialize_valid_json(self):
|
||||
line = '{"type": "query", "id": "123", "content": "test"}'
|
||||
msg = deserialize(line)
|
||||
assert msg.type == MessageType.QUERY
|
||||
assert msg.id == "123"
|
||||
assert msg.content == "test"
|
||||
|
||||
def test_deserialize_invalid_json(self):
|
||||
with pytest.raises(ValueError, match="Invalid JSON"):
|
||||
deserialize("not json")
|
||||
|
||||
def test_deserialize_unknown_type(self):
|
||||
with pytest.raises(ValueError, match="Unknown message type"):
|
||||
deserialize('{"type": "unknown_type", "content": "test"}')
|
||||
|
||||
def test_tool_call_roundtrip(self):
|
||||
msg = ProtocolMessage(
|
||||
type=MessageType.TOOL_CALL,
|
||||
tool_name="calculator",
|
||||
tool_args={"expression": "3*7"},
|
||||
)
|
||||
line = serialize(msg)
|
||||
restored = deserialize(line)
|
||||
assert restored.type == MessageType.TOOL_CALL
|
||||
assert restored.tool_name == "calculator"
|
||||
assert restored.tool_args == {"expression": "3*7"}
|
||||
@@ -1,106 +0,0 @@
|
||||
"""Tests for OpenClaw transport implementations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.agents.openclaw_protocol import MessageType, ProtocolMessage
|
||||
from openjarvis.agents.openclaw_transport import (
|
||||
HttpTransport,
|
||||
OpenClawTransport,
|
||||
SubprocessTransport,
|
||||
)
|
||||
|
||||
|
||||
class TestOpenClawTransportABC:
|
||||
def test_abc_cannot_instantiate(self):
|
||||
with pytest.raises(TypeError):
|
||||
OpenClawTransport()
|
||||
|
||||
def test_concrete_required_methods(self):
|
||||
class DummyTransport(OpenClawTransport):
|
||||
def send(self, msg):
|
||||
return msg
|
||||
|
||||
def health(self):
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
t = DummyTransport()
|
||||
assert t.health() is True
|
||||
|
||||
|
||||
class TestHttpTransport:
|
||||
def test_send(self):
|
||||
transport = HttpTransport(host="http://localhost:18789")
|
||||
msg = ProtocolMessage(type=MessageType.QUERY, content="Hello")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"type": "response",
|
||||
"id": "resp-1",
|
||||
"content": "World",
|
||||
}
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = transport.send(msg)
|
||||
assert result.type == MessageType.RESPONSE
|
||||
assert result.content == "World"
|
||||
mock_post.assert_called_once()
|
||||
|
||||
def test_health_ok(self):
|
||||
transport = HttpTransport()
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.status_code = 200
|
||||
with patch("httpx.get", return_value=mock_resp):
|
||||
assert transport.health() is True
|
||||
|
||||
def test_health_fail(self):
|
||||
transport = HttpTransport()
|
||||
with patch("httpx.get", side_effect=Exception("Connection refused")):
|
||||
assert transport.health() is False
|
||||
|
||||
def test_close(self):
|
||||
transport = HttpTransport()
|
||||
transport.close() # should not raise
|
||||
|
||||
|
||||
class TestSubprocessTransport:
|
||||
def test_send(self):
|
||||
transport = SubprocessTransport(node_path="node", script_path="test.js")
|
||||
|
||||
mock_proc = MagicMock()
|
||||
mock_proc.poll.return_value = None
|
||||
mock_proc.stdin = MagicMock()
|
||||
mock_proc.stdout = MagicMock()
|
||||
mock_proc.stdout.readline.return_value = (
|
||||
'{"type": "response", "id": "r1", "content": "reply"}\n'
|
||||
)
|
||||
|
||||
transport._process = mock_proc
|
||||
|
||||
msg = ProtocolMessage(type=MessageType.QUERY, content="test")
|
||||
result = transport.send(msg)
|
||||
assert result.type == MessageType.RESPONSE
|
||||
assert result.content == "reply"
|
||||
|
||||
def test_health_fail_no_node(self):
|
||||
transport = SubprocessTransport(node_path="/nonexistent/node")
|
||||
assert transport.health() is False
|
||||
|
||||
def test_close(self):
|
||||
transport = SubprocessTransport()
|
||||
mock_proc = MagicMock()
|
||||
transport._process = mock_proc
|
||||
transport.close()
|
||||
mock_proc.terminate.assert_called_once()
|
||||
assert transport._process is None
|
||||
|
||||
def test_close_no_process(self):
|
||||
transport = SubprocessTransport()
|
||||
transport.close() # should not raise
|
||||
@@ -0,0 +1,384 @@
|
||||
"""Tests for the RLM agent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.agents.rlm import RLMAgent
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
from openjarvis.core.types import ToolResult
|
||||
from openjarvis.tools._stubs import BaseTool, ToolSpec
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _CalcStub(BaseTool):
|
||||
tool_id = "calculator"
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="calculator",
|
||||
description="Math calculator.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"expression": {"type": "string"}},
|
||||
"required": ["expression"],
|
||||
},
|
||||
)
|
||||
|
||||
def execute(self, **params) -> ToolResult:
|
||||
expr = params.get("expression", "0")
|
||||
try:
|
||||
val = eval(expr) # noqa: S307
|
||||
except Exception as e:
|
||||
return ToolResult(tool_name="calculator", content=str(e), success=False)
|
||||
return ToolResult(tool_name="calculator", content=str(val), success=True)
|
||||
|
||||
|
||||
def _make_engine(content: str = "Final answer.") -> MagicMock:
|
||||
"""Engine that returns plain content (no code block)."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": content,
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
return engine
|
||||
|
||||
|
||||
def _make_engine_with_code(
|
||||
code: str,
|
||||
final_content: str = "Done.",
|
||||
) -> MagicMock:
|
||||
"""Engine that returns a python code block, then a final answer."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.side_effect = [
|
||||
{
|
||||
"content": f"```python\n{code}\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
{
|
||||
"content": final_content,
|
||||
"usage": {"prompt_tokens": 15, "completion_tokens": 5, "total_tokens": 20},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
]
|
||||
return engine
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRLMAgentRegistration:
|
||||
def test_registered(self):
|
||||
# Re-register after conftest clears all registries
|
||||
AgentRegistry.register_value("rlm", RLMAgent)
|
||||
assert AgentRegistry.contains("rlm")
|
||||
|
||||
def test_agent_id(self):
|
||||
engine = _make_engine()
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
assert agent.agent_id == "rlm"
|
||||
|
||||
|
||||
class TestRLMCodeExtraction:
|
||||
def test_extract_python_block(self):
|
||||
text = "Here is code:\n```python\nx = 1\n```\nDone."
|
||||
code = RLMAgent._extract_code(text)
|
||||
assert code == "x = 1"
|
||||
|
||||
def test_extract_bare_block(self):
|
||||
text = "Here is code:\n```\nx = 1\n```\nDone."
|
||||
code = RLMAgent._extract_code(text)
|
||||
assert code == "x = 1"
|
||||
|
||||
def test_no_block(self):
|
||||
text = "No code here, just text."
|
||||
code = RLMAgent._extract_code(text)
|
||||
assert code is None
|
||||
|
||||
def test_python_preferred_over_bare(self):
|
||||
text = "```python\nx = 1\n```\n\n```\ny = 2\n```"
|
||||
code = RLMAgent._extract_code(text)
|
||||
assert code == "x = 1"
|
||||
|
||||
|
||||
class TestRLMStripThink:
|
||||
def test_strip_think(self):
|
||||
text = "<think>thinking...</think>Answer here."
|
||||
result = RLMAgent._strip_think_tags(text)
|
||||
assert result == "Answer here."
|
||||
|
||||
def test_no_think_tags(self):
|
||||
text = "Just text."
|
||||
result = RLMAgent._strip_think_tags(text)
|
||||
assert result == "Just text."
|
||||
|
||||
|
||||
class TestRLMDirectAnswer:
|
||||
def test_no_code_block_returns_content(self):
|
||||
"""When model returns no code block, treat content as final answer."""
|
||||
engine = _make_engine("The answer is 42.")
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("What is the answer?")
|
||||
assert result.content == "The answer is 42."
|
||||
assert result.turns == 1
|
||||
assert result.tool_results == []
|
||||
|
||||
|
||||
class TestRLMFinalTermination:
|
||||
def test_final_terminates(self):
|
||||
"""FINAL() in code should terminate the agent."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "```python\nFINAL('hello world')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("Test")
|
||||
assert result.content == "hello world"
|
||||
assert len(result.tool_results) == 1
|
||||
assert result.tool_results[0].tool_name == "rlm_repl"
|
||||
|
||||
def test_final_var_terminates(self):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "```python\nresult = 42\nFINAL_VAR('result')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("Test")
|
||||
assert result.content == "42"
|
||||
|
||||
|
||||
class TestRLMContextInjection:
|
||||
def test_context_from_metadata(self):
|
||||
engine = _make_engine("The answer is 42.")
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
ctx = AgentContext(metadata={"context": "Some long document text."})
|
||||
result = agent.run("Summarize", context=ctx)
|
||||
assert result.content == "The answer is 42."
|
||||
|
||||
def test_context_from_memory_results(self):
|
||||
engine = _make_engine("Summary.")
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
ctx = AgentContext(memory_results=["chunk1", "chunk2"])
|
||||
result = agent.run("Summarize", context=ctx)
|
||||
assert result.content == "Summary."
|
||||
|
||||
|
||||
class TestRLMSubLMCalls:
|
||||
def test_sub_lm_called_from_repl(self):
|
||||
"""Verify that llm_query() inside REPL code calls engine.generate."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
# First call: root LM generates code that calls llm_query
|
||||
# Second call: sub-LM responds to llm_query
|
||||
# Third call: root LM gets REPL output, returns final (no code)
|
||||
engine.generate.side_effect = [
|
||||
{
|
||||
"content": "```python\nresult = llm_query('What is 2+2?')\nFINAL(result)\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
# Sub-LM response for llm_query
|
||||
{
|
||||
"content": "4",
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
]
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("Calculate")
|
||||
assert result.content == "4"
|
||||
# engine.generate should be called at least twice (root + sub)
|
||||
assert engine.generate.call_count >= 2
|
||||
|
||||
|
||||
class TestRLMMultiTurn:
|
||||
def test_multi_turn_loop(self):
|
||||
"""Agent should loop: generate code → execute → feed output → generate again."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.side_effect = [
|
||||
# Turn 1: code that sets a variable
|
||||
{
|
||||
"content": "```python\nx = 10\nprint(f'x = {x}')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
# Turn 2: code that uses the variable and terminates
|
||||
{
|
||||
"content": "```python\ny = x * 2\nFINAL(y)\n```",
|
||||
"usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
]
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("Calculate")
|
||||
assert result.content == "20"
|
||||
assert result.turns == 2
|
||||
assert len(result.tool_results) == 2
|
||||
|
||||
def test_max_turns_exceeded(self):
|
||||
"""Agent should stop after max_turns."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "```python\nprint('looping')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
agent = RLMAgent(engine, "test-model", max_turns=3)
|
||||
result = agent.run("Loop")
|
||||
assert result.turns == 3
|
||||
assert result.metadata.get("max_turns_exceeded") is True
|
||||
|
||||
def test_max_turns_with_partial_answer(self):
|
||||
"""When max turns exceeded but answer dict has value, use it."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "```python\nanswer['value'] = 'partial'\nprint('working')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
agent = RLMAgent(engine, "test-model", max_turns=2)
|
||||
result = agent.run("Work")
|
||||
assert result.content == "partial"
|
||||
assert result.metadata.get("max_turns_exceeded") is True
|
||||
|
||||
|
||||
class TestRLMEventBus:
|
||||
def test_agent_events(self):
|
||||
bus = EventBus(record_history=True)
|
||||
engine = _make_engine("Direct answer.")
|
||||
agent = RLMAgent(engine, "test-model", bus=bus)
|
||||
agent.run("Hello")
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.AGENT_TURN_START in event_types
|
||||
assert EventType.AGENT_TURN_END in event_types
|
||||
|
||||
def test_agent_events_with_code(self):
|
||||
bus = EventBus(record_history=True)
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "```python\nFINAL('done')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
agent = RLMAgent(engine, "test-model", bus=bus)
|
||||
agent.run("Test")
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.AGENT_TURN_START in event_types
|
||||
assert EventType.AGENT_TURN_END in event_types
|
||||
|
||||
|
||||
class TestRLMSubLMWithTools:
|
||||
def test_sub_lm_tool_resolution(self):
|
||||
"""When sub-LM returns tool_calls, agent resolves them."""
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.side_effect = [
|
||||
# Root LM: code that calls llm_query
|
||||
{
|
||||
"content": "```python\nresult = llm_query('Calculate 2+2')\nFINAL(result)\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
# Sub-LM: returns tool call
|
||||
{
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{"id": "sub_0", "name": "calculator", "arguments": '{"expression":"2+2"}'},
|
||||
],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8},
|
||||
"model": "test-model",
|
||||
"finish_reason": "tool_calls",
|
||||
},
|
||||
# Sub-LM follow-up after tool result
|
||||
{
|
||||
"content": "The answer is 4.",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
]
|
||||
agent = RLMAgent(engine, "test-model", tools=[_CalcStub()])
|
||||
result = agent.run("Calculate")
|
||||
assert result.content == "The answer is 4."
|
||||
|
||||
|
||||
class TestRLMBlockedCode:
|
||||
def test_blocked_code_returns_error(self):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.side_effect = [
|
||||
# Code with blocked pattern
|
||||
{
|
||||
"content": "```python\nos.system('ls')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
# After error feedback, model gives direct answer
|
||||
{
|
||||
"content": "I apologize, let me answer directly.",
|
||||
"usage": {"prompt_tokens": 15, "completion_tokens": 5, "total_tokens": 20},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
},
|
||||
]
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("Test")
|
||||
assert result.content == "I apologize, let me answer directly."
|
||||
# The blocked code should produce a failed tool result
|
||||
assert len(result.tool_results) == 1
|
||||
assert result.tool_results[0].success is False
|
||||
assert "Blocked" in result.tool_results[0].content
|
||||
|
||||
|
||||
class TestRLMReplResults:
|
||||
def test_repl_results_in_tool_results(self):
|
||||
engine = MagicMock()
|
||||
engine.engine_id = "mock"
|
||||
engine.generate.return_value = {
|
||||
"content": "```python\nprint('hello')\nFINAL('done')\n```",
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
|
||||
"model": "test-model",
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
agent = RLMAgent(engine, "test-model")
|
||||
result = agent.run("Test")
|
||||
assert len(result.tool_results) == 1
|
||||
assert result.tool_results[0].tool_name == "rlm_repl"
|
||||
assert "hello" in result.tool_results[0].content
|
||||
@@ -0,0 +1,206 @@
|
||||
"""Tests for the RLM REPL environment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.agents.rlm_repl import RLMRepl
|
||||
|
||||
|
||||
class TestRLMReplBasics:
|
||||
"""Basic execution and variable persistence."""
|
||||
|
||||
def test_variable_persistence(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("x = 42")
|
||||
repl.execute("y = x + 1")
|
||||
assert repl.get_variable("x") == 42
|
||||
assert repl.get_variable("y") == 43
|
||||
|
||||
def test_code_execution_stdout(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("print('hello')")
|
||||
assert "hello" in output
|
||||
|
||||
def test_function_definition(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("def double(n): return n * 2")
|
||||
repl.execute("result = double(5)")
|
||||
assert repl.get_variable("result") == 10
|
||||
|
||||
def test_multiline_code(self):
|
||||
repl = RLMRepl()
|
||||
code = "for i in range(3):\n print(i)"
|
||||
output = repl.execute(code)
|
||||
assert "0" in output
|
||||
assert "1" in output
|
||||
assert "2" in output
|
||||
|
||||
def test_set_and_get_variable(self):
|
||||
repl = RLMRepl()
|
||||
repl.set_variable("x", 99)
|
||||
assert repl.get_variable("x") == 99
|
||||
|
||||
def test_get_missing_variable(self):
|
||||
repl = RLMRepl()
|
||||
assert repl.get_variable("missing") is None
|
||||
|
||||
|
||||
class TestRLMReplSecurity:
|
||||
"""Security: blocked patterns and safe modules."""
|
||||
|
||||
def test_blocked_os_system(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("os.system('ls')")
|
||||
assert "Blocked" in output
|
||||
|
||||
def test_blocked_subprocess(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("import subprocess")
|
||||
assert "Blocked" in output
|
||||
|
||||
def test_blocked_open(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("f = open('/etc/passwd')")
|
||||
assert "Blocked" in output
|
||||
|
||||
def test_blocked_dunder_import(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("__import__('os')")
|
||||
assert "Blocked" in output
|
||||
|
||||
def test_blocked_socket(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("import socket")
|
||||
assert "Blocked" in output
|
||||
|
||||
def test_safe_modules_available(self):
|
||||
repl = RLMRepl()
|
||||
# json, re, math should be pre-injected
|
||||
assert repl.get_variable("json") is not None
|
||||
assert repl.get_variable("re") is not None
|
||||
assert repl.get_variable("math") is not None
|
||||
|
||||
def test_safe_module_usage(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("result = json.dumps({'a': 1})")
|
||||
assert repl.get_variable("result") == '{"a": 1}'
|
||||
|
||||
def test_math_module_usage(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("result = math.sqrt(16)")
|
||||
assert repl.get_variable("result") == 4.0
|
||||
|
||||
def test_security_check_returns_none_for_safe_code(self):
|
||||
repl = RLMRepl()
|
||||
assert repl.security_check("x = 1 + 2") is None
|
||||
|
||||
def test_security_check_returns_error_for_blocked(self):
|
||||
repl = RLMRepl()
|
||||
result = repl.security_check("os.system('rm -rf /')")
|
||||
assert result is not None
|
||||
assert "Blocked" in result
|
||||
|
||||
|
||||
class TestRLMReplTermination:
|
||||
"""FINAL, FINAL_VAR, and answer dict termination."""
|
||||
|
||||
def test_final_terminates(self):
|
||||
repl = RLMRepl()
|
||||
assert not repl.is_terminated
|
||||
repl.execute("FINAL('done')")
|
||||
assert repl.is_terminated
|
||||
assert repl.final_answer == "done"
|
||||
|
||||
def test_final_var_terminates(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("result = 42")
|
||||
repl.execute("FINAL_VAR('result')")
|
||||
assert repl.is_terminated
|
||||
assert repl.final_answer == 42
|
||||
|
||||
def test_answer_dict_terminates(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("answer['value'] = 'hello'")
|
||||
repl.execute("answer['ready'] = True")
|
||||
assert repl.is_terminated
|
||||
assert repl.final_answer == "hello"
|
||||
|
||||
def test_answer_dict_not_ready(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("answer['value'] = 'hello'")
|
||||
assert not repl.is_terminated
|
||||
|
||||
def test_final_with_complex_value(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("FINAL([1, 2, 3])")
|
||||
assert repl.is_terminated
|
||||
assert repl.final_answer == [1, 2, 3]
|
||||
|
||||
|
||||
class TestRLMReplCallbacks:
|
||||
"""llm_query and llm_batch callbacks."""
|
||||
|
||||
def test_llm_query_callback(self):
|
||||
calls = []
|
||||
|
||||
def mock_query(prompt):
|
||||
calls.append(prompt)
|
||||
return f"answer: {prompt}"
|
||||
|
||||
repl = RLMRepl(llm_query_fn=mock_query)
|
||||
repl.execute("result = llm_query('What is 2+2?')")
|
||||
assert len(calls) == 1
|
||||
assert calls[0] == "What is 2+2?"
|
||||
assert repl.get_variable("result") == "answer: What is 2+2?"
|
||||
|
||||
def test_llm_batch_callback(self):
|
||||
def mock_batch(prompts):
|
||||
return [f"answer: {p}" for p in prompts]
|
||||
|
||||
repl = RLMRepl(llm_batch_fn=mock_batch)
|
||||
repl.execute("results = llm_batch(['q1', 'q2'])")
|
||||
results = repl.get_variable("results")
|
||||
assert len(results) == 2
|
||||
assert results[0] == "answer: q1"
|
||||
assert results[1] == "answer: q2"
|
||||
|
||||
def test_no_callback_raises(self):
|
||||
"""If llm_query not injected, calling it raises NameError."""
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("llm_query('test')")
|
||||
assert "NameError" in output
|
||||
|
||||
|
||||
class TestRLMReplOutput:
|
||||
"""Output truncation and error handling."""
|
||||
|
||||
def test_output_truncation(self):
|
||||
repl = RLMRepl(max_output_chars=50)
|
||||
repl.execute("print('x' * 200)")
|
||||
output = repl.execute("print('y' * 200)")
|
||||
assert "truncated" in output
|
||||
|
||||
def test_syntax_error(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("def foo(")
|
||||
assert "SyntaxError" in output
|
||||
|
||||
def test_runtime_error(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("1 / 0")
|
||||
assert "ZeroDivisionError" in output
|
||||
|
||||
def test_name_error(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("print(undefined_var)")
|
||||
assert "NameError" in output
|
||||
|
||||
def test_error_doesnt_corrupt_state(self):
|
||||
repl = RLMRepl()
|
||||
repl.execute("x = 10")
|
||||
repl.execute("1 / 0") # Error
|
||||
assert repl.get_variable("x") == 10
|
||||
|
||||
def test_no_output_returns_empty(self):
|
||||
repl = RLMRepl()
|
||||
output = repl.execute("x = 1")
|
||||
assert output == ""
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Tests for the BlueBubblesChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.bluebubbles import BlueBubblesChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_bluebubbles():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("bluebubbles"):
|
||||
ChannelRegistry.register_value("bluebubbles", BlueBubblesChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("bluebubbles")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
assert ch.channel_id == "bluebubbles"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = BlueBubblesChannel()
|
||||
assert ch._url == ""
|
||||
assert ch._password == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_param(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
assert ch._url == "http://localhost:1234"
|
||||
assert ch._password == "test-pass"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"BLUEBUBBLES_URL": "http://env:1234", "BLUEBUBBLES_PASSWORD": "env-pass"}):
|
||||
ch = BlueBubblesChannel()
|
||||
assert ch._url == "http://env:1234"
|
||||
assert ch._password == "env-pass"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"BLUEBUBBLES_URL": "http://env:1234", "BLUEBUBBLES_PASSWORD": "env-pass"}):
|
||||
ch = BlueBubblesChannel(url="http://explicit:1234", password="explicit-pass")
|
||||
assert ch._url == "http://explicit:1234"
|
||||
assert ch._password == "explicit-pass"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("iMessage;+;chat123", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
assert call_args[1]["params"] == {"password": "test-pass"}
|
||||
payload = call_args[1]["json"]
|
||||
assert "chatGuid" in payload
|
||||
assert payload["chatGuid"] == "iMessage;+;chat123"
|
||||
assert "message" in payload
|
||||
assert payload["message"] == "Hello!"
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("iMessage;+;chat123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("iMessage;+;chat123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = BlueBubblesChannel()
|
||||
result = ch.send("iMessage;+;chat123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("iMessage;+;chat123", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
assert ch.list_channels() == ["bluebubbles"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_url_connect_error(self):
|
||||
ch = BlueBubblesChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = BlueBubblesChannel(url="http://localhost:1234", password="test-pass")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,195 @@
|
||||
"""Tests for channel configuration — nested sub-configs and TOML loading."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from openjarvis.core.config import (
|
||||
ChannelConfig,
|
||||
DiscordChannelConfig,
|
||||
EmailChannelConfig,
|
||||
SlackChannelConfig,
|
||||
TelegramChannelConfig,
|
||||
WebhookChannelConfig,
|
||||
load_config,
|
||||
)
|
||||
|
||||
|
||||
class TestChannelConfig:
|
||||
def test_defaults(self):
|
||||
cfg = ChannelConfig()
|
||||
assert cfg.enabled is False
|
||||
assert cfg.default_channel == ""
|
||||
assert cfg.default_agent == "simple"
|
||||
|
||||
def test_nested_defaults(self):
|
||||
cfg = ChannelConfig()
|
||||
assert isinstance(cfg.telegram, TelegramChannelConfig)
|
||||
assert isinstance(cfg.discord, DiscordChannelConfig)
|
||||
assert isinstance(cfg.slack, SlackChannelConfig)
|
||||
assert isinstance(cfg.webhook, WebhookChannelConfig)
|
||||
assert isinstance(cfg.email, EmailChannelConfig)
|
||||
|
||||
def test_telegram_defaults(self):
|
||||
cfg = TelegramChannelConfig()
|
||||
assert cfg.bot_token == ""
|
||||
assert cfg.allowed_chat_ids == ""
|
||||
assert cfg.parse_mode == "Markdown"
|
||||
|
||||
def test_discord_defaults(self):
|
||||
cfg = DiscordChannelConfig()
|
||||
assert cfg.bot_token == ""
|
||||
|
||||
def test_slack_defaults(self):
|
||||
cfg = SlackChannelConfig()
|
||||
assert cfg.bot_token == ""
|
||||
assert cfg.app_token == ""
|
||||
|
||||
def test_webhook_defaults(self):
|
||||
cfg = WebhookChannelConfig()
|
||||
assert cfg.url == ""
|
||||
assert cfg.secret == ""
|
||||
assert cfg.method == "POST"
|
||||
|
||||
def test_email_defaults(self):
|
||||
cfg = EmailChannelConfig()
|
||||
assert cfg.smtp_host == ""
|
||||
assert cfg.smtp_port == 587
|
||||
assert cfg.imap_host == ""
|
||||
assert cfg.imap_port == 993
|
||||
assert cfg.username == ""
|
||||
assert cfg.password == ""
|
||||
assert cfg.use_tls is True
|
||||
|
||||
|
||||
class TestTomlLoading:
|
||||
def _write_toml(self, content: str) -> Path:
|
||||
f = tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".toml", delete=False,
|
||||
)
|
||||
f.write(content)
|
||||
f.flush()
|
||||
f.close()
|
||||
return Path(f.name)
|
||||
|
||||
def test_load_channel_top_level(self):
|
||||
path = self._write_toml("""
|
||||
[channel]
|
||||
enabled = true
|
||||
default_channel = "telegram"
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.enabled is True
|
||||
assert cfg.channel.default_channel == "telegram"
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_load_channel_telegram(self):
|
||||
path = self._write_toml("""
|
||||
[channel]
|
||||
enabled = true
|
||||
default_channel = "telegram"
|
||||
|
||||
[channel.telegram]
|
||||
bot_token = "123:ABC"
|
||||
parse_mode = "HTML"
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.telegram.bot_token == "123:ABC"
|
||||
assert cfg.channel.telegram.parse_mode == "HTML"
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_load_channel_discord(self):
|
||||
path = self._write_toml("""
|
||||
[channel.discord]
|
||||
bot_token = "discord-token"
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.discord.bot_token == "discord-token"
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_load_channel_slack(self):
|
||||
path = self._write_toml("""
|
||||
[channel.slack]
|
||||
bot_token = "xoxb-slack"
|
||||
app_token = "xapp-slack"
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.slack.bot_token == "xoxb-slack"
|
||||
assert cfg.channel.slack.app_token == "xapp-slack"
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_load_channel_webhook(self):
|
||||
path = self._write_toml("""
|
||||
[channel.webhook]
|
||||
url = "https://example.com/hook"
|
||||
secret = "s3cr3t"
|
||||
method = "PUT"
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.webhook.url == "https://example.com/hook"
|
||||
assert cfg.channel.webhook.secret == "s3cr3t"
|
||||
assert cfg.channel.webhook.method == "PUT"
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_load_channel_email(self):
|
||||
path = self._write_toml("""
|
||||
[channel.email]
|
||||
smtp_host = "smtp.example.com"
|
||||
smtp_port = 465
|
||||
username = "user@example.com"
|
||||
password = "pass"
|
||||
use_tls = false
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.email.smtp_host == "smtp.example.com"
|
||||
assert cfg.channel.email.smtp_port == 465
|
||||
assert cfg.channel.email.username == "user@example.com"
|
||||
assert cfg.channel.email.password == "pass"
|
||||
assert cfg.channel.email.use_tls is False
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_backward_compat_no_default_channel(self):
|
||||
"""Old config without default_channel still works."""
|
||||
path = self._write_toml("""
|
||||
[channel]
|
||||
enabled = false
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.enabled is False
|
||||
assert cfg.channel.default_channel == ""
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
def test_multiple_channel_configs(self):
|
||||
path = self._write_toml("""
|
||||
[channel]
|
||||
enabled = true
|
||||
default_channel = "slack"
|
||||
|
||||
[channel.telegram]
|
||||
bot_token = "tg-token"
|
||||
|
||||
[channel.slack]
|
||||
bot_token = "slack-token"
|
||||
""")
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.channel.default_channel == "slack"
|
||||
assert cfg.channel.telegram.bot_token == "tg-token"
|
||||
assert cfg.channel.slack.bot_token == "slack-token"
|
||||
finally:
|
||||
path.unlink()
|
||||
@@ -74,17 +74,6 @@ class TestChannelRegistry:
|
||||
assert isinstance(instance, CreatableChannel)
|
||||
assert instance.channel_id == "creatable"
|
||||
|
||||
def test_openclaw_bridge_registered(self) -> None:
|
||||
"""After importing openclaw_bridge, 'openclaw' should be in the registry."""
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge
|
||||
|
||||
# Re-register since conftest clears registries
|
||||
if not ChannelRegistry.contains("openclaw"):
|
||||
ChannelRegistry.register_value("openclaw", OpenClawChannelBridge)
|
||||
|
||||
assert ChannelRegistry.contains("openclaw")
|
||||
assert ChannelRegistry.get("openclaw") is OpenClawChannelBridge
|
||||
|
||||
def test_duplicate_registration_raises(self) -> None:
|
||||
"""Registering the same key twice should raise ValueError."""
|
||||
ChannelRegistry.register_value("dup", object)
|
||||
|
||||
@@ -12,21 +12,15 @@ class TestChannelConfigDefaults:
|
||||
def test_channel_config_defaults(self) -> None:
|
||||
cfg = ChannelConfig()
|
||||
assert cfg.enabled is False
|
||||
assert cfg.gateway_url == "ws://127.0.0.1:18789/ws"
|
||||
assert cfg.default_agent == "simple"
|
||||
assert cfg.reconnect_interval == 5.0
|
||||
|
||||
def test_channel_config_custom(self) -> None:
|
||||
cfg = ChannelConfig(
|
||||
enabled=True,
|
||||
gateway_url="ws://custom:9999/ws",
|
||||
default_agent="orchestrator",
|
||||
reconnect_interval=10.0,
|
||||
)
|
||||
assert cfg.enabled is True
|
||||
assert cfg.gateway_url == "ws://custom:9999/ws"
|
||||
assert cfg.default_agent == "orchestrator"
|
||||
assert cfg.reconnect_interval == 10.0
|
||||
|
||||
|
||||
class TestChannelConfigInJarvisConfig:
|
||||
@@ -38,7 +32,7 @@ class TestChannelConfigInJarvisConfig:
|
||||
def test_jarvis_config_channel_defaults(self) -> None:
|
||||
cfg = JarvisConfig()
|
||||
assert cfg.channel.enabled is False
|
||||
assert cfg.channel.gateway_url == "ws://127.0.0.1:18789/ws"
|
||||
assert cfg.channel.default_channel == ""
|
||||
|
||||
|
||||
class TestLoadConfigWithChannel:
|
||||
@@ -47,18 +41,14 @@ class TestLoadConfigWithChannel:
|
||||
toml_content = textwrap.dedent("""\
|
||||
[channel]
|
||||
enabled = true
|
||||
gateway_url = "ws://my-gateway:7777/ws"
|
||||
default_agent = "orchestrator"
|
||||
reconnect_interval = 15.0
|
||||
""")
|
||||
config_file = tmp_path / "config.toml"
|
||||
config_file.write_text(toml_content)
|
||||
|
||||
cfg = load_config(path=config_file)
|
||||
assert cfg.channel.enabled is True
|
||||
assert cfg.channel.gateway_url == "ws://my-gateway:7777/ws"
|
||||
assert cfg.channel.default_agent == "orchestrator"
|
||||
assert cfg.channel.reconnect_interval == 15.0
|
||||
|
||||
def test_load_config_without_channel_section(self, tmp_path: Path) -> None:
|
||||
"""When no [channel] section, defaults should apply."""
|
||||
@@ -71,7 +61,7 @@ class TestLoadConfigWithChannel:
|
||||
|
||||
cfg = load_config(path=config_file)
|
||||
assert cfg.channel.enabled is False
|
||||
assert cfg.channel.gateway_url == "ws://127.0.0.1:18789/ws"
|
||||
assert cfg.channel.default_channel == ""
|
||||
|
||||
def test_load_config_partial_channel_section(self, tmp_path: Path) -> None:
|
||||
"""Partial [channel] section overlays only specified fields."""
|
||||
@@ -85,6 +75,4 @@ class TestLoadConfigWithChannel:
|
||||
cfg = load_config(path=config_file)
|
||||
assert cfg.channel.enabled is True
|
||||
# Non-specified fields keep defaults
|
||||
assert cfg.channel.gateway_url == "ws://127.0.0.1:18789/ws"
|
||||
assert cfg.channel.default_agent == "simple"
|
||||
assert cfg.channel.reconnect_interval == 5.0
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Tests for the DiscordChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.discord_channel import DiscordChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_discord():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("discord"):
|
||||
ChannelRegistry.register_value("discord", DiscordChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("discord")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = DiscordChannel(bot_token="test-token")
|
||||
assert ch.channel_id == "discord"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = DiscordChannel()
|
||||
assert ch._token == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_token(self):
|
||||
ch = DiscordChannel(bot_token="my-token")
|
||||
assert ch._token == "my-token"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"DISCORD_BOT_TOKEN": "env-token"}):
|
||||
ch = DiscordChannel()
|
||||
assert ch._token == "env-token"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"DISCORD_BOT_TOKEN": "env-token"}):
|
||||
ch = DiscordChannel(bot_token="explicit-token")
|
||||
assert ch._token == "explicit-token"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("987654321", "Hello Discord!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
url = call_args[0][0]
|
||||
assert "discord.com/api/v10/channels/987654321/messages" in url
|
||||
headers = call_args[1]["headers"]
|
||||
assert headers["Authorization"] == "Bot my-bot-token"
|
||||
payload = call_args[1]["json"]
|
||||
assert payload["content"] == "Hello Discord!"
|
||||
|
||||
def test_send_with_conversation_id(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
ch.send("987654321", "Reply!", conversation_id="msg-123")
|
||||
payload = mock_post.call_args[1]["json"]
|
||||
assert payload["message_reference"] == {"message_id": "msg-123"}
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 403
|
||||
mock_response.text = "Missing Permissions"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("987654321", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("987654321", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = DiscordChannel()
|
||||
result = ch.send("987654321", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = DiscordChannel(bot_token="my-bot-token", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("987654321", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
assert ch.list_channels() == ["discord"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_token_connect_error(self):
|
||||
ch = DiscordChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = DiscordChannel(bot_token="my-bot-token")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Tests for the EmailChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.email_channel import EmailChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_email():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("email"):
|
||||
ChannelRegistry.register_value("email", EmailChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("email")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com", username="user@example.com",
|
||||
)
|
||||
assert ch.channel_id == "email"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = EmailChannel()
|
||||
assert ch._smtp_host == ""
|
||||
assert ch._smtp_port == 587
|
||||
assert ch._imap_host == ""
|
||||
assert ch._imap_port == 993
|
||||
assert ch._username == ""
|
||||
assert ch._password == ""
|
||||
assert ch._use_tls is True
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_params(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
smtp_port=465,
|
||||
imap_host="imap.example.com",
|
||||
imap_port=143,
|
||||
username="user@example.com",
|
||||
password="pass123",
|
||||
use_tls=False,
|
||||
)
|
||||
assert ch._smtp_host == "smtp.example.com"
|
||||
assert ch._smtp_port == 465
|
||||
assert ch._imap_host == "imap.example.com"
|
||||
assert ch._imap_port == 143
|
||||
assert ch._username == "user@example.com"
|
||||
assert ch._password == "pass123"
|
||||
assert ch._use_tls is False
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {
|
||||
"EMAIL_USERNAME": "env@example.com",
|
||||
"EMAIL_PASSWORD": "env-pass",
|
||||
}):
|
||||
ch = EmailChannel()
|
||||
assert ch._username == "env@example.com"
|
||||
assert ch._password == "env-pass"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"EMAIL_USERNAME": "env@example.com"}):
|
||||
ch = EmailChannel(username="explicit@example.com")
|
||||
assert ch._username == "explicit@example.com"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success_tls(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
username="user@example.com",
|
||||
password="pass123",
|
||||
)
|
||||
|
||||
mock_smtp = MagicMock()
|
||||
with patch("smtplib.SMTP", return_value=mock_smtp) as mock_cls:
|
||||
mock_smtp.__enter__ = MagicMock(return_value=mock_smtp)
|
||||
mock_smtp.__exit__ = MagicMock(return_value=False)
|
||||
result = ch.send("recipient@example.com", "Hello!")
|
||||
assert result is True
|
||||
mock_cls.assert_called_once_with("smtp.example.com", 587)
|
||||
mock_smtp.starttls.assert_called_once()
|
||||
mock_smtp.login.assert_called_once_with("user@example.com", "pass123")
|
||||
mock_smtp.send_message.assert_called_once()
|
||||
|
||||
def test_send_success_no_tls(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
username="user@example.com",
|
||||
password="pass123",
|
||||
use_tls=False,
|
||||
)
|
||||
|
||||
mock_smtp = MagicMock()
|
||||
with patch("smtplib.SMTP", return_value=mock_smtp):
|
||||
mock_smtp.__enter__ = MagicMock(return_value=mock_smtp)
|
||||
mock_smtp.__exit__ = MagicMock(return_value=False)
|
||||
result = ch.send("recipient@example.com", "Hello!")
|
||||
assert result is True
|
||||
mock_smtp.starttls.assert_not_called()
|
||||
|
||||
def test_send_no_config(self):
|
||||
ch = EmailChannel()
|
||||
result = ch.send("recipient@example.com", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
username="user@example.com",
|
||||
password="pass123",
|
||||
)
|
||||
|
||||
with patch("smtplib.SMTP", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("recipient@example.com", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
username="user@example.com",
|
||||
password="pass123",
|
||||
bus=bus,
|
||||
)
|
||||
|
||||
mock_smtp = MagicMock()
|
||||
with patch("smtplib.SMTP", return_value=mock_smtp):
|
||||
mock_smtp.__enter__ = MagicMock(return_value=mock_smtp)
|
||||
mock_smtp.__exit__ = MagicMock(return_value=False)
|
||||
ch.send("recipient@example.com", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
def test_send_with_subject_metadata(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
username="user@example.com",
|
||||
password="pass123",
|
||||
)
|
||||
|
||||
mock_smtp = MagicMock()
|
||||
with patch("smtplib.SMTP", return_value=mock_smtp):
|
||||
mock_smtp.__enter__ = MagicMock(return_value=mock_smtp)
|
||||
mock_smtp.__exit__ = MagicMock(return_value=False)
|
||||
result = ch.send(
|
||||
"recipient@example.com",
|
||||
"Hello!",
|
||||
metadata={"subject": "Custom Subject"},
|
||||
)
|
||||
assert result is True
|
||||
sent_msg = mock_smtp.send_message.call_args[0][0]
|
||||
assert sent_msg["Subject"] == "Custom Subject"
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = EmailChannel(smtp_host="smtp.example.com", username="user@example.com")
|
||||
assert ch.list_channels() == ["email"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = EmailChannel(smtp_host="smtp.example.com", username="user@example.com")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_config_connect_error(self):
|
||||
ch = EmailChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestConnect:
|
||||
def test_connect_smtp_only(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com",
|
||||
username="user@example.com",
|
||||
)
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.CONNECTED
|
||||
# No IMAP, so no listener thread
|
||||
assert ch._listener_thread is None
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = EmailChannel(smtp_host="smtp.example.com", username="user@example.com")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = EmailChannel(
|
||||
smtp_host="smtp.example.com", username="user@example.com",
|
||||
)
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Tests for the FeishuChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.feishu import FeishuChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_feishu():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("feishu"):
|
||||
ChannelRegistry.register_value("feishu", FeishuChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("feishu")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
assert ch.channel_id == "feishu"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = FeishuChannel()
|
||||
assert ch._app_id == ""
|
||||
assert ch._app_secret == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_param(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
assert ch._app_id == "test-id"
|
||||
assert ch._app_secret == "test-secret"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"FEISHU_APP_ID": "env-id", "FEISHU_APP_SECRET": "env-secret"}):
|
||||
ch = FeishuChannel()
|
||||
assert ch._app_id == "env-id"
|
||||
assert ch._app_secret == "env-secret"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"FEISHU_APP_ID": "env-id", "FEISHU_APP_SECRET": "env-secret"}):
|
||||
ch = FeishuChannel(app_id="explicit-id", app_secret="explicit-secret")
|
||||
assert ch._app_id == "explicit-id"
|
||||
assert ch._app_secret == "explicit-secret"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
|
||||
token_resp = MagicMock()
|
||||
token_resp.status_code = 200
|
||||
token_resp.json.return_value = {"tenant_access_token": "fake-token"}
|
||||
|
||||
msg_resp = MagicMock()
|
||||
msg_resp.status_code = 200
|
||||
|
||||
with patch("httpx.post", side_effect=[token_resp, msg_resp]) as mock_post:
|
||||
result = ch.send("chat_id", "Hello!")
|
||||
assert result is True
|
||||
assert mock_post.call_count == 2
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
|
||||
token_resp = MagicMock()
|
||||
token_resp.status_code = 200
|
||||
token_resp.json.return_value = {"tenant_access_token": "fake-token"}
|
||||
|
||||
msg_resp = MagicMock()
|
||||
msg_resp.status_code = 400
|
||||
msg_resp.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", side_effect=[token_resp, msg_resp]):
|
||||
result = ch.send("chat_id", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("chat_id", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = FeishuChannel()
|
||||
result = ch.send("chat_id", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret", bus=bus)
|
||||
|
||||
token_resp = MagicMock()
|
||||
token_resp.status_code = 200
|
||||
token_resp.json.return_value = {"tenant_access_token": "fake-token"}
|
||||
|
||||
msg_resp = MagicMock()
|
||||
msg_resp.status_code = 200
|
||||
|
||||
with patch("httpx.post", side_effect=[token_resp, msg_resp]):
|
||||
ch.send("chat_id", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
assert ch.list_channels() == ["feishu"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_config_connect_error(self):
|
||||
ch = FeishuChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = FeishuChannel(app_id="test-id", app_secret="test-secret")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,139 @@
|
||||
"""Tests for the GoogleChatChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.google_chat import GoogleChatChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_google_chat():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("google_chat"):
|
||||
ChannelRegistry.register_value("google_chat", GoogleChatChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("google_chat")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
assert ch.channel_id == "google_chat"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = GoogleChatChannel()
|
||||
assert ch._webhook_url == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_url(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
assert ch._webhook_url == "https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {
|
||||
"GOOGLE_CHAT_WEBHOOK_URL": "https://chat.googleapis.com/v1/spaces/env/messages?key=env",
|
||||
}):
|
||||
ch = GoogleChatChannel()
|
||||
assert ch._webhook_url == "https://chat.googleapis.com/v1/spaces/env/messages?key=env"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {
|
||||
"GOOGLE_CHAT_WEBHOOK_URL": "https://chat.googleapis.com/v1/spaces/env/messages?key=env",
|
||||
}):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/explicit/messages?key=explicit")
|
||||
assert ch._webhook_url == "https://chat.googleapis.com/v1/spaces/explicit/messages?key=explicit"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("space", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("space", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("space", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_url(self):
|
||||
ch = GoogleChatChannel()
|
||||
result = ch.send("space", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = GoogleChatChannel(
|
||||
webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy",
|
||||
bus=bus,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("space", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
assert ch.list_channels() == ["google_chat"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_url_connect_error(self):
|
||||
ch = GoogleChatChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = GoogleChatChannel(webhook_url="https://chat.googleapis.com/v1/spaces/xxx/messages?key=yyy")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,139 @@
|
||||
"""Tests for the IRCChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.irc_channel import IRCChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_irc():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("irc"):
|
||||
ChannelRegistry.register_value("irc", IRCChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("irc")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
assert ch.channel_id == "irc"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = IRCChannel()
|
||||
assert ch._server == ""
|
||||
assert ch._nick == ""
|
||||
assert ch._password == ""
|
||||
assert ch._port == 6667
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_params(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
assert ch._server == "irc.example.com"
|
||||
assert ch._nick == "jarvis"
|
||||
assert ch._password == "pass123"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {
|
||||
"IRC_SERVER": "irc.env.com",
|
||||
"IRC_NICK": "envbot",
|
||||
"IRC_PASSWORD": "envpass",
|
||||
"IRC_PORT": "6697",
|
||||
}):
|
||||
ch = IRCChannel()
|
||||
assert ch._server == "irc.env.com"
|
||||
assert ch._nick == "envbot"
|
||||
assert ch._password == "envpass"
|
||||
assert ch._port == 6697
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {
|
||||
"IRC_SERVER": "irc.env.com",
|
||||
"IRC_NICK": "envbot",
|
||||
"IRC_PASSWORD": "envpass",
|
||||
}):
|
||||
ch = IRCChannel(server="irc.explicit.com", nick="explicit", password="explicit-pass")
|
||||
assert ch._server == "irc.explicit.com"
|
||||
assert ch._nick == "explicit"
|
||||
assert ch._password == "explicit-pass"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
|
||||
mock_sock = MagicMock()
|
||||
with patch("socket.socket", return_value=mock_sock):
|
||||
result = ch.send("#channel", "Hello!")
|
||||
assert result is True
|
||||
mock_sock.connect.assert_called_once()
|
||||
mock_sock.sendall.assert_called()
|
||||
|
||||
def test_send_failure_exception(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
|
||||
mock_sock = MagicMock()
|
||||
mock_sock.connect.side_effect = ConnectionError("refused")
|
||||
with patch("socket.socket", return_value=mock_sock):
|
||||
result = ch.send("#channel", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_config(self):
|
||||
ch = IRCChannel()
|
||||
result = ch.send("#channel", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123", bus=bus)
|
||||
|
||||
mock_sock = MagicMock()
|
||||
with patch("socket.socket", return_value=mock_sock):
|
||||
ch.send("#channel", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
assert ch.list_channels() == ["irc"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_server_connect_error(self):
|
||||
ch = IRCChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = IRCChannel(server="irc.example.com", nick="jarvis", password="pass123")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,140 @@
|
||||
"""Tests for the MatrixChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.matrix_channel import MatrixChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_matrix():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("matrix"):
|
||||
ChannelRegistry.register_value("matrix", MatrixChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("matrix")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
assert ch.channel_id == "matrix"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = MatrixChannel()
|
||||
assert ch._homeserver == ""
|
||||
assert ch._access_token == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_param(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
assert ch._homeserver == "https://matrix.example.com"
|
||||
assert ch._access_token == "test-token"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"MATRIX_HOMESERVER": "https://env.example.com", "MATRIX_ACCESS_TOKEN": "env-token"}):
|
||||
ch = MatrixChannel()
|
||||
assert ch._homeserver == "https://env.example.com"
|
||||
assert ch._access_token == "env-token"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"MATRIX_HOMESERVER": "https://env.example.com", "MATRIX_ACCESS_TOKEN": "env-token"}):
|
||||
ch = MatrixChannel(homeserver="https://explicit.example.com", access_token="explicit-token")
|
||||
assert ch._homeserver == "https://explicit.example.com"
|
||||
assert ch._access_token == "explicit-token"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.put", return_value=mock_response) as mock_put:
|
||||
result = ch.send("!room123:example.com", "Hello!")
|
||||
assert result is True
|
||||
mock_put.assert_called_once()
|
||||
call_args = mock_put.call_args
|
||||
url = call_args[0][0]
|
||||
assert "/_matrix/client/v3/rooms/" in url
|
||||
assert "!room123:example.com" in url
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.put", return_value=mock_response):
|
||||
result = ch.send("!room123:example.com", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
|
||||
with patch("httpx.put", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("!room123:example.com", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = MatrixChannel()
|
||||
result = ch.send("!room123:example.com", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.put", return_value=mock_response):
|
||||
ch.send("!room123:example.com", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
assert ch.list_channels() == ["matrix"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_homeserver_connect_error(self):
|
||||
ch = MatrixChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = MatrixChannel(homeserver="https://matrix.example.com", access_token="test-token")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,153 @@
|
||||
"""Tests for the MattermostChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.mattermost import MattermostChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_mattermost():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("mattermost"):
|
||||
ChannelRegistry.register_value("mattermost", MattermostChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("mattermost")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
assert ch.channel_id == "mattermost"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = MattermostChannel()
|
||||
assert ch._url == ""
|
||||
assert ch._token == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_param(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
assert ch._url == "https://mattermost.example.com"
|
||||
assert ch._token == "test-token"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"MATTERMOST_URL": "https://env.example.com", "MATTERMOST_TOKEN": "env-token"}):
|
||||
ch = MattermostChannel()
|
||||
assert ch._url == "https://env.example.com"
|
||||
assert ch._token == "env-token"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"MATTERMOST_URL": "https://env.example.com", "MATTERMOST_TOKEN": "env-token"}):
|
||||
ch = MattermostChannel(url="https://explicit.example.com", token="explicit-token")
|
||||
assert ch._url == "https://explicit.example.com"
|
||||
assert ch._token == "explicit-token"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("channel-id-123", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
url = call_args[0][0]
|
||||
assert url.endswith("/api/v4/posts")
|
||||
|
||||
def test_send_with_conversation_id(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("channel-id-123", "Hello!", conversation_id="root-123")
|
||||
assert result is True
|
||||
call_args = mock_post.call_args
|
||||
payload = call_args[1]["json"]
|
||||
assert "root_id" in payload
|
||||
assert payload["root_id"] == "root-123"
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("channel-id-123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("channel-id-123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = MattermostChannel()
|
||||
result = ch.send("channel-id-123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("channel-id-123", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
assert ch.list_channels() == ["mattermost"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_url_connect_error(self):
|
||||
ch = MattermostChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = MattermostChannel(url="https://mattermost.example.com", token="test-token")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -1,432 +0,0 @@
|
||||
"""Tests for the OpenClawChannelBridge implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelMessage, ChannelStatus
|
||||
from openjarvis.channels.openclaw_bridge import OpenClawChannelBridge
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_bridge():
|
||||
"""Re-register the OpenClawChannelBridge after registry clear."""
|
||||
if not ChannelRegistry.contains("openclaw"):
|
||||
ChannelRegistry.register_value("openclaw", OpenClawChannelBridge)
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_init_defaults(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
assert bridge._gateway_url == "ws://127.0.0.1:18789/ws"
|
||||
assert bridge._reconnect_interval == 5.0
|
||||
assert bridge._bus is None
|
||||
assert bridge._handlers == []
|
||||
assert bridge._status == ChannelStatus.DISCONNECTED
|
||||
assert bridge._ws is None
|
||||
|
||||
def test_init_custom_url(self) -> None:
|
||||
bridge = OpenClawChannelBridge(gateway_url="ws://custom:9999/ws")
|
||||
assert bridge._gateway_url == "ws://custom:9999/ws"
|
||||
|
||||
def test_init_custom_reconnect(self) -> None:
|
||||
bridge = OpenClawChannelBridge(reconnect_interval=10.0)
|
||||
assert bridge._reconnect_interval == 10.0
|
||||
|
||||
def test_init_with_bus(self) -> None:
|
||||
bus = EventBus()
|
||||
bridge = OpenClawChannelBridge(bus=bus)
|
||||
assert bridge._bus is bus
|
||||
|
||||
def test_channel_id(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
assert bridge.channel_id == "openclaw"
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_status_disconnected_initially(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
assert bridge.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
|
||||
class TestConnect:
|
||||
def test_connect_no_websockets(self) -> None:
|
||||
"""When websockets is not installed, connect uses HTTP fallback."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
with patch(
|
||||
"openjarvis.channels.openclaw_bridge.OpenClawChannelBridge.connect",
|
||||
wraps=bridge.connect,
|
||||
):
|
||||
# Simulate ImportError for websockets
|
||||
import builtins
|
||||
original_import = builtins.__import__
|
||||
|
||||
def mock_import(name, *args, **kwargs):
|
||||
if "websockets" in name:
|
||||
raise ImportError("No module named 'websockets'")
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=mock_import):
|
||||
bridge.connect()
|
||||
|
||||
assert bridge.status() == ChannelStatus.CONNECTED
|
||||
assert bridge._ws is None
|
||||
|
||||
def test_connect_websocket_success(self) -> None:
|
||||
"""When websockets is available, connect opens a WebSocket."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
mock_ws = MagicMock()
|
||||
|
||||
# Manually simulate successful connect
|
||||
bridge._ws = mock_ws
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
|
||||
assert bridge.status() == ChannelStatus.CONNECTED
|
||||
|
||||
def test_connect_connection_error(self) -> None:
|
||||
"""When connection fails, status is ERROR."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
|
||||
import builtins
|
||||
original_import = builtins.__import__
|
||||
|
||||
def mock_import(name, *args, **kwargs):
|
||||
if "websockets" in name:
|
||||
# Allow the import but make connect fail
|
||||
mod = MagicMock()
|
||||
mod.sync.client.connect.side_effect = ConnectionError("refused")
|
||||
return mod
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=mock_import):
|
||||
bridge.connect()
|
||||
|
||||
assert bridge.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
# Manually set to connected state
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
bridge.disconnect()
|
||||
assert bridge.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_disconnect_closes_ws(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
mock_ws = MagicMock()
|
||||
bridge._ws = mock_ws
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
bridge.disconnect()
|
||||
mock_ws.close.assert_called_once()
|
||||
assert bridge._ws is None
|
||||
assert bridge.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_disconnect_joins_listener_thread(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
mock_thread = MagicMock()
|
||||
bridge._listener_thread = mock_thread
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
bridge.disconnect()
|
||||
mock_thread.join.assert_called_once()
|
||||
assert bridge._listener_thread is None
|
||||
|
||||
def test_disconnect_when_already_disconnected(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
bridge.disconnect() # Should not raise
|
||||
assert bridge.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_http_fallback(self) -> None:
|
||||
"""When no WebSocket, send uses HTTP POST."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
bridge._ws = None
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = bridge.send("slack", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
assert call_args[0][0] == "http://127.0.0.1:18789/send"
|
||||
payload = call_args[1]["json"]
|
||||
assert payload["channel"] == "slack"
|
||||
assert payload["content"] == "Hello!"
|
||||
|
||||
def test_send_http_failure(self) -> None:
|
||||
"""When HTTP POST fails, send returns False."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
bridge._ws = None
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = bridge.send("slack", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_http_bad_status(self) -> None:
|
||||
"""When HTTP POST returns error status, send returns False."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
bridge._ws = None
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = bridge.send("slack", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_via_websocket(self) -> None:
|
||||
"""When WebSocket is available, send via WS."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
mock_ws = MagicMock()
|
||||
bridge._ws = mock_ws
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
|
||||
result = bridge.send("slack", "Hello!", conversation_id="conv-1")
|
||||
assert result is True
|
||||
mock_ws.send.assert_called_once()
|
||||
sent_data = json.loads(mock_ws.send.call_args[0][0])
|
||||
assert sent_data["channel"] == "slack"
|
||||
assert sent_data["content"] == "Hello!"
|
||||
assert sent_data["conversation_id"] == "conv-1"
|
||||
|
||||
def test_send_ws_failure_falls_back_to_http(self) -> None:
|
||||
"""When WebSocket send fails, fallback to HTTP."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
mock_ws = MagicMock()
|
||||
mock_ws.send.side_effect = Exception("WS error")
|
||||
bridge._ws = mock_ws
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = bridge.send("slack", "Hello!")
|
||||
assert result is True
|
||||
|
||||
def test_send_with_metadata(self) -> None:
|
||||
"""Metadata is included in the payload."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
mock_ws = MagicMock()
|
||||
bridge._ws = mock_ws
|
||||
|
||||
bridge.send("slack", "Hello!", metadata={"thread_ts": "123"})
|
||||
sent_data = json.loads(mock_ws.send.call_args[0][0])
|
||||
assert sent_data["metadata"] == {"thread_ts": "123"}
|
||||
|
||||
def test_send_publishes_event(self) -> None:
|
||||
"""Successful send publishes CHANNEL_MESSAGE_SENT event."""
|
||||
bus = EventBus(record_history=True)
|
||||
bridge = OpenClawChannelBridge(bus=bus)
|
||||
mock_ws = MagicMock()
|
||||
bridge._ws = mock_ws
|
||||
|
||||
bridge.send("slack", "Hello!")
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels_http(self) -> None:
|
||||
"""list_channels queries HTTP endpoint."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = ["slack", "discord", "telegram"]
|
||||
|
||||
with patch("httpx.get", return_value=mock_response) as mock_get:
|
||||
channels = bridge.list_channels()
|
||||
assert channels == ["slack", "discord", "telegram"]
|
||||
mock_get.assert_called_once()
|
||||
assert "/channels" in mock_get.call_args[0][0]
|
||||
|
||||
def test_list_channels_dict_response(self) -> None:
|
||||
"""list_channels handles dict response with 'channels' key."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"channels": ["slack", "discord"]}
|
||||
|
||||
with patch("httpx.get", return_value=mock_response):
|
||||
channels = bridge.list_channels()
|
||||
assert channels == ["slack", "discord"]
|
||||
|
||||
def test_list_channels_error(self) -> None:
|
||||
"""When HTTP GET fails, return empty list."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
|
||||
with patch("httpx.get", side_effect=ConnectionError("refused")):
|
||||
channels = bridge.list_channels()
|
||||
assert channels == []
|
||||
|
||||
def test_list_channels_unexpected_format(self) -> None:
|
||||
"""When response is not a list or dict with 'channels', return empty list."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"unexpected": "format"}
|
||||
|
||||
with patch("httpx.get", return_value=mock_response):
|
||||
channels = bridge.list_channels()
|
||||
assert channels == []
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
handler = MagicMock()
|
||||
bridge.on_message(handler)
|
||||
assert handler in bridge._handlers
|
||||
|
||||
def test_on_message_multiple_handlers(self) -> None:
|
||||
bridge = OpenClawChannelBridge()
|
||||
h1 = MagicMock()
|
||||
h2 = MagicMock()
|
||||
bridge.on_message(h1)
|
||||
bridge.on_message(h2)
|
||||
assert len(bridge._handlers) == 2
|
||||
assert h1 in bridge._handlers
|
||||
assert h2 in bridge._handlers
|
||||
|
||||
|
||||
class TestHttpUrlConversion:
|
||||
def test_ws_to_http(self) -> None:
|
||||
bridge = OpenClawChannelBridge(gateway_url="ws://host:1234/ws")
|
||||
assert bridge._http_url("/send") == "http://host:1234/send"
|
||||
|
||||
def test_wss_to_https(self) -> None:
|
||||
bridge = OpenClawChannelBridge(gateway_url="wss://host:1234/ws")
|
||||
assert bridge._http_url("/send") == "https://host:1234/send"
|
||||
|
||||
def test_http_url_channels(self) -> None:
|
||||
bridge = OpenClawChannelBridge(gateway_url="ws://127.0.0.1:18789/ws")
|
||||
assert bridge._http_url("/channels") == "http://127.0.0.1:18789/channels"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _listener_loop() tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_bridge_with_mock_ws(
|
||||
msg_json: str,
|
||||
*,
|
||||
bus: EventBus | None = None,
|
||||
) -> OpenClawChannelBridge:
|
||||
"""Return a bridge whose ``_ws.recv()`` returns *msg_json* once, then stops."""
|
||||
bridge = OpenClawChannelBridge(bus=bus)
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
|
||||
call_count = 0
|
||||
|
||||
def mock_recv(timeout=None):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
return msg_json
|
||||
bridge._stop_event.set()
|
||||
raise TimeoutError()
|
||||
|
||||
mock_ws = MagicMock()
|
||||
mock_ws.recv = mock_recv
|
||||
bridge._ws = mock_ws
|
||||
return bridge
|
||||
|
||||
|
||||
_SAMPLE_MSG = json.dumps({
|
||||
"channel": "slack",
|
||||
"sender": "user1",
|
||||
"content": "Hello!",
|
||||
"message_id": "msg-1",
|
||||
"conversation_id": "conv-1",
|
||||
"metadata": {},
|
||||
})
|
||||
|
||||
|
||||
class TestListenerLoop:
|
||||
def test_listener_receives_and_parses_message(self) -> None:
|
||||
"""Listener parses incoming JSON into a ChannelMessage."""
|
||||
bridge = _make_bridge_with_mock_ws(_SAMPLE_MSG)
|
||||
received: list[ChannelMessage] = []
|
||||
bridge.on_message(lambda msg: received.append(msg))
|
||||
|
||||
bridge._listener_loop()
|
||||
|
||||
assert len(received) == 1
|
||||
msg = received[0]
|
||||
assert isinstance(msg, ChannelMessage)
|
||||
assert msg.channel == "slack"
|
||||
assert msg.sender == "user1"
|
||||
assert msg.content == "Hello!"
|
||||
assert msg.message_id == "msg-1"
|
||||
|
||||
def test_listener_invokes_all_handlers(self) -> None:
|
||||
"""All registered handlers are called for each message."""
|
||||
bridge = _make_bridge_with_mock_ws(_SAMPLE_MSG)
|
||||
calls_a: list[ChannelMessage] = []
|
||||
calls_b: list[ChannelMessage] = []
|
||||
bridge.on_message(lambda msg: calls_a.append(msg))
|
||||
bridge.on_message(lambda msg: calls_b.append(msg))
|
||||
|
||||
bridge._listener_loop()
|
||||
|
||||
assert len(calls_a) == 1
|
||||
assert len(calls_b) == 1
|
||||
|
||||
def test_listener_publishes_event(self) -> None:
|
||||
"""Listener publishes CHANNEL_MESSAGE_RECEIVED on the event bus."""
|
||||
bus = EventBus(record_history=True)
|
||||
bridge = _make_bridge_with_mock_ws(_SAMPLE_MSG, bus=bus)
|
||||
|
||||
bridge._listener_loop()
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_RECEIVED in event_types
|
||||
recv_event = next(
|
||||
e for e in bus.history
|
||||
if e.event_type == EventType.CHANNEL_MESSAGE_RECEIVED
|
||||
)
|
||||
assert recv_event.data["channel"] == "slack"
|
||||
assert recv_event.data["sender"] == "user1"
|
||||
|
||||
def test_listener_stops_on_stop_event(self) -> None:
|
||||
"""Listener returns immediately when _stop_event is already set."""
|
||||
bridge = OpenClawChannelBridge()
|
||||
bridge._status = ChannelStatus.CONNECTED
|
||||
bridge._ws = MagicMock()
|
||||
bridge._stop_event.set()
|
||||
|
||||
handler = MagicMock()
|
||||
bridge.on_message(handler)
|
||||
|
||||
# Should return immediately without calling any handler
|
||||
import threading
|
||||
t = threading.Thread(target=bridge._listener_loop)
|
||||
t.start()
|
||||
t.join(timeout=3.0)
|
||||
assert not t.is_alive(), "_listener_loop did not exit in time"
|
||||
handler.assert_not_called()
|
||||
|
||||
def test_listener_catches_handler_exceptions(self) -> None:
|
||||
"""A failing handler does not prevent subsequent handlers from running."""
|
||||
bridge = _make_bridge_with_mock_ws(_SAMPLE_MSG)
|
||||
second_calls: list[ChannelMessage] = []
|
||||
|
||||
def bad_handler(msg: ChannelMessage) -> None:
|
||||
raise ValueError("boom")
|
||||
|
||||
bridge.on_message(bad_handler)
|
||||
bridge.on_message(lambda msg: second_calls.append(msg))
|
||||
|
||||
bridge._listener_loop()
|
||||
|
||||
assert len(second_calls) == 1
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Tests for the SignalChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.signal_channel import SignalChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_signal():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("signal"):
|
||||
ChannelRegistry.register_value("signal", SignalChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("signal")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
assert ch.channel_id == "signal"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = SignalChannel()
|
||||
assert ch._api_url == ""
|
||||
assert ch._phone_number == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_params(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
assert ch._api_url == "http://localhost:8080"
|
||||
assert ch._phone_number == "+1234567890"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {
|
||||
"SIGNAL_API_URL": "http://env-signal:8080",
|
||||
"SIGNAL_PHONE_NUMBER": "+9876543210",
|
||||
}):
|
||||
ch = SignalChannel()
|
||||
assert ch._api_url == "http://env-signal:8080"
|
||||
assert ch._phone_number == "+9876543210"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {
|
||||
"SIGNAL_API_URL": "http://env-signal:8080",
|
||||
"SIGNAL_PHONE_NUMBER": "+9876543210",
|
||||
}):
|
||||
ch = SignalChannel(api_url="http://explicit:8080", phone_number="+1111111111")
|
||||
assert ch._api_url == "http://explicit:8080"
|
||||
assert ch._phone_number == "+1111111111"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("+0987654321", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("+0987654321", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("+0987654321", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_config(self):
|
||||
ch = SignalChannel()
|
||||
result = ch.send("+0987654321", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("+0987654321", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
assert ch.list_channels() == ["signal"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_config_connect_error(self):
|
||||
ch = SignalChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = SignalChannel(api_url="http://localhost:8080", phone_number="+1234567890")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Tests for the SlackChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.slack import SlackChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_slack():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("slack"):
|
||||
ChannelRegistry.register_value("slack", SlackChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("slack")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
assert ch.channel_id == "slack"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = SlackChannel()
|
||||
assert ch._token == ""
|
||||
assert ch._app_token == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_token(self):
|
||||
ch = SlackChannel(bot_token="xoxb-my-token")
|
||||
assert ch._token == "xoxb-my-token"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"SLACK_BOT_TOKEN": "xoxb-env"}):
|
||||
ch = SlackChannel()
|
||||
assert ch._token == "xoxb-env"
|
||||
|
||||
def test_app_token_env_var(self):
|
||||
with patch.dict(os.environ, {"SLACK_APP_TOKEN": "xapp-env"}):
|
||||
ch = SlackChannel()
|
||||
assert ch._app_token == "xapp-env"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"SLACK_BOT_TOKEN": "xoxb-env"}):
|
||||
ch = SlackChannel(bot_token="xoxb-explicit")
|
||||
assert ch._token == "xoxb-explicit"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"ok": True}
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("C1234567890", "Hello Slack!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
url = call_args[0][0]
|
||||
assert "slack.com/api/chat.postMessage" in url
|
||||
headers = call_args[1]["headers"]
|
||||
assert headers["Authorization"] == "Bearer xoxb-test"
|
||||
payload = call_args[1]["json"]
|
||||
assert payload["channel"] == "C1234567890"
|
||||
assert payload["text"] == "Hello Slack!"
|
||||
|
||||
def test_send_with_thread(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"ok": True}
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
ch.send("C123", "Reply", conversation_id="1234567890.123456")
|
||||
payload = mock_post.call_args[1]["json"]
|
||||
assert payload["thread_ts"] == "1234567890.123456"
|
||||
|
||||
def test_send_api_error(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"ok": False, "error": "channel_not_found"}
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("C123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_http_failure(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("C123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("C123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = SlackChannel()
|
||||
result = ch.send("C123", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = SlackChannel(bot_token="xoxb-test", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"ok": True}
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("C123", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
assert ch.list_channels() == ["slack"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_token_connect_error(self):
|
||||
ch = SlackChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = SlackChannel(bot_token="xoxb-test")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,140 @@
|
||||
"""Tests for the TeamsChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.teams import TeamsChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_teams():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("teams"):
|
||||
ChannelRegistry.register_value("teams", TeamsChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("teams")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
assert ch.channel_id == "teams"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = TeamsChannel()
|
||||
assert ch._app_id == ""
|
||||
assert ch._app_password == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_param(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
assert ch._app_id == "test-id"
|
||||
assert ch._app_password == "test-pass"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"TEAMS_APP_ID": "env-id", "TEAMS_APP_PASSWORD": "env-pass"}):
|
||||
ch = TeamsChannel()
|
||||
assert ch._app_id == "env-id"
|
||||
assert ch._app_password == "env-pass"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"TEAMS_APP_ID": "env-id", "TEAMS_APP_PASSWORD": "env-pass"}):
|
||||
ch = TeamsChannel(app_id="explicit-id", app_password="explicit-pass")
|
||||
assert ch._app_id == "explicit-id"
|
||||
assert ch._app_password == "explicit-pass"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("general", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
url = call_args[0][0]
|
||||
assert "/v3/conversations/" in url
|
||||
assert "general" in url
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("general", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("general", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_config(self):
|
||||
ch = TeamsChannel()
|
||||
result = ch.send("general", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("general", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
assert ch.list_channels() == ["teams"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_config_connect_error(self):
|
||||
ch = TeamsChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = TeamsChannel(app_id="test-id", app_password="test-pass")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Tests for the TelegramChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.telegram import TelegramChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_telegram():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("telegram"):
|
||||
ChannelRegistry.register_value("telegram", TelegramChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("telegram")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = TelegramChannel(bot_token="test-token")
|
||||
assert ch.channel_id == "telegram"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = TelegramChannel()
|
||||
assert ch._token == ""
|
||||
assert ch._parse_mode == "Markdown"
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_token(self):
|
||||
ch = TelegramChannel(bot_token="my-token")
|
||||
assert ch._token == "my-token"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {"TELEGRAM_BOT_TOKEN": "env-token"}):
|
||||
ch = TelegramChannel()
|
||||
assert ch._token == "env-token"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {"TELEGRAM_BOT_TOKEN": "env-token"}):
|
||||
ch = TelegramChannel(bot_token="explicit-token")
|
||||
assert ch._token == "explicit-token"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("12345678", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
url = call_args[0][0]
|
||||
assert "api.telegram.org" in url
|
||||
assert "bot123:ABC" in url
|
||||
assert "sendMessage" in url
|
||||
payload = call_args[1]["json"]
|
||||
assert payload["chat_id"] == "12345678"
|
||||
assert payload["text"] == "Hello!"
|
||||
assert payload["parse_mode"] == "Markdown"
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("12345678", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("12345678", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = TelegramChannel()
|
||||
result = ch.send("12345678", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = TelegramChannel(bot_token="123:ABC", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("12345678", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
assert ch.list_channels() == ["telegram"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_token_connect_error(self):
|
||||
ch = TelegramChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = TelegramChannel(bot_token="123:ABC")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Tests for the WebChatChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.webchat import WebChatChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_webchat():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("webchat"):
|
||||
ChannelRegistry.register_value("webchat", WebChatChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("webchat")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = WebChatChannel()
|
||||
assert ch.channel_id == "webchat"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = WebChatChannel()
|
||||
assert ch._messages == []
|
||||
assert ch._handlers == []
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = WebChatChannel()
|
||||
result = ch.send("user", "Hello!")
|
||||
assert result is True
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = WebChatChannel(bus=bus)
|
||||
ch.send("user", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
def test_get_messages(self):
|
||||
ch = WebChatChannel()
|
||||
ch.send("user1", "Hello!")
|
||||
ch.send("user2", "World!")
|
||||
ch.send("user1", "Again!")
|
||||
messages = ch.get_messages()
|
||||
assert len(messages) == 3
|
||||
assert messages[0].content == "Hello!"
|
||||
assert messages[1].content == "World!"
|
||||
assert messages[2].content == "Again!"
|
||||
|
||||
def test_clear_messages(self):
|
||||
ch = WebChatChannel()
|
||||
ch.send("user1", "Hello!")
|
||||
ch.send("user2", "World!")
|
||||
assert len(ch.get_messages()) == 2
|
||||
ch.clear_messages()
|
||||
assert len(ch.get_messages()) == 0
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = WebChatChannel()
|
||||
assert ch.list_channels() == ["webchat"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = WebChatChannel()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_connected_after_connect(self):
|
||||
ch = WebChatChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.CONNECTED
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = WebChatChannel()
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = WebChatChannel()
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Tests for the WebhookChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.webhook import WebhookChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_webhook():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("webhook"):
|
||||
ChannelRegistry.register_value("webhook", WebhookChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("webhook")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
assert ch.channel_id == "webhook"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = WebhookChannel()
|
||||
assert ch._url == ""
|
||||
assert ch._secret == ""
|
||||
assert ch._method == "POST"
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_params(self):
|
||||
ch = WebhookChannel(
|
||||
url="https://example.com/hook",
|
||||
secret="s3cr3t",
|
||||
method="PUT",
|
||||
)
|
||||
assert ch._url == "https://example.com/hook"
|
||||
assert ch._secret == "s3cr3t"
|
||||
assert ch._method == "PUT"
|
||||
|
||||
|
||||
class TestConnect:
|
||||
def test_connect_with_url(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.CONNECTED
|
||||
|
||||
def test_connect_no_url(self):
|
||||
ch = WebhookChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.request", return_value=mock_response) as mock_req:
|
||||
result = ch.send("target", "Hello!")
|
||||
assert result is True
|
||||
mock_req.assert_called_once()
|
||||
call_args = mock_req.call_args
|
||||
assert call_args[0][0] == "POST"
|
||||
assert call_args[0][1] == "https://example.com/hook"
|
||||
payload = call_args[1]["json"]
|
||||
assert payload["channel"] == "target"
|
||||
assert payload["content"] == "Hello!"
|
||||
|
||||
def test_send_put_method(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook", method="PUT")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.request", return_value=mock_response) as mock_req:
|
||||
ch.send("target", "Hello!")
|
||||
assert mock_req.call_args[0][0] == "PUT"
|
||||
|
||||
def test_send_with_secret(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook", secret="s3cr3t")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.request", return_value=mock_response) as mock_req:
|
||||
ch.send("target", "Hello!")
|
||||
headers = mock_req.call_args[1]["headers"]
|
||||
assert headers["X-Webhook-Secret"] == "s3cr3t"
|
||||
|
||||
def test_send_without_secret(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.request", return_value=mock_response) as mock_req:
|
||||
ch.send("target", "Hello!")
|
||||
headers = mock_req.call_args[1]["headers"]
|
||||
assert "X-Webhook-Secret" not in headers
|
||||
|
||||
def test_send_with_metadata(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.request", return_value=mock_response) as mock_req:
|
||||
ch.send("target", "Hello!", metadata={"key": "value"})
|
||||
payload = mock_req.call_args[1]["json"]
|
||||
assert payload["metadata"] == {"key": "value"}
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
|
||||
with patch("httpx.request", return_value=mock_response):
|
||||
result = ch.send("target", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
|
||||
with patch("httpx.request", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("target", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_url(self):
|
||||
ch = WebhookChannel()
|
||||
result = ch.send("target", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = WebhookChannel(url="https://example.com/hook", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.request", return_value=mock_response):
|
||||
ch.send("target", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
assert ch.list_channels() == ["webhook"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = WebhookChannel(url="https://example.com/hook")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -0,0 +1,146 @@
|
||||
"""Tests for the WhatsAppChannel adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
from openjarvis.channels.whatsapp import WhatsAppChannel
|
||||
from openjarvis.core.events import EventBus, EventType
|
||||
from openjarvis.core.registry import ChannelRegistry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _register_whatsapp():
|
||||
"""Re-register after any registry clear."""
|
||||
if not ChannelRegistry.contains("whatsapp"):
|
||||
ChannelRegistry.register_value("whatsapp", WhatsAppChannel)
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
def test_registry_key(self):
|
||||
assert ChannelRegistry.contains("whatsapp")
|
||||
|
||||
def test_channel_id(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
assert ch.channel_id == "whatsapp"
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_defaults(self):
|
||||
ch = WhatsAppChannel()
|
||||
assert ch._token == ""
|
||||
assert ch._phone_number_id == ""
|
||||
assert ch._status == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_constructor_token(self):
|
||||
ch = WhatsAppChannel(access_token="my-token", phone_number_id="12345")
|
||||
assert ch._token == "my-token"
|
||||
assert ch._phone_number_id == "12345"
|
||||
|
||||
def test_env_var_fallback(self):
|
||||
with patch.dict(os.environ, {
|
||||
"WHATSAPP_ACCESS_TOKEN": "env-token",
|
||||
"WHATSAPP_PHONE_NUMBER_ID": "env-id",
|
||||
}):
|
||||
ch = WhatsAppChannel()
|
||||
assert ch._token == "env-token"
|
||||
assert ch._phone_number_id == "env-id"
|
||||
|
||||
def test_constructor_overrides_env(self):
|
||||
with patch.dict(os.environ, {
|
||||
"WHATSAPP_ACCESS_TOKEN": "env-token",
|
||||
"WHATSAPP_PHONE_NUMBER_ID": "env-id",
|
||||
}):
|
||||
ch = WhatsAppChannel(access_token="explicit-token", phone_number_id="explicit-id")
|
||||
assert ch._token == "explicit-token"
|
||||
assert ch._phone_number_id == "explicit-id"
|
||||
|
||||
|
||||
class TestSend:
|
||||
def test_send_success(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response) as mock_post:
|
||||
result = ch.send("+1234567890", "Hello!")
|
||||
assert result is True
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
url = call_args[0][0]
|
||||
assert "graph.facebook.com" in url
|
||||
assert "12345" in url
|
||||
|
||||
def test_send_failure(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "Bad Request"
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
result = ch.send("+1234567890", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_exception(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
|
||||
with patch("httpx.post", side_effect=ConnectionError("refused")):
|
||||
result = ch.send("+1234567890", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_no_token(self):
|
||||
ch = WhatsAppChannel()
|
||||
result = ch.send("+1234567890", "Hello!")
|
||||
assert result is False
|
||||
|
||||
def test_send_publishes_event(self):
|
||||
bus = EventBus(record_history=True)
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345", bus=bus)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
|
||||
with patch("httpx.post", return_value=mock_response):
|
||||
ch.send("+1234567890", "Hello!")
|
||||
|
||||
event_types = [e.event_type for e in bus.history]
|
||||
assert EventType.CHANNEL_MESSAGE_SENT in event_types
|
||||
|
||||
|
||||
class TestListChannels:
|
||||
def test_list_channels(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
assert ch.list_channels() == ["whatsapp"]
|
||||
|
||||
|
||||
class TestStatus:
|
||||
def test_disconnected_initially(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
|
||||
def test_no_token_connect_error(self):
|
||||
ch = WhatsAppChannel()
|
||||
ch.connect()
|
||||
assert ch.status() == ChannelStatus.ERROR
|
||||
|
||||
|
||||
class TestOnMessage:
|
||||
def test_on_message(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
handler = MagicMock()
|
||||
ch.on_message(handler)
|
||||
assert handler in ch._handlers
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_disconnect(self):
|
||||
ch = WhatsAppChannel(access_token="test-token", phone_number_id="12345")
|
||||
ch._status = ChannelStatus.CONNECTED
|
||||
ch.disconnect()
|
||||
assert ch.status() == ChannelStatus.DISCONNECTED
|
||||
@@ -31,7 +31,6 @@ def _mock_engine(content="Hello from engine"):
|
||||
def _register_agents():
|
||||
"""Re-register agents after registry clear."""
|
||||
from openjarvis.agents.custom import CustomAgent
|
||||
from openjarvis.agents.openclaw import OpenClawAgent
|
||||
from openjarvis.agents.orchestrator import OrchestratorAgent
|
||||
from openjarvis.agents.simple import SimpleAgent
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
@@ -40,7 +39,6 @@ def _register_agents():
|
||||
("simple", SimpleAgent),
|
||||
("orchestrator", OrchestratorAgent),
|
||||
("custom", CustomAgent),
|
||||
("openclaw", OpenClawAgent),
|
||||
]:
|
||||
if not AgentRegistry.contains(name):
|
||||
AgentRegistry.register_value(name, cls)
|
||||
|
||||
@@ -15,9 +15,9 @@ def _patch_channel(
|
||||
send_return=True,
|
||||
status_return=ChannelStatus.DISCONNECTED,
|
||||
):
|
||||
"""Return patches for load_config and OpenClawChannelBridge."""
|
||||
"""Return patches for load_config and _get_channel."""
|
||||
cfg = mock.MagicMock()
|
||||
cfg.channel.gateway_url = "ws://127.0.0.1:18789/ws"
|
||||
cfg.channel.default_channel = ""
|
||||
|
||||
bridge_instance = mock.MagicMock()
|
||||
bridge_instance.list_channels.return_value = list_channels or []
|
||||
@@ -27,11 +27,11 @@ def _patch_channel(
|
||||
config_patch = mock.patch(
|
||||
"openjarvis.core.config.load_config", return_value=cfg,
|
||||
)
|
||||
bridge_patch = mock.patch(
|
||||
"openjarvis.channels.openclaw_bridge.OpenClawChannelBridge",
|
||||
get_channel_patch = mock.patch(
|
||||
"openjarvis.cli.channel_cmd._get_channel",
|
||||
return_value=bridge_instance,
|
||||
)
|
||||
return config_patch, bridge_patch, bridge_instance
|
||||
return config_patch, get_channel_patch, bridge_instance
|
||||
|
||||
|
||||
class TestChannelHelp:
|
||||
@@ -45,24 +45,26 @@ class TestChannelHelp:
|
||||
|
||||
class TestChannelList:
|
||||
def test_list_with_channels(self) -> None:
|
||||
config_p, bridge_p, _ = _patch_channel(list_channels=["slack", "discord"])
|
||||
with config_p, bridge_p:
|
||||
config_p, getch_p, _ = _patch_channel(
|
||||
list_channels=["slack", "discord"],
|
||||
)
|
||||
with config_p, getch_p:
|
||||
result = CliRunner().invoke(cli, ["channel", "list"])
|
||||
assert result.exit_code == 0
|
||||
assert "slack" in result.output
|
||||
assert "discord" in result.output
|
||||
|
||||
def test_list_no_channels(self) -> None:
|
||||
config_p, bridge_p, _ = _patch_channel(list_channels=[])
|
||||
with config_p, bridge_p:
|
||||
config_p, getch_p, _ = _patch_channel(list_channels=[])
|
||||
with config_p, getch_p:
|
||||
result = CliRunner().invoke(cli, ["channel", "list"])
|
||||
assert result.exit_code == 0
|
||||
assert "No channels available" in result.output
|
||||
|
||||
def test_list_connection_error(self) -> None:
|
||||
config_p, bridge_p, inst = _patch_channel()
|
||||
config_p, getch_p, inst = _patch_channel()
|
||||
inst.list_channels.side_effect = ConnectionError("refused")
|
||||
with config_p, bridge_p:
|
||||
with config_p, getch_p:
|
||||
result = CliRunner().invoke(cli, ["channel", "list"])
|
||||
assert result.exit_code == 0
|
||||
assert "Failed" in result.output or "refused" in result.output
|
||||
@@ -70,25 +72,30 @@ class TestChannelList:
|
||||
|
||||
class TestChannelSend:
|
||||
def test_send_success(self) -> None:
|
||||
config_p, bridge_p, _ = _patch_channel(send_return=True)
|
||||
with config_p, bridge_p:
|
||||
result = CliRunner().invoke(cli, ["channel", "send", "slack", "Hello!"])
|
||||
config_p, getch_p, _ = _patch_channel(send_return=True)
|
||||
with config_p, getch_p:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["channel", "send", "slack", "Hello!"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "Message sent" in result.output
|
||||
|
||||
def test_send_failure(self) -> None:
|
||||
config_p, bridge_p, _ = _patch_channel(send_return=False)
|
||||
with config_p, bridge_p:
|
||||
result = CliRunner().invoke(cli, ["channel", "send", "slack", "Hello!"])
|
||||
config_p, getch_p, _ = _patch_channel(send_return=False)
|
||||
with config_p, getch_p:
|
||||
result = CliRunner().invoke(
|
||||
cli, ["channel", "send", "slack", "Hello!"],
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
assert "Failed to send" in result.output
|
||||
|
||||
|
||||
class TestChannelStatus:
|
||||
def test_status_shows_info(self) -> None:
|
||||
config_p, bridge_p, _ = _patch_channel(status_return=ChannelStatus.DISCONNECTED)
|
||||
with config_p, bridge_p:
|
||||
config_p, getch_p, _ = _patch_channel(
|
||||
status_return=ChannelStatus.DISCONNECTED,
|
||||
)
|
||||
with config_p, getch_p:
|
||||
result = CliRunner().invoke(cli, ["channel", "status"])
|
||||
assert result.exit_code == 0
|
||||
assert "disconnected" in result.output
|
||||
assert "127.0.0.1" in result.output
|
||||
|
||||
@@ -127,9 +127,7 @@ class TestChannelConfig:
|
||||
def test_channel_config_defaults(self) -> None:
|
||||
cc = ChannelConfig()
|
||||
assert cc.enabled is False
|
||||
assert cc.gateway_url == "ws://127.0.0.1:18789/ws"
|
||||
assert cc.default_agent == "simple"
|
||||
assert cc.reconnect_interval == 5.0
|
||||
|
||||
def test_channel_config_on_jarvis_config(self) -> None:
|
||||
cfg = JarvisConfig()
|
||||
@@ -138,11 +136,11 @@ class TestChannelConfig:
|
||||
def test_channel_config_loads_from_toml(self, tmp_path: Path) -> None:
|
||||
toml_file = tmp_path / "config.toml"
|
||||
toml_file.write_text(
|
||||
'[channel]\nenabled = true\ngateway_url = "ws://custom:9999/ws"\n'
|
||||
'[channel]\nenabled = true\ndefault_channel = "telegram"\n'
|
||||
)
|
||||
cfg = load_config(toml_file)
|
||||
assert cfg.channel.enabled is True
|
||||
assert cfg.channel.gateway_url == "ws://custom:9999/ws"
|
||||
assert cfg.channel.default_channel == "telegram"
|
||||
|
||||
def test_channel_config_in_default_toml(self) -> None:
|
||||
output = generate_default_toml(HardwareInfo())
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
"""Tests for the LiteLLM engine backend."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
from openjarvis.core.registry import EngineRegistry
|
||||
from openjarvis.core.types import Message, Role
|
||||
from openjarvis.engine.litellm import LiteLLMEngine
|
||||
|
||||
|
||||
class TestLiteLLMEngineHealth:
|
||||
def test_health_importable(self) -> None:
|
||||
fake_litellm = mock.MagicMock()
|
||||
with mock.patch.dict("sys.modules", {"litellm": fake_litellm}):
|
||||
engine = LiteLLMEngine()
|
||||
assert engine.health() is True
|
||||
|
||||
def test_health_not_importable(self) -> None:
|
||||
with mock.patch.dict("sys.modules", {"litellm": None}):
|
||||
engine = LiteLLMEngine()
|
||||
assert engine.health() is False
|
||||
|
||||
|
||||
class TestLiteLLMEngineGenerate:
|
||||
def test_generate(self) -> None:
|
||||
fake_usage = SimpleNamespace(
|
||||
prompt_tokens=10, completion_tokens=5, total_tokens=15
|
||||
)
|
||||
fake_choice = SimpleNamespace(
|
||||
message=SimpleNamespace(content="Hello!", tool_calls=None),
|
||||
finish_reason="stop",
|
||||
)
|
||||
fake_resp = SimpleNamespace(
|
||||
choices=[fake_choice], usage=fake_usage, model="gpt-4o"
|
||||
)
|
||||
|
||||
fake_litellm = mock.MagicMock()
|
||||
fake_litellm.completion.return_value = fake_resp
|
||||
fake_litellm.completion_cost.return_value = 0.001
|
||||
|
||||
with mock.patch.dict("sys.modules", {"litellm": fake_litellm}):
|
||||
engine = LiteLLMEngine()
|
||||
result = engine.generate(
|
||||
[Message(role=Role.USER, content="Hi")], model="gpt-4o"
|
||||
)
|
||||
|
||||
assert result["content"] == "Hello!"
|
||||
assert result["usage"]["prompt_tokens"] == 10
|
||||
assert result["usage"]["completion_tokens"] == 5
|
||||
assert result["usage"]["total_tokens"] == 15
|
||||
assert result["model"] == "gpt-4o"
|
||||
assert result["finish_reason"] == "stop"
|
||||
assert result["cost_usd"] == 0.001
|
||||
|
||||
def test_generate_with_tools(self) -> None:
|
||||
fake_tool_call = SimpleNamespace(
|
||||
id="call_123",
|
||||
function=SimpleNamespace(
|
||||
name="calculator",
|
||||
arguments='{"expression": "2+2"}',
|
||||
),
|
||||
)
|
||||
fake_usage = SimpleNamespace(
|
||||
prompt_tokens=20, completion_tokens=10, total_tokens=30
|
||||
)
|
||||
fake_choice = SimpleNamespace(
|
||||
message=SimpleNamespace(
|
||||
content="",
|
||||
tool_calls=[fake_tool_call],
|
||||
),
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
fake_resp = SimpleNamespace(
|
||||
choices=[fake_choice], usage=fake_usage, model="gpt-4o"
|
||||
)
|
||||
|
||||
fake_litellm = mock.MagicMock()
|
||||
fake_litellm.completion.return_value = fake_resp
|
||||
fake_litellm.completion_cost.return_value = 0.002
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "calculator",
|
||||
"description": "Evaluate math",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
with mock.patch.dict("sys.modules", {"litellm": fake_litellm}):
|
||||
engine = LiteLLMEngine()
|
||||
result = engine.generate(
|
||||
[Message(role=Role.USER, content="What is 2+2?")],
|
||||
model="gpt-4o",
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
assert "tool_calls" in result
|
||||
assert len(result["tool_calls"]) == 1
|
||||
tc = result["tool_calls"][0]
|
||||
assert tc["id"] == "call_123"
|
||||
assert tc["type"] == "function"
|
||||
assert tc["function"]["name"] == "calculator"
|
||||
assert tc["function"]["arguments"] == '{"expression": "2+2"}'
|
||||
|
||||
def test_generate_with_api_base(self) -> None:
|
||||
fake_usage = SimpleNamespace(
|
||||
prompt_tokens=5, completion_tokens=3, total_tokens=8
|
||||
)
|
||||
fake_choice = SimpleNamespace(
|
||||
message=SimpleNamespace(content="Hi!", tool_calls=None),
|
||||
finish_reason="stop",
|
||||
)
|
||||
fake_resp = SimpleNamespace(
|
||||
choices=[fake_choice], usage=fake_usage, model="custom-model"
|
||||
)
|
||||
|
||||
fake_litellm = mock.MagicMock()
|
||||
fake_litellm.completion.return_value = fake_resp
|
||||
fake_litellm.completion_cost.return_value = 0.0
|
||||
|
||||
with mock.patch.dict("sys.modules", {"litellm": fake_litellm}):
|
||||
engine = LiteLLMEngine(api_base="http://localhost:8080")
|
||||
engine.generate(
|
||||
[Message(role=Role.USER, content="Hi")], model="custom-model"
|
||||
)
|
||||
|
||||
call_kwargs = fake_litellm.completion.call_args
|
||||
assert call_kwargs[1]["api_base"] == "http://localhost:8080"
|
||||
|
||||
def test_generate_cost_error_fallback(self) -> None:
|
||||
fake_usage = SimpleNamespace(
|
||||
prompt_tokens=10, completion_tokens=5, total_tokens=15
|
||||
)
|
||||
fake_choice = SimpleNamespace(
|
||||
message=SimpleNamespace(content="Hello!", tool_calls=None),
|
||||
finish_reason="stop",
|
||||
)
|
||||
fake_resp = SimpleNamespace(
|
||||
choices=[fake_choice], usage=fake_usage, model="unknown/model"
|
||||
)
|
||||
|
||||
fake_litellm = mock.MagicMock()
|
||||
fake_litellm.completion.return_value = fake_resp
|
||||
fake_litellm.completion_cost.side_effect = Exception("Unknown model")
|
||||
|
||||
with mock.patch.dict("sys.modules", {"litellm": fake_litellm}):
|
||||
engine = LiteLLMEngine()
|
||||
result = engine.generate(
|
||||
[Message(role=Role.USER, content="Hi")], model="unknown/model"
|
||||
)
|
||||
|
||||
assert result["cost_usd"] == 0.0
|
||||
|
||||
|
||||
class TestLiteLLMEngineStream:
|
||||
def test_stream(self) -> None:
|
||||
chunk1 = SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content="Hel"))]
|
||||
)
|
||||
chunk2 = SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content="lo!"))]
|
||||
)
|
||||
chunk3 = SimpleNamespace(
|
||||
choices=[SimpleNamespace(delta=SimpleNamespace(content=None))]
|
||||
)
|
||||
|
||||
fake_litellm = mock.MagicMock()
|
||||
fake_litellm.completion.return_value = iter([chunk1, chunk2, chunk3])
|
||||
|
||||
with mock.patch.dict("sys.modules", {"litellm": fake_litellm}):
|
||||
engine = LiteLLMEngine()
|
||||
import asyncio
|
||||
|
||||
async def collect() -> list[str]:
|
||||
tokens: list[str] = []
|
||||
async for token in engine.stream(
|
||||
[Message(role=Role.USER, content="Hi")], model="gpt-4o"
|
||||
):
|
||||
tokens.append(token)
|
||||
return tokens
|
||||
|
||||
tokens = asyncio.run(collect())
|
||||
|
||||
assert tokens == ["Hel", "lo!"]
|
||||
|
||||
|
||||
class TestLiteLLMEngineListModels:
|
||||
def test_list_models_default(self) -> None:
|
||||
engine = LiteLLMEngine()
|
||||
assert engine.list_models() == []
|
||||
|
||||
def test_list_models_with_default_model(self) -> None:
|
||||
engine = LiteLLMEngine(default_model="anthropic/claude-sonnet-4-20250514")
|
||||
assert engine.list_models() == ["anthropic/claude-sonnet-4-20250514"]
|
||||
|
||||
|
||||
class TestLiteLLMEngineRegistry:
|
||||
def test_registry_key(self) -> None:
|
||||
EngineRegistry.register_value("litellm", LiteLLMEngine)
|
||||
assert EngineRegistry.contains("litellm")
|
||||
cls = EngineRegistry.get("litellm")
|
||||
assert cls is LiteLLMEngine
|
||||
@@ -143,7 +143,7 @@ class TestLlamaCppGenerate:
|
||||
return tokens
|
||||
|
||||
import asyncio
|
||||
tokens = asyncio.get_event_loop().run_until_complete(collect())
|
||||
tokens = asyncio.run(collect())
|
||||
assert tokens == ["Hi", " there"]
|
||||
|
||||
|
||||
|
||||
@@ -136,7 +136,7 @@ class TestOllamaGenerate:
|
||||
return tokens
|
||||
|
||||
import asyncio
|
||||
tokens = asyncio.get_event_loop().run_until_complete(collect())
|
||||
tokens = asyncio.run(collect())
|
||||
assert "Hello" in tokens
|
||||
|
||||
|
||||
|
||||
@@ -144,7 +144,7 @@ class TestVLLMGenerate:
|
||||
return tokens
|
||||
|
||||
import asyncio
|
||||
tokens = asyncio.get_event_loop().run_until_complete(collect())
|
||||
tokens = asyncio.run(collect())
|
||||
assert tokens == ["Hello", " world"]
|
||||
|
||||
|
||||
|
||||
@@ -122,10 +122,10 @@ class TestICLUpdaterPolicy:
|
||||
result = policy.update(store)
|
||||
assert len(result["examples"]) <= 5
|
||||
|
||||
def test_is_tool_policy(self):
|
||||
from openjarvis.learning._stubs import ToolLearningPolicy
|
||||
assert issubclass(ICLUpdaterPolicy, ToolLearningPolicy)
|
||||
assert ICLUpdaterPolicy.target == "tools"
|
||||
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()
|
||||
|
||||
@@ -8,7 +8,6 @@ from openjarvis.learning._stubs import (
|
||||
AgentLearningPolicy,
|
||||
IntelligenceLearningPolicy,
|
||||
LearningPolicy,
|
||||
ToolLearningPolicy,
|
||||
)
|
||||
|
||||
|
||||
@@ -25,20 +24,12 @@ class TestLearningPolicyABC:
|
||||
with pytest.raises(TypeError):
|
||||
AgentLearningPolicy()
|
||||
|
||||
def test_cannot_instantiate_tools(self):
|
||||
with pytest.raises(TypeError):
|
||||
ToolLearningPolicy()
|
||||
|
||||
def test_target_intelligence(self):
|
||||
assert IntelligenceLearningPolicy.target == "intelligence"
|
||||
|
||||
def test_target_agent(self):
|
||||
assert AgentLearningPolicy.target == "agent"
|
||||
|
||||
def test_target_tools(self):
|
||||
assert ToolLearningPolicy.target == "tools"
|
||||
|
||||
def test_hierarchy(self):
|
||||
assert issubclass(IntelligenceLearningPolicy, LearningPolicy)
|
||||
assert issubclass(AgentLearningPolicy, LearningPolicy)
|
||||
assert issubclass(ToolLearningPolicy, LearningPolicy)
|
||||
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from openjarvis.learning.sft_policy import SFTPolicy
|
||||
from openjarvis.learning.sft_policy import SFTRouterPolicy
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -26,9 +26,9 @@ class _MockTraceStore:
|
||||
return self._traces
|
||||
|
||||
|
||||
class TestSFTPolicy:
|
||||
class TestSFTRouterPolicy:
|
||||
def test_empty_traces(self):
|
||||
policy = SFTPolicy()
|
||||
policy = SFTRouterPolicy()
|
||||
store = _MockTraceStore([])
|
||||
result = policy.update(store)
|
||||
assert result["updated"] is False
|
||||
@@ -42,7 +42,7 @@ class TestSFTPolicy:
|
||||
)
|
||||
for i in range(6)
|
||||
]
|
||||
policy = SFTPolicy(min_samples=5)
|
||||
policy = SFTRouterPolicy(min_samples=5)
|
||||
store = _MockTraceStore(traces)
|
||||
result = policy.update(store)
|
||||
assert result["updated"] is True
|
||||
@@ -54,29 +54,29 @@ class TestSFTPolicy:
|
||||
_MockTrace(query="def foo(): pass", model="code-model", outcome="success")
|
||||
for _ in range(3)
|
||||
]
|
||||
policy = SFTPolicy(min_samples=5)
|
||||
policy = SFTRouterPolicy(min_samples=5)
|
||||
store = _MockTraceStore(traces)
|
||||
result = policy.update(store)
|
||||
assert result["updated"] is False
|
||||
|
||||
def test_classify_code(self):
|
||||
assert SFTPolicy._classify_query("def hello(): pass") == "code"
|
||||
assert SFTRouterPolicy._classify_query("def hello(): pass") == "code"
|
||||
|
||||
def test_classify_math(self):
|
||||
assert SFTPolicy._classify_query("solve the integral") == "math"
|
||||
assert SFTRouterPolicy._classify_query("solve the integral") == "math"
|
||||
|
||||
def test_classify_short(self):
|
||||
assert SFTPolicy._classify_query("hello world") == "short"
|
||||
assert SFTRouterPolicy._classify_query("hello world") == "short"
|
||||
|
||||
def test_classify_general(self):
|
||||
query = "tell me about " + " ".join(["something"] * 20)
|
||||
assert SFTPolicy._classify_query(query) == "general"
|
||||
assert SFTRouterPolicy._classify_query(query) == "general"
|
||||
|
||||
def test_policy_map_property(self):
|
||||
policy = SFTPolicy()
|
||||
policy = SFTRouterPolicy()
|
||||
assert policy.policy_map == {}
|
||||
|
||||
def test_is_intelligence_policy(self):
|
||||
from openjarvis.learning._stubs import IntelligenceLearningPolicy
|
||||
assert issubclass(SFTPolicy, IntelligenceLearningPolicy)
|
||||
assert SFTPolicy.target == "intelligence"
|
||||
assert issubclass(SFTRouterPolicy, IntelligenceLearningPolicy)
|
||||
assert SFTRouterPolicy.target == "intelligence"
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
"""Tests for /v1/channels endpoints."""
|
||||
"""Tests for /v1/channels endpoints.
|
||||
|
||||
Requires the ``[server]`` optional extra (fastapi, uvicorn, pydantic).
|
||||
Skipped automatically when those packages are not installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -6,7 +10,9 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus
|
||||
pytest.importorskip("fastapi", reason="openjarvis[server] not installed")
|
||||
|
||||
from openjarvis.channels._stubs import ChannelStatus # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
||||
@@ -559,59 +559,6 @@ class TestBenchmarkSuiteRunAll:
|
||||
assert "benchmark_name" in obj
|
||||
|
||||
|
||||
class TestOpenClawProtocolRoundtrip:
|
||||
"""OpenClaw protocol serialize/deserialize roundtrip."""
|
||||
|
||||
def test_roundtrip(self):
|
||||
from openjarvis.agents.openclaw_protocol import (
|
||||
MessageType,
|
||||
ProtocolMessage,
|
||||
deserialize,
|
||||
serialize,
|
||||
)
|
||||
|
||||
msg = ProtocolMessage(
|
||||
type=MessageType.QUERY,
|
||||
content="integration test",
|
||||
metadata={"source": "test"},
|
||||
)
|
||||
line = serialize(msg)
|
||||
restored = deserialize(line)
|
||||
assert restored.type == MessageType.QUERY
|
||||
assert restored.content == "integration test"
|
||||
|
||||
|
||||
class TestOpenClawAgentE2E:
|
||||
"""OpenClawAgent runs with mock transport end-to-end."""
|
||||
|
||||
def test_agent_e2e(self):
|
||||
from openjarvis.agents._stubs import AgentResult
|
||||
from openjarvis.agents.openclaw import OpenClawAgent
|
||||
from openjarvis.agents.openclaw_protocol import MessageType, ProtocolMessage
|
||||
from openjarvis.agents.openclaw_transport import OpenClawTransport
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
class _MockTransport(OpenClawTransport):
|
||||
def send(self, msg):
|
||||
return ProtocolMessage(
|
||||
type=MessageType.RESPONSE,
|
||||
content="E2E response",
|
||||
)
|
||||
|
||||
def health(self):
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
if not AgentRegistry.contains("openclaw"):
|
||||
AgentRegistry.register_value("openclaw", OpenClawAgent)
|
||||
agent = OpenClawAgent(transport=_MockTransport())
|
||||
result = agent.run("Integration test query")
|
||||
assert isinstance(result, AgentResult)
|
||||
assert result.content == "E2E response"
|
||||
|
||||
|
||||
class TestFullPipeline:
|
||||
"""Full pipeline: SDK → agent → engine → telemetry."""
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for orchestrator learning infrastructure."""
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Tests for orchestrator RL environment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.core.types import ToolResult
|
||||
from openjarvis.learning.orchestrator.environment import (
|
||||
OrchestratorEnvironment,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import OrchestratorAction
|
||||
from openjarvis.tools._stubs import BaseTool, ToolSpec
|
||||
|
||||
# -- Mock tool ---------------------------------------------------------------
|
||||
|
||||
|
||||
class _MockCalculator(BaseTool):
|
||||
tool_id = "calculator"
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="calculator",
|
||||
description="mock calculator",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "math expression",
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def execute(self, **params) -> ToolResult:
|
||||
expr = params.get("expression", "")
|
||||
try:
|
||||
result = str(eval(expr)) # noqa: S307
|
||||
except Exception as e:
|
||||
return ToolResult(
|
||||
tool_name="calculator",
|
||||
content=f"Error: {e}",
|
||||
success=False,
|
||||
)
|
||||
return ToolResult(
|
||||
tool_name="calculator",
|
||||
content=result,
|
||||
success=True,
|
||||
)
|
||||
|
||||
|
||||
# -- Tests -------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestOrchestratorEnvironment:
|
||||
def test_reset_creates_clean_state(self):
|
||||
env = OrchestratorEnvironment(tools=[_MockCalculator()])
|
||||
state = env.reset("What is 2+2?")
|
||||
assert state.initial_prompt == "What is 2+2?"
|
||||
assert state.num_turns() == 0
|
||||
assert state.final_answer is None
|
||||
|
||||
def test_step_executes_tool(self):
|
||||
env = OrchestratorEnvironment(tools=[_MockCalculator()])
|
||||
state = env.reset("q")
|
||||
action = OrchestratorAction(
|
||||
thought="calc",
|
||||
tool_name="calculator",
|
||||
tool_input="2+2",
|
||||
)
|
||||
state, obs = env.step(state, action)
|
||||
assert state.num_turns() == 1
|
||||
assert obs.latency_seconds >= 0
|
||||
|
||||
def test_is_done_on_final_answer(self):
|
||||
env = OrchestratorEnvironment(tools=[_MockCalculator()])
|
||||
state = env.reset("q")
|
||||
action = OrchestratorAction(
|
||||
thought="done",
|
||||
tool_name="calculator",
|
||||
tool_input="2+2",
|
||||
is_final_answer=True,
|
||||
)
|
||||
state, obs = env.step(state, action)
|
||||
assert env.is_done(state) is True
|
||||
|
||||
def test_is_done_on_max_turns(self):
|
||||
env = OrchestratorEnvironment(
|
||||
tools=[_MockCalculator()], max_turns=2
|
||||
)
|
||||
state = env.reset("q")
|
||||
for _ in range(2):
|
||||
action = OrchestratorAction(
|
||||
thought="go", tool_name="calculator", tool_input="1+1"
|
||||
)
|
||||
state, obs = env.step(state, action)
|
||||
assert env.is_done(state) is True
|
||||
|
||||
def test_invalid_tool_raises(self):
|
||||
env = OrchestratorEnvironment(tools=[_MockCalculator()])
|
||||
state = env.reset("q")
|
||||
action = OrchestratorAction(
|
||||
thought="t", tool_name="nonexistent", tool_input="x"
|
||||
)
|
||||
with pytest.raises(ValueError, match="not available"):
|
||||
env.step(state, action)
|
||||
|
||||
def test_max_turns_exceeded_raises(self):
|
||||
env = OrchestratorEnvironment(
|
||||
tools=[_MockCalculator()], max_turns=1
|
||||
)
|
||||
state = env.reset("q")
|
||||
action = OrchestratorAction(
|
||||
thought="go", tool_name="calculator", tool_input="1+1"
|
||||
)
|
||||
state, _ = env.step(state, action)
|
||||
with pytest.raises(ValueError, match="exceeded"):
|
||||
env.step(state, action)
|
||||
|
||||
def test_get_available_tools(self):
|
||||
env = OrchestratorEnvironment(tools=[_MockCalculator()])
|
||||
assert env.get_available_tools() == ["calculator"]
|
||||
|
||||
def test_not_done_initially(self):
|
||||
env = OrchestratorEnvironment(tools=[_MockCalculator()])
|
||||
state = env.reset("q")
|
||||
assert env.is_done(state) is False
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Tests for orchestrator GRPO trainer."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.orchestrator.grpo_trainer import (
|
||||
OrchestratorGRPOConfig,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import Episode
|
||||
|
||||
|
||||
class TestOrchestratorGRPOConfig:
|
||||
def test_defaults(self):
|
||||
cfg = OrchestratorGRPOConfig()
|
||||
assert cfg.model_name == "Qwen/Qwen3-1.7B"
|
||||
assert cfg.num_epochs == 10
|
||||
assert cfg.batch_size == 16
|
||||
assert cfg.num_samples_per_prompt == 8
|
||||
assert cfg.kl_coef == 0.0001
|
||||
assert cfg.clip_ratio == 0.2
|
||||
assert cfg.gradient_checkpointing is True
|
||||
assert cfg.use_8bit_ref is True
|
||||
|
||||
def test_custom_values(self):
|
||||
cfg = OrchestratorGRPOConfig(
|
||||
model_name="test-model",
|
||||
num_samples_per_prompt=4,
|
||||
kl_coef=0.01,
|
||||
)
|
||||
assert cfg.model_name == "test-model"
|
||||
assert cfg.num_samples_per_prompt == 4
|
||||
assert cfg.kl_coef == 0.01
|
||||
|
||||
def test_default_tools(self):
|
||||
cfg = OrchestratorGRPOConfig()
|
||||
assert "calculator" in cfg.available_tools
|
||||
assert "think" in cfg.available_tools
|
||||
|
||||
|
||||
class TestGroupAdvantageNormalization:
|
||||
"""Test the math behind group-relative advantage normalization."""
|
||||
|
||||
def test_uniform_rewards_zero_advantages(self):
|
||||
rewards = [0.5, 0.5, 0.5, 0.5]
|
||||
mean_r = sum(rewards) / len(rewards)
|
||||
std_r = (sum((r - mean_r) ** 2 for r in rewards) / len(rewards)) ** 0.5
|
||||
# All same → std=0 → advantages should be 0
|
||||
assert std_r < 1e-8
|
||||
advantages = [0.0] * len(rewards)
|
||||
assert all(a == 0.0 for a in advantages)
|
||||
|
||||
def test_varied_rewards_normalized(self):
|
||||
rewards = [0.1, 0.3, 0.5, 0.7]
|
||||
mean_r = sum(rewards) / len(rewards)
|
||||
std_r = (sum((r - mean_r) ** 2 for r in rewards) / len(rewards)) ** 0.5
|
||||
advantages = [(r - mean_r) / std_r for r in rewards]
|
||||
# Mean of advantages should be ~0
|
||||
assert abs(sum(advantages) / len(advantages)) < 1e-6
|
||||
# Std should be ~1
|
||||
adv_mean = sum(advantages) / len(advantages)
|
||||
adv_std = (
|
||||
sum((a - adv_mean) ** 2 for a in advantages) / len(advantages)
|
||||
) ** 0.5
|
||||
assert abs(adv_std - 1.0) < 1e-6
|
||||
|
||||
def test_best_gets_positive_advantage(self):
|
||||
rewards = [0.1, 0.9, 0.3, 0.5]
|
||||
mean_r = sum(rewards) / len(rewards)
|
||||
std_r = (sum((r - mean_r) ** 2 for r in rewards) / len(rewards)) ** 0.5
|
||||
advantages = [(r - mean_r) / std_r for r in rewards]
|
||||
# Index 1 (reward=0.9) should have highest advantage
|
||||
assert advantages[1] == max(advantages)
|
||||
assert advantages[1] > 0
|
||||
|
||||
|
||||
class TestRewardIntegration:
|
||||
def test_episode_reward(self):
|
||||
from openjarvis.learning.orchestrator.reward import (
|
||||
MultiObjectiveReward,
|
||||
Normalizers,
|
||||
RewardWeights,
|
||||
)
|
||||
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
ep = Episode(
|
||||
task_id="t",
|
||||
initial_prompt="q",
|
||||
correct=True,
|
||||
total_energy_joules=10.0,
|
||||
total_cost_usd=0.01,
|
||||
total_latency_seconds=1.0,
|
||||
max_power_watts=50.0,
|
||||
)
|
||||
r = reward_fn.compute(ep)
|
||||
assert isinstance(r, float)
|
||||
assert r > 0 # correct episode
|
||||
|
||||
|
||||
class TestGRPORegistration:
|
||||
def test_registered_in_learning_registry(self):
|
||||
import openjarvis.learning.orchestrator.grpo_trainer # noqa: F401
|
||||
from openjarvis.core.registry import LearningRegistry
|
||||
assert LearningRegistry.contains("orchestrator_grpo")
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Tests for OrchestratorAgent structured mode."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.agents.orchestrator import OrchestratorAgent
|
||||
from openjarvis.core.types import ToolResult
|
||||
from openjarvis.engine._stubs import InferenceEngine
|
||||
from openjarvis.tools._stubs import BaseTool, ToolSpec
|
||||
|
||||
# -- Mocks -------------------------------------------------------------------
|
||||
|
||||
|
||||
class _MockEngine(InferenceEngine):
|
||||
"""Engine that returns pre-scripted responses."""
|
||||
|
||||
engine_id = "mock"
|
||||
|
||||
def __init__(self, responses: list[str]) -> None:
|
||||
self._responses = list(responses)
|
||||
self._call_idx = 0
|
||||
|
||||
def generate(self, messages, **kwargs) -> dict:
|
||||
if self._call_idx < len(self._responses):
|
||||
content = self._responses[self._call_idx]
|
||||
self._call_idx += 1
|
||||
else:
|
||||
content = "FINAL_ANSWER: fallback"
|
||||
return {"content": content}
|
||||
|
||||
def stream(self, messages, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def list_models(self):
|
||||
return []
|
||||
|
||||
def health(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _MockTool(BaseTool):
|
||||
tool_id = "calculator"
|
||||
|
||||
@property
|
||||
def spec(self) -> ToolSpec:
|
||||
return ToolSpec(
|
||||
name="calculator",
|
||||
description="A calculator",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"expression": {"type": "string"},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def execute(self, **params) -> ToolResult:
|
||||
expr = params.get("expression", "")
|
||||
return ToolResult(
|
||||
tool_name="calculator", content=str(expr), success=True
|
||||
)
|
||||
|
||||
|
||||
# -- Tests -------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStructuredMode:
|
||||
def test_thought_tool_input_then_final_answer(self):
|
||||
"""Test full structured loop: TOOL call -> FINAL_ANSWER."""
|
||||
engine = _MockEngine([
|
||||
"THOUGHT: I need to calculate\nTOOL: calculator\nINPUT: 2+2",
|
||||
"THOUGHT: Got result\nFINAL_ANSWER: 4",
|
||||
])
|
||||
agent = OrchestratorAgent(
|
||||
engine=engine,
|
||||
model="test",
|
||||
tools=[_MockTool()],
|
||||
mode="structured",
|
||||
)
|
||||
result = agent.run("What is 2+2?")
|
||||
assert result.content == "4"
|
||||
assert result.turns == 2
|
||||
assert len(result.tool_results) == 1
|
||||
|
||||
def test_direct_final_answer(self):
|
||||
"""Test that FINAL_ANSWER on first turn works."""
|
||||
engine = _MockEngine([
|
||||
"THOUGHT: Easy\nFINAL_ANSWER: Paris",
|
||||
])
|
||||
agent = OrchestratorAgent(
|
||||
engine=engine,
|
||||
model="test",
|
||||
tools=[_MockTool()],
|
||||
mode="structured",
|
||||
)
|
||||
result = agent.run("Capital of France?")
|
||||
assert result.content == "Paris"
|
||||
assert result.turns == 1
|
||||
assert len(result.tool_results) == 0
|
||||
|
||||
def test_no_tool_no_final_treats_as_answer(self):
|
||||
"""If neither TOOL nor FINAL_ANSWER found, treat content as answer."""
|
||||
engine = _MockEngine(["Just a plain response with no format"])
|
||||
agent = OrchestratorAgent(
|
||||
engine=engine,
|
||||
model="test",
|
||||
tools=[_MockTool()],
|
||||
mode="structured",
|
||||
)
|
||||
result = agent.run("Hello")
|
||||
assert result.content == "Just a plain response with no format"
|
||||
assert result.turns == 1
|
||||
|
||||
def test_max_turns(self):
|
||||
"""Test that max_turns terminates the loop."""
|
||||
# Always returns a tool call, never a final answer
|
||||
responses = [
|
||||
"THOUGHT: calc\nTOOL: calculator\nINPUT: 1+1",
|
||||
] * 5
|
||||
engine = _MockEngine(responses)
|
||||
agent = OrchestratorAgent(
|
||||
engine=engine,
|
||||
model="test",
|
||||
tools=[_MockTool()],
|
||||
mode="structured",
|
||||
max_turns=3,
|
||||
)
|
||||
result = agent.run("loop forever")
|
||||
assert result.turns == 3
|
||||
assert result.metadata.get("max_turns_exceeded") is True
|
||||
|
||||
def test_custom_system_prompt(self):
|
||||
"""Test that custom system_prompt is used."""
|
||||
engine = _MockEngine(["FINAL_ANSWER: ok"])
|
||||
agent = OrchestratorAgent(
|
||||
engine=engine,
|
||||
model="test",
|
||||
tools=[_MockTool()],
|
||||
mode="structured",
|
||||
system_prompt="Custom prompt here",
|
||||
)
|
||||
result = agent.run("test")
|
||||
assert result.content == "ok"
|
||||
|
||||
|
||||
class TestFunctionCallingModeUnchanged:
|
||||
"""Ensure function_calling mode still works as before."""
|
||||
|
||||
def test_default_mode(self):
|
||||
engine = _MockEngine(["Hello back"])
|
||||
agent = OrchestratorAgent(
|
||||
engine=engine,
|
||||
model="test",
|
||||
)
|
||||
# Default mode should be function_calling
|
||||
assert agent._mode == "function_calling"
|
||||
result = agent.run("Hello")
|
||||
assert result.content == "Hello back"
|
||||
|
||||
|
||||
class TestParseStructuredResponse:
|
||||
def test_parse_thought_tool_input(self):
|
||||
parsed = OrchestratorAgent._parse_structured_response(
|
||||
"THOUGHT: reasoning\nTOOL: calculator\nINPUT: 2+2"
|
||||
)
|
||||
assert parsed["thought"] == "reasoning"
|
||||
assert parsed["tool"] == "calculator"
|
||||
assert parsed["input"] == "2+2"
|
||||
assert parsed["final_answer"] == ""
|
||||
|
||||
def test_parse_final_answer(self):
|
||||
parsed = OrchestratorAgent._parse_structured_response(
|
||||
"THOUGHT: done\nFINAL_ANSWER: 42"
|
||||
)
|
||||
assert parsed["final_answer"] == "42"
|
||||
assert parsed["thought"] == "done"
|
||||
|
||||
def test_parse_empty(self):
|
||||
parsed = OrchestratorAgent._parse_structured_response("")
|
||||
assert parsed["thought"] == ""
|
||||
assert parsed["tool"] == ""
|
||||
assert parsed["final_answer"] == ""
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Tests for orchestrator policy model."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.learning.orchestrator.policy_model import (
|
||||
OrchestratorPolicyModel,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
EpisodeState,
|
||||
OrchestratorAction,
|
||||
OrchestratorObservation,
|
||||
)
|
||||
|
||||
|
||||
class TestParseOutput:
|
||||
"""Test _parse_output without loading a real model."""
|
||||
|
||||
def _model(self) -> OrchestratorPolicyModel:
|
||||
return OrchestratorPolicyModel()
|
||||
|
||||
def test_valid_thought_tool_input(self):
|
||||
m = self._model()
|
||||
text = (
|
||||
"THOUGHT: I need to calculate 2+2\n"
|
||||
"TOOL: calculator\n"
|
||||
"INPUT: 2+2"
|
||||
)
|
||||
po = m._parse_output(text, ["calculator", "think"])
|
||||
assert po.thought == "I need to calculate 2+2"
|
||||
assert po.tool_name == "calculator"
|
||||
assert po.tool_input == "2+2"
|
||||
assert po.is_final_answer is False
|
||||
|
||||
def test_final_answer(self):
|
||||
m = self._model()
|
||||
text = (
|
||||
"THOUGHT: I have the result\n"
|
||||
"FINAL_ANSWER: 42"
|
||||
)
|
||||
po = m._parse_output(text, ["calculator"])
|
||||
assert po.is_final_answer is True
|
||||
assert po.tool_input == "42"
|
||||
|
||||
def test_final_answer_with_space(self):
|
||||
m = self._model()
|
||||
text = "FINAL ANSWER: the result is 7"
|
||||
po = m._parse_output(text, ["calculator"])
|
||||
assert po.is_final_answer is True
|
||||
|
||||
def test_missing_fields_fallback(self):
|
||||
m = self._model()
|
||||
text = "just some random output"
|
||||
po = m._parse_output(text, ["calculator", "think"])
|
||||
# Should fallback to first available tool
|
||||
assert po.tool_name == "calculator"
|
||||
assert po.thought == "No thought provided"
|
||||
|
||||
def test_invalid_tool_name_fallback(self):
|
||||
m = self._model()
|
||||
text = "THOUGHT: reason\nTOOL: nonexistent_tool\nINPUT: hello"
|
||||
po = m._parse_output(text, ["calculator", "think"])
|
||||
assert po.tool_name == "calculator" # fallback
|
||||
|
||||
def test_case_insensitive_tool_match(self):
|
||||
m = self._model()
|
||||
text = "THOUGHT: reason\nTOOL: Calculator\nINPUT: 5+5"
|
||||
po = m._parse_output(text, ["calculator", "think"])
|
||||
assert po.tool_name == "calculator"
|
||||
|
||||
def test_empty_tools_list(self):
|
||||
m = self._model()
|
||||
text = "THOUGHT: reason\nTOOL: calc\nINPUT: 1"
|
||||
po = m._parse_output(text, [])
|
||||
assert po.tool_name == "unknown"
|
||||
|
||||
|
||||
class TestBuildPrompt:
|
||||
def test_includes_task(self):
|
||||
m = OrchestratorPolicyModel()
|
||||
state = EpisodeState(initial_prompt="What is 2+2?")
|
||||
prompt = m._build_prompt(state, ["calculator"])
|
||||
assert "What is 2+2?" in prompt
|
||||
|
||||
def test_includes_tools(self):
|
||||
m = OrchestratorPolicyModel()
|
||||
state = EpisodeState(initial_prompt="q")
|
||||
prompt = m._build_prompt(state, ["calculator", "think"])
|
||||
assert "calculator" in prompt
|
||||
assert "think" in prompt
|
||||
|
||||
def test_includes_history(self):
|
||||
m = OrchestratorPolicyModel()
|
||||
state = EpisodeState(initial_prompt="q")
|
||||
action = OrchestratorAction(
|
||||
thought="use calc", tool_name="calculator", tool_input="2+2"
|
||||
)
|
||||
obs = OrchestratorObservation(content="4")
|
||||
state.add_turn(action, obs)
|
||||
prompt = m._build_prompt(state, ["calculator"])
|
||||
assert "Turn 1:" in prompt
|
||||
assert "use calc" in prompt
|
||||
|
||||
def test_format_instructions(self):
|
||||
m = OrchestratorPolicyModel()
|
||||
state = EpisodeState(initial_prompt="q")
|
||||
prompt = m._build_prompt(state, ["calculator"])
|
||||
assert "THOUGHT:" in prompt
|
||||
assert "TOOL:" in prompt
|
||||
assert "INPUT:" in prompt
|
||||
|
||||
|
||||
class TestPredictActionRequiresModel:
|
||||
def test_raises_without_model(self):
|
||||
m = OrchestratorPolicyModel()
|
||||
state = EpisodeState(initial_prompt="q")
|
||||
with pytest.raises(RuntimeError, match="Cannot generate"):
|
||||
m.predict_action(state, ["calculator"])
|
||||
|
||||
|
||||
class TestRepr:
|
||||
def test_repr(self):
|
||||
m = OrchestratorPolicyModel()
|
||||
r = repr(m)
|
||||
assert "OrchestratorPolicyModel" in r
|
||||
assert "None" in r
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Tests for orchestrator prompt registry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.orchestrator.prompt_registry import (
|
||||
TOOL_DESCRIPTIONS,
|
||||
build_system_prompt,
|
||||
)
|
||||
|
||||
|
||||
class TestToolDescriptions:
|
||||
def test_calculator_present(self):
|
||||
assert "calculator" in TOOL_DESCRIPTIONS
|
||||
|
||||
def test_think_present(self):
|
||||
assert "think" in TOOL_DESCRIPTIONS
|
||||
|
||||
def test_each_has_category(self):
|
||||
for name, info in TOOL_DESCRIPTIONS.items():
|
||||
assert "category" in info, f"{name} missing category"
|
||||
|
||||
def test_each_has_description(self):
|
||||
for name, info in TOOL_DESCRIPTIONS.items():
|
||||
assert "description" in info, f"{name} missing description"
|
||||
assert len(info["description"]) > 10
|
||||
|
||||
|
||||
class TestBuildSystemPrompt:
|
||||
def test_includes_tool_descriptions(self):
|
||||
prompt = build_system_prompt(["calculator", "think"])
|
||||
assert "calculator" in prompt.lower()
|
||||
assert "think" in prompt.lower()
|
||||
|
||||
def test_includes_response_format(self):
|
||||
prompt = build_system_prompt(["calculator"])
|
||||
assert "THOUGHT:" in prompt
|
||||
assert "TOOL:" in prompt
|
||||
assert "INPUT:" in prompt
|
||||
assert "FINAL_ANSWER:" in prompt
|
||||
|
||||
def test_includes_guide_sections(self):
|
||||
prompt = build_system_prompt(
|
||||
["calculator", "think", "code_interpreter", "web_search"]
|
||||
)
|
||||
assert "MATH PROBLEMS:" in prompt
|
||||
assert "CODING TASKS:" in prompt
|
||||
assert "REASONING/LOGIC:" in prompt
|
||||
|
||||
def test_with_single_tool(self):
|
||||
prompt = build_system_prompt(["calculator"])
|
||||
assert "calculator" in prompt
|
||||
|
||||
def test_default_all_tools(self):
|
||||
prompt = build_system_prompt()
|
||||
for name in TOOL_DESCRIPTIONS:
|
||||
assert name in prompt
|
||||
|
||||
def test_unknown_tool_handled(self):
|
||||
prompt = build_system_prompt(["calculator", "custom_tool_xyz"])
|
||||
assert "custom_tool_xyz" in prompt
|
||||
|
||||
def test_with_memory_tools(self):
|
||||
prompt = build_system_prompt(
|
||||
["calculator", "memory_search", "memory_store"]
|
||||
)
|
||||
assert "memory_search" in prompt
|
||||
|
||||
def test_with_llm_tool(self):
|
||||
prompt = build_system_prompt(["calculator", "llm"])
|
||||
assert "llm" in prompt.lower()
|
||||
assert "REASONING" in prompt
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Tests for orchestrator multi-objective reward."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.learning.orchestrator.reward import (
|
||||
AdaptiveRewardWeights,
|
||||
MultiObjectiveReward,
|
||||
Normalizers,
|
||||
RewardWeights,
|
||||
)
|
||||
from openjarvis.learning.orchestrator.types import (
|
||||
Episode,
|
||||
)
|
||||
|
||||
|
||||
class TestRewardWeights:
|
||||
def test_default_sum(self):
|
||||
w = RewardWeights()
|
||||
total = w.alpha + w.beta_cost + w.beta_energy + w.gamma_latency + w.gamma_power
|
||||
assert abs(total - 1.0) < 0.01
|
||||
|
||||
def test_invalid_sum_raises(self):
|
||||
with pytest.raises(ValueError, match="sum to 1.0"):
|
||||
RewardWeights(alpha=0.9, beta_cost=0.5)
|
||||
|
||||
def test_custom_weights(self):
|
||||
w = RewardWeights(
|
||||
alpha=0.5,
|
||||
beta_cost=0.1,
|
||||
beta_energy=0.1,
|
||||
gamma_latency=0.2,
|
||||
gamma_power=0.1,
|
||||
)
|
||||
assert w.alpha == 0.5
|
||||
|
||||
|
||||
class TestMultiObjectiveReward:
|
||||
def _make_episode(self, correct: bool = True) -> Episode:
|
||||
ep = Episode(
|
||||
task_id="t",
|
||||
initial_prompt="q",
|
||||
ground_truth="4",
|
||||
final_answer="4" if correct else "5",
|
||||
correct=correct,
|
||||
total_energy_joules=50.0,
|
||||
total_cost_usd=0.05,
|
||||
total_latency_seconds=15.0,
|
||||
max_power_watts=100.0,
|
||||
)
|
||||
return ep
|
||||
|
||||
def test_correct_episode_positive(self):
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
ep = self._make_episode(correct=True)
|
||||
r = reward_fn.compute(ep)
|
||||
assert r > 0, "Correct episode should have positive reward"
|
||||
|
||||
def test_incorrect_episode_negative(self):
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
ep = self._make_episode(correct=False)
|
||||
r = reward_fn.compute(ep)
|
||||
assert r < 0, "Incorrect episode should have negative reward"
|
||||
|
||||
def test_correct_better_than_incorrect(self):
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
correct = reward_fn.compute(self._make_episode(correct=True))
|
||||
incorrect = reward_fn.compute(self._make_episode(correct=False))
|
||||
assert correct > incorrect
|
||||
|
||||
def test_compute_with_breakdown(self):
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
ep = self._make_episode(correct=True)
|
||||
bd = reward_fn.compute_with_breakdown(ep)
|
||||
assert "total_reward" in bd
|
||||
assert "accuracy_reward" in bd
|
||||
assert bd["accuracy_reward"] == 1.0
|
||||
assert bd["cost_penalty"] > 0
|
||||
assert bd["energy_penalty"] > 0
|
||||
assert "ipj" in bd
|
||||
|
||||
def test_compute_batch(self):
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
episodes = [
|
||||
self._make_episode(correct=True),
|
||||
self._make_episode(correct=False),
|
||||
]
|
||||
rewards = reward_fn.compute_batch(episodes)
|
||||
assert len(rewards) == 2
|
||||
assert rewards[0] > rewards[1]
|
||||
|
||||
def test_zero_cost_episode(self):
|
||||
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
ep = Episode(
|
||||
task_id="t",
|
||||
initial_prompt="q",
|
||||
correct=True,
|
||||
total_energy_joules=0.0,
|
||||
total_cost_usd=0.0,
|
||||
total_latency_seconds=0.0,
|
||||
max_power_watts=0.0,
|
||||
)
|
||||
r = reward_fn.compute(ep)
|
||||
# Only accuracy component, no penalties
|
||||
assert r == pytest.approx(RewardWeights().alpha)
|
||||
|
||||
|
||||
class TestAdaptiveRewardWeights:
|
||||
def test_at_zero_progress(self):
|
||||
adaptive = AdaptiveRewardWeights(total_steps=10000)
|
||||
w = adaptive.get_weights(0)
|
||||
total = w.alpha + w.beta_cost + w.beta_energy + w.gamma_latency + w.gamma_power
|
||||
assert abs(total - 1.0) < 0.01
|
||||
# At step 0, alpha should be close to initial (highest)
|
||||
assert w.alpha > 0.5
|
||||
|
||||
def test_at_fifty_percent(self):
|
||||
adaptive = AdaptiveRewardWeights(total_steps=10000)
|
||||
w = adaptive.get_weights(5000)
|
||||
total = w.alpha + w.beta_cost + w.beta_energy + w.gamma_latency + w.gamma_power
|
||||
assert abs(total - 1.0) < 0.01
|
||||
|
||||
def test_at_hundred_percent(self):
|
||||
adaptive = AdaptiveRewardWeights(total_steps=10000)
|
||||
w = adaptive.get_weights(10000)
|
||||
total = w.alpha + w.beta_cost + w.beta_energy + w.gamma_latency + w.gamma_power
|
||||
assert abs(total - 1.0) < 0.01
|
||||
# At step 10000, alpha should be lower
|
||||
w0 = adaptive.get_weights(0)
|
||||
assert w.alpha < w0.alpha
|
||||
|
||||
def test_alpha_decreases(self):
|
||||
adaptive = AdaptiveRewardWeights(total_steps=1000)
|
||||
w_start = adaptive.get_weights(0)
|
||||
w_end = adaptive.get_weights(1000)
|
||||
assert w_start.alpha > w_end.alpha
|
||||
|
||||
def test_beyond_total_steps_clamped(self):
|
||||
adaptive = AdaptiveRewardWeights(total_steps=100)
|
||||
w = adaptive.get_weights(200)
|
||||
total = w.alpha + w.beta_cost + w.beta_energy + w.gamma_latency + w.gamma_power
|
||||
assert abs(total - 1.0) < 0.01
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user