Schema and AI-generated columns

"""Build a typed schema, evolve it, then attach AI generators."""

from vectorshift.table import (
    ColumnSpec,
    CompoundFilter,
    FilterCondition,
    NumberFormat,
    NumberKind,
    SelectOption,
    SingleSelectFormat,
    StringFormat,
    Table,
    TableFilterOperator,
    TimestampFormat,
)

# --- Create with seeded columns ---
t = Table.new(
    name="Vendor scorecard",
    columns=[
        ColumnSpec(name="vendor", format=StringFormat()),
        ColumnSpec(
            name="amount",
            format=NumberFormat(number_kind=NumberKind.CURRENCY, decimals=2, symbol="$"),
        ),
        ColumnSpec(
            name="status",
            format=SingleSelectFormat(
                options=[
                    SelectOption(id="open", label="Open", color="#22aa22"),
                    SelectOption(id="done", label="Done", color="#888888"),
                ]
            ),
        ),
    ],
)
print(f"1. Created Table id={t.id} with 3 typed columns:")
for c in t.columns:
    print(f"     - {c.name:10} kind={c.kind}")

# --- Add, rename, then drop columns post-hoc ---
t.add_column(name="due", format=TimestampFormat(format="yyyy-MM-dd"))
print(f"2. Added 'due' (timestamp). Columns: {[c.name for c in t.columns]}")

t.rename_column("status", "current_status")
print(f"3. Renamed 'status' -> 'current_status'.")

# --- Insert some typed rows ---
t.insert_rows([
    {"vendor": "Acme",  "amount": 1200.50, "current_status": "open"},
    {"vendor": "Beta",  "amount": 49.99,   "current_status": "open"},
    {"vendor": "Gamma", "amount": 980.00,  "current_status": "done"},
])

# --- Read with a filter on the string column ---
filt = CompoundFilter.of(
    FilterCondition("vendor", TableFilterOperator.CONTAINS, "a"),
).order_by("vendor", desc=False)
page = t.read_rows(filters=filt)
print(f"4. Filtered (vendor contains 'a'): {page['total']} rows")
for row in page["rows"]:
    v = row["values"]
    print(f"     - {v['vendor']:6} amount={v.get('amount')} status={v.get('current_status')}")

t.delete_column("due")
print(f"5. Deleted 'due'. Columns: {[c.name for c in t.columns]}")

from vectorshift.pipeline import Pipeline
from vectorshift.table import PipelineGenerator

# Fetch a pipeline you've built and deployed via the SDK or platform UI.
# It should accept `vendor_name: str` and `raw_notes: str` inputs and emit
# a `summary: str` output.
summarizer = Pipeline.fetch(name="vendor-summarizer")

# Add a 'notes' column to capture free-text context for each vendor,
# plus a 'summary' column that the pipeline fills.
t.add_column("notes", StringFormat())
t.add_column(
    "summary",
    StringFormat(),
    generation=PipelineGenerator(
        pipeline=summarizer,
        # Keys are the pipeline's input names; values are the table-column
        # names read for each row.
        input_mapping={"vendor_name": "vendor", "raw_notes": "notes"},
        output_name="summary",
    ),
)

# Populate notes on the existing rows.
from vectorshift.table import RelationalOperator
t.update_rows({"notes": "fast shipping, premium pricing"},
              filters=CompoundFilter.of(FilterCondition("vendor", TableFilterOperator.EQ, "Acme")))
t.update_rows({"notes": "slow but cheap"},
              filters=CompoundFilter.of(FilterCondition("vendor", TableFilterOperator.EQ, "Beta")))
t.update_rows({"notes": "reliable mid-tier vendor"},
              filters=CompoundFilter.of(FilterCondition("vendor", TableFilterOperator.EQ, "Gamma")))

# Fill the 'summary' column for every row — synchronously.
task = t.run_and_wait(columns=["summary"], timeout=120)
print(f"6. summary fill: status={task.status} rows_processed={task.rows_processed}")

from vectorshift.agent import Agent
from vectorshift.knowledge_base import KnowledgeBase
from vectorshift.table import AgentGenerator

reviewer = Agent.fetch(name="vendor-risk-reviewer")
policies = KnowledgeBase.fetch(name="vendor-policies")

t.add_column(
    "risk",
    StringFormat(),
    generation=AgentGenerator(
        agent=reviewer,
        knowledge_base=policies,
        instructions="Classify the vendor as low / medium / high risk based on the notes; cite the relevant policy.",
        input_mapping={"vendor_notes": "notes"},
    ),
)

task = t.run_and_wait(columns=["risk"], timeout=240)
print(f"7. risk fill:    status={task.status} rows_processed={task.rows_processed}")

# Cleanup
t.delete()
print("8. Deleted Table.\n\nDone.")