mirror of
https://github.com/aaPanel/aaPanel.git
synced 2026-08-29 03:09:46 +02:00
522 lines
19 KiB
Python
522 lines
19 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import uuid
|
|
import requests
|
|
from typing import List, Dict, Any
|
|
import openai
|
|
import numpy as np
|
|
|
|
panelPath = os.getenv('BT_PANEL')
|
|
if not panelPath: panelPath = '/www/server/panel'
|
|
if panelPath not in sys.path:
|
|
os.chdir(panelPath)
|
|
sys.path.insert(0, panelPath + "/class/")
|
|
try:
|
|
import public
|
|
from public import lang
|
|
except ImportError:
|
|
public = None
|
|
lang = lambda x: x
|
|
|
|
|
|
class SimpleVectorDB:
|
|
def __init__(self, persist_dir: str):
|
|
self.persist_dir = persist_dir
|
|
self.persist_file = os.path.join(persist_dir, "vector_store.json")
|
|
self.data: List[Dict[str, Any]] = [] # stores {"id":, "text":, "metadata":, "embedding":}
|
|
|
|
if not os.path.exists(persist_dir):
|
|
try:
|
|
os.makedirs(persist_dir, exist_ok=True)
|
|
except OSError as e:
|
|
public.print_log(f"[ERROR] Failed to create vector DB directory {persist_dir}: {str(e)}")
|
|
raise
|
|
|
|
self.load()
|
|
|
|
def load(self):
|
|
if os.path.exists(self.persist_file):
|
|
try:
|
|
with open(self.persist_file, 'r', encoding='utf-8') as f:
|
|
self.data = json.load(f)
|
|
except Exception as e:
|
|
public.print_log(f"[ERROR] Failed to load vector store {self.persist_file}: {str(e)}")
|
|
self.data = []
|
|
|
|
def save(self):
|
|
try:
|
|
with open(self.persist_file, 'w', encoding='utf-8') as f:
|
|
json.dump(self.data, f, ensure_ascii=False)
|
|
except Exception as e:
|
|
public.print_log(f"[ERROR] Failed to save vector store {self.persist_file}: {str(e)}")
|
|
pass
|
|
|
|
def add(self, documents: List[str], metadatas: List[Dict], ids: List[str], embeddings: List[List[float]]):
|
|
for doc, meta, doc_id, emb in zip(documents, metadatas, ids, embeddings):
|
|
self.data = [d for d in self.data if d['id'] != doc_id]
|
|
|
|
self.data.append({
|
|
"id": doc_id,
|
|
"text": doc,
|
|
"metadata": meta,
|
|
"embedding": emb
|
|
})
|
|
self.save()
|
|
|
|
def search(self, query_embedding: List[float], n_results: int = 5, score: float = 0.5, where: Dict = None,
|
|
metric: str = "cosine") -> List[Dict]:
|
|
"""
|
|
Search for similar documents.
|
|
|
|
Args:
|
|
query_embedding: The query vector.
|
|
n_results: Number of results to return.
|
|
score: Minimum score threshold.
|
|
where: Metadata filter.
|
|
metric: Distance metric to use. Options: "cosine" (default), "euclidean", "dot".
|
|
- "cosine": Cosine similarity (normalized dot product). Range [-1, 1].
|
|
- "euclidean": 1 / (1 + L2_distance). Range (0, 1].
|
|
- "dot": Dot product. Range (-inf, inf).
|
|
"""
|
|
if not self.data:
|
|
return []
|
|
|
|
candidates = self.data
|
|
if where:
|
|
filtered = []
|
|
for item in candidates:
|
|
match = True
|
|
for k, v in where.items():
|
|
if item['metadata'].get(k) != v:
|
|
match = False
|
|
break
|
|
if match:
|
|
filtered.append(item)
|
|
candidates = filtered
|
|
|
|
if not candidates:
|
|
return []
|
|
|
|
results = []
|
|
|
|
query_vec = np.array(query_embedding)
|
|
cand_vecs = np.array([c['embedding'] for c in candidates])
|
|
|
|
if metric == "cosine":
|
|
query_norm = np.linalg.norm(query_vec)
|
|
if query_norm == 0: query_norm = 1e-9
|
|
|
|
cand_norms = np.linalg.norm(cand_vecs, axis=1)
|
|
cand_norms[cand_norms == 0] = 1e-9
|
|
|
|
dot_products = np.dot(cand_vecs, query_vec)
|
|
scores = dot_products / (cand_norms * query_norm)
|
|
|
|
elif metric == "euclidean":
|
|
# Euclidean distance: smaller is closer.
|
|
# Convert to similarity score: 1 / (1 + distance)
|
|
dists = np.linalg.norm(cand_vecs - query_vec, axis=1)
|
|
scores = 1 / (1 + dists)
|
|
|
|
elif metric == "dot":
|
|
scores = np.dot(cand_vecs, query_vec)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown metric: {metric}")
|
|
|
|
# Sort by score descending
|
|
top_k_indices = np.argsort(scores)[-n_results:][::-1]
|
|
|
|
for idx in top_k_indices:
|
|
item = candidates[idx]
|
|
score_ = float(scores[idx])
|
|
if score_ >= score:
|
|
results.append({
|
|
"id": item['id'],
|
|
"text": item['text'],
|
|
"metadata": item['metadata'],
|
|
"score": score_
|
|
})
|
|
|
|
return results
|
|
|
|
|
|
class RAGService:
|
|
def __init__(
|
|
self,
|
|
session_dir: str,
|
|
openai_api_key: str,
|
|
openai_base_url: str,
|
|
embedding_api_key: str = None,
|
|
embedding_base_url: str = None,
|
|
embedding_model_name: str = "text-embedding-v4",
|
|
small_model_name: str = None,
|
|
rag_retrieval_count: int = 10,
|
|
rag_final_count: int = 5,
|
|
default_headers: Dict = None
|
|
):
|
|
self.session_dir = session_dir
|
|
self.openai_api_key = openai_api_key
|
|
self.openai_base_url = openai_base_url
|
|
self.embedding_api_key = embedding_api_key or openai_api_key or ""
|
|
self.embedding_base_url = embedding_base_url or openai_base_url or ""
|
|
self.embedding_model_name = embedding_model_name
|
|
self.small_model_name = small_model_name
|
|
self.rag_retrieval_count = rag_retrieval_count
|
|
self.rag_final_count = rag_final_count
|
|
|
|
self.vector_db = SimpleVectorDB(persist_dir=session_dir)
|
|
# self.openai_client = openai.OpenAI(
|
|
# api_key=openai_api_key,
|
|
# base_url=openai_base_url,
|
|
# default_headers=default_headers
|
|
# )
|
|
|
|
self.embedding_client = openai.OpenAI(
|
|
api_key=self.embedding_api_key,
|
|
base_url=self.embedding_base_url,
|
|
default_headers=default_headers
|
|
)
|
|
# else:
|
|
# self.embedding_client = self.openai_client
|
|
|
|
self.small_model_name = small_model_name
|
|
|
|
def close(self):
|
|
self.embedding_client.close()
|
|
|
|
def get_embedding(self, text: str) -> List[float]:
|
|
try:
|
|
model = self.embedding_model_name or "text-embedding-3-small"
|
|
text = text.replace("\n", " ")
|
|
|
|
response = self.embedding_client.embeddings.create(
|
|
input=[text],
|
|
model=model
|
|
)
|
|
return response.data[0].embedding
|
|
except Exception as e:
|
|
return []
|
|
|
|
def generate_context_prefix(self, user_text: str, ai_text: str) -> str:
|
|
# 暂时禁用小模型精简功能
|
|
return "[Context: General]"
|
|
# try:
|
|
# prompt = (
|
|
# f"请分析以下对话片段，并提供一个非常简要的背景标签 (例如:[Context: Nginx Config]).\\n"
|
|
# f"Dialogue:\\nUser: {user_text}\\nAI: {ai_text}\\n\\nContext Label:"
|
|
# )
|
|
# model_name = self.small_model_name
|
|
# response = self.openai_client.chat.completions.create(
|
|
# model=model_name,
|
|
# messages=[{"role": "user", "content": prompt}],
|
|
# max_tokens=20
|
|
# )
|
|
# return response.choices[0].message.content.strip()
|
|
# except Exception as e:
|
|
# return "[Context: General]"
|
|
|
|
def add_memory(self, user_msg: Dict, ai_msg: Dict, session_id: str):
|
|
user_content = user_msg.get("content", "")
|
|
if isinstance(user_content, list):
|
|
user_text = "\n".join([item.get("text", "") for item in user_content if
|
|
isinstance(item, dict) and item.get("type") == "text"])
|
|
else:
|
|
user_text = str(user_content)
|
|
|
|
ai_content = ai_msg.get("content", "")
|
|
if isinstance(ai_content, list):
|
|
ai_text = "\n".join(
|
|
[item.get("text", "") for item in ai_content if isinstance(item, dict) and item.get("type") == "text"])
|
|
else:
|
|
ai_text = str(ai_content)
|
|
|
|
prefix = self.generate_context_prefix(user_text, ai_text)
|
|
full_text = f"{prefix} User: {user_text} -> AI: {ai_text}"
|
|
|
|
embedding = self.get_embedding(full_text)
|
|
if not embedding:
|
|
return
|
|
|
|
doc_id = f"{user_msg['id']}_{ai_msg['id']}"
|
|
|
|
self.vector_db.add(
|
|
documents=[full_text],
|
|
metadatas=[{
|
|
"session_id": session_id,
|
|
"timestamp": ai_msg.get("timestamp", int(time.time())),
|
|
"user_msg_id": user_msg["id"],
|
|
"ai_msg_id": ai_msg["id"],
|
|
"type": "conversation_pair"
|
|
}],
|
|
ids=[doc_id],
|
|
embeddings=[embedding]
|
|
)
|
|
|
|
def add_document(self, text: str, metadata: Dict = None):
|
|
"""
|
|
Add a generic document to the knowledge base.
|
|
"""
|
|
if metadata is None:
|
|
metadata = {}
|
|
|
|
embedding = self.get_embedding(text)
|
|
if not embedding:
|
|
return
|
|
|
|
doc_id = str(uuid.uuid4())
|
|
self.vector_db.add(
|
|
documents=[text],
|
|
metadatas=[metadata],
|
|
ids=[doc_id],
|
|
embeddings=[embedding]
|
|
)
|
|
|
|
def search(self, query: str, session_id: str = None, score: float = 0.35, scope: str = "session",
|
|
exclude_ids: List[str] = None, full_text: bool = True, session_history: List[Dict[str, Any]] = None,
|
|
enable_rag_judgment: bool = None, ) -> List[str]:
|
|
"""
|
|
从知识库中检索与查询最相关的上下文信息。
|
|
- query: 用户查询文本
|
|
- session_id: 当前会话ID（仅在scope为"session"时使用
|
|
- score: 最小相似度分数阈值（范围0-1，默认0.2）
|
|
- scope: 检索范围，"session"（仅当前会话）或"global"（所有会话）
|
|
- exclude_ids: 要排除的消息ID列表（例如，当前对话中的消息ID，避免检索到自己）
|
|
- full_text: 是否返回完整文本（包含时间戳等元信息）还是仅返回原始文本
|
|
"""
|
|
if exclude_ids is None:
|
|
exclude_ids = []
|
|
|
|
where_clause = {}
|
|
if scope == "session" and session_id:
|
|
where_clause = {"session_id": session_id}
|
|
|
|
query_embedding = self.get_embedding(query)
|
|
if not query_embedding:
|
|
return []
|
|
|
|
results = self.vector_db.search(
|
|
query_embedding=query_embedding,
|
|
n_results=self.rag_retrieval_count,
|
|
score=score,
|
|
where=where_clause if where_clause else None
|
|
)
|
|
|
|
if not results:
|
|
return []
|
|
|
|
candidates = []
|
|
for res in results:
|
|
meta = res['metadata']
|
|
if meta.get('user_msg_id', '') in exclude_ids or meta.get('ai_msg_id', '') in exclude_ids:
|
|
continue
|
|
|
|
candidates.append(res)
|
|
|
|
candidates = candidates[:self.rag_final_count]
|
|
if not full_text:
|
|
return candidates
|
|
|
|
final_context = []
|
|
for c in candidates:
|
|
ts = c['metadata'].get('timestamp')
|
|
if ts:
|
|
time_str = str(ts)
|
|
final_context.append(f"[Time: {time_str}] {c['text']}")
|
|
else:
|
|
final_context.append(f"{c['text']}")
|
|
return final_context
|
|
|
|
def close(self):
|
|
self.embedding_client.close()
|
|
|
|
|
|
class ExternalRAGService:
|
|
def __init__(
|
|
self,
|
|
appid: str = "",
|
|
index_id: str = "main",
|
|
rag_final_count: int = 5,
|
|
enable_rag_judgment: bool = False,
|
|
api_key: str = "",
|
|
base_url: str = "",
|
|
retrieve_url: str = "",
|
|
model_name: str = "fast",
|
|
default_headers: Dict[str, str] = None
|
|
):
|
|
self.appid = appid
|
|
self.rag_final_count = rag_final_count
|
|
self.index_id = index_id
|
|
self.enable_rag_judgment = enable_rag_judgment
|
|
self.retrieve_url = retrieve_url
|
|
|
|
self.judgment_api_key = api_key
|
|
self.judgment_base_url = base_url
|
|
self.judgment_model_name = model_name
|
|
self.judgment_default_headers = default_headers or {}
|
|
|
|
if public is not None:
|
|
user_info = public.get_user_info()
|
|
self.uid = str(user_info.get('uid', ''))
|
|
self.access_key = user_info.get('access_key', '')
|
|
else:
|
|
self.uid = ''
|
|
self.access_key = ''
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
def _should_use_rag(self, user_input: str, session_history: List[Dict[str, Any]] = None) -> Dict[str, Any]:
|
|
"""
|
|
判断是否需要 RAG 检索（内部方法）
|
|
- user_input: 当前用户输入
|
|
- session_history: 最近对话历史（用于结合上下文判断）
|
|
"""
|
|
prompt = """You are a RAG retrieval judgment assistant. Based on user input and conversation context, determine whether relevant information needs to be retrieved from the external knowledge base.
|
|
|
|
Judgment criteria:
|
|
- If the question involves specific domain knowledge, technical documentation, product information, FAQs, etc., retrieval should be performed
|
|
- If it is simple casual chat, emotional exchange, or general knowledge, no retrieval needed
|
|
- If the context clearly points to a specific domain question, retrieval should be performed
|
|
|
|
Return in JSON format:
|
|
{
|
|
"use_rag": true/false,
|
|
"confidence": 0.0-1.0,
|
|
"reason": "judgment reason"
|
|
}"""
|
|
|
|
try:
|
|
from ..chat_client.single_agent import SingleAgent
|
|
single_agent = SingleAgent(
|
|
api_key=self.judgment_api_key,
|
|
base_url=self.judgment_base_url,
|
|
model_name=self.judgment_model_name,
|
|
default_headers=self.judgment_default_headers,
|
|
temperature=0.1
|
|
)
|
|
|
|
notime = time.time()
|
|
if session_history:
|
|
formatted_history = []
|
|
for msg in session_history[-6:]:
|
|
role = msg.get("role", "user")
|
|
if role not in ("user", "assistant"):
|
|
continue
|
|
if role == "user":
|
|
content = msg.get("content", "")
|
|
if isinstance(content, list):
|
|
content = " ".join([c.get("text", "") for c in content if c.get("type") == "text"])
|
|
formatted_history.append({"role": role, "content": str(content)})
|
|
else:
|
|
content = msg.get("content", "")
|
|
if isinstance(content, list):
|
|
content = " ".join([c.get("text", "") for c in content if c.get("type") == "text"])
|
|
if not content:
|
|
content = msg.get("reasoning_content", "")
|
|
if not content:
|
|
continue
|
|
formatted_history.append({"role": role, "content": str(content)})
|
|
# logging.info(f"RAG Judgment Session History:{formatted_history}")
|
|
result = single_agent.chat(
|
|
prompt=prompt,
|
|
messages=formatted_history,
|
|
input_text=f"Current user input: {user_input}",
|
|
json_response=True,
|
|
temperature=0.1,
|
|
extra_body={"enable_thinking": False}
|
|
)
|
|
else:
|
|
result = single_agent.chat(
|
|
prompt=prompt,
|
|
input_text=user_input,
|
|
json_response=True,
|
|
temperature=0.1,
|
|
extra_body={"enable_thinking": False}
|
|
)
|
|
|
|
single_agent.close()
|
|
# logging.info(f"RAG Judgment Result:{time.time()-notime} {result}")
|
|
if result["success"]:
|
|
return {
|
|
"use_rag": result["data"].get("use_rag", False),
|
|
"confidence": result["data"].get("confidence", 0.0),
|
|
"reason": result["data"].get("reason", "")
|
|
}
|
|
else:
|
|
return {
|
|
"use_rag": True,
|
|
"confidence": 0.0,
|
|
"reason": result.get("error", lang("Judgment failed, defaulting to retrieve"))
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"use_rag": True,
|
|
"confidence": 0.0,
|
|
"reason": lang(f"Judgment exception: {str(e)}, defaulting to retrieve")
|
|
}
|
|
|
|
def search(self, query: str, score: float = 0.2, scope: str = "session", full_text: bool = True,
|
|
enable_rag_judgment: bool = None, session_history: List[Dict[str, Any]] = None) -> List[str]:
|
|
"""
|
|
从外部知识库检索与查询最相关的上下文信息。
|
|
- query: 用户查询文本
|
|
- score: 最小相似度分数阈值（范围0-1，默认0.2）
|
|
- enable_rag_judgment: 覆盖类级别的判断开关
|
|
- session_history: 最近对话历史（用于 RAG 判断时结合上下文）
|
|
"""
|
|
if enable_rag_judgment is None:
|
|
enable_rag_judgment = self.enable_rag_judgment
|
|
|
|
if enable_rag_judgment:
|
|
rag_judgment = self._should_use_rag(query, session_history)
|
|
if not rag_judgment.get("use_rag", True):
|
|
return []
|
|
try:
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"uid": self.uid,
|
|
"access-key": self.access_key,
|
|
"appid": self.appid
|
|
}
|
|
|
|
payload = {
|
|
"query": query,
|
|
"dense_similarity_top_k": 100,
|
|
"enable_reranking": True,
|
|
"rerank_min_score": 0.5,
|
|
"index_id": self.index_id,
|
|
"rerank_top_n": self.rag_final_count
|
|
}
|
|
response = requests.post(
|
|
self.retrieve_url,
|
|
headers=headers,
|
|
json=payload,
|
|
timeout=30
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
return []
|
|
|
|
result = response.json()
|
|
|
|
if not result.get("status") or "data" not in result:
|
|
return []
|
|
|
|
nodes = result["data"].get("Nodes", [])
|
|
|
|
filtered_docs = []
|
|
for node in nodes:
|
|
node_score = node.get("Score", 0)
|
|
if node_score >= score:
|
|
text = node.get("Text", "")
|
|
if text:
|
|
filtered_docs.append(text)
|
|
|
|
return filtered_docs[:self.rag_final_count]
|
|
|
|
except Exception as e:
|
|
return []
|