321 lines
10 KiB
Python
321 lines
10 KiB
Python
"""
|
||
Elasticsearch 连接和工具模块
|
||
用于聊天记录存储和向量搜索
|
||
"""
|
||
|
||
from elasticsearch import Elasticsearch, helpers
|
||
from datetime import datetime
|
||
from typing import List, Dict, Any, Optional
|
||
import logging
|
||
import json
|
||
import openai
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ==================== 配置 ====================
|
||
|
||
# ES 配置
|
||
ES_CONFIG = {
|
||
"host": "http://222.128.1.157:19200",
|
||
"index_chat_history": "agent_chat_history", # 聊天记录索引
|
||
}
|
||
|
||
# Embedding 配置
|
||
EMBEDDING_CONFIG = {
|
||
"api_key": "dummy",
|
||
"base_url": "http://222.128.1.157:18008/v1",
|
||
"model": "qwen3-embedding-8b",
|
||
"dims": 4096, # 向量维度
|
||
}
|
||
|
||
# ==================== ES 客户端 ====================
|
||
|
||
class ESClient:
|
||
"""Elasticsearch 客户端封装"""
|
||
|
||
def __init__(self):
|
||
self.es = Elasticsearch([ES_CONFIG["host"]], request_timeout=60)
|
||
self.chat_index = ES_CONFIG["index_chat_history"]
|
||
|
||
# 初始化 OpenAI 客户端用于 embedding
|
||
self.embedding_client = openai.OpenAI(
|
||
api_key=EMBEDDING_CONFIG["api_key"],
|
||
base_url=EMBEDDING_CONFIG["base_url"],
|
||
)
|
||
self.embedding_model = EMBEDDING_CONFIG["model"]
|
||
|
||
# 初始化索引
|
||
self.create_chat_history_index()
|
||
|
||
def create_chat_history_index(self):
|
||
"""创建聊天记录索引"""
|
||
if self.es.indices.exists(index=self.chat_index):
|
||
logger.info(f"索引 {self.chat_index} 已存在")
|
||
return
|
||
|
||
mappings = {
|
||
"properties": {
|
||
"session_id": {"type": "keyword"}, # 会话ID
|
||
"user_id": {"type": "keyword"}, # 用户ID
|
||
"user_nickname": {"type": "text"}, # 用户昵称
|
||
"user_avatar": {"type": "keyword"}, # 用户头像URL
|
||
"message_type": {"type": "keyword"}, # user / assistant
|
||
"message": {"type": "text"}, # 消息内容
|
||
"message_embedding": { # 消息向量
|
||
"type": "dense_vector",
|
||
"dims": EMBEDDING_CONFIG["dims"],
|
||
"index": True,
|
||
"similarity": "cosine"
|
||
},
|
||
"plan": {"type": "text"}, # 执行计划（仅 assistant）
|
||
"steps": {"type": "text"}, # 执行步骤（仅 assistant）
|
||
"timestamp": {"type": "date"}, # 时间戳
|
||
"created_at": {"type": "date"}, # 创建时间
|
||
}
|
||
}
|
||
|
||
self.es.indices.create(index=self.chat_index, body={"mappings": mappings})
|
||
logger.info(f"创建索引: {self.chat_index}")
|
||
|
||
def generate_embedding(self, text: str) -> List[float]:
|
||
"""生成文本向量"""
|
||
try:
|
||
if not text or len(text.strip()) == 0:
|
||
return []
|
||
|
||
# 截断过长文本
|
||
text = text[:16000] if len(text) > 16000 else text
|
||
|
||
response = self.embedding_client.embeddings.create(
|
||
model=self.embedding_model,
|
||
input=[text]
|
||
)
|
||
return response.data[0].embedding
|
||
except Exception as e:
|
||
logger.error(f"Embedding 生成失败: {e}")
|
||
return []
|
||
|
||
def save_chat_message(
|
||
self,
|
||
session_id: str,
|
||
user_id: str,
|
||
user_nickname: str,
|
||
user_avatar: str,
|
||
message_type: str, # "user" or "assistant"
|
||
message: str,
|
||
plan: Optional[str] = None,
|
||
steps: Optional[str] = None,
|
||
) -> str:
|
||
"""
|
||
保存聊天消息
|
||
|
||
Args:
|
||
session_id: 会话ID
|
||
user_id: 用户ID
|
||
user_nickname: 用户昵称
|
||
user_avatar: 用户头像URL
|
||
message_type: 消息类型 (user/assistant)
|
||
message: 消息内容
|
||
plan: 执行计划（可选）
|
||
steps: 执行步骤（可选）
|
||
|
||
Returns:
|
||
文档ID
|
||
"""
|
||
try:
|
||
# 生成向量
|
||
embedding = self.generate_embedding(message)
|
||
|
||
doc = {
|
||
"session_id": session_id,
|
||
"user_id": user_id,
|
||
"user_nickname": user_nickname,
|
||
"user_avatar": user_avatar,
|
||
"message_type": message_type,
|
||
"message": message,
|
||
"message_embedding": embedding if embedding else None,
|
||
"plan": plan,
|
||
"steps": steps,
|
||
"timestamp": datetime.now(),
|
||
"created_at": datetime.now(),
|
||
}
|
||
|
||
result = self.es.index(index=self.chat_index, body=doc)
|
||
logger.info(f"保存聊天记录: {result['_id']}")
|
||
return result["_id"]
|
||
|
||
except Exception as e:
|
||
logger.error(f"保存聊天记录失败: {e}")
|
||
raise
|
||
|
||
def get_chat_sessions(self, user_id: str, limit: int = 50) -> List[Dict[str, Any]]:
|
||
"""
|
||
获取用户的聊天会话列表
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
limit: 返回数量
|
||
|
||
Returns:
|
||
会话列表，每个会话包含：session_id, last_message, last_timestamp
|
||
"""
|
||
try:
|
||
# 聚合查询：按 session_id 分组，获取每个会话的最后一条消息
|
||
query = {
|
||
"query": {
|
||
"term": {"user_id": user_id}
|
||
},
|
||
"aggs": {
|
||
"sessions": {
|
||
"terms": {
|
||
"field": "session_id",
|
||
"size": limit,
|
||
"order": {"last_message": "desc"}
|
||
},
|
||
"aggs": {
|
||
"last_message": {
|
||
"max": {"field": "timestamp"}
|
||
},
|
||
"last_message_content": {
|
||
"top_hits": {
|
||
"size": 1,
|
||
"sort": [{"timestamp": {"order": "desc"}}],
|
||
"_source": ["message", "timestamp", "message_type"]
|
||
}
|
||
}
|
||
}
|
||
}
|
||
},
|
||
"size": 0
|
||
}
|
||
|
||
result = self.es.search(index=self.chat_index, body=query)
|
||
|
||
sessions = []
|
||
for bucket in result["aggregations"]["sessions"]["buckets"]:
|
||
session_data = bucket["last_message_content"]["hits"]["hits"][0]["_source"]
|
||
sessions.append({
|
||
"session_id": bucket["key"],
|
||
"last_message": session_data["message"],
|
||
"last_timestamp": session_data["timestamp"],
|
||
"message_count": bucket["doc_count"],
|
||
})
|
||
|
||
return sessions
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取会话列表失败: {e}")
|
||
return []
|
||
|
||
def get_chat_history(
|
||
self,
|
||
session_id: str,
|
||
limit: int = 100
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
获取指定会话的聊天历史
|
||
|
||
Args:
|
||
session_id: 会话ID
|
||
limit: 返回数量
|
||
|
||
Returns:
|
||
聊天记录列表
|
||
"""
|
||
try:
|
||
query = {
|
||
"query": {
|
||
"term": {"session_id": session_id}
|
||
},
|
||
"sort": [{"timestamp": {"order": "asc"}}],
|
||
"size": limit
|
||
}
|
||
|
||
result = self.es.search(index=self.chat_index, body=query)
|
||
|
||
messages = []
|
||
for hit in result["hits"]["hits"]:
|
||
doc = hit["_source"]
|
||
messages.append({
|
||
"message_type": doc["message_type"],
|
||
"message": doc["message"],
|
||
"plan": doc.get("plan"),
|
||
"steps": doc.get("steps"),
|
||
"timestamp": doc["timestamp"],
|
||
})
|
||
|
||
return messages
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取聊天历史失败: {e}")
|
||
return []
|
||
|
||
def search_chat_history(
|
||
self,
|
||
user_id: str,
|
||
query_text: str,
|
||
top_k: int = 10
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
向量搜索聊天历史
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
query_text: 查询文本
|
||
top_k: 返回数量
|
||
|
||
Returns:
|
||
相关聊天记录列表
|
||
"""
|
||
try:
|
||
# 生成查询向量
|
||
query_embedding = self.generate_embedding(query_text)
|
||
if not query_embedding:
|
||
return []
|
||
|
||
# 向量搜索
|
||
query = {
|
||
"query": {
|
||
"bool": {
|
||
"must": [
|
||
{"term": {"user_id": user_id}},
|
||
{
|
||
"script_score": {
|
||
"query": {"match_all": {}},
|
||
"script": {
|
||
"source": "cosineSimilarity(params.query_vector, 'message_embedding') + 1.0",
|
||
"params": {"query_vector": query_embedding}
|
||
}
|
||
}
|
||
}
|
||
]
|
||
}
|
||
},
|
||
"size": top_k
|
||
}
|
||
|
||
result = self.es.search(index=self.chat_index, body=query)
|
||
|
||
messages = []
|
||
for hit in result["hits"]["hits"]:
|
||
doc = hit["_source"]
|
||
messages.append({
|
||
"session_id": doc["session_id"],
|
||
"message_type": doc["message_type"],
|
||
"message": doc["message"],
|
||
"timestamp": doc["timestamp"],
|
||
"score": hit["_score"],
|
||
})
|
||
|
||
return messages
|
||
|
||
except Exception as e:
|
||
logger.error(f"向量搜索失败: {e}")
|
||
return []
|
||
|
||
|
||
# ==================== 全局实例 ====================
|
||
|
||
# 创建全局 ES 客户端
|
||
es_client = ESClient()
|