| from llm_with_prompt import LLM |
| from typing import List, Dict, Any, Optional |
| import yaml |
| import os |
|
|
| current_dir = os.path.dirname(os.path.abspath(__file__)) |
| import logging |
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
| |
| |
| |
| config = yaml.load(open(os.path.join(current_dir, "..", "config.yaml")), Loader=yaml.FullLoader)[0] |
| model_config = config["model_config"] |
|
|
| USER_PROMPT_TEMPLATE = """ |
| Câu hỏi của sinh viên: |
| {query} |
| |
| Dựa trên các thông tin sau, hãy đưa ra câu trả lời ngắn gọn, chính xác, trích từ nội dung khi cần: |
| {context} |
| """ |
|
|
|
|
| def build_user_prompt(query: str, chunks_retrieved: List[Dict[str, Any]]) -> str: |
| """ |
| Build prompt for the model to generate answer. |
| """ |
| context = [] |
| for i, chunk in enumerate(chunks_retrieved): |
| context.append(f"Chunk {i+1}**: \n {chunk['doc_id']}: \n {chunk['chunk_for_embedding']}") |
| context = "\n\n".join(context) |
| return USER_PROMPT_TEMPLATE.format(query=query, context=context) |
|
|
| class ResponseGenerator: |
| def __init__(self): |
| self.gen_answer_agent = LLM(model_name=model_config["model_name"], prompt_id="gen_answer", version="1.0") |
| |
| def generate(self, query, chunks_retrieved: List[Dict[str, Any]]) -> Dict[str, Any]: |
| """ |
| chunks_retrieved: list of chunk dicts (ranked). We'll take top self.top_k_chunks for prompt building. |
| Returns structured dict: answer_json + grounding metrics + raw model output |
| """ |
| user_prompt = build_user_prompt(query, chunks_retrieved) |
| try: |
| raw_answer = self.gen_answer_agent.generate_response(user_prompt) |
| except Exception as e: |
| logger.exception("Error: %s", e) |
| return self._get_error_message(e) |
| return raw_answer, user_prompt |
| |
| def _get_error_message(self, e: Exception) -> Dict[str, Any]: |
| return { |
| "answer": "Error: " + str(e), |
| "claims": [], |
| "citations": [], |
| "grounding": {"total_claims": 0, "supported_claims": 0, "rate": 0.0}, |
| "raw_model": e |
| } |
| |
| def _safe_not_found(self, raw_model: Optional[Any] = None) -> Dict[str, Any]: |
| return { |
| "answer": "Không tìm thấy quy định", |
| "claims": [], |
| "citations": [], |
| "grounding": {"total_claims": 0, "supported_claims": 0, "rate": 0.0}, |
| "raw_model": raw_model |
| } |
|
|
| if __name__ == "__main__": |
| chunks_retrieved = [ |
| { |
| "id": "11.QD::CH3::A5::C1", |
| "doc_id": "1563/QĐ-ĐHTL", |
| "chunk_text": "Nếu sinh viên nghỉ học dưới 10 ngày tại kỳ chính thì không phải nộp học phí.", |
| "chunk_for_embedding": None, |
| "metadata": {"chapter":3, "article":5, "identifiers":["1563/QĐ-ĐHTL"]} |
| }, |
| { |
| "id": "11.QD::CH3::A5::C2", |
| "doc_id": "1563/QĐ-ĐHTL", |
| "chunk_text": "Nếu nghỉ từ 10 ngày đến dưới 20 ngày tại kỳ chính phải nộp 30% học phí.", |
| "metadata": {"chapter":3, "article":5} |
| } |
| ] |
| query = "Em quên đóng học phí quá hạn có bị cấm bảo vệ khóa luận không ạ ?" |
| gen = ResponseGenerator() |
| out = gen.generate(query, chunks_retrieved) |
