src/query_rewriter.py

from groq import Groq
import os
from dotenv import load_dotenv

load_dotenv()
_client = Groq(api_key=os.environ["GROQ_API_KEY"])

REWRITE_MODEL = "llama-3.3-70b-versatile"


def rewrite_query(query: str) -> str:
    """Paraphrase câu hỏi để tăng khả năng tìm kiếm."""
    prompt = f"""Bạn là công cụ cải thiện câu truy vấn tìm kiếm tiếng Việt.
Hãy viết lại câu hỏi sau thành một câu truy vấn ngắn gọn, súc tích hơn với các từ khóa quan trọng.
Chỉ trả về câu truy vấn mới, không giải thích.

Câu hỏi gốc: {query}
Câu truy vấn mới:"""

    response = _client.chat.completions.create(
        model=REWRITE_MODEL,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=100,
        temperature=0.3,
    )
    rewritten = response.choices[0].message.content.strip()
    return rewritten if rewritten else query