query.py

import argparse
import os
import re
import sys
from config import TIME_PERIOD_MAP, TIME_PERIOD_CUISINES_MAP, CUISINES
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from opensearchpy import OpenSearchException
from opensearch_client import get_opensearch_client
from requirement_parser import UserIntent, get_user_intent
from sentence_transformers import SentenceTransformer
from typing import Any


def find(keywords, string: str | None, find_all: bool = False):
    if not string:
        return [] if find_all else None
    pattern = "|".join(map(re.escape, keywords))
    if find_all:
        result = re.findall(pattern, string)
        return result if result else []
    result = re.match(pattern, string)
    if result:
        return result.group()
    return None


def normalize_intent(intent: UserIntent) -> UserIntent:
    cuisines = []
    for cuisine in intent.cuisines:
        if cuisine not in CUISINES:
            intent.keywords.append(cuisine)
        else:
            cuisines.append(cuisine)
    intent.cuisines = cuisines

    cuisines = find(TIME_PERIOD_MAP.keys(), intent.cuisines, True)
    normalized_time_period = find(TIME_PERIOD_MAP.keys(), intent.time_period)
    if normalized_time_period:
        cuisines.extend(TIME_PERIOD_CUISINES_MAP[normalized_time_period])
        intent.time_period = normalized_time_period

    if cuisines:
        intent.cuisines = " ".join(sorted(set(cuisines)))

    if intent.price_min and intent.price_max and intent.price_min == intent.price_max:
        intent.price_min *= 0.5
        intent.price_max *= 1.5
    elif intent.price_min and not intent.price_max:
        intent.price_max *= 1.5
    elif not intent.price_min and intent.price_max:
        intent.price_min *= 0.5

    return intent


def build_keyword_query(intent: UserIntent) -> dict[str, Any]:
    should_clauses: list[dict[str, Any]] = []
    filter_clauses: list[dict[str, Any]] = []
    cuisines = find(TIME_PERIOD_MAP.keys(), intent.cuisines, True)

    intent.time_period = find(TIME_PERIOD_MAP.keys(), intent.time_period)
    if intent.time_period:
        cuisines.extend(TIME_PERIOD_CUISINES_MAP[intent.time_period])

    if cuisines:
        should_clauses.append({'terms': {'cuisines': cuisines}})

    if intent.price_min and intent.price_max:
        filter_clauses.append({
            "nested": {
                "path": "dishes",
                "query": {
                    'range': {
                        'dishes.price': {
                            'gte': intent.price_min,
                            'lte': intent.price_max
                        }
                    }
                }
            }})

    if intent.keywords:
        keywords = ' '.join(intent.keywords)
        should_clauses.append({
            'multi_match': {
                'query': keywords,
                'fields': ['name^2', 'cuisines^2']
            }
        })
        should_clauses.append({
            "nested": {
                "path": "dishes",
                "query": {
                    'multi_match': {
                        'query': keywords,
                        'fields': ["dishes.name", "dishes.description"]
                    }
                },
                "inner_hits": {
                    "_source": ["dishes.name", "dishes.description", "dishes.price"],
                    "size": 1
                }
            }
        })

    result = {
        "bool": {
            "should": should_clauses,
            "filter": filter_clauses,
        }
    }
    return result


def get_embedding(passage: str) -> list[list[float]]:
    model = SentenceTransformer('intfloat/multilingual-e5-small')
    return model.encode(f"passage: {passage}", normalize_embeddings=True).tolist()


def build_semantic_like_query(intent: UserIntent, k) -> dict[str, Any]:
    query = {
        "nested": {
            "path": "dishes",
            "query": {
                "knn": {
                    "dishes.embedding": {
                        "vector": get_embedding(' '.join(intent.keywords)),
                        "k": k,
                        "expand_nested_docs": True
                    }
                }
            },
            "inner_hits": {
                "_source": ["dishes.name", "dishes.description", "dishes.price"],
                "size": 1
            }
        }
    }
    return query


def search_once(query_body: dict[str, Any], k: int) -> list[dict[str, Any]]:
    client = get_opensearch_client()
    response = client.search(
        index='stores',
        body={
            "_source": {"excludes": ["dishes"]},
            "size": k,
            "query": query_body
        }
    )
    return response.get("hits", {}).get("hits", [])


def print_search_result(result: list[dict[str, Any]]) -> None:
    for i, item in enumerate(result, start=1):
        print(f"{i}. {item['_source']['name']} | score={item['_score']:.4f}")


def weighted_fuse(
        result_sets: list[list[dict[str, Any]]],
        weights: list[float]
) -> list[dict[str, Any]]:
    fused: dict[str, dict[str, Any]] = {}
    if not result_sets:
        return []

    total_weight = sum(weights) if weights else 0.0
    if total_weight <= 0:
        weights = [1.0 / len(result_sets)] * len(result_sets)
    else:
        weights = [w / total_weight for w in weights]

    for result_set, weight in zip(result_sets, weights):
        scores = [float(hit.get("_score") or 0.0) for hit in result_set]
        if scores:
            min_score = min(scores)
            max_score = max(scores)
        else:
            min_score = 0.0
            max_score = 0.0

        for hit, score in zip(result_set, scores):
            if max_score > min_score:
                normalized = (score - min_score) / (max_score - min_score)
            else:
                normalized = 1.0 if score > 0.0 else 0.0

            doc_id = str(hit.get("_id"))
            src = hit.get("_source", {})
            inner_hits = hit.get("inner_hits", {})

            if doc_id not in fused:
                fused[doc_id] = {
                    "id": doc_id,
                    "score": 0.0,
                    "_source": src,
                    "inner_hits": inner_hits
                }
            else:
                if not fused[doc_id].get("inner_hits") and inner_hits:
                    fused[doc_id]["inner_hits"] = inner_hits

            fused[doc_id]["score"] += weight * normalized

    reranked = list(fused.values())
    reranked.sort(key=lambda x: x["score"], reverse=True)
    return reranked


def build_llm_intro(user_query: str, top_n: list[dict[str, Any]]) -> str:
    if not top_n:
        return "找不到符合條件的店家。"

    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "你是美食推薦助手。請根據使用者的問題，從提供的資料中選出 3 間店，並簡短介紹每間店的 1 個餐點"
            ),
            (
                "human",
                "使用者問題：{query}\n"
                "候選資料(JSON)：{stores}\n"
            ),
        ]
    )

    llm = ChatOpenAI(openai_api_key=os.getenv('OPENAI_API_KEY'), openai_api_base=os.getenv('OPENAI_API_BASE'))
    chain = prompt | llm
    return str(chain.invoke({"query": user_query, "stores": top_n}).content)


def get_must_try_dishes(query: str, k: int):
    intent = get_user_intent(query)
    print(intent)
    keyword_hits = search_once(build_keyword_query(intent), k)
    print("keyword hits:")
    print_search_result(keyword_hits)
    semantic_hits = search_once(build_semantic_like_query(intent, k), k)
    print("semantic hits:")
    print_search_result(semantic_hits)

    fused = weighted_fuse([keyword_hits, semantic_hits], [0.8, 0.2])

    top_n = fused[:k]
    result = []
    for store in top_n:
        dishes = []
        for dish in store.get('inner_hits', {}).get('dishes', {}).get('hits', {}).get('hits', []):
            dishes.append({
                'name': dish['_source']['name'],
                'description': dish['_source']['description']
            })
        result.append({
            'cuisines': store['_source']['cuisines'],
            'name': store['_source']['name'],
            'url': store['_source']['url'],
            'dishes': dishes,
        })
    intro = build_llm_intro(query, result)

    return {
        "intent": intent.model_dump(),
        "top_n": top_n,
        "intro": intro,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="餐點推薦（Hybrid Retrieval with OpenSearch only）")
    parser.add_argument("--query", default="下午茶喝什麼飲料好？仙草不錯，大概 50 元左右", help="使用者問題")
    parser.add_argument("--k", default=6, help="K 值")
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    try:
        result = get_must_try_dishes(args.query, 3 if args.k < 3 else args.k)
    except OpenSearchException as exc:
        print(f"[ERROR] OpenSearch 操作失敗: {exc}", file=sys.stderr)
        return 1

    print(result["intro"])
    print(f"\n--- top{args.k} raw ---")
    for i, item in enumerate(result["top_n"], start=1):
        src = item["_source"]
        print(f"{i}. {src.get('name')} | score={item['score']:.4f}")

    return 0


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