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.
- llm_waterfall/LICENSE +21 -0
- llm_waterfall/__init__.py +53 -0
- llm_waterfall/adapters/__init__.py +36 -0
- llm_waterfall/adapters/anthropic.py +105 -0
- llm_waterfall/adapters/aws_mantle.py +47 -0
- llm_waterfall/adapters/azure_openai.py +71 -0
- llm_waterfall/adapters/base.py +51 -0
- llm_waterfall/adapters/bedrock.py +309 -0
- llm_waterfall/adapters/openai.py +130 -0
- llm_waterfall/classify.py +184 -0
- llm_waterfall/pricing.py +110 -0
- llm_waterfall/py.typed +0 -0
- llm_waterfall/types.py +295 -0
- llm_waterfall/waterfall.py +255 -0
- wmo/__init__.py +38 -0
- wmo/agents/__init__.py +7 -0
- wmo/agents/default.py +29 -0
- wmo/agents/meta.py +55 -0
- wmo/agents/optimizer.py +55 -0
- wmo/agents/project.py +928 -0
- wmo/cli/__init__.py +5 -0
- wmo/cli/agent_session.py +1123 -0
- wmo/cli/app.py +2489 -0
- wmo/cli/e2b_cmds.py +212 -0
- wmo/cli/eval_closed_loop.py +207 -0
- wmo/cli/harness_app.py +1147 -0
- wmo/cli/harness_distill.py +659 -0
- wmo/cli/hosted_session.py +880 -0
- wmo/cli/ingest_cmd.py +165 -0
- wmo/cli/model_roles.py +82 -0
- wmo/cli/platform_cmds.py +372 -0
- wmo/cli/route_app.py +274 -0
- wmo/cli/session_state.py +243 -0
- wmo/cli/ui.py +1107 -0
- wmo/cli/workspace_sync.py +504 -0
- wmo/config/__init__.py +60 -0
- wmo/config/card.py +129 -0
- wmo/config/config.py +367 -0
- wmo/config/dotenv.py +67 -0
- wmo/config/settings.py +128 -0
- wmo/config/store.py +177 -0
- wmo/conftest.py +19 -0
- wmo/connect/__init__.py +88 -0
- wmo/connect/apps.py +78 -0
- wmo/connect/brave.py +284 -0
- wmo/connect/connector.py +79 -0
- wmo/connect/credentials.py +164 -0
- wmo/connect/github.py +321 -0
- wmo/connect/google.py +627 -0
- wmo/connect/notion.py +790 -0
- wmo/connect/oauth.py +461 -0
- wmo/connect/slack.py +555 -0
- wmo/connect/store.py +199 -0
- wmo/connect/types.py +156 -0
- wmo/core/__init__.py +21 -0
- wmo/core/parsing.py +281 -0
- wmo/core/render.py +271 -0
- wmo/core/text.py +40 -0
- wmo/core/types.py +116 -0
- wmo/distill/__init__.py +14 -0
- wmo/distill/agents.py +140 -0
- wmo/distill/config.py +1006 -0
- wmo/distill/cost.py +437 -0
- wmo/distill/data.py +921 -0
- wmo/distill/deadlines.py +254 -0
- wmo/distill/fake_tinker.py +734 -0
- wmo/distill/gate.py +122 -0
- wmo/distill/loop.py +3499 -0
- wmo/distill/renderers.py +399 -0
- wmo/distill/rendering.py +620 -0
- wmo/distill/rollouts.py +726 -0
- wmo/distill/samples.py +195 -0
- wmo/distill/store.py +829 -0
- wmo/distill/teacher.py +714 -0
- wmo/distill/tokens.py +535 -0
- wmo/distill/tracking.py +552 -0
- wmo/distill/tripwire.py +411 -0
- wmo/distill/xtoken/byte_offsets.py +152 -0
- wmo/distill/xtoken/chunks.py +457 -0
- wmo/distill/xtoken/prompt_logprobs.py +475 -0
- wmo/distill/xtoken/teacher_render.py +346 -0
- wmo/engine/__init__.py +28 -0
- wmo/engine/autoconfig.py +367 -0
- wmo/engine/build.py +346 -0
- wmo/engine/demo.py +77 -0
- wmo/engine/eval_suites.py +245 -0
- wmo/engine/grounding.py +491 -0
- wmo/engine/knowledge.py +291 -0
- wmo/engine/loader.py +36 -0
- wmo/engine/play.py +92 -0
- wmo/engine/prompts.py +99 -0
- wmo/engine/replay.py +443 -0
- wmo/engine/reporting.py +58 -0
- wmo/engine/workspace.py +468 -0
- wmo/engine/world_model.py +568 -0
- wmo/env/__init__.py +22 -0
- wmo/env/base.py +121 -0
- wmo/env/closed_loop.py +229 -0
- wmo/env/episode.py +107 -0
- wmo/env/llm_agent.py +93 -0
- wmo/env/scenarios.py +73 -0
- wmo/evals/__init__.py +52 -0
- wmo/evals/agreement.py +110 -0
- wmo/evals/base.py +45 -0
- wmo/evals/closed_loop.py +480 -0
- wmo/evals/failover.py +96 -0
- wmo/evals/gold.py +127 -0
- wmo/evals/grid.py +394 -0
- wmo/evals/grid_plot.py +205 -0
- wmo/evals/harbor/__init__.py +27 -0
- wmo/evals/harbor/agent.py +573 -0
- wmo/evals/harbor/ctrf.py +171 -0
- wmo/evals/harbor/e2b_environment.py +587 -0
- wmo/evals/harbor/e2b_template_policy.py +144 -0
- wmo/evals/harbor/scorer.py +875 -0
- wmo/evals/harbor/tasks.py +140 -0
- wmo/evals/open_loop.py +194 -0
- wmo/evals/tasks.py +53 -0
- wmo/harness/__init__.py +51 -0
- wmo/harness/code_runtime.py +288 -0
- wmo/harness/create.py +1191 -0
- wmo/harness/delta.py +220 -0
- wmo/harness/doc.py +556 -0
- wmo/harness/e2b_ledger.py +342 -0
- wmo/harness/e2b_reap.py +476 -0
- wmo/harness/e2b_sandbox.py +350 -0
- wmo/harness/environment.py +35 -0
- wmo/harness/live_session.py +543 -0
- wmo/harness/mutate.py +343 -0
- wmo/harness/pi_e2b.py +1710 -0
- wmo/harness/pi_entry/entry.ts +268 -0
- wmo/harness/pi_entry/runner_frames.ts +92 -0
- wmo/harness/pi_entry/runner_live.ts +587 -0
- wmo/harness/pi_entry/runner_service.ts +270 -0
- wmo/harness/pi_entry/runner_stdio.ts +374 -0
- wmo/harness/pi_entry/runner_termination.ts +142 -0
- wmo/harness/pi_local.py +262 -0
- wmo/harness/pi_runtime.py +495 -0
- wmo/harness/pi_vendor.py +65 -0
- wmo/harness/population.py +509 -0
- wmo/harness/project_proposer.py +569 -0
- wmo/harness/proposer.py +977 -0
- wmo/harness/runner_link.py +619 -0
- wmo/harness/runtime.py +389 -0
- wmo/harness/scoring.py +247 -0
- wmo/harness/skills.py +116 -0
- wmo/harness/source_tree.py +319 -0
- wmo/harness/store.py +176 -0
- wmo/harness/tools.py +105 -0
- wmo/harness/vendor/manifest.sha256 +58 -0
- wmo/harness/vendor/pi-agent/CHANGELOG.md +556 -0
- wmo/harness/vendor/pi-agent/LICENSE +21 -0
- wmo/harness/vendor/pi-agent/README.md +488 -0
- wmo/harness/vendor/pi-agent/VENDOR.md +39 -0
- wmo/harness/vendor/pi-agent/docs/agent-harness.md +486 -0
- wmo/harness/vendor/pi-agent/docs/durable-harness.md +212 -0
- wmo/harness/vendor/pi-agent/docs/hooks.md +445 -0
- wmo/harness/vendor/pi-agent/docs/models.md +966 -0
- wmo/harness/vendor/pi-agent/docs/observability.md +376 -0
- wmo/harness/vendor/pi-agent/package.json +60 -0
- wmo/harness/vendor/pi-agent/src/agent-loop.ts +748 -0
- wmo/harness/vendor/pi-agent/src/agent.ts +575 -0
- wmo/harness/vendor/pi-agent/src/harness/agent-harness.ts +1029 -0
- wmo/harness/vendor/pi-agent/src/harness/compaction/branch-summarization.ts +261 -0
- wmo/harness/vendor/pi-agent/src/harness/compaction/compaction.ts +747 -0
- wmo/harness/vendor/pi-agent/src/harness/compaction/utils.ts +144 -0
- wmo/harness/vendor/pi-agent/src/harness/env/nodejs.ts +550 -0
- wmo/harness/vendor/pi-agent/src/harness/messages.ts +164 -0
- wmo/harness/vendor/pi-agent/src/harness/prompt-templates.ts +267 -0
- wmo/harness/vendor/pi-agent/src/harness/session/jsonl-repo.ts +177 -0
- wmo/harness/vendor/pi-agent/src/harness/session/jsonl-storage.ts +293 -0
- wmo/harness/vendor/pi-agent/src/harness/session/memory-repo.ts +50 -0
- wmo/harness/vendor/pi-agent/src/harness/session/memory-storage.ts +131 -0
- wmo/harness/vendor/pi-agent/src/harness/session/repo-utils.ts +51 -0
- wmo/harness/vendor/pi-agent/src/harness/session/session.ts +267 -0
- wmo/harness/vendor/pi-agent/src/harness/session/uuid.ts +54 -0
- wmo/harness/vendor/pi-agent/src/harness/skills.ts +375 -0
- wmo/harness/vendor/pi-agent/src/harness/system-prompt.ts +34 -0
- wmo/harness/vendor/pi-agent/src/harness/types.ts +836 -0
- wmo/harness/vendor/pi-agent/src/harness/utils/shell-output.ts +135 -0
- wmo/harness/vendor/pi-agent/src/harness/utils/truncate.ts +344 -0
- wmo/harness/vendor/pi-agent/src/index.ts +44 -0
- wmo/harness/vendor/pi-agent/src/node.ts +2 -0
- wmo/harness/vendor/pi-agent/src/proxy.ts +367 -0
- wmo/harness/vendor/pi-agent/src/types.ts +428 -0
- wmo/harness/vendor/pi-agent/test/agent-loop.test.ts +1351 -0
- wmo/harness/vendor/pi-agent/test/agent.test.ts +699 -0
- wmo/harness/vendor/pi-agent/test/e2e.test.ts +404 -0
- wmo/harness/vendor/pi-agent/test/harness/agent-harness-stream.test.ts +213 -0
- wmo/harness/vendor/pi-agent/test/harness/agent-harness.test.ts +608 -0
- wmo/harness/vendor/pi-agent/test/harness/compaction.test.ts +655 -0
- wmo/harness/vendor/pi-agent/test/harness/nodejs-env.test.ts +321 -0
- wmo/harness/vendor/pi-agent/test/harness/prompt-templates.test.ts +90 -0
- wmo/harness/vendor/pi-agent/test/harness/repo.test.ts +68 -0
- wmo/harness/vendor/pi-agent/test/harness/resource-formatting.test.ts +24 -0
- wmo/harness/vendor/pi-agent/test/harness/session-test-utils.ts +55 -0
- wmo/harness/vendor/pi-agent/test/harness/session-uuid.test.ts +50 -0
- wmo/harness/vendor/pi-agent/test/harness/session.test.ts +156 -0
- wmo/harness/vendor/pi-agent/test/harness/skills.test.ts +116 -0
- wmo/harness/vendor/pi-agent/test/harness/storage.test.ts +299 -0
- wmo/harness/vendor/pi-agent/test/harness/system-prompt.test.ts +66 -0
- wmo/harness/vendor/pi-agent/test/harness/truncate.test.ts +169 -0
- wmo/harness/vendor/pi-agent/test/scratch/simple.ts +72 -0
- wmo/harness/vendor/pi-agent/test/utils/calculate.ts +32 -0
- wmo/harness/vendor/pi-agent/test/utils/get-current-time.ts +46 -0
- wmo/harness/vendor/pi-agent/tsconfig.build.json +13 -0
- wmo/harness/vendor/pi-agent/vitest.config.ts +19 -0
- wmo/harness/vendor/pi-agent/vitest.harness.config.ts +28 -0
- wmo/harness/vendor/vendor_pi.sh +59 -0
- wmo/harness/workspace_patch.py +270 -0
- wmo/ingest/__init__.py +47 -0
- wmo/ingest/adapter.py +72 -0
- wmo/ingest/base.py +114 -0
- wmo/ingest/braintrust.py +339 -0
- wmo/ingest/detect.py +126 -0
- wmo/ingest/langfuse.py +291 -0
- wmo/ingest/langsmith.py +444 -0
- wmo/ingest/mastra.py +330 -0
- wmo/ingest/messages.py +170 -0
- wmo/ingest/normalize.py +679 -0
- wmo/ingest/otel_genai.py +69 -0
- wmo/ingest/otel_writer.py +100 -0
- wmo/ingest/phoenix.py +150 -0
- wmo/ingest/postgres.py +246 -0
- wmo/ingest/posthog.py +320 -0
- wmo/ingest/quality.py +28 -0
- wmo/ingest/stream.py +209 -0
- wmo/ingest/testdata/sample_otlp.json +60 -0
- wmo/ingest/testdata/sample_spans.jsonl +3 -0
- wmo/optimize/__init__.py +25 -0
- wmo/optimize/base.py +143 -0
- wmo/optimize/gepa.py +806 -0
- wmo/optimize/judge.py +262 -0
- wmo/optimize/judge_quality.py +359 -0
- wmo/optimize/knn.py +468 -0
- wmo/optimize/numeric.py +152 -0
- wmo/optimize/outcomes.py +103 -0
- wmo/optimize/policy.py +669 -0
- wmo/optimize/report.py +231 -0
- wmo/optimize/reward.py +129 -0
- wmo/optimize/routing.py +373 -0
- wmo/platform/__init__.py +6 -0
- wmo/platform/auth.py +115 -0
- wmo/platform/client.py +551 -0
- wmo/platform/credentials.py +126 -0
- wmo/platform/transfer.py +158 -0
- wmo/providers/__init__.py +40 -0
- wmo/providers/_bedrock_chat.py +155 -0
- wmo/providers/_openai_common.py +182 -0
- wmo/providers/_responses_common.py +472 -0
- wmo/providers/anthropic.py +134 -0
- wmo/providers/azure_openai.py +296 -0
- wmo/providers/base.py +300 -0
- wmo/providers/bedrock.py +312 -0
- wmo/providers/models.py +205 -0
- wmo/providers/openai.py +143 -0
- wmo/providers/openai_responses.py +240 -0
- wmo/providers/pool.py +170 -0
- wmo/providers/registry.py +73 -0
- wmo/providers/retry.py +151 -0
- wmo/providers/tinker.py +936 -0
- wmo/providers/waterfall.py +336 -0
- wmo/research/__init__.py +81 -0
- wmo/research/ablation.py +133 -0
- wmo/research/concurrency_plot.py +523 -0
- wmo/research/concurrency_run.py +240 -0
- wmo/research/concurrency_scaling.py +270 -0
- wmo/research/gepa_scaling.py +274 -0
- wmo/research/pipeline.py +198 -0
- wmo/research/scaling_split.py +82 -0
- wmo/research/scenario_fidelity.py +198 -0
- wmo/research/scenario_recovery.py +92 -0
- wmo/research/seed_stability.py +90 -0
- wmo/research/trace_scaling.py +348 -0
- wmo/retrieval/__init__.py +6 -0
- wmo/retrieval/embedders.py +105 -0
- wmo/retrieval/leakfree.py +52 -0
- wmo/retrieval/retriever.py +173 -0
- wmo/scenarios/__init__.py +58 -0
- wmo/scenarios/builder.py +152 -0
- wmo/scenarios/mining/__init__.py +27 -0
- wmo/scenarios/mining/clustering.py +171 -0
- wmo/scenarios/mining/facets.py +226 -0
- wmo/scenarios/mining/selection.py +220 -0
- wmo/scenarios/synthesis/__init__.py +6 -0
- wmo/scenarios/synthesis/scenario_set.py +63 -0
- wmo/scenarios/synthesis/synthesizer.py +85 -0
- wmo/scenarios/verification/__init__.py +17 -0
- wmo/scenarios/verification/judge.py +97 -0
- wmo/scenarios/verification/verify.py +135 -0
- wmo/serving/__init__.py +5 -0
- wmo/serving/builds.py +451 -0
- wmo/serving/chat.py +878 -0
- wmo/serving/endpoint_config.py +64 -0
- wmo/serving/savings.py +250 -0
- wmo/serving/server.py +553 -0
- wmo/serving/traces_source.py +206 -0
- wmo/telemetry.py +213 -0
- wmo/tracking/__init__.py +36 -0
- wmo/tracking/clock.py +24 -0
- wmo/tracking/metered.py +125 -0
- wmo/tracking/pricing.py +99 -0
- wmo/tracking/store.py +31 -0
- wmo/tracking/tracker.py +149 -0
- world_model_optimizer-0.2.0.dist-info/METADATA +203 -0
- world_model_optimizer-0.2.0.dist-info/RECORD +308 -0
- world_model_optimizer-0.2.0.dist-info/WHEEL +4 -0
- 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
|
+
)
|
wmo/connect/__init__.py
ADDED
|
@@ -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())
|