RAG/utils/code_intent.py

"""
代码意图识别模块
用于识别用户问题是否与代码相关，并进行详细的意图分类
支持多种代码问题类型的识别，为后续检索和Prompt生成提供指导
"""
import sys
import os

# 添加项目根目录到 Python 模块搜索路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if project_root not in sys.path:
    sys.path.insert(0, project_root)

import json
import asyncio
from typing import Dict, Any, Optional, List
from enum import Enum
from functools import lru_cache
from loguru import logger
from config import settings
from llama_index.llms.ollama import Ollama
from utils.prompt import INTENT_DETECTION_TEMPLATE


class CodeIntentCategory(Enum):
    """代码意图分类枚举 - 细粒度分类体系"""
    
    # 代码解释与逻辑类 (Existing Code Focus)
    LOGIC_EXPLANATION = "logic_explanation"  # 解释既有代码的底层逻辑
    ENTITY_INTRODUCTION = "entity_introduction"  # 介绍具体的代码实体（函数定义、类属性、API参数）
    CODE_STRUCTURE = "code_structure"  # 询问项目组织
    
    # 代码生成与实现类 (New Code Focus)
    CODE_GENERATION = "code_generation"  # 请求从零编写完整代码或功能块
    BOILERPLATE_IMPLEMENTATION = "boilerplate_implementation"  # 请求提供标准算法/模板
    
    # 调试、优化与理论类
    ERROR_DEBUGGING = "error_debugging"  # 排查 Bug 或异常
    CODE_OPTIMIZATION = "code_optimization"  # 改进既有代码的性能或质量
    ALGORITHM_THEORY = "algorithm_theory"  # 算法原理或复杂度分析
    
    # 非代码类
    GENERAL_TECHNICAL = "general_technical"  # 通用技术咨询
    NON_TECHNICAL = "non_technical"  # 非技术问题
    UNKNOWN = "unknown"  # 未知类型


class PromptTemplateType(Enum):
    """Prompt模板类型枚举"""
    CODE_EXPLANATION = "code_explanation"  # 代码解释模板（逻辑解释、实体介绍、代码结构）
    CODE_GENERATION = "code_generation"  # 代码生成模板（代码生成、模板实现）
    CODE_DEBUGGING = "code_debugging"  # 代码调试模板（错误调试）
    CODE_OPTIMIZATION = "code_optimization"  # 代码优化模板（代码优化）
    ALGORITHM_EXPLANATION = "algorithm_explanation"  # 算法解释模板（算法理论）
    GENERAL_QA = "general_qa"  # 通用问答模板（通用技术咨询、非技术问题）


class CodeIntentResult:
    """代码意图识别结果"""
    
    def __init__(
        self,
        is_code_related: bool,
        category: CodeIntentCategory,
        confidence: float,
        prompt_template_type: PromptTemplateType,
        keywords: List[str],
        reasoning: str,
        requires_code_context: bool,
        suggested_search_terms: List[str],
    ):
        self.is_code_related = is_code_related  # 是否与代码相关（True/False）
        self.category = category  # 代码意图分类
        self.confidence = confidence  # 置信度分数，范围0-1之间
        self.prompt_template_type = prompt_template_type  # Prompt模板类型
        self.keywords = keywords  # 相关关键词列表
        self.reasoning = reasoning  # 解释或理由
        self.requires_code_context = requires_code_context  # 是否需要代码上下文（True/False）
        self.suggested_search_terms = suggested_search_terms  # 建议搜索条款列表
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "is_code_related": self.is_code_related,
            "category": self.category.value,
            "confidence": self.confidence,
            "prompt_template_type": self.prompt_template_type.value,
            "keywords": self.keywords,
            "reasoning": self.reasoning,
            "requires_code_context": self.requires_code_context,
            "suggested_search_terms": self.suggested_search_terms,
        }
    
    def to_json(self) -> str:
        """转换为JSON格式"""
        return json.dumps(self.to_dict(), ensure_ascii=False, indent=2)


class CodeIntentDetector:
    """代码意图检测器"""
    
    def __init__(self, llm: Optional[Ollama] = None):
        """
        初始化代码意图检测器
        
        Args:
            llm: LLM实例，如果为None则使用默认配置
        """
        if llm is None:
            self.llm = Ollama(
                model=settings.OLLAMA_MODEL,
                base_url=settings.OLLAMA_BASE_URL,
                temperature=0.1,  # 低温度确保输出稳定
                request_timeout=1200.0
            )
        else:
            self.llm = llm
        
        logger.info("代码意图检测器初始化完成")
    
    def _build_classification_prompt(self, query: str, history: Optional[str] = None) -> str:
        """
        构建分类Prompt
        
        Args:
            query: 用户问题
            history: 对话历史
            
        Returns:
            str: 分类Prompt
        """
        history_str = "无" if not history else history
        prompt = INTENT_DETECTION_TEMPLATE.format(
            history_str=history_str,
            query=query
        )
        return prompt
    
    def _parse_llm_response(self, response_text: str) -> Optional[Dict[str, Any]]:
        """
        解析LLM响应
        
        Args:
            response_text: LLM响应文本
            
        Returns:
            解析后的字典，如果解析失败返回None
        """
        try:
            response_text = response_text.strip()
            
            # 尝试提取JSON部分
            json_start = response_text.find('{')
            json_end = response_text.rfind('}')
            
            if json_start == -1 or json_end == -1:
                logger.warning(f"未找到JSON格式响应: {response_text}")
                return None
            
            json_str = response_text[json_start:json_end + 1]
            result = json.loads(json_str)
            
            # 验证必需字段并使用默认值
            required_fields = {
                'is_code_related': False,
                'category': 'unknown',
                'confidence': 0.5,
                'keywords': [],
                'reasoning': '',
                'requires_code_context': False,
                'suggested_search_terms': []
            }
            
            for field, default_value in required_fields.items():
                if field not in result:
                    logger.warning(f"缺少必需字段: {field}, 使用默认值: {default_value}")
                    result[field] = default_value
            
            return result
            
        except json.JSONDecodeError as e:
            logger.error(f"JSON解析失败: {e}, 响应: {response_text}")
            return None
        except Exception as e:
            logger.error(f"解析响应失败: {e}")
            return None
    
    def _map_to_enums(
        self,
        parsed_result: Dict[str, Any]
    ) -> tuple[CodeIntentCategory, PromptTemplateType]:
        """
        将解析结果映射到枚举类型
        
        Args:
            parsed_result: 解析后的结果字典
            
        Returns:
            (CodeIntentCategory, PromptTemplateType)
        """
        category_str = parsed_result.get('category', 'unknown')
        
        try:
            category = CodeIntentCategory(category_str)
        except ValueError:
            logger.warning(f"未知的分类: {category_str}, 使用默认值")
            category = CodeIntentCategory.UNKNOWN
        
        # 根据分类确定Prompt模板类型
        if not parsed_result.get('is_code_related', False):
            prompt_template_type = PromptTemplateType.GENERAL_QA
        
        # 代码解释与逻辑类 -> CODE_EXPLANATION
        elif category in [CodeIntentCategory.LOGIC_EXPLANATION, CodeIntentCategory.ENTITY_INTRODUCTION,
                       CodeIntentCategory.CODE_STRUCTURE]:
            prompt_template_type = PromptTemplateType.CODE_EXPLANATION
        
        # 代码生成与实现类 -> CODE_GENERATION
        elif category in [CodeIntentCategory.CODE_GENERATION, CodeIntentCategory.BOILERPLATE_IMPLEMENTATION]:
            prompt_template_type = PromptTemplateType.CODE_GENERATION
        
        # 调试、优化与理论类
        elif category == CodeIntentCategory.ERROR_DEBUGGING:
            prompt_template_type = PromptTemplateType.CODE_DEBUGGING
        elif category == CodeIntentCategory.CODE_OPTIMIZATION:
            prompt_template_type = PromptTemplateType.CODE_OPTIMIZATION
        elif category == CodeIntentCategory.ALGORITHM_THEORY:
            prompt_template_type = PromptTemplateType.ALGORITHM_EXPLANATION
        
        # 非代码问题 -> GENERAL_QA
        elif category in [CodeIntentCategory.GENERAL_TECHNICAL, CodeIntentCategory.NON_TECHNICAL,
                       CodeIntentCategory.UNKNOWN]:
            prompt_template_type = PromptTemplateType.GENERAL_QA
        
        else:
            prompt_template_type = PromptTemplateType.GENERAL_QA

        return category, prompt_template_type
    
    async def detect_intent_async(self, query: str, history: Optional[str] = None) -> CodeIntentResult:
        """
        异步检测代码意图
        
        Args:
            query: 用户问题
            history: 对话历史
            
        Returns:
            CodeIntentResult: 意图识别结果
        """
        try:
            # 构建Prompt
            prompt = self._build_classification_prompt(query, history)
            # 调用LLM
            response = await self.llm.acomplete(prompt)
            response_text = response.text
            logger.debug(f"LLM响应: {response_text}")
            
            # 解析响应
            parsed_result = self._parse_llm_response(response_text)
            
            if parsed_result is None:
                logger.error("解析LLM响应失败，使用默认分类")
                return self._get_default_result(query)
            
            # 映射到枚举类型
            category, prompt_template_type = self._map_to_enums(parsed_result)
            
            # 构建结果对象
            result = CodeIntentResult(
                is_code_related=parsed_result.get('is_code_related', False),
                category=category,
                confidence=parsed_result.get('confidence', 0.5),
                prompt_template_type=prompt_template_type,
                keywords=parsed_result.get('keywords', []),
                reasoning=parsed_result.get('reasoning', ''),
                requires_code_context=parsed_result.get('requires_code_context', False),
                suggested_search_terms=parsed_result.get('suggested_search_terms', [])
            )
            
            logger.info(f"代码意图检测完成: {result.to_json()}")
            return result
            
        except Exception as e:
            logger.error(f"代码意图检测失败: {e}")
            return self._get_default_result(query)
    
    def detect_intent(self, query: str) -> CodeIntentResult:
        """
        同步检测代码意图（包装异步方法）
        
        Args:
            query: 用户问题
            
        Returns:
            CodeIntentResult: 意图识别结果
        """
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        
        return loop.run_until_complete(self.detect_intent_async(query))
    
    def _get_default_result(self, query: str) -> CodeIntentResult:
        """
        获取默认结果（当检测失败时使用）
        
        Args:
            query: 用户问题
            
        Returns:
            CodeIntentResult: 默认结果
        """
        logger.warning(f"使用默认意图识别结果: {query}")
        
        return CodeIntentResult(
            is_code_related=False,
            category=CodeIntentCategory.UNKNOWN,
            confidence=0.0,
            prompt_template_type=PromptTemplateType.GENERAL_QA,
            keywords=[],
            reasoning="意图识别失败，使用默认分类",
            requires_code_context=False,
            suggested_search_terms=[query]
        )


@lru_cache(maxsize=100)
def detect_code_intent_cached(query: str) -> CodeIntentResult:
    """
    带缓存的代码意图检测（同步版本）
    
    Args:
        query: 用户问题
        
    Returns:
        CodeIntentResult: 意图识别结果
    """
    detector = CodeIntentDetector()
    return detector.detect_intent(query)


async def detect_code_intent_async_cached(query: str) -> CodeIntentResult:
    """
    带缓存的代码意图检测（异步版本）
    
    Args:
        query: 用户问题
        
    Returns:
        CodeIntentResult: 意图识别结果
    """
    detector = CodeIntentDetector()
    return await detector.detect_intent_async(query)


def create_intent_detector(llm: Optional[Ollama] = None) -> CodeIntentDetector:
    """
    创建代码意图检测器实例
    
    Args:
        llm: LLM实例，如果为None则使用默认配置
        
    Returns:
        CodeIntentDetector: 检测器实例
    """
    return CodeIntentDetector(llm)


if __name__ == "__main__":
    import asyncio
    
    async def test_intent_detection():
        """测试代码意图检测"""
        detector = CodeIntentDetector()
        
        test_queries = [
            "kth_number怎么实现的",
            "kth_number的时间复杂度是多少",
            "kth_number的测试用例有哪些",
            "kth_number的代码质量如何",
            "trieste实现的采集函数有哪些",
            "这个函数是做什么的？",
            "如何优化这段代码？",
            "Python中如何实现异步？",
            "今天天气怎么样？",
            "为什么会报这个错误？",
            "项目结构是怎样的？",
            "如何调用这个API？"
        ]
        
        for query in test_queries:
            print(f"\n{'='*60}")
            print(f"问题: {query}")
            print('='*60)
            
            result = await detector.detect_intent_async(query)
            print(result.to_json())
    
    asyncio.run(test_intent_detection())