lambdagent/agentexample/qaagent67lambda/scripts/search_engine.py

#!/usr/bin/env python3
"""
RAG search engine v2: uses optimized pickle index
Loads in ~6s instead of ~50s, uses 96% less disk
"""
import json
import os
import re
import math
import pickle
import struct
from collections import Counter
from pathlib import Path


try:
    from config import cfg as _app_cfg
    _HAS_CONFIG = True
except ImportError:
    _HAS_CONFIG = False

_index_cache = {}


def _get_config():
    if _HAS_CONFIG:
        return {
            'base_dir': _app_cfg.base_dir,
            'index_file': _app_cfg.index_file,
            'keyword_file': _app_cfg.keyword_file,
        }
    # Fallback: environment variable
    base = Path(os.environ.get('KNOWLEDGE_DIR', str(Path(__file__).parent)))
    return {
        'base_dir': base,
        'index_file': base / 'rag_index.json',
        'keyword_file': base / 'rag_index.keywords.json',
    }


def _load_index():
    if 'chunks' not in _index_cache:
        cfg = _get_config()
        base = cfg['base_dir']

        # Try optimized pickle first, fall back to JSON
        pkl_index = base / 'rag_index.pkl'
        pkl_kw = base / 'rag_keywords.pkl'

        if pkl_index.exists() and pkl_kw.exists():
            with open(pkl_index, 'rb') as f:
                data = pickle.load(f)
            with open(pkl_kw, 'rb') as f:
                kw_raw = pickle.load(f)

            chunks = data['chunks']  # list of (source, chunk_id, text) tuples
            use_bytes = data.get('use_bytes', False)

            # Build lookup
            chunk_lookup = {}
            for i, (source, chunk_id, text) in enumerate(chunks):
                cid = f"{source}#chunk{chunk_id}"
                chunk_lookup[i] = {
                    'id': cid, 'source': source,
                    'chunk_id': chunk_id, 'text': text, 'length': len(text),
                }

            # Convert bytes postings back to int lists on demand
            _index_cache['chunks'] = chunk_lookup
            _index_cache['total'] = len(chunks)
            _index_cache['keywords'] = kw_raw
            _index_cache['use_bytes'] = use_bytes
        else:
            # Fallback to JSON
            json_index = base / 'rag_index.json'
            json_kw = base / 'rag_index.keywords.json'
            with open(json_index, 'r', encoding='utf-8') as f:
                data = json.load(f)
            _index_cache['chunks'] = {c['id']: c for c in data['chunks']}
            _index_cache['total'] = data['total_chunks']
            with open(json_kw, 'r', encoding='utf-8') as f:
                _index_cache['keywords'] = json.load(f)
            _index_cache['use_bytes'] = False

    return _index_cache


def _decode_postings(raw, use_bytes):
    """Decode posting list from bytes or list"""
    if use_bytes and isinstance(raw, bytes):
        n = len(raw) // 2
        return list(struct.unpack(f"<{n}H", raw))
    return raw


def tokenize(text):
    text = text.lower()
    words = []
    chinese = re.findall(r'[\u4e00-\u9fff]+', text)
    for cc in chinese:
        for i in range(len(cc) - 1):
            words.append(cc[i:i+2])
        if len(cc) >= 3:
            for i in range(len(cc) - 2):
                words.append(cc[i:i+3])
    words.extend(re.findall(r'[a-z]+', text))
    return words


def search(query, top_k=5):
    idx = _load_index()
    keywords = idx['keywords']
    chunks = idx['chunks']
    total = idx['total']
    use_bytes = idx.get('use_bytes', False)

    query_tokens = tokenize(query)
    if not query_tokens:
        return []

    scores = Counter()
    for token in query_tokens:
        if token not in keywords:
            continue
        posting_raw = keywords[token]
        posting_list = _decode_postings(posting_raw, use_bytes)
        idf = math.log((total - len(posting_list) + 0.5) / (len(posting_list) + 0.5) + 1)
        for idx_num in posting_list:
            scores[idx_num] += idf

    top = scores.most_common(top_k)
    results = []
    for idx_num, score in top:
        chunk = chunks.get(idx_num)
        if chunk:
            results.append({
                'source': chunk['source'],
                'chunk_id': chunk['chunk_id'],
                'score': round(score, 3),
                'text': chunk['text'],
            })
    return results


if __name__ == '__main__':
    import sys, time
    cfg = _get_config()
    print(f'Config: base={cfg["base_dir"]}')

    t0 = time.time()
    _load_index()
    t1 = time.time()
    print(f'Index loaded in {t1-t0:.1f}s')

    query = sys.argv[1] if len(sys.argv) > 1 else '船舶进出港报告'
    t0 = time.time()
    results = search(query, top_k=3)
    t1 = time.time()
    print(f'Search "{query}" in {(t1-t0)*1000:.0f}ms')
    for i, r in enumerate(results):
        print(f'\n[{i+1}] {r["source"][:50]} (score={r["score"]})')
        print(f'    {r["text"][:150]}...')