from flask import Flask, render_template, request, jsonify
from functools import lru_cache
import math
import os
from dotenv import load_dotenv
os.environ['CUDA_VISIBLE_DEVICES'] = "2"
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher
import torch
import base64
import io

app = Flask(__name__)

dataset_name = "multiqa_text"

searcher = Searcher(index=f"/data1/liuyaoyang/Papers/icml2025/multi_rag/RAG/Search-in-the-Chain/ColBERT/experiments/{dataset_name}/indexes/{dataset_name}.nbits=2")


def api_search_query(query, sub_query, k=100):
    # pids, ranks, scores = searcher.search_with_mask(query, mask, k=100)
    pids, ranks, scores = searcher.search_with_sentence(query, sub_query[0], k=100)

    pids, ranks, scores = pids[:100], ranks[:100], scores[:100]
    passages = [searcher.collection[pid] for pid in pids]
    probs = [math.exp(score) for score in scores]
    probs = [prob / sum(probs) for prob in probs]
    topk = []
    for pid, rank, score, prob in zip(pids, ranks, scores, probs):
        text = searcher.collection[pid]            
        d = {'text': text, 'pid': pid, 'rank': rank, 'score': score, 'prob': prob}
        topk.append(d)
    topk = list(sorted(topk, key=lambda p: (-1 * p['score'], p['pid'])))
    
    return {
        "topk": topk,
    }

@app.route('/api/search', methods=['POST'])
def handle_search():
    # 获取JSON数据
    data = request.get_json()
    try:
        # 解析参数
        query = data['query']
        sub_query = data['sub_query']
        k = data.get('k', 100)  # 默认值100

        result = api_search_query(query, sub_query, k)
        return jsonify(result)
    
    except Exception as e:
        print(e)
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run("0.0.0.0", 8220)