llm/scripts/search_1c_code_vectors.py

from __future__ import annotations

import argparse
import json
from typing import Any

from embed_1c_code_vectors import code_embedding_model_label, embed_pending_code_vectors
from embed_1c_semantic_cache import DEFAULT_ADAPTER_URL, adapter_call
from rag_embedding_providers import LOCAL_HASHING_MODEL, LOCAL_HASHING_PROVIDER, embed_texts, provider_metadata


DEFAULT_QWEN3_CODE_RETRIEVAL_INSTRUCTION = (
    "Given a natural-language software task, retrieve the relevant 1C Enterprise "
    "BSL source-code fragment that implements or explains it"
)


def prepare_query_embedding_text(
    query: str,
    *,
    embedding_model: str,
    query_instruction: str | None,
) -> tuple[str, str]:
    instruction = query_instruction
    if instruction is None and "qwen3-embedding" in str(embedding_model or "").strip().lower():
        instruction = DEFAULT_QWEN3_CODE_RETRIEVAL_INSTRUCTION
    clean_instruction = str(instruction or "").strip()
    if not clean_instruction:
        return query, ""
    return f"Instruct: {clean_instruction}\nQuery:{query}", clean_instruction


def search_code_vectors(
    *,
    adapter_url: str,
    base_id: str,
    query: str,
    limit: int = 10,
    scan_limit: int = 2000,
    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",
    query_instruction: str | None = None,
    embed_pending: bool = False,
    embed_limit: int = 100,
    embed_batch_size: int = 16,
    timeout_seconds: int = 180,
) -> dict[str, Any]:
    refresh = None
    if embed_pending:
        refresh = embed_pending_code_vectors(
            adapter_url=adapter_url,
            base_id=base_id,
            limit=embed_limit,
            batch_size=embed_batch_size,
            embedding_provider=embedding_provider,
            embedding_model=embedding_model,
            dimensions=dimensions,
            embedding_base_url=embedding_base_url,
            embedding_api_key_env=embedding_api_key_env,
            timeout_seconds=timeout_seconds,
        )
    query_embedding_text, applied_instruction = prepare_query_embedding_text(
        query,
        embedding_model=embedding_model,
        query_instruction=query_instruction,
    )
    query_embedding = embed_texts(
        [query_embedding_text],
        provider=embedding_provider,
        model=embedding_model,
        dimensions=dimensions,
        base_url=embedding_base_url,
        api_key_env=embedding_api_key_env,
        timeout_seconds=timeout_seconds,
    )[0]
    stored_model = code_embedding_model_label(
        provider=embedding_provider,
        model=embedding_model,
        dimensions=len(query_embedding),
    )
    result = adapter_call(
        adapter_url,
        "metadata.code_vector.search",
        {
            "base_id": base_id,
            "query": query,
            "query_embedding": query_embedding,
            "embedding_model": stored_model,
            "limit": limit,
            "scan_limit": scan_limit,
            "verify": True,
            "strict": True,
        },
        timeout_seconds=timeout_seconds,
    )
    provider = provider_metadata(
        provider=embedding_provider,
        model=embedding_model,
        dimensions=len(query_embedding),
        base_url=embedding_base_url,
    )
    result["client_embedding"] = {
        "provider": provider.get("embedding_provider"),
        "model": embedding_model,
        "stored_embedding_model": stored_model,
        "dimensions": len(query_embedding),
        "query_instruction": applied_instruction or None,
    }
    if refresh is not None:
        result["embedding_refresh"] = {
            "status": refresh.get("status"),
            "counts": refresh.get("counts") or {},
        }
    return result


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Search BSL code with a local or OpenAI-compatible embedding model."
    )
    parser.add_argument("query")
    parser.add_argument("--adapter-url", default=DEFAULT_ADAPTER_URL)
    parser.add_argument("--base-id", required=True)
    parser.add_argument("--limit", type=int, default=10)
    parser.add_argument("--scan-limit", type=int, default=2000)
    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(
        "--query-instruction",
        default=None,
        help=(
            "Instruction prepended only to the query embedding. "
            "Qwen3 Embedding gets a 1C-code retrieval instruction automatically; "
            "pass an empty value to disable it."
        ),
    )
    parser.add_argument("--embed-pending", action="store_true")
    parser.add_argument("--embed-limit", type=int, default=100)
    parser.add_argument("--embed-batch-size", type=int, default=16)
    parser.add_argument("--timeout-seconds", type=int, default=180)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    result = search_code_vectors(
        adapter_url=args.adapter_url,
        base_id=args.base_id,
        query=args.query,
        limit=args.limit,
        scan_limit=args.scan_limit,
        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,
        query_instruction=args.query_instruction,
        embed_pending=args.embed_pending,
        embed_limit=args.embed_limit,
        embed_batch_size=args.embed_batch_size,
        timeout_seconds=args.timeout_seconds,
    )
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(
            f"status={result.get('status')} matches={len(result.get('matches') or [])} "
            f"model={(result.get('client_embedding') or {}).get('stored_embedding_model')}"
        )
        for position, match in enumerate(result.get("matches") or [], start=1):
            print(
                f"{position}. score={float(match.get('score') or 0):.4f} "
                f"object={match.get('object_ref')} routine={(match.get('chunk') or {}).get('routine_name')}"
            )
    return 0 if result.get("status") in {"ok", "not_found"} else 1


if __name__ == "__main__":
    raise SystemExit(main())