File size: 6,062 Bytes
638c806 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | from __future__ import annotations
import importlib
import inspect
import json
import sqlite3
import sys
from pathlib import Path
from typing import Any
import pytest
@pytest.fixture()
def temp_db(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:
db_path = tmp_path / "promptboost-test.db"
monkeypatch.setenv("PROMPTBOOST_DB_PATH", str(db_path))
return db_path
@pytest.fixture()
def promptboost_db_path(temp_db: Path) -> Path:
return temp_db
def maybe_import(module_name: str):
try:
return importlib.import_module(module_name)
except ImportError:
return None
def import_any(*module_names: str):
for module_name in module_names:
module = maybe_import(module_name)
if module is not None:
return module
raise ImportError(f"Could not import any module: {module_names}")
def resolve_callable(module_names: tuple[str, ...], func_names: tuple[str, ...]):
for module_name in module_names:
module = maybe_import(module_name)
if module is None:
continue
for func_name in func_names:
candidate = getattr(module, func_name, None)
if callable(candidate):
return candidate
raise LookupError(f"Could not resolve callable {func_names} in {module_names}")
def call_with_supported_kwargs(func, **kwargs):
sig = inspect.signature(func)
if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in sig.parameters.values()):
return func(**{key: value for key, value in kwargs.items() if value is not None})
supported = {
key: value
for key, value in kwargs.items()
if key in sig.parameters and value is not None
}
return func(**supported)
def as_mapping(value: Any) -> dict[str, Any]:
if isinstance(value, dict):
return value
if hasattr(value, "model_dump"):
return value.model_dump()
if hasattr(value, "__dict__"):
return dict(value.__dict__)
return {}
def text_blob(value: Any) -> str:
if isinstance(value, str):
return value.lower()
if hasattr(value, "model_dump"):
value = value.model_dump()
try:
return json.dumps(value, sort_keys=True, default=str).lower()
except TypeError:
return str(value).lower()
def unload_promptboost_modules() -> None:
for module_name in list(sys.modules):
if module_name == "promptboost" or module_name.startswith("promptboost."):
sys.modules.pop(module_name, None)
def install_fake_adapter(monkeypatch: pytest.MonkeyPatch, boosted_prompt: str) -> None:
def fake_boost(
raw_prompt: str = "",
prompt: str = "",
text: str = "",
target_harness: str = "codex",
harness: str | None = None,
target_model: str = "test-model",
model: str | None = None,
**_: Any,
) -> dict[str, Any]:
from promptboost.config import promptboost_db_path
source = raw_prompt or prompt or text
chosen_harness = harness or target_harness
chosen_model = model or target_model
db_path = promptboost_db_path()
db_path.parent.mkdir(parents=True, exist_ok=True)
with sqlite3.connect(db_path) as conn:
conn.execute(
"""
CREATE TABLE IF NOT EXISTS boost_runs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
raw_prompt TEXT NOT NULL,
detected_task_type TEXT NOT NULL,
target_harness TEXT NOT NULL,
target_model TEXT NOT NULL,
adapter_rules_json TEXT NOT NULL,
boosted_prompt TEXT NOT NULL,
no_new_facts_passed INTEGER NOT NULL,
warnings_json TEXT NOT NULL
)
"""
)
cur = conn.execute(
"""
INSERT INTO boost_runs (
raw_prompt,
detected_task_type,
target_harness,
target_model,
adapter_rules_json,
boosted_prompt,
no_new_facts_passed,
warnings_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(
source,
"general_agent_task",
chosen_harness,
chosen_model,
"{}",
boosted_prompt,
1,
"[]",
),
)
return {
"run_id": cur.lastrowid,
"raw_prompt": source,
"boosted_prompt": boosted_prompt,
"target_harness": chosen_harness,
"target_model": chosen_model,
}
for module_name in ("promptboost", "promptboost.core", "promptboost.boost"):
module = maybe_import(module_name)
if module is not None:
monkeypatch.setattr(module, "boost", fake_boost, raising=False)
monkeypatch.setattr(module, "boost_prompt", fake_boost, raising=False)
def sqlite_user_tables(path: Path) -> list[str]:
with sqlite3.connect(path) as conn:
rows = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
).fetchall()
return [row[0] for row in rows]
def sqlite_contains_text(path: Path, needle: str) -> bool:
with sqlite3.connect(path) as conn:
for table in sqlite_user_tables(path):
columns = conn.execute(f"PRAGMA table_info({table})").fetchall()
text_cols = [col[1] for col in columns if "TEXT" in col[2].upper()]
for col in text_cols:
rows = conn.execute(f"SELECT 1 FROM {table} WHERE {col} LIKE ? LIMIT 1", (f"%{needle}%",)).fetchall()
if rows:
return True
return False
|