470 lines
20 KiB
Python
470 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
|
||
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.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 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. 处理查询（整合意图识别、filter生成和查询转换）
|
||
logger.info(f"开始查询处理: {query}")
|
||
process_result = self.query_processor.process_query(query, history)
|
||
|
||
# 提取结果
|
||
intent_result = process_result['intent']
|
||
filters = process_result['filters']
|
||
transformed_queries = [process_result['transformed']['rewritten']] + process_result['transformed']['sub_queries']
|
||
|
||
logger.info(f"代码意图识别结果: {intent_result.category}")
|
||
logger.info(f"过滤条件: {filters}")
|
||
logger.info(f"查询转换结果: {transformed_queries}")
|
||
|
||
# 4. 使用融合检索，对每个转换后的查询进行检索
|
||
logger.info("使用融合检索策略")
|
||
all_hybrid_results = []
|
||
for transformed_query in transformed_queries:
|
||
hybrid_results = await self.vector_store_manager.ahybrid_search(
|
||
query=transformed_query,
|
||
top_k=settings.TOP_K,
|
||
filters=filters
|
||
)
|
||
all_hybrid_results.extend(hybrid_results)
|
||
|
||
# 去重并按得分排序
|
||
seen_doc_ids = set()
|
||
unique_hybrid_results = []
|
||
for doc_id, score, metadata in all_hybrid_results:
|
||
if doc_id not in seen_doc_ids:
|
||
seen_doc_ids.add(doc_id)
|
||
unique_hybrid_results.append((doc_id, score, metadata))
|
||
|
||
# 按得分排序
|
||
unique_hybrid_results.sort(key=lambda x: x[1], reverse=True)
|
||
|
||
# 限制结果数量
|
||
hybrid_results = unique_hybrid_results[:top_k or settings.TOP_K]
|
||
# 美化输出融合检索结果
|
||
logger.info("融合检索结果:")
|
||
for i, (doc_id, score, metadata) in enumerate(hybrid_results):
|
||
func_name = metadata.get('func_name', 'N/A')
|
||
file_path = metadata.get('file_path', 'N/A')
|
||
lang = metadata.get('lang', 'N/A')
|
||
logger.info(f" [{i+1}] 相似度: {score:.4f}")
|
||
logger.info(f" 函数: {func_name}")
|
||
logger.info(f" 文件: {file_path}")
|
||
logger.info(f" 语言: {lang}")
|
||
logger.info(" " + "-" * 50)
|
||
# 根据文档ID获取完整的节点信息
|
||
retrieved_nodes = []
|
||
for doc_id, score, metadata in hybrid_results:
|
||
# 从向量存储中获取文档内容
|
||
doc_chunks = self.vector_store_manager.get_document_by_id(doc_id)
|
||
for chunk in doc_chunks:
|
||
# 创建节点对象
|
||
from llama_index.core.schema import TextNode
|
||
node = TextNode(
|
||
text=chunk['text'],
|
||
node_id=chunk['id'],
|
||
metadata=chunk['metadata']
|
||
)
|
||
retrieved_nodes.append(node)
|
||
|
||
# 3. 构建上下文
|
||
context_parts = []
|
||
max_nodes = top_k or settings.TOP_K
|
||
|
||
# 构建原始的检索结果列表，包含 text 和 metadata
|
||
retrieved_results = []
|
||
for i, node in enumerate(retrieved_nodes[:max_nodes], 1):
|
||
text = node.text if hasattr(node, 'text') else str(node)
|
||
text = text.strip()
|
||
|
||
# 获取原始 metadata
|
||
metadata = getattr(node, 'metadata', {})
|
||
|
||
# 保存原始的 text 和 metadata
|
||
retrieved_results.append({
|
||
'id': i,
|
||
'text': text,
|
||
'metadata': metadata
|
||
})
|
||
|
||
# 构建简单的上下文字符串，只包含基本信息
|
||
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 "未找到相关参考信息"
|
||
logger.info(f"上下文: {context_str}")
|
||
# 4. 生成优化的Prompt
|
||
logger.info("生成优化的Prompt")
|
||
if intent_result.is_code_related:
|
||
# 对于代码相关问题，使用代码专用Prompt
|
||
filled_prompt = generate_dynamic_code_prompt(
|
||
user_query=query,
|
||
intent_result=intent_result.to_dict(),
|
||
code_context=context_str,
|
||
retrieved_results=retrieved_results,
|
||
conversation_history=history
|
||
)
|
||
logger.info(f"使用代码专用Prompt，类型: {intent_result.category}")
|
||
else:
|
||
# 对于非代码问题，使用通用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")
|
||
|
||
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) # slight delay to yield control
|
||
|
||
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. 处理查询（整合意图识别、filter生成和查询转换）
|
||
logger.info(f"开始查询处理: {query[:50]}...")
|
||
process_result = self.query_processor.process_query(query, history)
|
||
|
||
# 提取结果
|
||
intent_result = process_result['intent']
|
||
filters = process_result['filters']
|
||
transformed_queries = [process_result['transformed']['rewritten']] + process_result['transformed']['sub_queries']
|
||
|
||
logger.info(f"代码意图识别结果: {intent_result.category}")
|
||
logger.info(f"过滤条件: {filters}")
|
||
logger.info(f"查询转换完成，生成了 {len(transformed_queries)} 个转换后的查询")
|
||
|
||
# 2. 使用融合检索，对每个转换后的查询进行检索
|
||
logger.info("使用融合检索策略")
|
||
all_hybrid_results = []
|
||
for transformed_query in transformed_queries:
|
||
hybrid_results = await self.vector_store_manager.ahybrid_search(
|
||
query=transformed_query,
|
||
top_k=top_k or settings.TOP_K,
|
||
filters=filters
|
||
)
|
||
all_hybrid_results.extend(hybrid_results)
|
||
|
||
# 去重并按得分排序
|
||
seen_doc_ids = set()
|
||
unique_hybrid_results = []
|
||
for doc_id, score, metadata in all_hybrid_results:
|
||
if doc_id not in seen_doc_ids:
|
||
seen_doc_ids.add(doc_id)
|
||
unique_hybrid_results.append((doc_id, score, metadata))
|
||
|
||
# 按得分排序
|
||
unique_hybrid_results.sort(key=lambda x: x[1], reverse=True)
|
||
|
||
# 限制结果数量
|
||
hybrid_results = unique_hybrid_results[:top_k or settings.TOP_K]
|
||
|
||
# 根据文档ID获取完整的节点信息
|
||
retrieved_nodes = []
|
||
for doc_id, score, metadata in hybrid_results:
|
||
# 从向量存储中获取文档内容
|
||
doc_chunks = self.vector_store_manager.get_document_by_id(doc_id)
|
||
for chunk in doc_chunks:
|
||
# 创建节点对象
|
||
from llama_index.core.schema import TextNode
|
||
node = TextNode(
|
||
text=chunk['text'],
|
||
node_id=chunk['id'],
|
||
metadata=chunk['metadata']
|
||
)
|
||
retrieved_nodes.append(node)
|
||
|
||
# 3. 构建原始的检索结果列表，包含 text 和 metadata
|
||
retrieved_results = []
|
||
for i, node in enumerate(retrieved_nodes[:top_k or settings.TOP_K], 1):
|
||
text = node.text if hasattr(node, 'text') else str(node)
|
||
text = text.strip()
|
||
|
||
# 获取原始 metadata
|
||
metadata = getattr(node, 'metadata', {})
|
||
|
||
# 保存原始的 text 和 metadata
|
||
retrieved_results.append({
|
||
'id': i,
|
||
'text': text,
|
||
'metadata': metadata
|
||
})
|
||
|
||
# 构建简单的上下文字符串，只包含基本信息
|
||
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 "未找到相关参考信息"
|
||
|
||
# 4. 生成优化的Prompt
|
||
logger.info("生成优化的Prompt")
|
||
if intent_result.is_code_related:
|
||
# 对于代码相关问题，使用代码专用Prompt
|
||
filled_prompt = generate_dynamic_code_prompt(
|
||
user_query=query,
|
||
intent_result=intent_result.to_dict(),
|
||
code_context=context_str,
|
||
retrieved_results=retrieved_results,
|
||
conversation_history=history
|
||
)
|
||
logger.info(f"使用代码专用Prompt，类型: {intent_result.category}")
|
||
else:
|
||
# 对于非代码问题，使用通用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")
|
||
|
||
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)}"
|
||
|