llm/scripts/search_1c_rag_hybrid.py

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from common import read_json, search_lexical_index
from rag_profiles import resolve_rag_profile
from search_1c_rag_vector import DEFAULT_CORPUS, DEFAULT_INDEX as DEFAULT_VECTOR_INDEX
from search_1c_rag_vector import search_vector_index


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_LEXICAL_INDEX = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_index.json"


def document_key(document: dict[str, Any]) -> str:
    return str(document.get("id") or f"{document.get('source_path')}#{document.get('chunk_index')}")


def reciprocal_rank(position: int, *, k: int = 60) -> float:
    return 1.0 / float(k + position)


def hybrid_search(
    query: str,
    *,
    lexical_index_path: Path,
    vector_index_path: Path,
    corpus_path: Path,
    limit: int,
    candidate_limit: int,
    source_types: list[str] | None,
    lexical_weight: float = 1.0,
    vector_weight: float = 1.0,
    embedding_base_url: str = "",
    embedding_api_key_env: str = "OPENAI_API_KEY",
) -> dict[str, Any]:
    lexical_index = read_json(lexical_index_path)
    lexical_results = search_lexical_index(
        lexical_index,
        query=query,
        limit=candidate_limit,
        candidate_limit=candidate_limit,
        min_score=0.0,
        source_types=source_types,
    )
    vector_response = search_vector_index(
        vector_index_path,
        query,
        limit=candidate_limit,
        candidate_limit=candidate_limit,
        min_score=-1.0,
        source_types=source_types,
        corpus_path=corpus_path,
        embedding_base_url=embedding_base_url,
        embedding_api_key_env=embedding_api_key_env,
    )
    fused: dict[str, dict[str, Any]] = {}
    for position, item in enumerate(lexical_results, start=1):
        key = document_key(item["document"])
        fused.setdefault(key, {"document": item["document"], "score": 0.0, "channels": {}})
        fused[key]["score"] += lexical_weight * reciprocal_rank(position)
        fused[key]["channels"]["lexical"] = {"rank": position, "score": item["score"]}
    for position, item in enumerate(vector_response.get("results") or [], start=1):
        key = document_key(item["document"])
        fused.setdefault(key, {"document": item["document"], "score": 0.0, "channels": {}})
        fused[key]["score"] += vector_weight * reciprocal_rank(position)
        fused[key]["channels"]["vector"] = {"rank": position, "score": item["score"]}
    results = sorted(fused.values(), key=lambda item: item["score"], reverse=True)[:limit]
    return {
        "schema": "onec_rag_hybrid_search.v1",
        "status": "ok",
        "query": query,
        "indexes": {"lexical": str(lexical_index_path), "vector": str(vector_index_path)},
        "vector_freshness": vector_response.get("freshness"),
        "vector_meta": vector_response.get("meta"),
        "results": results,
        "counts": {
            "results": len(results),
            "lexical_candidates": len(lexical_results),
            "vector_candidates": len(vector_response.get("results") or []),
            "fused_candidates": len(fused),
        },
    }


def print_result(result: dict, index: int) -> None:
    document = result["document"]
    channels = ", ".join(f"{name}#{data['rank']}" for name, data in sorted((result.get("channels") or {}).items()))
    content = (document.get("content") or "").strip().replace("\n", " ")
    if len(content) > 500:
        content = content[:497].rstrip() + "..."
    print(f"{index}. score={result['score']:.4f} channels={channels}")
    print(f"   source={document.get('source_path')} chunk={document.get('chunk_index')}")
    print(f"   title={document.get('title')}")
    print(f"   content={content}")


def main() -> int:
    parser = argparse.ArgumentParser(description="Hybrid search over 1C RAG lexical and vector indexes.")
    parser.add_argument("query")
    parser.add_argument("--lexical-index", type=Path, default=DEFAULT_LEXICAL_INDEX)
    parser.add_argument("--vector-index", type=Path, default=DEFAULT_VECTOR_INDEX)
    parser.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS)
    parser.add_argument("--profile", default="auto")
    parser.add_argument("--limit", type=int)
    parser.add_argument("--candidate-limit", type=int)
    parser.add_argument("--source-type", action="append", dest="source_types")
    parser.add_argument("--lexical-weight", type=float, default=1.0)
    parser.add_argument("--vector-weight", type=float, default=1.0)
    parser.add_argument("--embedding-base-url", default="")
    parser.add_argument("--embedding-api-key-env", default="OPENAI_API_KEY")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    profile = resolve_rag_profile(args.profile, args.query)
    source_types = args.source_types if args.source_types is not None else profile["source_types"]
    result = hybrid_search(
        args.query,
        lexical_index_path=args.lexical_index,
        vector_index_path=args.vector_index,
        corpus_path=args.corpus,
        limit=args.limit or int(profile["limit"]),
        candidate_limit=args.candidate_limit or int(profile["candidate_limit"]),
        source_types=source_types,
        lexical_weight=args.lexical_weight,
        vector_weight=args.vector_weight,
        embedding_base_url=args.embedding_base_url,
        embedding_api_key_env=args.embedding_api_key_env,
    )
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    elif result["results"]:
        print(f"vector_freshness={(result.get('vector_freshness') or {}).get('status')}")
        for position, item in enumerate(result["results"], start=1):
            print_result(item, position)
    else:
        print("No matches.")
    return 0 if (result.get("vector_freshness") or {}).get("status") != "stale" else 1


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