Initial project import
This commit is contained in:
@@ -0,0 +1,231 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user