from __future__ import annotations """AI Quỳnh QS Interactive Training & Experience Memory Engine.""" import asyncio import json import logging import os import re from collections.abc import AsyncGenerator from typing import Any from app.common.sse import format_sse_event from .learning_prompts import build_quynh_fallback_response, build_quynh_system_prompt logger = logging.getLogger("dscons.takeoff.learning_memory") class TakeoffLearningMemoryEngine: """Engine handling interactive training and real-time reasoning for AI Quỳnh QS.""" def __init__(self, erp_client: Any, llm_client: Any): self._erp_client = erp_client self._llm_client = llm_client async def chat_train_ai_quynh( self, takeoff_id: str, user_message: str, current_user_name: str = "Kỹ Sư Trưởng QS", auto_apply: bool = False, ) -> dict[str, Any]: takeoff = self._erp_client.get_drawing_takeoff(takeoff_id) if not takeoff: raise ValueError("Không tìm thấy hồ sơ bóc tách bản vẽ") current_items = self._erp_client.list_drawing_takeoff_items(takeoff_id) or [] drawing_type = takeoff.get("drawing_type") or "general" self._erp_client.save_takeoff_chat_message( takeoff_id=takeoff_id, sender_role="user", sender_name=current_user_name, message_text=user_message, ) items_summary = "\n".join( [ f"- STT {i.get('item_order', idx + 1)}: [{i.get('wbs_code', '')}] [{i.get('norm_code', '')}] {i.get('item_name')} | Nhóm: {i.get('category')} | Công thức: {i.get('dimension_formula')} | ĐVT: {i.get('unit')} | KL: {i.get('quantity')} | Đơn giá: {i.get('unit_price_vnd'):,} đ | Thành tiền: {i.get('total_amount_vnd'):,} đ" for idx, i in enumerate(current_items) ] ) system_prompt = build_quynh_system_prompt() user_content = ( f"BẢN VẼ: {takeoff.get('drawing_code')} - {takeoff.get('drawing_title')}\n" f"LOẠI CÔNG TRÌNH: {drawing_type}\n" f"DANH MỤC BOQ HIỆN TẠI:\n{items_summary}\n\n" f'Ý KIẾN CHỈ ĐẠO CỦA KỸ SƯ ({current_user_name}):\n"{user_message}"\n\n' f"Hãy tiếp thu, hiệu chỉnh danh mục BoQ và đúc kết quy tắc kinh nghiệm theo đúng định dạng JSON." ) parsed_response = None try: llm_timeout = 2.0 if os.environ.get("PYTEST_CURRENT_TEST") else 60.0 resp = await asyncio.wait_for( self._llm_client.create_chat_completion( [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_content}, ], model="google/gemini-3.1-pro-preview", agent_code="quynh", temperature=0.2, max_tokens=32768, extra_body={ "thinking_config": {"thinking_budget": 16384}, "reasoning_effort": "high", }, ), timeout=llm_timeout, ) raw_text = resp["choices"][0]["message"].get("content") or "" json_match = re.search(r"(\{.*\})", raw_text, re.DOTALL) if json_match: parsed_response = json.loads(json_match.group(1)) except Exception as e: logger.warning("[TAKEOFF] LLM chat training fallback: %s", e) if not parsed_response: parsed_response = build_quynh_fallback_response( user_message=user_message, current_user_name=current_user_name, current_items=current_items, drawing_type=drawing_type, ) reply_text = ( parsed_response.get("reply_message") or "Em Quỳnh QS đã tiếp thu và ghi nhận hiệu chỉnh." ) suggested_items = parsed_response.get("suggested_corrections") or [] extracted_rules = parsed_response.get("extracted_rules") or [] sanitized_items = [] for itm in suggested_items: s_itm = { str(k): (str(v) if k in ["id", "takeoff_id"] else v) for k, v in itm.items() } sanitized_items.append(s_itm) try: from app.modules.takeoff.application import get_drawing_takeoff_service takeoff_service = get_drawing_takeoff_service() sanitized_items = takeoff_service.apply_learned_rules_and_state_prices( sanitized_items, drawing_type=drawing_type, takeoff_id=takeoff_id, ) except Exception as norm_err: logger.debug("[LEARNING] Non-critical normalization notice: %s", norm_err) saved_rules = [] for r in extracted_rules: saved_r = self._erp_client.save_takeoff_learned_rule( drawing_type=r.get("drawing_type") or drawing_type, category=r.get("category") or "all", error_pattern=r.get("error_pattern") or user_message[:150], correction_rule=r.get("correction_rule") or user_message, learned_from_user=current_user_name, takeoff_id=takeoff_id, ) saved_rules.append(saved_r) assistant_msg = self._erp_client.save_takeoff_chat_message( takeoff_id=takeoff_id, sender_role="assistant", sender_name="AI Quỳnh (Kỹ Sư Trưởng QS)", message_text=reply_text, suggested_corrections=sanitized_items, extracted_rules=extracted_rules, applied=auto_apply, ) if auto_apply and sanitized_items: self._erp_client.apply_chat_suggested_corrections( takeoff_id=takeoff_id, suggested_items=sanitized_items, message_id=str(assistant_msg["id"]), ) return { "status": "success", "reply_message": reply_text, "ai_response": reply_text, "message_id": str(assistant_msg["id"]), "suggested_corrections": sanitized_items, "suggested_boq_items": sanitized_items, "extracted_rules": extracted_rules, "learned_rule_created": saved_rules[0] if saved_rules else None, "saved_rules_count": len(saved_rules), "applied": auto_apply, "applied_to_takeoff": auto_apply, } async def chat_train_ai_quynh_stream( self, takeoff_id: str, user_message: str, current_user_name: str = "Kỹ Sư Trưởng QS", auto_apply: bool = False, ) -> AsyncGenerator[str, None]: yield format_sse_event( "thinking", { "status": "thinking", "text": "AI Quỳnh đang đối soát bản vẽ kỹ thuật, tra cứu định mức xây dựng 2026 và phân tích chỉ đạo...", }, ) await asyncio.sleep(0.08) res = await self.chat_train_ai_quynh( takeoff_id=takeoff_id, user_message=user_message, current_user_name=current_user_name, auto_apply=auto_apply, ) reply_text = res.get("reply_message") or "" words = reply_text.split(" ") chunk_size = 3 for i in range(0, len(words), chunk_size): chunk = " ".join(words[i : i + chunk_size]) if i + chunk_size < len(words): chunk += " " yield format_sse_event("delta", {"chunk": chunk}) await asyncio.sleep(0.03) yield format_sse_event("complete", res)