419 lines
16 KiB
Python
419 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
查询处理器
|
||
整合意图识别、filter生成和查询转换功能，使用一次LLM调用生成所有结果
|
||
"""
|
||
import json
|
||
from typing import Dict, Any, Optional, List
|
||
from loguru import logger
|
||
from config import settings
|
||
from llama_index.llms.ollama import Ollama
|
||
|
||
|
||
class CodeIntentCategory:
|
||
"""代码意图分类枚举 - 细粒度分类体系"""
|
||
|
||
# 代码解释与逻辑类 (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:
|
||
"""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: str,
|
||
confidence: float,
|
||
prompt_template_type: str,
|
||
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,
|
||
"confidence": self.confidence,
|
||
"prompt_template_type": self.prompt_template_type,
|
||
"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 QueryProcessor:
|
||
"""
|
||
查询处理器
|
||
整合意图识别、filter生成和查询转换功能
|
||
"""
|
||
|
||
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_integrated_prompt(self, query: str, history: Optional[str] = None) -> str:
|
||
"""
|
||
构建集成Prompt，一次调用生成所有结果
|
||
|
||
Args:
|
||
query: 用户问题
|
||
history: 对话历史
|
||
|
||
Returns:
|
||
str: 集成Prompt
|
||
"""
|
||
from utils.prompt.integrated_query_processing import INTEGRATED_QUERY_PROCESSING_TEMPLATE
|
||
|
||
history_str = "无" if not history else history
|
||
prompt = INTEGRATED_QUERY_PROCESSING_TEMPLATE.format(
|
||
history_str=history_str,
|
||
query=query
|
||
)
|
||
|
||
return prompt
|
||
|
||
def _map_to_prompt_template_type(self, category: str) -> str:
|
||
"""
|
||
根据分类确定Prompt模板类型
|
||
|
||
Args:
|
||
category: 意图分类
|
||
|
||
Returns:
|
||
str: Prompt模板类型
|
||
"""
|
||
# 代码解释与逻辑类 -> CODE_EXPLANATION
|
||
if category in [CodeIntentCategory.LOGIC_EXPLANATION, CodeIntentCategory.ENTITY_INTRODUCTION, CodeIntentCategory.CODE_STRUCTURE]:
|
||
return PromptTemplateType.CODE_EXPLANATION
|
||
|
||
# 代码生成与实现类 -> CODE_GENERATION
|
||
elif category in [CodeIntentCategory.CODE_GENERATION, CodeIntentCategory.BOILERPLATE_IMPLEMENTATION]:
|
||
return PromptTemplateType.CODE_GENERATION
|
||
|
||
# 调试、优化与理论类
|
||
elif category == CodeIntentCategory.ERROR_DEBUGGING:
|
||
return PromptTemplateType.CODE_DEBUGGING
|
||
elif category == CodeIntentCategory.CODE_OPTIMIZATION:
|
||
return PromptTemplateType.CODE_OPTIMIZATION
|
||
elif category == CodeIntentCategory.ALGORITHM_THEORY:
|
||
return PromptTemplateType.ALGORITHM_EXPLANATION
|
||
|
||
# 非代码问题 -> GENERAL_QA
|
||
elif category in [CodeIntentCategory.GENERAL_TECHNICAL, CodeIntentCategory.NON_TECHNICAL, CodeIntentCategory.UNKNOWN]:
|
||
return PromptTemplateType.GENERAL_QA
|
||
|
||
else:
|
||
return PromptTemplateType.GENERAL_QA
|
||
|
||
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]
|
||
logger.debug(f"提取的JSON字符串: {json_str}")
|
||
result = json.loads(json_str)
|
||
|
||
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 process_query(self, query: str, history: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
处理查询，一次调用生成所有结果
|
||
|
||
Args:
|
||
query: 用户问题
|
||
history: 对话历史
|
||
|
||
Returns:
|
||
包含意图识别、filter生成和查询转换结果的字典
|
||
"""
|
||
try:
|
||
# 构建集成Prompt
|
||
prompt = self._build_integrated_prompt(query, history)
|
||
|
||
# 调用LLM
|
||
response = self.llm.complete(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)
|
||
|
||
# 验证并处理结果
|
||
result = {
|
||
"intent": None,
|
||
"filters": {},
|
||
"transformed": {
|
||
"rewritten": query,
|
||
"backward": query,
|
||
"sub_queries": [query]
|
||
}
|
||
}
|
||
|
||
# 处理意图识别结果
|
||
try:
|
||
if "intent" in parsed_result:
|
||
intent_data = parsed_result["intent"]
|
||
# 确保所有必需字段都存在
|
||
intent_data.setdefault("is_code_related", False)
|
||
intent_data.setdefault("category", CodeIntentCategory.UNKNOWN)
|
||
intent_data.setdefault("confidence", 0.5)
|
||
intent_data.setdefault("keywords", [])
|
||
intent_data.setdefault("reasoning", "")
|
||
intent_data.setdefault("requires_code_context", False)
|
||
intent_data.setdefault("suggested_search_terms", [])
|
||
|
||
# 确定Prompt模板类型
|
||
prompt_template_type = self._map_to_prompt_template_type(intent_data["category"])
|
||
|
||
# 构建CodeIntentResult对象
|
||
intent_result = CodeIntentResult(
|
||
is_code_related=intent_data["is_code_related"],
|
||
category=intent_data["category"],
|
||
confidence=intent_data["confidence"],
|
||
prompt_template_type=prompt_template_type,
|
||
keywords=intent_data["keywords"],
|
||
reasoning=intent_data["reasoning"],
|
||
requires_code_context=intent_data["requires_code_context"],
|
||
suggested_search_terms=intent_data["suggested_search_terms"]
|
||
)
|
||
result["intent"] = intent_result
|
||
except Exception as e:
|
||
logger.error(f"处理意图识别结果失败: {e}")
|
||
|
||
# 处理过滤条件
|
||
try:
|
||
if "filters" in parsed_result:
|
||
filters = parsed_result["filters"]
|
||
if isinstance(filters, dict):
|
||
# 验证并过滤结果，确保只包含有效的metadata key
|
||
valid_keys = ['func_id', 'func_name', 'class_name', 'file_path', 'lang', 'params', 'return_type', 'docstring', 'start_line', 'end_line', 'repo_id', 'branch', 'func_body']
|
||
filtered_filters = {}
|
||
for key, value in filters.items():
|
||
if key in valid_keys and value:
|
||
# 确保value是字符串类型
|
||
if isinstance(value, str):
|
||
filtered_filters[key] = value
|
||
result["filters"] = filtered_filters
|
||
except Exception as e:
|
||
logger.error(f"处理过滤条件失败: {e}")
|
||
|
||
# 处理查询转换结果
|
||
try:
|
||
if "transformed" in parsed_result:
|
||
transformed = parsed_result["transformed"]
|
||
if isinstance(transformed, dict):
|
||
result["transformed"].update({
|
||
"rewritten": transformed.get("rewritten", query),
|
||
"backward": transformed.get("backward", query),
|
||
"sub_queries": transformed.get("sub_queries", [query])
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"处理查询转换结果失败: {e}")
|
||
|
||
logger.info("查询处理完成")
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"查询处理失败: {e}")
|
||
return self._get_default_result(query)
|
||
|
||
def _get_default_result(self, query: str) -> Dict[str, Any]:
|
||
"""
|
||
获取默认结果（当处理失败时使用）
|
||
|
||
Args:
|
||
query: 用户问题
|
||
|
||
Returns:
|
||
默认结果
|
||
"""
|
||
logger.warning(f"使用默认查询处理结果: {query}")
|
||
|
||
# 构建默认的意图识别结果
|
||
default_intent = 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]
|
||
)
|
||
|
||
return {
|
||
"intent": default_intent,
|
||
"filters": {},
|
||
"transformed": {
|
||
"rewritten": query,
|
||
"backward": query,
|
||
"sub_queries": [query]
|
||
}
|
||
}
|
||
|
||
|
||
# 辅助函数
|
||
def create_query_processor(llm: Optional[Ollama] = None) -> QueryProcessor:
|
||
"""
|
||
创建查询处理器实例
|
||
|
||
Args:
|
||
llm: LLM实例，如果为None则使用默认配置
|
||
|
||
Returns:
|
||
QueryProcessor: 处理器实例
|
||
"""
|
||
return QueryProcessor(llm)
|
||
|
||
|
||
class MetadataFilter:
|
||
"""
|
||
Metadata过滤工具类
|
||
"""
|
||
|
||
@staticmethod
|
||
def apply_filter(metadata: Dict[str, Any], filters: Dict[str, Any]) -> bool:
|
||
"""
|
||
应用过滤条件到metadata
|
||
|
||
Args:
|
||
metadata: 文档的metadata
|
||
filters: 过滤条件
|
||
|
||
Returns:
|
||
bool: 如果metadata符合过滤条件返回True，否则返回False
|
||
"""
|
||
if not filters:
|
||
return True
|
||
|
||
for key, value in filters.items():
|
||
if key not in metadata:
|
||
return False
|
||
|
||
metadata_value = metadata[key]
|
||
if isinstance(metadata_value, str) and isinstance(value, str):
|
||
# 对于字符串类型，使用大小写不敏感的模糊匹配
|
||
if value.lower() not in metadata_value.lower():
|
||
return False
|
||
else:
|
||
if metadata_value != value:
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
if __name__ == "__main__":
|
||
"""测试查询处理器"""
|
||
processor = QueryProcessor()
|
||
|
||
test_queries = [
|
||
"在 data_structures 目录下，二叉搜索树（Binary Search Tree）的删除操作依赖于哪些辅助方法来寻找后继节点？",
|
||
"如何实现快速排序算法？",
|
||
"kth_number的时间复杂度是多少？",
|
||
"今天天气怎么样？"
|
||
]
|
||
|
||
for query in test_queries:
|
||
print(f"\n{'='*60}")
|
||
print(f"问题: {query}")
|
||
print('='*60)
|
||
|
||
result = processor.process_query(query)
|
||
|
||
print("意图识别结果:")
|
||
if result["intent"]:
|
||
print(result["intent"].to_json())
|
||
|
||
print("\n过滤条件:")
|
||
import json
|
||
print(json.dumps(result["filters"], ensure_ascii=False, indent=2))
|
||
|
||
print("\n查询转换结果:")
|
||
print(json.dumps(result["transformed"], ensure_ascii=False, indent=2))
|