world-model-optimizer 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (308) hide show
  1. llm_waterfall/LICENSE +21 -0
  2. llm_waterfall/__init__.py +53 -0
  3. llm_waterfall/adapters/__init__.py +36 -0
  4. llm_waterfall/adapters/anthropic.py +105 -0
  5. llm_waterfall/adapters/aws_mantle.py +47 -0
  6. llm_waterfall/adapters/azure_openai.py +71 -0
  7. llm_waterfall/adapters/base.py +51 -0
  8. llm_waterfall/adapters/bedrock.py +309 -0
  9. llm_waterfall/adapters/openai.py +130 -0
  10. llm_waterfall/classify.py +184 -0
  11. llm_waterfall/pricing.py +110 -0
  12. llm_waterfall/py.typed +0 -0
  13. llm_waterfall/types.py +295 -0
  14. llm_waterfall/waterfall.py +255 -0
  15. wmo/__init__.py +38 -0
  16. wmo/agents/__init__.py +7 -0
  17. wmo/agents/default.py +29 -0
  18. wmo/agents/meta.py +55 -0
  19. wmo/agents/optimizer.py +55 -0
  20. wmo/agents/project.py +928 -0
  21. wmo/cli/__init__.py +5 -0
  22. wmo/cli/agent_session.py +1123 -0
  23. wmo/cli/app.py +2489 -0
  24. wmo/cli/e2b_cmds.py +212 -0
  25. wmo/cli/eval_closed_loop.py +207 -0
  26. wmo/cli/harness_app.py +1147 -0
  27. wmo/cli/harness_distill.py +659 -0
  28. wmo/cli/hosted_session.py +880 -0
  29. wmo/cli/ingest_cmd.py +165 -0
  30. wmo/cli/model_roles.py +82 -0
  31. wmo/cli/platform_cmds.py +372 -0
  32. wmo/cli/route_app.py +274 -0
  33. wmo/cli/session_state.py +243 -0
  34. wmo/cli/ui.py +1107 -0
  35. wmo/cli/workspace_sync.py +504 -0
  36. wmo/config/__init__.py +60 -0
  37. wmo/config/card.py +129 -0
  38. wmo/config/config.py +367 -0
  39. wmo/config/dotenv.py +67 -0
  40. wmo/config/settings.py +128 -0
  41. wmo/config/store.py +177 -0
  42. wmo/conftest.py +19 -0
  43. wmo/connect/__init__.py +88 -0
  44. wmo/connect/apps.py +78 -0
  45. wmo/connect/brave.py +284 -0
  46. wmo/connect/connector.py +79 -0
  47. wmo/connect/credentials.py +164 -0
  48. wmo/connect/github.py +321 -0
  49. wmo/connect/google.py +627 -0
  50. wmo/connect/notion.py +790 -0
  51. wmo/connect/oauth.py +461 -0
  52. wmo/connect/slack.py +555 -0
  53. wmo/connect/store.py +199 -0
  54. wmo/connect/types.py +156 -0
  55. wmo/core/__init__.py +21 -0
  56. wmo/core/parsing.py +281 -0
  57. wmo/core/render.py +271 -0
  58. wmo/core/text.py +40 -0
  59. wmo/core/types.py +116 -0
  60. wmo/distill/__init__.py +14 -0
  61. wmo/distill/agents.py +140 -0
  62. wmo/distill/config.py +1006 -0
  63. wmo/distill/cost.py +437 -0
  64. wmo/distill/data.py +921 -0
  65. wmo/distill/deadlines.py +254 -0
  66. wmo/distill/fake_tinker.py +734 -0
  67. wmo/distill/gate.py +122 -0
  68. wmo/distill/loop.py +3499 -0
  69. wmo/distill/renderers.py +399 -0
  70. wmo/distill/rendering.py +620 -0
  71. wmo/distill/rollouts.py +726 -0
  72. wmo/distill/samples.py +195 -0
  73. wmo/distill/store.py +829 -0
  74. wmo/distill/teacher.py +714 -0
  75. wmo/distill/tokens.py +535 -0
  76. wmo/distill/tracking.py +552 -0
  77. wmo/distill/tripwire.py +411 -0
  78. wmo/distill/xtoken/byte_offsets.py +152 -0
  79. wmo/distill/xtoken/chunks.py +457 -0
  80. wmo/distill/xtoken/prompt_logprobs.py +475 -0
  81. wmo/distill/xtoken/teacher_render.py +346 -0
  82. wmo/engine/__init__.py +28 -0
  83. wmo/engine/autoconfig.py +367 -0
  84. wmo/engine/build.py +346 -0
  85. wmo/engine/demo.py +77 -0
  86. wmo/engine/eval_suites.py +245 -0
  87. wmo/engine/grounding.py +491 -0
  88. wmo/engine/knowledge.py +291 -0
  89. wmo/engine/loader.py +36 -0
  90. wmo/engine/play.py +92 -0
  91. wmo/engine/prompts.py +99 -0
  92. wmo/engine/replay.py +443 -0
  93. wmo/engine/reporting.py +58 -0
  94. wmo/engine/workspace.py +468 -0
  95. wmo/engine/world_model.py +568 -0
  96. wmo/env/__init__.py +22 -0
  97. wmo/env/base.py +121 -0
  98. wmo/env/closed_loop.py +229 -0
  99. wmo/env/episode.py +107 -0
  100. wmo/env/llm_agent.py +93 -0
  101. wmo/env/scenarios.py +73 -0
  102. wmo/evals/__init__.py +52 -0
  103. wmo/evals/agreement.py +110 -0
  104. wmo/evals/base.py +45 -0
  105. wmo/evals/closed_loop.py +480 -0
  106. wmo/evals/failover.py +96 -0
  107. wmo/evals/gold.py +127 -0
  108. wmo/evals/grid.py +394 -0
  109. wmo/evals/grid_plot.py +205 -0
  110. wmo/evals/harbor/__init__.py +27 -0
  111. wmo/evals/harbor/agent.py +573 -0
  112. wmo/evals/harbor/ctrf.py +171 -0
  113. wmo/evals/harbor/e2b_environment.py +587 -0
  114. wmo/evals/harbor/e2b_template_policy.py +144 -0
  115. wmo/evals/harbor/scorer.py +875 -0
  116. wmo/evals/harbor/tasks.py +140 -0
  117. wmo/evals/open_loop.py +194 -0
  118. wmo/evals/tasks.py +53 -0
  119. wmo/harness/__init__.py +51 -0
  120. wmo/harness/code_runtime.py +288 -0
  121. wmo/harness/create.py +1191 -0
  122. wmo/harness/delta.py +220 -0
  123. wmo/harness/doc.py +556 -0
  124. wmo/harness/e2b_ledger.py +342 -0
  125. wmo/harness/e2b_reap.py +476 -0
  126. wmo/harness/e2b_sandbox.py +350 -0
  127. wmo/harness/environment.py +35 -0
  128. wmo/harness/live_session.py +543 -0
  129. wmo/harness/mutate.py +343 -0
  130. wmo/harness/pi_e2b.py +1710 -0
  131. wmo/harness/pi_entry/entry.ts +268 -0
  132. wmo/harness/pi_entry/runner_frames.ts +92 -0
  133. wmo/harness/pi_entry/runner_live.ts +587 -0
  134. wmo/harness/pi_entry/runner_service.ts +270 -0
  135. wmo/harness/pi_entry/runner_stdio.ts +374 -0
  136. wmo/harness/pi_entry/runner_termination.ts +142 -0
  137. wmo/harness/pi_local.py +262 -0
  138. wmo/harness/pi_runtime.py +495 -0
  139. wmo/harness/pi_vendor.py +65 -0
  140. wmo/harness/population.py +509 -0
  141. wmo/harness/project_proposer.py +569 -0
  142. wmo/harness/proposer.py +977 -0
  143. wmo/harness/runner_link.py +619 -0
  144. wmo/harness/runtime.py +389 -0
  145. wmo/harness/scoring.py +247 -0
  146. wmo/harness/skills.py +116 -0
  147. wmo/harness/source_tree.py +319 -0
  148. wmo/harness/store.py +176 -0
  149. wmo/harness/tools.py +105 -0
  150. wmo/harness/vendor/manifest.sha256 +58 -0
  151. wmo/harness/vendor/pi-agent/CHANGELOG.md +556 -0
  152. wmo/harness/vendor/pi-agent/LICENSE +21 -0
  153. wmo/harness/vendor/pi-agent/README.md +488 -0
  154. wmo/harness/vendor/pi-agent/VENDOR.md +39 -0
  155. wmo/harness/vendor/pi-agent/docs/agent-harness.md +486 -0
  156. wmo/harness/vendor/pi-agent/docs/durable-harness.md +212 -0
  157. wmo/harness/vendor/pi-agent/docs/hooks.md +445 -0
  158. wmo/harness/vendor/pi-agent/docs/models.md +966 -0
  159. wmo/harness/vendor/pi-agent/docs/observability.md +376 -0
  160. wmo/harness/vendor/pi-agent/package.json +60 -0
  161. wmo/harness/vendor/pi-agent/src/agent-loop.ts +748 -0
  162. wmo/harness/vendor/pi-agent/src/agent.ts +575 -0
  163. wmo/harness/vendor/pi-agent/src/harness/agent-harness.ts +1029 -0
  164. wmo/harness/vendor/pi-agent/src/harness/compaction/branch-summarization.ts +261 -0
  165. wmo/harness/vendor/pi-agent/src/harness/compaction/compaction.ts +747 -0
  166. wmo/harness/vendor/pi-agent/src/harness/compaction/utils.ts +144 -0
  167. wmo/harness/vendor/pi-agent/src/harness/env/nodejs.ts +550 -0
  168. wmo/harness/vendor/pi-agent/src/harness/messages.ts +164 -0
  169. wmo/harness/vendor/pi-agent/src/harness/prompt-templates.ts +267 -0
  170. wmo/harness/vendor/pi-agent/src/harness/session/jsonl-repo.ts +177 -0
  171. wmo/harness/vendor/pi-agent/src/harness/session/jsonl-storage.ts +293 -0
  172. wmo/harness/vendor/pi-agent/src/harness/session/memory-repo.ts +50 -0
  173. wmo/harness/vendor/pi-agent/src/harness/session/memory-storage.ts +131 -0
  174. wmo/harness/vendor/pi-agent/src/harness/session/repo-utils.ts +51 -0
  175. wmo/harness/vendor/pi-agent/src/harness/session/session.ts +267 -0
  176. wmo/harness/vendor/pi-agent/src/harness/session/uuid.ts +54 -0
  177. wmo/harness/vendor/pi-agent/src/harness/skills.ts +375 -0
  178. wmo/harness/vendor/pi-agent/src/harness/system-prompt.ts +34 -0
  179. wmo/harness/vendor/pi-agent/src/harness/types.ts +836 -0
  180. wmo/harness/vendor/pi-agent/src/harness/utils/shell-output.ts +135 -0
  181. wmo/harness/vendor/pi-agent/src/harness/utils/truncate.ts +344 -0
  182. wmo/harness/vendor/pi-agent/src/index.ts +44 -0
  183. wmo/harness/vendor/pi-agent/src/node.ts +2 -0
  184. wmo/harness/vendor/pi-agent/src/proxy.ts +367 -0
  185. wmo/harness/vendor/pi-agent/src/types.ts +428 -0
  186. wmo/harness/vendor/pi-agent/test/agent-loop.test.ts +1351 -0
  187. wmo/harness/vendor/pi-agent/test/agent.test.ts +699 -0
  188. wmo/harness/vendor/pi-agent/test/e2e.test.ts +404 -0
  189. wmo/harness/vendor/pi-agent/test/harness/agent-harness-stream.test.ts +213 -0
  190. wmo/harness/vendor/pi-agent/test/harness/agent-harness.test.ts +608 -0
  191. wmo/harness/vendor/pi-agent/test/harness/compaction.test.ts +655 -0
  192. wmo/harness/vendor/pi-agent/test/harness/nodejs-env.test.ts +321 -0
  193. wmo/harness/vendor/pi-agent/test/harness/prompt-templates.test.ts +90 -0
  194. wmo/harness/vendor/pi-agent/test/harness/repo.test.ts +68 -0
  195. wmo/harness/vendor/pi-agent/test/harness/resource-formatting.test.ts +24 -0
  196. wmo/harness/vendor/pi-agent/test/harness/session-test-utils.ts +55 -0
  197. wmo/harness/vendor/pi-agent/test/harness/session-uuid.test.ts +50 -0
  198. wmo/harness/vendor/pi-agent/test/harness/session.test.ts +156 -0
  199. wmo/harness/vendor/pi-agent/test/harness/skills.test.ts +116 -0
  200. wmo/harness/vendor/pi-agent/test/harness/storage.test.ts +299 -0
  201. wmo/harness/vendor/pi-agent/test/harness/system-prompt.test.ts +66 -0
  202. wmo/harness/vendor/pi-agent/test/harness/truncate.test.ts +169 -0
  203. wmo/harness/vendor/pi-agent/test/scratch/simple.ts +72 -0
  204. wmo/harness/vendor/pi-agent/test/utils/calculate.ts +32 -0
  205. wmo/harness/vendor/pi-agent/test/utils/get-current-time.ts +46 -0
  206. wmo/harness/vendor/pi-agent/tsconfig.build.json +13 -0
  207. wmo/harness/vendor/pi-agent/vitest.config.ts +19 -0
  208. wmo/harness/vendor/pi-agent/vitest.harness.config.ts +28 -0
  209. wmo/harness/vendor/vendor_pi.sh +59 -0
  210. wmo/harness/workspace_patch.py +270 -0
  211. wmo/ingest/__init__.py +47 -0
  212. wmo/ingest/adapter.py +72 -0
  213. wmo/ingest/base.py +114 -0
  214. wmo/ingest/braintrust.py +339 -0
  215. wmo/ingest/detect.py +126 -0
  216. wmo/ingest/langfuse.py +291 -0
  217. wmo/ingest/langsmith.py +444 -0
  218. wmo/ingest/mastra.py +330 -0
  219. wmo/ingest/messages.py +170 -0
  220. wmo/ingest/normalize.py +679 -0
  221. wmo/ingest/otel_genai.py +69 -0
  222. wmo/ingest/otel_writer.py +100 -0
  223. wmo/ingest/phoenix.py +150 -0
  224. wmo/ingest/postgres.py +246 -0
  225. wmo/ingest/posthog.py +320 -0
  226. wmo/ingest/quality.py +28 -0
  227. wmo/ingest/stream.py +209 -0
  228. wmo/ingest/testdata/sample_otlp.json +60 -0
  229. wmo/ingest/testdata/sample_spans.jsonl +3 -0
  230. wmo/optimize/__init__.py +25 -0
  231. wmo/optimize/base.py +143 -0
  232. wmo/optimize/gepa.py +806 -0
  233. wmo/optimize/judge.py +262 -0
  234. wmo/optimize/judge_quality.py +359 -0
  235. wmo/optimize/knn.py +468 -0
  236. wmo/optimize/numeric.py +152 -0
  237. wmo/optimize/outcomes.py +103 -0
  238. wmo/optimize/policy.py +669 -0
  239. wmo/optimize/report.py +231 -0
  240. wmo/optimize/reward.py +129 -0
  241. wmo/optimize/routing.py +373 -0
  242. wmo/platform/__init__.py +6 -0
  243. wmo/platform/auth.py +115 -0
  244. wmo/platform/client.py +551 -0
  245. wmo/platform/credentials.py +126 -0
  246. wmo/platform/transfer.py +158 -0
  247. wmo/providers/__init__.py +40 -0
  248. wmo/providers/_bedrock_chat.py +155 -0
  249. wmo/providers/_openai_common.py +182 -0
  250. wmo/providers/_responses_common.py +472 -0
  251. wmo/providers/anthropic.py +134 -0
  252. wmo/providers/azure_openai.py +296 -0
  253. wmo/providers/base.py +300 -0
  254. wmo/providers/bedrock.py +312 -0
  255. wmo/providers/models.py +205 -0
  256. wmo/providers/openai.py +143 -0
  257. wmo/providers/openai_responses.py +240 -0
  258. wmo/providers/pool.py +170 -0
  259. wmo/providers/registry.py +73 -0
  260. wmo/providers/retry.py +151 -0
  261. wmo/providers/tinker.py +936 -0
  262. wmo/providers/waterfall.py +336 -0
  263. wmo/research/__init__.py +81 -0
  264. wmo/research/ablation.py +133 -0
  265. wmo/research/concurrency_plot.py +523 -0
  266. wmo/research/concurrency_run.py +240 -0
  267. wmo/research/concurrency_scaling.py +270 -0
  268. wmo/research/gepa_scaling.py +274 -0
  269. wmo/research/pipeline.py +198 -0
  270. wmo/research/scaling_split.py +82 -0
  271. wmo/research/scenario_fidelity.py +198 -0
  272. wmo/research/scenario_recovery.py +92 -0
  273. wmo/research/seed_stability.py +90 -0
  274. wmo/research/trace_scaling.py +348 -0
  275. wmo/retrieval/__init__.py +6 -0
  276. wmo/retrieval/embedders.py +105 -0
  277. wmo/retrieval/leakfree.py +52 -0
  278. wmo/retrieval/retriever.py +173 -0
  279. wmo/scenarios/__init__.py +58 -0
  280. wmo/scenarios/builder.py +152 -0
  281. wmo/scenarios/mining/__init__.py +27 -0
  282. wmo/scenarios/mining/clustering.py +171 -0
  283. wmo/scenarios/mining/facets.py +226 -0
  284. wmo/scenarios/mining/selection.py +220 -0
  285. wmo/scenarios/synthesis/__init__.py +6 -0
  286. wmo/scenarios/synthesis/scenario_set.py +63 -0
  287. wmo/scenarios/synthesis/synthesizer.py +85 -0
  288. wmo/scenarios/verification/__init__.py +17 -0
  289. wmo/scenarios/verification/judge.py +97 -0
  290. wmo/scenarios/verification/verify.py +135 -0
  291. wmo/serving/__init__.py +5 -0
  292. wmo/serving/builds.py +451 -0
  293. wmo/serving/chat.py +878 -0
  294. wmo/serving/endpoint_config.py +64 -0
  295. wmo/serving/savings.py +250 -0
  296. wmo/serving/server.py +553 -0
  297. wmo/serving/traces_source.py +206 -0
  298. wmo/telemetry.py +213 -0
  299. wmo/tracking/__init__.py +36 -0
  300. wmo/tracking/clock.py +24 -0
  301. wmo/tracking/metered.py +125 -0
  302. wmo/tracking/pricing.py +99 -0
  303. wmo/tracking/store.py +31 -0
  304. wmo/tracking/tracker.py +149 -0
  305. world_model_optimizer-0.2.0.dist-info/METADATA +203 -0
  306. world_model_optimizer-0.2.0.dist-info/RECORD +308 -0
  307. world_model_optimizer-0.2.0.dist-info/WHEEL +4 -0
  308. world_model_optimizer-0.2.0.dist-info/entry_points.txt +2 -0
wmo/config/store.py ADDED
@@ -0,0 +1,177 @@
1
+ """Named world models on disk.
2
+
3
+ The selected project root (`.wmo/` by default) holds named world models under `models/<name>/`.
4
+ Each model directory is a self-contained artifact in the layout `ArtifactPaths` already understands
5
+ (config.toml, prompts/, index/, metrics.json). The store turns names into directories, lists what
6
+ is available, and reads a small summary for `wmo list`.
7
+
8
+ .wmo/ <- writable: where `wmo build` writes
9
+ models/
10
+ tau2-airline/ <- one artifact (config.toml, prompts/, index/, metrics.json)
11
+ retail-bench/ <- another
12
+
13
+ "Filesystem as DB": loading a model is just reading its folder. The default root is the writable
14
+ `.wmo/` directory used by `wmo build`; callers can pass another root such as `examples/<task>` to
15
+ read intentional prebuilt example artifacts from that root's `models/` directory.
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ import json
21
+ import re
22
+ import tomllib
23
+ from pathlib import Path
24
+
25
+ from pydantic import BaseModel, JsonValue
26
+
27
+ from wmo.config.card import ModelCard, load_card
28
+ from wmo.config.config import ARTIFACT_DIR, ArtifactPaths, HarnessConfig
29
+ from wmo.providers.models import resolve_provider_model
30
+
31
+ # The implicit model name used when the user does not pass `--name`.
32
+ DEFAULT_MODEL_NAME = "default"
33
+
34
+ # A safe, filesystem-friendly model name: no path separators, traversal, or leading dot.
35
+ _NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
36
+
37
+
38
+ def validate_name(name: str) -> str:
39
+ """Return `name` if it is a safe single path segment, else raise a friendly ValueError."""
40
+ if not _NAME_RE.match(name) or name in {".", ".."} or "/" in name or "\\" in name:
41
+ raise ValueError(
42
+ f"invalid world model name {name!r}: use letters, digits, '.', '_', '-' "
43
+ "(must start with a letter or digit, no path separators)"
44
+ )
45
+ return name
46
+
47
+
48
+ def normalize_name(name: str) -> str:
49
+ """Dash-join whitespace: 'tau bench' -> 'tau-bench'. Validity stays validate_name's call."""
50
+ return re.sub(r"\s+", "-", name.strip())
51
+
52
+
53
+ class ModelInfo(BaseModel):
54
+ """A one-line summary of a built world model, for `wmo list`."""
55
+
56
+ name: str
57
+ serve_provider: str
58
+ serve_model: str
59
+ held_out_accuracy: float | None = None
60
+ rollouts_used: int | None = None
61
+ frontier_size: int | None = None
62
+
63
+
64
+ class WorldModelStore:
65
+ """Resolves and enumerates named world models under a selected project root."""
66
+
67
+ def __init__(self, root: str | Path = ARTIFACT_DIR) -> None:
68
+ self.root = Path(root)
69
+
70
+ @property
71
+ def models_dir(self) -> Path:
72
+ """The selected root's models dir (`.wmo/models/` by default)."""
73
+ return self.root / "models"
74
+
75
+ def model_dir(self, name: str) -> Path:
76
+ """The artifact directory for `name` (where a build writes; may not exist)."""
77
+ return self.models_dir / validate_name(name)
78
+
79
+ def dir_for(self, name: str) -> Path | None:
80
+ """The artifact dir holding `name`'s config, or None."""
81
+ model_dir = self.model_dir(name)
82
+ if ArtifactPaths(model_dir).config.exists():
83
+ return model_dir
84
+ return None
85
+
86
+ def exists(self, name: str) -> bool:
87
+ return self.dir_for(name) is not None
88
+
89
+ def _names_in(self, models_dir: Path) -> set[str]:
90
+ """Names of every artifact (a dir containing config.toml) directly under `models_dir`."""
91
+ if not models_dir.exists():
92
+ return set()
93
+ return {d.name for d in models_dir.iterdir() if d.is_dir() and (d / "config.toml").exists()}
94
+
95
+ def list_names(self) -> list[str]:
96
+ """Sorted names of every locally built model."""
97
+ return sorted(self._names_in(self.models_dir))
98
+
99
+ def resolve(self, name: str | None) -> Path:
100
+ """Resolve `name` to a model's artifact dir for read commands (serve/demo/play).
101
+
102
+ With an explicit `name`, require it to exist. With `name=None`, fall back to the single
103
+ available model if there is exactly one; else raise, listing the choices.
104
+ """
105
+ if name is not None:
106
+ found = self.dir_for(name)
107
+ if found is None:
108
+ available = self.list_names()
109
+ hint = f" (have: {', '.join(available)})" if available else ""
110
+ raise FileNotFoundError(
111
+ f"no world model named {name!r} under {self.models_dir}{hint}; "
112
+ "run `wmo build --name <name>` first"
113
+ )
114
+ return found
115
+
116
+ names = self.list_names()
117
+ if not names:
118
+ raise FileNotFoundError(
119
+ f"no world models built under {self.models_dir}; "
120
+ "run `wmo build --name <name>` first"
121
+ )
122
+ if len(names) > 1:
123
+ raise ValueError(
124
+ f"multiple world models built ({', '.join(names)}); pass --name to choose one"
125
+ )
126
+ resolved = self.dir_for(names[0])
127
+ assert resolved is not None # name came from list_names(), so it resolves
128
+ return resolved
129
+
130
+ def card(self, name: str) -> ModelCard | None:
131
+ """The model's `card.json` metadata, or None when the model carries no card."""
132
+ model_dir = self.dir_for(name)
133
+ if model_dir is None:
134
+ raise FileNotFoundError(f"no world model named {name!r} under {self.models_dir}")
135
+ return load_card(model_dir)
136
+
137
+ def info(self, name: str) -> ModelInfo:
138
+ """Read a model's config + metrics into a summary (for `wmo list`)."""
139
+ model_dir = self.dir_for(name)
140
+ if model_dir is None:
141
+ raise FileNotFoundError(f"no world model named {name!r}")
142
+ paths = ArtifactPaths(model_dir)
143
+ with paths.config.open("rb") as fh:
144
+ config = HarnessConfig.model_validate(tomllib.load(fh))
145
+ accuracy: float | None = None
146
+ rollouts: int | None = None
147
+ if paths.metrics.exists():
148
+ metrics = json.loads(paths.metrics.read_text(encoding="utf-8"))
149
+ accuracy = _as_float(metrics.get("held_out_accuracy"))
150
+ rollouts = _as_int(metrics.get("rollouts_used"))
151
+ frontier_size: int | None = None
152
+ if paths.frontier.exists():
153
+ frontier = json.loads(paths.frontier.read_text(encoding="utf-8"))
154
+ if isinstance(frontier, list):
155
+ frontier_size = len(frontier)
156
+ serve = config.serve_provider_config()
157
+ model_type = serve.model_type or resolve_provider_model(serve.kind, serve.model).model_type
158
+ return ModelInfo(
159
+ name=name,
160
+ serve_provider=serve.kind.value,
161
+ serve_model=model_type,
162
+ held_out_accuracy=accuracy,
163
+ rollouts_used=rollouts,
164
+ frontier_size=frontier_size,
165
+ )
166
+
167
+ def list_info(self) -> list[ModelInfo]:
168
+ return [self.info(name) for name in self.list_names()]
169
+
170
+
171
+ def _as_float(value: JsonValue) -> float | None:
172
+ # bool is an int subclass; exclude it so a stray `true` doesn't read as 1.0.
173
+ return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
174
+
175
+
176
+ def _as_int(value: JsonValue) -> int | None:
177
+ return int(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
wmo/conftest.py ADDED
@@ -0,0 +1,19 @@
1
+ """Suite-wide fixtures.
2
+
3
+ The default failover chain lives in a developer-local, gitignored `.wmo/fallback.toml`; tests must
4
+ behave identically whether or not the developer running them has one, so the default lookup path is
5
+ pointed at a nonexistent file for every test. Chain tests pass an explicit `path=`.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from pathlib import Path
11
+
12
+ import pytest
13
+
14
+
15
+ @pytest.fixture(autouse=True)
16
+ def _no_local_fallback_chain(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
17
+ monkeypatch.setattr(
18
+ "wmo.providers.waterfall.FALLBACK_CONFIG_PATH", tmp_path / "no-fallback.toml"
19
+ )
@@ -0,0 +1,88 @@
1
+ """Context connectors: fetch content from a service into normalized `ContextItem`s.
2
+
3
+ This is a library, not a CLI. A host (the platform's connector tools) supplies a per-service
4
+ access token and calls a connector's `pull`:
5
+
6
+ from wmo.connect import ConnectorAuth, PullQuery, get_connector
7
+
8
+ items = get_connector("github").pull(
9
+ ConnectorAuth(kind="token", access_token=token),
10
+ PullQuery(query="repo:me/proj is:open", limit=20),
11
+ )
12
+
13
+ A `ContextConnector` owns one service: `pull` (normalized `ContextItem`s) plus `verify` (a cheap
14
+ identity check) and `connect` (an interactive auth flow, for callers that want to acquire tokens
15
+ themselves). Token/OAuth acquisition is the caller's responsibility. Connectors register
16
+ themselves on import and are looked up by name (`get_connector`) or listed (`list_connectors`),
17
+ mirroring `wmo.ingest`.
18
+ """
19
+
20
+ # Import each built-in connector module for its registration side effect so `get_connector(...)`
21
+ # works on package import, mirroring wmo/ingest/__init__.py. All are SDK-free at import
22
+ # time: they talk httpx directly, and notion imports its optional `mcp` SDK (the `connectors`
23
+ # extra) lazily inside the MCP code paths, so a bare install still imports this package.
24
+ from wmo.connect import brave as brave # noqa: F401
25
+ from wmo.connect import github as github # noqa: F401
26
+ from wmo.connect import google as google # noqa: F401
27
+ from wmo.connect import notion as notion # noqa: F401
28
+ from wmo.connect import slack as slack # noqa: F401
29
+ from wmo.connect.apps import EMBEDDED_APPS, get_app
30
+ from wmo.connect.connector import (
31
+ ConnectUI,
32
+ ContextConnector,
33
+ get_connector,
34
+ list_connectors,
35
+ register_connector,
36
+ )
37
+ from wmo.connect.credentials import (
38
+ connectors_path,
39
+ delete_connector_auth,
40
+ list_connected,
41
+ load_connector_auth,
42
+ resolve_env_token,
43
+ save_connector_auth,
44
+ token_env_var,
45
+ token_env_vars,
46
+ )
47
+ from wmo.connect.oauth import (
48
+ OAuthApp,
49
+ ensure_fresh,
50
+ pkce_challenge,
51
+ refresh_auth,
52
+ run_device_flow,
53
+ run_loopback_flow,
54
+ )
55
+ from wmo.connect.store import BundleManifest, ContextStore, render_markdown
56
+ from wmo.connect.types import ConnectError, ConnectorAuth, ContextItem, ItemKind, PullQuery
57
+
58
+ __all__ = [
59
+ "EMBEDDED_APPS",
60
+ "BundleManifest",
61
+ "ConnectError",
62
+ "ConnectUI",
63
+ "ConnectorAuth",
64
+ "ContextConnector",
65
+ "ContextItem",
66
+ "ContextStore",
67
+ "ItemKind",
68
+ "OAuthApp",
69
+ "PullQuery",
70
+ "connectors_path",
71
+ "delete_connector_auth",
72
+ "ensure_fresh",
73
+ "get_app",
74
+ "get_connector",
75
+ "list_connected",
76
+ "list_connectors",
77
+ "load_connector_auth",
78
+ "pkce_challenge",
79
+ "refresh_auth",
80
+ "register_connector",
81
+ "render_markdown",
82
+ "resolve_env_token",
83
+ "run_device_flow",
84
+ "run_loopback_flow",
85
+ "save_connector_auth",
86
+ "token_env_var",
87
+ "token_env_vars",
88
+ ]
wmo/connect/apps.py ADDED
@@ -0,0 +1,78 @@
1
+ """The shared OAuth app registry: embedded endpoint defaults, env-var client credentials.
2
+
3
+ Endpoint configuration per provider is embedded here; client credentials never are. The repo
4
+ policy is all-or-nothing: since some providers' credentials cannot ship (Google and Slack
5
+ token exchanges demand a client secret), none do, and every OAuth connector requires its
6
+ `WMO_<NAME>_CLIENT_ID` (plus `WMO_<NAME>_CLIENT_SECRET` where the provider demands one) when,
7
+ and only when, that connector is used. `get_app` resolves the env credentials over the
8
+ embedded endpoints. A planned hosted path will supply registered apps platform-side instead.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import os
14
+
15
+ from wmo.connect.oauth import OAuthApp
16
+ from wmo.connect.types import ConnectError
17
+
18
+ # Full endpoint config per provider; client ids/secrets come from the env, never the repo.
19
+ EMBEDDED_APPS: dict[str, OAuthApp] = {
20
+ "github": OAuthApp(
21
+ name="github",
22
+ # Expects a GitHub App client id ($WMO_GITHUB_CLIENT_ID): GitHub Apps carry their
23
+ # permissions in the app registration, repo-scoped at install time, so no scopes are
24
+ # sent in the flow, and the device grant needs no client secret.
25
+ client_id="",
26
+ auth_url="https://github.com/login/oauth/authorize",
27
+ token_url="https://github.com/login/oauth/access_token",
28
+ device_url="https://github.com/login/device/code",
29
+ ),
30
+ "google": OAuthApp(
31
+ name="google",
32
+ # Google's Desktop-client token exchange requires the client secret even with PKCE
33
+ # ($WMO_GOOGLE_CLIENT_ID + $WMO_GOOGLE_CLIENT_SECRET). The planned zero-setup path
34
+ # points token_url at a hosted stateless exchange holding a secret platform-side.
35
+ client_id="",
36
+ auth_url="https://accounts.google.com/o/oauth2/v2/auth",
37
+ token_url="https://oauth2.googleapis.com/token",
38
+ extra_auth_params={"access_type": "offline", "prompt": "consent"},
39
+ ),
40
+ "slack": OAuthApp(
41
+ name="slack",
42
+ client_id="",
43
+ auth_url="https://slack.com/oauth/v2/authorize",
44
+ token_url="https://slack.com/api/oauth.v2.access",
45
+ ),
46
+ }
47
+
48
+
49
+ def get_app(name: str) -> OAuthApp:
50
+ """Resolve the OAuth app for `name`, layering env overrides over the embedded defaults.
51
+
52
+ Env vars (name upper-cased, hyphens to underscores): `WMO_<NAME>_CLIENT_ID` and
53
+ `WMO_<NAME>_CLIENT_SECRET`. Set-but-empty values are treated as unset.
54
+
55
+ Raises:
56
+ ConnectError: When `name` is unknown, or when neither layer provides a client id.
57
+ """
58
+ if name not in EMBEDDED_APPS:
59
+ known = ", ".join(sorted(EMBEDDED_APPS))
60
+ raise ConnectError(f"no OAuth app configuration for {name!r}; known apps: {known}")
61
+ prefix = f"WMO_{name.upper().replace('-', '_')}"
62
+ app = EMBEDDED_APPS[name]
63
+ updates: dict[str, str] = {}
64
+ client_id = os.environ.get(f"{prefix}_CLIENT_ID")
65
+ if client_id:
66
+ updates["client_id"] = client_id
67
+ client_secret = os.environ.get(f"{prefix}_CLIENT_SECRET")
68
+ if client_secret:
69
+ updates["client_secret"] = client_secret
70
+ if updates:
71
+ app = app.model_copy(update=updates)
72
+ if not app.client_id:
73
+ raise ConnectError(
74
+ f"no OAuth client id available for {name!r}: set ${prefix}_CLIENT_ID (and "
75
+ f"${prefix}_CLIENT_SECRET if the provider requires one) to use your own OAuth "
76
+ "app; see docs/reference/connect-library.md"
77
+ )
78
+ return app
wmo/connect/brave.py ADDED
@@ -0,0 +1,284 @@
1
+ """Brave Search context connector: API-key auth, web search results pulled as PAGE items.
2
+
3
+ `connect` prefers an env-injected key (the generic `WMO_BRAVE_TOKEN` override, then the
4
+ `BRAVE_SEARCH_API_KEY` deployments already carry for the grounding engine) and falls back to a
5
+ pasted key; `verify` runs one minimal search; `pull` queries the web search endpoint (`--query`
6
+ required, `--target` becomes a `site:` filter) and fetches each result page through the
7
+ grounding engine's SSRF-guarded fetcher, stripped to readable text. A failed or guarded-out
8
+ page fetch degrades that item's body to the search snippet, never aborting the pull.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import logging
14
+ from datetime import UTC, datetime
15
+
16
+ import httpx
17
+
18
+ from wmo.connect.connector import ConnectUI, register_connector
19
+ from wmo.connect.credentials import resolve_env_token
20
+ from wmo.connect.types import (
21
+ ConnectError,
22
+ ConnectorAuth,
23
+ ContextItem,
24
+ ItemKind,
25
+ PullQuery,
26
+ capped,
27
+ opt_str,
28
+ strip_html,
29
+ transport_errors,
30
+ )
31
+ from wmo.core.types import JsonObject
32
+ from wmo.engine.grounding import FetchFn, http_get
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+ _API_BASE = "https://api.search.brave.com"
37
+ _API_HOST = "api.search.brave.com"
38
+ _SEARCH_PATH = "/res/v1/web/search"
39
+ _TIMEOUT_SECONDS = 30.0
40
+ _SOURCE = "brave"
41
+
42
+ _DASHBOARD_URL = "https://api-dashboard.search.brave.com/"
43
+
44
+ # Brave's per-request page-size ceiling for the web search endpoint.
45
+ _PAGE_MAX = 20
46
+
47
+ # Hard cap on results per pull, below PullQuery's default limit of 100: Brave result relevance
48
+ # degrades sharply at deep offsets (and the API caps `offset` at 9), so a pull stops at 50
49
+ # results (3 pages of 20) no matter how large --limit is.
50
+ MAX_RESULTS = 50
51
+
52
+
53
+ def _raise_for_response(response: httpx.Response, *, doing: str) -> None:
54
+ """Turn Brave error statuses into ConnectErrors that say what to do next."""
55
+ status = response.status_code
56
+ if status < 400:
57
+ return
58
+ if status in (401, 403):
59
+ raise ConnectError(
60
+ f"brave search rejected the API key during {doing} (HTTP {status}); check "
61
+ f"$BRAVE_SEARCH_API_KEY (or the stored key) against your subscription at "
62
+ f"{_DASHBOARD_URL}; the key is invalid or expired and a valid key must be supplied"
63
+ )
64
+ if status == 429:
65
+ wait = response.headers.get("Retry-After", "1")
66
+ raise ConnectError(
67
+ f"brave search rate-limited {doing} (HTTP 429): wait {wait}s (Retry-After), then "
68
+ "re-run the command; the free tier allows 1 request/second"
69
+ )
70
+ raise ConnectError(
71
+ f"brave search {doing} failed (HTTP {status}): {response.text[:200]}; "
72
+ "check the query and retry"
73
+ )
74
+
75
+
76
+ def _result_rows(response: httpx.Response) -> list[JsonObject]:
77
+ """The `web.results` rows of one search response ([] on any other shape)."""
78
+ try:
79
+ raw = response.json()
80
+ except ValueError:
81
+ return []
82
+ if not isinstance(raw, dict):
83
+ return []
84
+ web = raw.get("web")
85
+ results = web.get("results") if isinstance(web, dict) else None
86
+ if not isinstance(results, list):
87
+ return []
88
+ return [row for row in results if isinstance(row, dict)]
89
+
90
+
91
+ def _iso_date(value: str, *, field: str) -> str:
92
+ """An ISO-8601 date or datetime as the YYYY-MM-DD form Brave's freshness range expects.
93
+
94
+ Raises:
95
+ ConnectError: When the value is not ISO-8601.
96
+ """
97
+ try:
98
+ parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
99
+ except ValueError:
100
+ raise ConnectError(
101
+ f"could not parse {field}={value!r} as ISO-8601; use YYYY-MM-DD or a full timestamp"
102
+ ) from None
103
+ return parsed.date().isoformat()
104
+
105
+
106
+ def _freshness(since: str | None, until: str | None) -> str | None:
107
+ """The Brave `freshness` range for the query bounds, or None when there is no clean mapping.
108
+
109
+ Both bounds map to `YYYY-MM-DDtoYYYY-MM-DD`; `since` alone ranges to today. `until` alone
110
+ is dropped (Brave has no until-only freshness form).
111
+ """
112
+ if not since:
113
+ if until:
114
+ logger.debug("brave freshness has no until-only form; ignoring until=%s", until)
115
+ return None
116
+ start = _iso_date(since, field="since")
117
+ end = _iso_date(until, field="until") if until else datetime.now(UTC).date().isoformat()
118
+ return f"{start}to{end}"
119
+
120
+
121
+ def _created_at(row: JsonObject) -> str | None:
122
+ """`page_age` (usually ISO) or `age`, normalized to ISO-8601 when parseable, else None."""
123
+ for field in ("page_age", "age"):
124
+ value = opt_str(row.get(field))
125
+ if not value:
126
+ continue
127
+ try:
128
+ parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
129
+ except ValueError:
130
+ continue # relative forms like "2 days ago" carry no usable timestamp
131
+ return parsed.isoformat()
132
+ return None
133
+
134
+
135
+ class BraveConnector:
136
+ """Brave Search connector: keyed web searches with result pages fetched as bodies.
137
+
138
+ Args:
139
+ transport: Injected httpx transport for the search API calls (tests pass
140
+ `httpx.MockTransport`); None means the real network.
141
+ fetch: The page-body fetcher; defaults to the grounding engine's SSRF-guarded
142
+ `http_get` (http(s)-only, public addresses only, redirects re-checked). Tests
143
+ inject a plain function.
144
+ """
145
+
146
+ name = _SOURCE
147
+ label = "Brave Search"
148
+
149
+ def __init__(
150
+ self,
151
+ transport: httpx.BaseTransport | None = None,
152
+ fetch: FetchFn = http_get,
153
+ ) -> None:
154
+ self._transport = transport
155
+ self._fetch = fetch
156
+
157
+ def connect(self, ui: ConnectUI) -> ConnectorAuth:
158
+ """Use the env-injected key when set (never persisted), else prompt for a pasted key.
159
+
160
+ Raises:
161
+ ConnectError: When no key is entered or Brave rejects the key.
162
+ """
163
+ resolved = resolve_env_token(self.name)
164
+ if resolved is not None:
165
+ env_var, token = resolved
166
+ auth = ConnectorAuth(kind="token", access_token=token)
167
+ account = self.verify(auth)
168
+ ui.info(f"using the ${env_var} key ({account})")
169
+ return auth.model_copy(update={"account": account})
170
+ key = ui.prompt_secret(f"Brave Search API key (free at {_DASHBOARD_URL})").strip()
171
+ if not key:
172
+ raise ConnectError(
173
+ f"no Brave Search API key entered; create a free key at {_DASHBOARD_URL} "
174
+ "and supply it as the connection credential"
175
+ )
176
+ auth = ConnectorAuth(kind="token", access_token=key)
177
+ account = self.verify(auth)
178
+ ui.info(f"Brave Search key verified ({account})")
179
+ return auth.model_copy(update={"account": account})
180
+
181
+ def verify(self, auth: ConnectorAuth) -> str:
182
+ """Check the key with one minimal search (`q=wmo, count=1`).
183
+
184
+ Raises:
185
+ ConnectError: When Brave rejects the key (with dashboard guidance) or the call
186
+ fails at the transport level.
187
+ """
188
+ with self._client(auth) as client, transport_errors(_API_HOST):
189
+ response = client.get(_SEARCH_PATH, params={"q": "wmo", "count": "1"})
190
+ _raise_for_response(response, doing="the Brave Search key check")
191
+ return "Brave Search (key valid)"
192
+
193
+ def pull(self, auth: ConnectorAuth, query: PullQuery) -> list[ContextItem]:
194
+ """Search the web and normalize each result into a PAGE item with its fetched body.
195
+
196
+ `query.query` is required (it is the web search itself); `query.target` scopes results
197
+ to one site via a prepended `site:` filter; `query.since`/`query.until` map to Brave's
198
+ `freshness` range. Results paginate by offset in fixed pages (`count` stays constant
199
+ across pages because Brave skips `count * offset` results, so a varying count would
200
+ misalign them) up to `query.limit`, hard-capped at `MAX_RESULTS`.
201
+
202
+ Raises:
203
+ ConnectError: On a missing query, an unparseable since/until, a rejected key, or
204
+ a rate limit.
205
+ """
206
+ if not query.query:
207
+ raise ConnectError(
208
+ "brave pull needs search terms: pass --query '<terms>' "
209
+ "(add --target <domain> to scope results to one site)"
210
+ )
211
+ q = f"site:{query.target} {query.query}" if query.target else query.query
212
+ limit = min(query.limit, MAX_RESULTS)
213
+ if limit <= 0:
214
+ return []
215
+ freshness = _freshness(query.since, query.until)
216
+ count = min(_PAGE_MAX, limit)
217
+ rows: list[JsonObject] = []
218
+ with self._client(auth) as client, transport_errors(_API_HOST):
219
+ offset = 0
220
+ while len(rows) < limit:
221
+ params = {"q": q, "count": str(count), "offset": str(offset)}
222
+ if freshness:
223
+ params["freshness"] = freshness
224
+ response = client.get(_SEARCH_PATH, params=params)
225
+ _raise_for_response(response, doing="the Brave web search")
226
+ page = _result_rows(response)
227
+ rows.extend(page)
228
+ if len(page) < count:
229
+ break # a short page means Brave has no further results
230
+ offset += 1
231
+ items = [self._page_item(row, rank) for rank, row in enumerate(rows[:limit], start=1)]
232
+ logger.debug("pulled %d brave results for %r", len(items), q)
233
+ return items
234
+
235
+ def _client(self, auth: ConnectorAuth) -> httpx.Client:
236
+ """A search API client carrying the subscription-token header."""
237
+ return httpx.Client(
238
+ base_url=_API_BASE,
239
+ headers={
240
+ "Accept": "application/json",
241
+ "X-Subscription-Token": auth.access_token,
242
+ },
243
+ timeout=_TIMEOUT_SECONDS,
244
+ transport=self._transport,
245
+ )
246
+
247
+ def _page_item(self, row: JsonObject, rank: int) -> ContextItem:
248
+ """Normalize one search result, fetching its page for the body."""
249
+ url = opt_str(row.get("url"))
250
+ snippet = opt_str(row.get("description")) or ""
251
+ body, fetch_error = self._page_body(url, snippet)
252
+ metadata: JsonObject = {"rank": rank, "snippet": snippet}
253
+ if fetch_error:
254
+ metadata["fetch_error"] = fetch_error
255
+ return ContextItem(
256
+ id=url or f"result-{rank}",
257
+ source=self.name,
258
+ kind=ItemKind.PAGE,
259
+ title=opt_str(row.get("title")) or url or "(untitled)",
260
+ body=body,
261
+ url=url,
262
+ created_at=_created_at(row),
263
+ metadata=metadata,
264
+ )
265
+
266
+ def _page_body(self, url: str | None, snippet: str) -> tuple[str, str | None]:
267
+ """(body, fetch_error) for one result: the page's readable text, else the snippet.
268
+
269
+ The fetch goes through the injected SSRF-guarded fetcher, so a result URL resolving
270
+ into private address space is refused there; any refusal or transport failure degrades
271
+ to the search snippet with the reason recorded, never aborting the pull.
272
+ """
273
+ if not url:
274
+ return snippet, "result has no url"
275
+ try:
276
+ raw = self._fetch(url, {"Accept": "text/html", "User-Agent": "wmo-connector"})
277
+ except Exception as exc: # noqa: BLE001 - any failed/guarded fetch degrades to the snippet
278
+ reason = str(exc)[:200] or type(exc).__name__
279
+ logger.debug("brave page fetch failed for %s: %s", url, reason)
280
+ return snippet, reason
281
+ return capped(strip_html(raw)), None
282
+
283
+
284
+ register_connector(BraveConnector())