282 lines
12 KiB
Python
282 lines
12 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
|
||
|
||
|
||
# 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 = 120.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,
|
||
)
|
||
|
||
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:
|
||
# Create query engine with streaming mode
|
||
retriever = self.vector_store_manager.get_retriever(top_k=top_k or settings.TOP_K)
|
||
# query index
|
||
retrieved_nodes = await retriever.aretrieve(query)
|
||
|
||
# 2. 构建上下文
|
||
context_parts = []
|
||
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()
|
||
if len(text) > 400:
|
||
text = text[:400] + "..."
|
||
context_parts.append(f"【参考信息{i}】{text}")
|
||
|
||
context_str = "\n\n".join(context_parts) if context_parts else "未找到相关参考信息"
|
||
|
||
if history is not None:
|
||
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)
|
||
|
||
stream_response = await self.llm.astream_complete(
|
||
prompt=filled_prompt
|
||
)
|
||
|
||
full_response = ""
|
||
think_filter = OptimizedDeltaThinkFilter()
|
||
|
||
async for chunk in stream_response:
|
||
# 提取文本内容
|
||
# text_chunk = self.extract_text_from_chunk(chunk)
|
||
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)}字符")
|
||
|
||
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:
|
||
# Create query engine with streaming mode
|
||
retriever = self.vector_store_manager.get_retriever(top_k=top_k or settings.TOP_K)
|
||
# query index
|
||
retrieved_nodes = await retriever.aretrieve(query)
|
||
|
||
# 2. 构建上下文
|
||
context_parts = []
|
||
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()
|
||
if len(text) > 400:
|
||
text = text[:400] + "..."
|
||
context_parts.append(f"【参考信息{i}】{text}")
|
||
|
||
context_str = "\n\n".join(context_parts) if context_parts else "未找到相关参考信息"
|
||
|
||
if history is not None:
|
||
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)
|
||
|
||
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)}"
|
||
|