Spaces:
Running
Running
File size: 1,576 Bytes
8e8639a |
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 |
from __future__ import annotations
from .types import ContextBudget, ContextPacket, SchemaPack, SchemaTable
from .parse import parse_sqlite_schema_preview
from .budgeter import apply_budget
DEFAULT_CONSTRAINTS = [
"SELECT_ONLY",
"NO_DDL_DML",
"NO_ATTACH_PRAGMA",
"SINGLE_STATEMENT",
"LIMIT_REQUIRED_IF_MISSING",
]
class ContextEngineer:
def __init__(
self,
*,
budget: ContextBudget,
constraints: list[str] | None = None,
) -> None:
self.budget = budget
self.constraints = constraints or DEFAULT_CONSTRAINTS
def build(self, *, schema_preview: str) -> ContextPacket:
raw_tables = parse_sqlite_schema_preview(schema_preview)
tables_sorted = sorted(raw_tables.keys())
tables = {t: SchemaTable(columns=raw_tables[t], fks={}) for t in tables_sorted}
pack = SchemaPack(tables=tables, version="v1")
tables_before = len(pack.tables)
columns_before = sum(len(t.columns) for t in pack.tables.values())
packed, reason = apply_budget(pack, self.budget)
tables_after = len(packed.tables)
columns_after = sum(len(t.columns) for t in packed.tables.values())
return ContextPacket(
schema_pack=packed,
constraints=self.constraints,
db_hints=None,
budget=self.budget,
tables_before=tables_before,
columns_before=columns_before,
tables_after=tables_after,
columns_after=columns_after,
budget_reason=reason,
)
|