232 lines
8.3 KiB
Python
232 lines
8.3 KiB
Python
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())
|