from __future__ import annotations import argparse import json from typing import Any from embed_1c_semantic_cache import DEFAULT_ADAPTER_URL, adapter_call, batched, embedding_model_label from rag_embedding_providers import LOCAL_HASHING_MODEL, LOCAL_HASHING_PROVIDER, embed_texts, provider_metadata def code_embedding_model_label(*, provider: str, model: str, dimensions: int) -> str: label = embedding_model_label(provider=provider, model=model) normalized_provider = str(provider or "").strip().lower().replace("_", "-") if normalized_provider in {"openai-compatible", "openai"} and int(dimensions or 0) > 0: return f"{label}@d{int(dimensions)}" return label def embed_pending_code_vectors( *, adapter_url: str, base_id: str, limit: int = 100, batch_size: int = 16, embedding_provider: str = LOCAL_HASHING_PROVIDER, embedding_model: str = LOCAL_HASHING_MODEL, dimensions: int = 384, embedding_base_url: str = "", embedding_api_key_env: str = "OPENAI_API_KEY", chunk_kinds: tuple[str, ...] | list[str] = ("routine",), max_text_chars: int = 4000, dry_run: bool = False, timeout_seconds: int = 180, ) -> dict[str, Any]: stored_model = code_embedding_model_label( provider=embedding_provider, model=embedding_model, dimensions=dimensions, ) pending = adapter_call( adapter_url, "metadata.code_vector.pending", { "base_id": base_id, "embedding_model": stored_model, "limit": int(limit or 100), "chunk_kinds": list(dict.fromkeys(str(value).strip().lower() for value in chunk_kinds if str(value).strip())), "max_text_chars": int(max_text_chars), }, timeout_seconds=timeout_seconds, ) if pending.get("status") != "ok": return { "schema": "onec_code_vector_embedding_worker.v1", "status": pending.get("status") or "error", "error": pending.get("error"), "pending": pending, } chunks = [item for item in pending.get("chunks") or [] if isinstance(item, dict)] upserts: list[dict[str, Any]] = [] skipped: list[dict[str, Any]] = [] for batch in batched(chunks, max(int(batch_size or 1), 1)): texts = [str(item.get("text") or "") for item in batch] vectors = embed_texts( texts, provider=embedding_provider, model=embedding_model, dimensions=dimensions, base_url=embedding_base_url, api_key_env=embedding_api_key_env, timeout_seconds=timeout_seconds, ) for item, vector in zip(batch, vectors): chunk_id = str(item.get("chunk_id") or "") text_sha1 = str(item.get("text_sha1") or "") if not chunk_id or not text_sha1 or not vector: skipped.append( { "chunk_id": chunk_id or None, "reason": "missing_chunk_id_text_sha1_or_embedding", } ) continue if dry_run: upserts.append( { "status": "dry_run", "chunk_id": chunk_id, "text_sha1": text_sha1, "dimensions": len(vector), } ) continue result = adapter_call( adapter_url, "metadata.code_vector.embedding.upsert", { "base_id": base_id, "chunk_id": chunk_id, "text_sha1": text_sha1, "embedding_model": stored_model, "embedding": vector, }, timeout_seconds=timeout_seconds, ) upserts.append( { "status": result.get("status"), "error": result.get("error"), "chunk_id": chunk_id, "text_sha1": text_sha1, "dimensions": result.get("dimensions") or len(vector), } ) observed_dimensions = next( ( int(item.get("dimensions") or 0) for item in upserts if int(item.get("dimensions") or 0) > 0 ), int(dimensions), ) provider = provider_metadata( provider=embedding_provider, model=embedding_model, dimensions=observed_dimensions, base_url=embedding_base_url, ) return { "schema": "onec_code_vector_embedding_worker.v1", "status": "ok", "base_id": base_id, "adapter_url": adapter_url, "dry_run": bool(dry_run), "embedding": { "provider": provider.get("embedding_provider"), "model": embedding_model, "stored_embedding_model": stored_model, "dimensions": observed_dimensions, "chunk_kinds": list(chunk_kinds), "max_text_chars": int(max_text_chars), }, "counts": { "pending": len(chunks), "processed": len(upserts), "stored": len([item for item in upserts if item.get("status") == "ok"]), "conflicts": len([item for item in upserts if item.get("status") == "conflict"]), "skipped": len(skipped), "errors": len( [ item for item in upserts if item.get("status") not in {"ok", "dry_run", "conflict"} ] ), }, "upserts": upserts, **({"skipped": skipped} if skipped else {}), } def main() -> int: parser = argparse.ArgumentParser( description="Embed pending BSL code chunks from the 1C adapter local code index." ) parser.add_argument("--adapter-url", default=DEFAULT_ADAPTER_URL) parser.add_argument("--base-id", required=True) parser.add_argument("--limit", type=int, default=100) parser.add_argument("--batch-size", type=int, default=16) parser.add_argument( "--embedding-provider", default=LOCAL_HASHING_PROVIDER, choices=[LOCAL_HASHING_PROVIDER, "openai-compatible"], ) parser.add_argument("--embedding-model", default=LOCAL_HASHING_MODEL) parser.add_argument("--dimensions", type=int, default=384) parser.add_argument("--embedding-base-url", default="") parser.add_argument("--embedding-api-key-env", default="OPENAI_API_KEY") parser.add_argument( "--chunk-kind", action="append", choices=["routine", "module"], default=None, help="Chunk kind to embed; repeat to include both. Defaults to routine.", ) parser.add_argument( "--max-text-chars", type=int, default=4000, help="Skip oversized chunks in this pass. Defaults to 4000 characters.", ) parser.add_argument("--timeout-seconds", type=int, default=180) parser.add_argument("--dry-run", action="store_true") parser.add_argument("--json", action="store_true") args = parser.parse_args() result = embed_pending_code_vectors( adapter_url=args.adapter_url, base_id=args.base_id, limit=args.limit, batch_size=args.batch_size, embedding_provider=args.embedding_provider, embedding_model=args.embedding_model, dimensions=args.dimensions, embedding_base_url=args.embedding_base_url, embedding_api_key_env=args.embedding_api_key_env, chunk_kinds=tuple(args.chunk_kind or ["routine"]), max_text_chars=args.max_text_chars, dry_run=args.dry_run, timeout_seconds=args.timeout_seconds, ) if args.json: print(json.dumps(result, ensure_ascii=False, indent=2)) else: counts = result.get("counts") or {} print( "code vector embeddings: " f"pending={counts.get('pending')} processed={counts.get('processed')} " f"stored={counts.get('stored')} conflicts={counts.get('conflicts')} " f"errors={counts.get('errors')}" ) return ( 0 if result.get("status") == "ok" and int((result.get("counts") or {}).get("errors") or 0) == 0 else 1 ) if __name__ == "__main__": raise SystemExit(main())