479 lines
20 KiB
Python
479 lines
20 KiB
Python
"""
|
||
RAG engine for query processing and response generation
|
||
"""
|
||
from llama_index.core.query_engine import RetrieverQueryEngine
|
||
from llama_index.core.response_synthesizers import ResponseMode
|
||
from llama_index.core.base.response.schema import StreamingResponse
|
||
from llama_index.llms.ollama import Ollama
|
||
from llama_index.core import PromptTemplate
|
||
from typing import AsyncIterator, Optional, Tuple, Dict, Any
|
||
import asyncio
|
||
import httpx
|
||
from loguru import logger
|
||
from config import settings
|
||
from .vector_store import VectorStoreManager
|
||
from .chunk_handler import OptimizedDeltaThinkFilter
|
||
from utils.query_processor import QueryProcessor
|
||
from utils.code_prompt_manager import generate_dynamic_code_prompt
|
||
|
||
|
||
# Single module-level prompt string for easy editing in one place
|
||
QA_PROMPT_STR_HISTORY = """
|
||
# 角色设定
|
||
你是一个专业的AI助手，具备丰富知识且乐于助人。请基于提供的参考信息和对话历史，以专业、准确且友好的方式回答问题。
|
||
|
||
# 对话历史
|
||
{history}
|
||
|
||
# 参考信息
|
||
{context_str}
|
||
|
||
# 当前用户问题
|
||
{query_str}
|
||
|
||
# 回答策略
|
||
请按以下优先级回答问题：
|
||
1. **优先使用参考信息**：如果参考信息充分相关，直接基于参考信息回答
|
||
2. **补充通用知识**：如果参考信息部分相关但不足，结合你的通用知识进行补充回答
|
||
3. **明确说明限制**：如果参考信息完全不相关或为空，明确说明后将基于通用知识回答
|
||
4. **完全无信息的情况**：如果参考信息为空或完全无关，直接基于你的知识库提供有帮助的回答
|
||
|
||
# 重要原则
|
||
1. **绝不拒绝回答**：即使参考信息为空，也必须提供有帮助的回答
|
||
2. **透明诚实**：明确说明信息的来源和局限性
|
||
3. **保持价值**：确保回答对用户有实际帮助
|
||
4. **历史连贯**：保持与对话历史的一致性
|
||
5. **结构清晰**：使用适当的格式增强可读性
|
||
|
||
现在请基于以上指导原则回答用户问题：
|
||
"""
|
||
|
||
QA_PROMPT_STR_NO_HISTORY = """
|
||
# 角色设定
|
||
你是一个专业的AI助手，具备丰富知识且乐于助人。请基于提供的参考信息和对话历史，以专业、准确且友好的方式回答问题。
|
||
|
||
# 参考信息
|
||
{context_str}
|
||
|
||
# 当前用户问题
|
||
{query_str}
|
||
|
||
# 回答策略
|
||
请按以下优先级回答问题：
|
||
1. **优先使用参考信息**：如果参考信息充分相关，直接基于参考信息回答
|
||
2. **补充通用知识**：如果参考信息部分相关但不足，结合你的通用知识进行补充回答
|
||
3. **明确说明限制**：如果参考信息完全不相关或为空，明确说明后将基于通用知识回答
|
||
4. **完全无信息的情况**：如果参考信息为空或完全无关，直接基于你的知识库提供有帮助的回答
|
||
|
||
# 重要原则
|
||
1. **绝不拒绝回答**：即使参考信息为空，也必须提供有帮助的回答
|
||
2. **透明诚实**：明确说明信息的来源和局限性
|
||
3. **保持价值**：确保回答对用户有实际帮助
|
||
4. **历史连贯**：保持与对话历史的一致性
|
||
5. **结构清晰**：使用适当的格式增强可读性
|
||
|
||
现在请基于以上指导原则回答用户问题：
|
||
"""
|
||
|
||
# Default PromptTemplate object
|
||
QA_PROMPT_HISTORY = PromptTemplate(QA_PROMPT_STR_HISTORY)
|
||
QA_PROMPT_NO_HISTORY = PromptTemplate(QA_PROMPT_STR_NO_HISTORY)
|
||
|
||
|
||
class RAGEngine:
|
||
"""Main RAG engine for query processing"""
|
||
|
||
@staticmethod
|
||
def check_ollama_connection() -> Tuple[bool, str]:
|
||
"""
|
||
Check if Ollama server is accessible and connection can be established
|
||
|
||
Returns:
|
||
Tuple of (is_connected: bool, error_message: str)
|
||
If connected, error_message will be empty string
|
||
"""
|
||
try:
|
||
logger.info(f"Checking Ollama connection to {settings.OLLAMA_BASE_URL}...")
|
||
|
||
# Test connection by calling Ollama API
|
||
with httpx.Client(timeout=10.0) as client:
|
||
response = client.get(f"{settings.OLLAMA_BASE_URL}/api/tags")
|
||
if response.status_code == 200:
|
||
models = response.json().get("models", [])
|
||
model_names = [m.get("name", "unknown") for m in models]
|
||
logger.info(f"✓ Ollama connection successful (found {len(models)} model(s): {', '.join(model_names[:3])}{'...' if len(model_names) > 3 else ''})")
|
||
return True, ""
|
||
else:
|
||
error_message = f"Ollama API returned status {response.status_code}: {response.text}"
|
||
logger.error(f"✗ Ollama connection failed: {error_message}")
|
||
return False, error_message
|
||
except httpx.ConnectError as e:
|
||
error_message = f"Cannot connect to Ollama server at {settings.OLLAMA_BASE_URL}. " \
|
||
f"Please check if Ollama server is running and accessible."
|
||
logger.error(f"✗ Ollama connection failed: {error_message}")
|
||
return False, error_message
|
||
except httpx.TimeoutException:
|
||
error_message = f"Connection to Ollama server at {settings.OLLAMA_BASE_URL} timed out. " \
|
||
f"Please check if Ollama server is running and accessible."
|
||
logger.error(f"✗ Ollama connection failed: {error_message}")
|
||
return False, error_message
|
||
except Exception as e:
|
||
error_message = f"Unexpected error while checking Ollama connection: {str(e)}"
|
||
logger.error(f"✗ Ollama connection check failed: {error_message}")
|
||
return False, error_message
|
||
|
||
def __init__(
|
||
self,
|
||
vector_store_manager: VectorStoreManager,
|
||
prompt_template: Optional[PromptTemplate] = None,
|
||
system_prompt: Optional[str] = None,
|
||
temperature: float = 0.7,
|
||
request_timeout: float = 1200.0,
|
||
):
|
||
self.vector_store_manager = vector_store_manager
|
||
# Configurable LLM / prompt parameters
|
||
self._user_prompt_template = prompt_template
|
||
self._system_prompt = system_prompt
|
||
self._temperature = temperature
|
||
self._request_timeout = request_timeout
|
||
|
||
# Check Ollama connection before initializing
|
||
is_connected, error_message = self.check_ollama_connection()
|
||
if not is_connected:
|
||
error_msg = (
|
||
f"Error: Cannot connect to Ollama server.\n"
|
||
f" {error_message}\n"
|
||
f"Connection info: {settings.OLLAMA_BASE_URL}\n"
|
||
f"Please check:\n"
|
||
f" 1. Ollama server is running\n"
|
||
f" 2. Ollama server is accessible from this host\n"
|
||
f" 3. OLLAMA_BASE_URL is correctly configured\n"
|
||
f" 4. Firewall rules allow connection to Ollama port"
|
||
)
|
||
logger.error(error_msg)
|
||
raise RuntimeError(error_msg)
|
||
|
||
self.llm = Ollama(
|
||
model=settings.OLLAMA_MODEL,
|
||
base_url=settings.OLLAMA_BASE_URL,
|
||
temperature=self._temperature,
|
||
request_timeout=self._request_timeout,
|
||
)
|
||
|
||
# 初始化小型模型用于代码问题判断
|
||
self.small_llm = Ollama(
|
||
model="qwen3:0.6b",
|
||
base_url=settings.OLLAMA_BASE_URL,
|
||
temperature=0.3,
|
||
request_timeout=30.0,
|
||
)
|
||
|
||
# 初始化查询处理器
|
||
self.query_processor = QueryProcessor(llm=self.llm)
|
||
|
||
def extract_text_from_chunk(self, chunk) -> Optional[str]:
|
||
if hasattr(chunk, 'delta'):
|
||
# CompletionChunk类型
|
||
return chunk.delta
|
||
elif hasattr(chunk, 'text'):
|
||
# 某些版本用text
|
||
return chunk.text
|
||
elif isinstance(chunk, str):
|
||
return chunk
|
||
else:
|
||
# 尝试转换为字符串
|
||
try:
|
||
str(chunk)
|
||
except:
|
||
pass
|
||
|
||
async def is_code_related(self, query: str, history: str) -> Dict[str, Any]:
|
||
"""
|
||
使用小型模型判断是否是代码相关问题，并提取可能的filter
|
||
|
||
Args:
|
||
query: 用户查询字符串
|
||
history: 对话历史字符串
|
||
|
||
Returns:
|
||
dict: 包含 confidence (float) 和 filters (dict)
|
||
"""
|
||
try:
|
||
# 获取需要提取的metadata字段
|
||
filter_fields = settings.FILTER_METADATA_FIELDS.split(',')
|
||
filter_fields_str = ', '.join(filter_fields)
|
||
|
||
prompt = f"""你是一个分类器，需要分析用户的问题。
|
||
|
||
用户问题：{query}
|
||
|
||
对话历史：{history}
|
||
|
||
请分析这个问题并返回JSON格式的分析结果：
|
||
{{
|
||
"confidence": 0.0-1.0之间的置信度，1表示完全确定是代码相关问题，0表示完全确定不是代码相关问题,
|
||
"filters": {{}} 或 {{"字段名": "从问题中提取的值"}}，只有当问题中明确提到"xxx文件/xxx函数/xxx类"时才提取
|
||
}}
|
||
|
||
提取规则：
|
||
- 只有当用户明确说明了"xxx文件"、"xxx函数"、"xxx类"时才提取对应的metadata
|
||
- class_name: 用户提到具体类名时提取，如"User类"、"ArrayList"
|
||
- func_name: 用户提到具体函数名时提取，如"main函数"、"delete方法"
|
||
- file_path: 用户提到具体文件时提取，如"utils.py"、"config.json"
|
||
|
||
请仅返回JSON，不要添加任何其他内容。"""
|
||
|
||
response = await self.small_llm.acomplete(prompt=prompt)
|
||
response_text = response.text.strip()
|
||
|
||
logger.info(f"小型模型分析结果: {response_text}")
|
||
|
||
# 尝试解析JSON
|
||
import json
|
||
json_start = response_text.find('{')
|
||
json_end = response_text.rfind('}')
|
||
if json_start != -1 and json_end != -1:
|
||
json_str = response_text[json_start:json_end + 1]
|
||
result = json.loads(json_str)
|
||
confidence = float(result.get('confidence', 0.5))
|
||
filters = result.get('filters', {})
|
||
# 过滤掉空的filter
|
||
filters = {k: v for k, v in filters.items() if v}
|
||
logger.info(f"解析成功: confidence={confidence}, filters={filters}")
|
||
return {'confidence': confidence, 'filters': filters}
|
||
else:
|
||
logger.warning(f"无法解析JSON，使用默认结果")
|
||
return {'confidence': 0.5, 'filters': {}}
|
||
|
||
except Exception as e:
|
||
logger.error(f"判断代码问题时出错: {e}")
|
||
# 出错时返回中间值，不过滤
|
||
return {'confidence': 0.5, 'filters': {}}
|
||
|
||
async def query_stream(self, query: str, history: str, top_k: Optional[int] = None) -> AsyncIterator[str]:
|
||
"""
|
||
Query the RAG system and stream the response
|
||
|
||
Args:
|
||
query: User query string
|
||
history: Chat history string
|
||
top_k: Number of documents to retrieve (optional)
|
||
|
||
Yields:
|
||
Response text chunks
|
||
"""
|
||
try:
|
||
# 1. 不区分是否代码相关问题，直接在所有collection中检索
|
||
logger.info("在所有collection中进行向量检索")
|
||
|
||
# 2. 使用纯向量检索
|
||
logger.info("使用纯向量检索")
|
||
k = top_k or settings.TOP_K
|
||
|
||
retriever = self.vector_store_manager.get_retriever(
|
||
top_k=k * 2,
|
||
filters=None,
|
||
collection_key=None
|
||
)
|
||
|
||
if isinstance(retriever, list):
|
||
all_nodes = []
|
||
for key, r in retriever:
|
||
nodes = r.retrieve(query)
|
||
for node in nodes:
|
||
actual_node = node.node if hasattr(node, 'node') else node
|
||
if hasattr(actual_node, 'metadata'):
|
||
actual_node.metadata['_collection_key'] = key
|
||
all_nodes.extend(nodes)
|
||
vector_nodes = all_nodes
|
||
else:
|
||
vector_nodes = retriever.retrieve(query)
|
||
|
||
# 去重并按得分排序
|
||
seen_doc_ids = set()
|
||
unique_results = []
|
||
for node in vector_nodes:
|
||
doc_id = getattr(node, 'id_', None) or getattr(node, 'node_id', None)
|
||
if doc_id and doc_id not in seen_doc_ids:
|
||
seen_doc_ids.add(doc_id)
|
||
unique_results.append(node)
|
||
|
||
unique_results.sort(key=lambda x: getattr(x, 'score', 0), reverse=True)
|
||
vector_nodes = unique_results[:k]
|
||
|
||
logger.info(f"检索到 {len(vector_nodes)} 个结果")
|
||
|
||
# 4. 构建检索结果
|
||
retrieved_results = []
|
||
for i, node in enumerate(vector_nodes, 1):
|
||
metadata = getattr(node, 'metadata', {})
|
||
if 'func_body' in metadata:
|
||
text = metadata['func_body']
|
||
else:
|
||
text = node.text if hasattr(node, 'text') else str(node)
|
||
text = text.strip()
|
||
|
||
retrieved_results.append({
|
||
'id': i,
|
||
'text': text,
|
||
'metadata': metadata
|
||
})
|
||
|
||
# 5. 构建上下文
|
||
context_parts = []
|
||
for result in retrieved_results:
|
||
metadata = result['metadata']
|
||
file_path = metadata.get('file_path', '')
|
||
func_name = metadata.get('func_name', '')
|
||
|
||
context_part = f"【参考信息{result['id']}】"
|
||
if file_path:
|
||
context_part += f"（来源：{file_path}）"
|
||
if func_name:
|
||
context_part += f"\n函数：{func_name}"
|
||
context_part += f"\n{result['text']}"
|
||
context_parts.append(context_part)
|
||
|
||
context_str = "\n\n".join(context_parts) if context_parts else "未找到相关参考信息"
|
||
|
||
# 6. 生成Prompt - 根据是否有历史对话选择不同模板
|
||
logger.info("生成Prompt")
|
||
if history:
|
||
qa_prompt = QA_PROMPT_HISTORY
|
||
filled_prompt = qa_prompt.format(history=history, context_str=context_str, query_str=query)
|
||
else:
|
||
qa_prompt = QA_PROMPT_NO_HISTORY
|
||
filled_prompt = qa_prompt.format(context_str=context_str, query_str=query)
|
||
|
||
logger.info("根据历史对话选择prompt模板")
|
||
|
||
# 7. 流式生成回答
|
||
stream_response = await self.llm.astream_complete(prompt=filled_prompt)
|
||
|
||
full_response = ""
|
||
think_filter = OptimizedDeltaThinkFilter()
|
||
|
||
async for chunk in stream_response:
|
||
delta, full_text, has_output = think_filter.process_delta_robust(chunk)
|
||
|
||
if delta is not None:
|
||
full_response += delta
|
||
yield delta.encode('utf-8')
|
||
await asyncio.sleep(0.001)
|
||
|
||
logger.info(f"响应完成，长度: {len(full_response)}字符")
|
||
print(full_response)
|
||
except Exception as e:
|
||
logger.error(f"Error in RAG query: {e}")
|
||
|
||
async def query(self, query: str, history: str, top_k: Optional[int] = None) -> str:
|
||
"""
|
||
Query the RAG system and return complete response
|
||
|
||
Args:
|
||
query: User query string
|
||
history: Chat history string
|
||
top_k: Number of documents to retrieve (optional)
|
||
|
||
Returns:
|
||
Complete response string
|
||
"""
|
||
try:
|
||
# 1. 不区分是否代码相关问题，直接在所有collection中检索
|
||
logger.info("在所有collection中进行向量检索")
|
||
|
||
# 2. 使用纯向量检索
|
||
logger.info("使用纯向量检索")
|
||
k = top_k or settings.TOP_K
|
||
|
||
# 获取retriever
|
||
retriever = self.vector_store_manager.get_retriever(
|
||
top_k=k * 2, # 获取更多结果用于去重
|
||
filters=None,
|
||
collection_key=None
|
||
)
|
||
|
||
# 执行检索
|
||
if isinstance(retriever, list):
|
||
# 多个collection
|
||
all_nodes = []
|
||
for key, r in retriever:
|
||
nodes = r.retrieve(query)
|
||
for node in nodes:
|
||
actual_node = node.node if hasattr(node, 'node') else node
|
||
if hasattr(actual_node, 'metadata'):
|
||
actual_node.metadata['_collection_key'] = key
|
||
all_nodes.extend(nodes)
|
||
vector_nodes = all_nodes
|
||
else:
|
||
vector_nodes = retriever.retrieve(query)
|
||
|
||
# 去重并按得分排序
|
||
seen_doc_ids = set()
|
||
unique_results = []
|
||
for node in vector_nodes:
|
||
doc_id = getattr(node, 'id_', None) or getattr(node, 'node_id', None)
|
||
if doc_id and doc_id not in seen_doc_ids:
|
||
seen_doc_ids.add(doc_id)
|
||
unique_results.append(node)
|
||
|
||
# 按得分排序，取top_k
|
||
unique_results.sort(key=lambda x: getattr(x, 'score', 0), reverse=True)
|
||
vector_nodes = unique_results[:k]
|
||
|
||
logger.info(f"检索到 {len(vector_nodes)} 个结果")
|
||
|
||
# 4. 获取文档内容
|
||
retrieved_results = []
|
||
for i, node in enumerate(vector_nodes, 1):
|
||
# 优先使用metadata中的func_body字段
|
||
metadata = getattr(node, 'metadata', {})
|
||
if 'func_body' in metadata:
|
||
text = metadata['func_body']
|
||
else:
|
||
text = node.text if hasattr(node, 'text') else str(node)
|
||
text = text.strip()
|
||
|
||
retrieved_results.append({
|
||
'id': i,
|
||
'text': text,
|
||
'metadata': metadata
|
||
})
|
||
|
||
# 5. 构建上下文
|
||
context_parts = []
|
||
for result in retrieved_results:
|
||
metadata = result['metadata']
|
||
file_path = metadata.get('file_path', '')
|
||
func_name = metadata.get('func_name', '')
|
||
|
||
context_part = f"【参考信息{result['id']}】"
|
||
if file_path:
|
||
context_part += f"（来源：{file_path}）"
|
||
if func_name:
|
||
context_part += f"\n函数：{func_name}"
|
||
context_part += f"\n{result['text']}"
|
||
context_parts.append(context_part)
|
||
|
||
context_str = "\n\n".join(context_parts) if context_parts else "未找到相关参考信息"
|
||
|
||
# 6. 生成Prompt - 根据是否有历史对话选择不同模板
|
||
logger.info("生成Prompt")
|
||
if history:
|
||
qa_prompt = QA_PROMPT_HISTORY
|
||
filled_prompt = qa_prompt.format(history=history, context_str=context_str, query_str=query)
|
||
else:
|
||
qa_prompt = QA_PROMPT_NO_HISTORY
|
||
filled_prompt = qa_prompt.format(context_str=context_str, query_str=query)
|
||
|
||
logger.info("根据历史对话选择prompt模板")
|
||
|
||
# 7. 调用LLM生成回答
|
||
response = await self.llm.acomplete(prompt=filled_prompt)
|
||
|
||
return response.text
|
||
except Exception as e:
|
||
logger.error(f"Error in RAG query: {e}")
|
||
return f"Error: {str(e)}"
|
||
|