fix: write JSON/JSONL output as real lists, reject unsupported formats

The non-Parquet branch still used json.dumps(result, default=str), which
stringifies numpy arrays to their repr() — the same corrupt-output bug
#994 reports, just for .json/.jsonl extensions instead of .parquet.
embed index reads .json/.jsonl via pd.read_json(lines=...) and detects a
vector column by isinstance(val, (list, np.ndarray)); a repr() string
fails that check, so generate→index still breaks for JSON outputs.

- .json/.jsonl now use pandas to_json(orient='records') with real lists
- Unsupported extensions (.txt, .csv, etc.) now raise ClickException
  instead of silently writing JSON text, matching embed index behavior
- Error message corrected: pyarrow is now a core dep, not an extra
This commit is contained in:
Varun Sahni
2026-08-16 15:49:39 +05:30
parent 616f5ca9b9
commit 4b6cc09585
+23 -14
View File
@@ -1709,32 +1709,41 @@ def embed_generate(cli_ctx: CLIContext, input_path: str, model: str,
raise click.ClickException(f"Embeddings module not available: {exc}") from exc
if output:
output_path = Path(output)
suffix = output_path.suffix.lower()
try:
import numpy as np
import pandas as pd
if output_path.suffix.lower() == ".parquet":
arr = np.asarray(result)
if arr.ndim == 1:
arr = arr[np.newaxis, :]
arr = np.asarray(result)
if arr.ndim == 1:
arr = arr[np.newaxis, :]
if arr.ndim != 2:
raise click.ClickException(
f"embed generate --output expects a 1-D or 2-D array, "
f"got {arr.ndim}-D (shape {arr.shape})"
)
rows = [list(row) for row in arr]
if suffix == ".parquet":
# Schema: single 'embedding' column (list[float] per row).
# embed index detects vector columns via
# isinstance(df[c].iloc[0], (list, np.ndarray)).
# Row indices serve as ids: embed index will see ids=None
# but vectors will index correctly regardless.
df = pd.DataFrame({
"embedding": [list(row) for row in arr],
})
df.index.name = "id"
df.index = [str(i) for i in range(len(arr))]
df = pd.DataFrame({"embedding": rows})
df.to_parquet(output_path, index=False)
elif suffix in (".json", ".jsonl"):
df = pd.DataFrame({"embedding": rows})
df.to_json(
output_path,
orient="records",
lines=(suffix == ".jsonl"),
)
else:
output_path.write_text(
json.dumps(result, default=str), encoding="utf-8"
raise click.ClickException(
f"Unsupported output format '{suffix}'. "
"Use .parquet, .json, or .jsonl"
)
except ImportError as exc:
raise click.ClickException(
f"Missing dependency for --output: {exc}. "
f"Install pyarrow/pandas with: pip install semantica[ingest-parquet]"
"Install pyarrow with: pip install pyarrow"
) from exc
_ok(cli_ctx, f"Wrote {output}")
elif _is_json(cli_ctx, local_json):