This commit is contained in:
@@ -10,7 +10,6 @@ import asyncio
|
|||||||
from ..main_graph.utils.main_graph_builder import build_react_main_graph
|
from ..main_graph.utils.main_graph_builder import build_react_main_graph
|
||||||
from ..main_graph.tools.graph_tools import AVAILABLE_TOOLS, TOOLS_BY_NAME
|
from ..main_graph.tools.graph_tools import AVAILABLE_TOOLS, TOOLS_BY_NAME
|
||||||
from ..main_graph.config import set_stream_writer
|
from ..main_graph.config import set_stream_writer
|
||||||
from ..model_services.chat_services import get_all_chat_services, LocalVLLMChatProvider
|
|
||||||
from ..main_graph.utils.rag_initializer import init_rag_tool
|
from ..main_graph.utils.rag_initializer import init_rag_tool
|
||||||
from ..core.intent_classifier import get_intent_classifier
|
from ..core.intent_classifier import get_intent_classifier
|
||||||
from ..logger import info, warning, error
|
from ..logger import info, warning, error
|
||||||
@@ -33,18 +32,10 @@ class AIAgentService:
|
|||||||
async def initialize(self):
|
async def initialize(self):
|
||||||
# 0. 初始化 Mem0 客户端
|
# 0. 初始化 Mem0 客户端
|
||||||
from ..memory.mem0_client import Mem0Client
|
from ..memory.mem0_client import Mem0Client
|
||||||
# 创建一个临时的 LLM 用于 Mem0(用第一个可用的)
|
self.mem0_client = Mem0Client()
|
||||||
chat_services = get_all_chat_services()
|
|
||||||
temp_llm = None
|
|
||||||
if chat_services:
|
|
||||||
temp_llm = list(chat_services.values())[0]
|
|
||||||
self.mem0_client = Mem0Client(temp_llm)
|
|
||||||
|
|
||||||
# 1. 初始化 RAG 工具(如果需要)
|
# 1. 初始化 RAG 工具(如果需要)
|
||||||
def create_local_llm():
|
rag_tool = await init_rag_tool()
|
||||||
provider = LocalVLLMChatProvider()
|
|
||||||
return provider.get_service()
|
|
||||||
rag_tool = await init_rag_tool(create_local_llm)
|
|
||||||
if rag_tool:
|
if rag_tool:
|
||||||
self.tools.append(rag_tool)
|
self.tools.append(rag_tool)
|
||||||
self.tools_by_name[rag_tool.name] = rag_tool
|
self.tools_by_name[rag_tool.name] = rag_tool
|
||||||
|
|||||||
@@ -20,12 +20,11 @@ def is_initialized() -> bool:
|
|||||||
return _initialized
|
return _initialized
|
||||||
|
|
||||||
|
|
||||||
async def init_rag_tool(local_llm_creator, force: bool = False):
|
async def init_rag_tool(force: bool = False):
|
||||||
"""
|
"""
|
||||||
初始化 RAG 工具(注册到模块级变量)
|
初始化 RAG 工具(注册到模块级变量,内部获取所需服务)
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
local_llm_creator: 返回 LLM 实例的函数
|
|
||||||
force: 是否强制重新初始化
|
force: 是否强制重新初始化
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -39,20 +38,22 @@ async def init_rag_tool(local_llm_creator, force: bool = False):
|
|||||||
return _rag_tool
|
return _rag_tool
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
from app.model_services.chat_services import get_chat_service
|
||||||
|
|
||||||
info("🔄 正在初始化 RAG 检索系统...")
|
info("🔄 正在初始化 RAG 检索系统...")
|
||||||
embeddings = get_embedding_service()
|
embeddings = get_embedding_service()
|
||||||
retriever = create_parent_hybrid_retriever(
|
retriever = create_parent_hybrid_retriever(
|
||||||
collection_name="rag_documents",
|
collection_name="rag_documents",
|
||||||
search_k=5,
|
search_k=5,
|
||||||
embeddings=embeddings
|
embeddings=embeddings,
|
||||||
)
|
)
|
||||||
rewrite_llm = local_llm_creator()
|
rewrite_llm = get_chat_service()
|
||||||
|
|
||||||
rag_tool = create_rag_tool(
|
rag_tool = create_rag_tool(
|
||||||
retriever=retriever,
|
retriever=retriever,
|
||||||
llm=rewrite_llm,
|
llm=rewrite_llm,
|
||||||
num_queries=3,
|
num_queries=3,
|
||||||
rerank_top_n=5
|
rerank_top_n=5,
|
||||||
)
|
)
|
||||||
|
|
||||||
_rag_tool = rag_tool
|
_rag_tool = rag_tool
|
||||||
|
|||||||
@@ -1,33 +1,36 @@
|
|||||||
from app.config import (
|
|
||||||
LLM_API_KEY, ZHIPUAI_API_KEY,
|
|
||||||
VLLM_BASE_URL, QDRANT_URL, QDRANT_COLLECTION_NAME, QDRANT_API_KEY,
|
|
||||||
LLAMACPP_EMBEDDING_URL, LLAMACPP_API_KEY,
|
|
||||||
ZHIPU_EMBEDDING_MODEL, ZHIPU_API_BASE
|
|
||||||
)
|
|
||||||
from ..model_services import get_embedding_service
|
|
||||||
from app.logger import info, warning, error
|
|
||||||
import time
|
|
||||||
"""
|
"""
|
||||||
Mem0 记忆层客户端封装模块
|
Mem0 记忆层客户端封装模块
|
||||||
负责 Mem0 的初始化、检索和存储
|
负责 Mem0 的初始化、检索和存储
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from typing import Optional, List, Dict
|
import time
|
||||||
|
from typing import Optional, List
|
||||||
|
|
||||||
from mem0 import AsyncMemory
|
from mem0 import AsyncMemory
|
||||||
|
|
||||||
|
from app.config import (
|
||||||
|
LLM_API_KEY,
|
||||||
|
ZHIPUAI_API_KEY,
|
||||||
|
VLLM_BASE_URL,
|
||||||
|
QDRANT_URL,
|
||||||
|
QDRANT_COLLECTION_NAME,
|
||||||
|
QDRANT_API_KEY,
|
||||||
|
LLAMACPP_EMBEDDING_URL,
|
||||||
|
LLAMACPP_API_KEY,
|
||||||
|
ZHIPU_EMBEDDING_MODEL,
|
||||||
|
ZHIPU_API_BASE,
|
||||||
|
)
|
||||||
|
from app.logger import info, warning, error
|
||||||
|
from app.model_services import get_embedding_service
|
||||||
|
from app.model_services.chat_services import get_chat_service
|
||||||
|
|
||||||
|
|
||||||
class Mem0Client:
|
class Mem0Client:
|
||||||
"""Mem0 异步客户端封装类"""
|
"""Mem0 异步客户端封装类"""
|
||||||
|
|
||||||
def __init__(self, llm_instance):
|
def __init__(self):
|
||||||
"""
|
"""初始化 Mem0 客户端(内部获取所需服务)"""
|
||||||
初始化 Mem0 客户端
|
|
||||||
|
|
||||||
Args:
|
|
||||||
llm_instance: LangChain LLM 实例(用于事实提取)
|
|
||||||
"""
|
|
||||||
self.llm = llm_instance
|
|
||||||
self.mem0: Optional[AsyncMemory] = None
|
self.mem0: Optional[AsyncMemory] = None
|
||||||
self._initialized = False
|
self._initialized = False
|
||||||
|
|
||||||
@@ -35,7 +38,7 @@ class Mem0Client:
|
|||||||
"""异步初始化 Mem0 客户端,并进行实际连接测试"""
|
"""异步初始化 Mem0 客户端,并进行实际连接测试"""
|
||||||
if self._initialized:
|
if self._initialized:
|
||||||
return
|
return
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 获取可用的 embedding 服务并确定维度
|
# 获取可用的 embedding 服务并确定维度
|
||||||
info("🔄 正在获取嵌入服务...")
|
info("🔄 正在获取嵌入服务...")
|
||||||
@@ -43,14 +46,16 @@ class Mem0Client:
|
|||||||
test_embedding = embeddings.embed_query("test")
|
test_embedding = embeddings.embed_query("test")
|
||||||
embedding_dim = len(test_embedding)
|
embedding_dim = len(test_embedding)
|
||||||
info(f"✅ 嵌入服务可用,向量维度: {embedding_dim}")
|
info(f"✅ 嵌入服务可用,向量维度: {embedding_dim}")
|
||||||
|
|
||||||
# 构建 embedder 配置 - 改进的方法
|
# 构建 embedder 配置
|
||||||
# 检查本地 provider
|
from app.model_services.embedding_services import (
|
||||||
from ..model_services.embedding_services import LocalLlamaCppEmbeddingProvider, ZhipuEmbeddingProvider
|
LocalLlamaCppEmbeddingProvider,
|
||||||
|
ZhipuEmbeddingProvider,
|
||||||
|
)
|
||||||
|
|
||||||
embedder_config = None
|
embedder_config = None
|
||||||
local_provider = LocalLlamaCppEmbeddingProvider()
|
local_provider = LocalLlamaCppEmbeddingProvider()
|
||||||
|
|
||||||
if local_provider.is_available():
|
if local_provider.is_available():
|
||||||
info("✅ 使用本地 llama.cpp 作为 mem0 embedder")
|
info("✅ 使用本地 llama.cpp 作为 mem0 embedder")
|
||||||
embedder_config = {
|
embedder_config = {
|
||||||
@@ -59,22 +64,20 @@ class Mem0Client:
|
|||||||
"model": "Qwen3-Embedding-0.6B-Q8_0",
|
"model": "Qwen3-Embedding-0.6B-Q8_0",
|
||||||
"api_key": LLAMACPP_API_KEY or "dummy-key",
|
"api_key": LLAMACPP_API_KEY or "dummy-key",
|
||||||
"openai_base_url": LLAMACPP_EMBEDDING_URL,
|
"openai_base_url": LLAMACPP_EMBEDDING_URL,
|
||||||
}
|
},
|
||||||
}
|
}
|
||||||
else:
|
else:
|
||||||
# 检查智谱
|
# 检查智谱
|
||||||
zhipu_provider = ZhipuEmbeddingProvider()
|
zhipu_provider = ZhipuEmbeddingProvider()
|
||||||
if zhipu_provider.is_available():
|
if zhipu_provider.is_available():
|
||||||
info("✅ 使用智谱 API 作为 mem0 embedder")
|
info("✅ 使用智谱 API 作为 mem0 embedder")
|
||||||
# 使用自定义 embedder 或者 openai 兼容方式
|
|
||||||
# 注意:这里我们使用一个特殊的配置方法
|
|
||||||
embedder_config = {
|
embedder_config = {
|
||||||
"provider": "openai",
|
"provider": "openai",
|
||||||
"config": {
|
"config": {
|
||||||
"model": ZHIPU_EMBEDDING_MODEL,
|
"model": ZHIPU_EMBEDDING_MODEL,
|
||||||
"api_key": ZHIPUAI_API_KEY,
|
"api_key": ZHIPUAI_API_KEY,
|
||||||
"openai_base_url": ZHIPU_API_BASE,
|
"openai_base_url": ZHIPU_API_BASE,
|
||||||
}
|
},
|
||||||
}
|
}
|
||||||
else:
|
else:
|
||||||
# 都不可用,使用 dummy 配置并警告
|
# 都不可用,使用 dummy 配置并警告
|
||||||
@@ -83,12 +86,17 @@ class Mem0Client:
|
|||||||
"provider": "openai",
|
"provider": "openai",
|
||||||
"config": {
|
"config": {
|
||||||
"model": "text-embedding-ada-002",
|
"model": "text-embedding-ada-002",
|
||||||
"api_key": "dummy-key",
|
"api_key": "***",
|
||||||
"openai_base_url": "http://localhost:8080/v1",
|
"openai_base_url": "http://localhost:8080/v1",
|
||||||
}
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
# Mem0 配置 - 简化配置,先确保能启动
|
# 获取 LLM 服务(内部获取)
|
||||||
|
info("🔄 正在获取 LLM 服务...")
|
||||||
|
chat_llm = get_chat_service()
|
||||||
|
info("✅ LLM 服务获取成功")
|
||||||
|
|
||||||
|
# Mem0 配置
|
||||||
info("🔄 正在构建 Mem0 配置...")
|
info("🔄 正在构建 Mem0 配置...")
|
||||||
config = {
|
config = {
|
||||||
"vector_store": {
|
"vector_store": {
|
||||||
@@ -98,7 +106,7 @@ class Mem0Client:
|
|||||||
"api_key": QDRANT_API_KEY,
|
"api_key": QDRANT_API_KEY,
|
||||||
"collection_name": QDRANT_COLLECTION_NAME,
|
"collection_name": QDRANT_COLLECTION_NAME,
|
||||||
"embedding_model_dims": embedding_dim,
|
"embedding_model_dims": embedding_dim,
|
||||||
}
|
},
|
||||||
},
|
},
|
||||||
"llm": {
|
"llm": {
|
||||||
"provider": "openai",
|
"provider": "openai",
|
||||||
@@ -108,31 +116,30 @@ class Mem0Client:
|
|||||||
"openai_base_url": VLLM_BASE_URL or ZHIPU_API_BASE,
|
"openai_base_url": VLLM_BASE_URL or ZHIPU_API_BASE,
|
||||||
"temperature": 0.1,
|
"temperature": 0.1,
|
||||||
"max_tokens": 2000,
|
"max_tokens": 2000,
|
||||||
}
|
},
|
||||||
},
|
},
|
||||||
"embedder": embedder_config,
|
"embedder": embedder_config,
|
||||||
"version": "v1.1"
|
"version": "v1.1",
|
||||||
}
|
}
|
||||||
|
|
||||||
info("🔄 正在初始化 Mem0 实例...")
|
info("🔄 正在初始化 Mem0 实例...")
|
||||||
self.mem0 = AsyncMemory.from_config(config)
|
self.mem0 = AsyncMemory.from_config(config)
|
||||||
info("✅ Mem0 配置加载成功")
|
info("✅ Mem0 配置加载成功")
|
||||||
|
|
||||||
# 尝试进行连接测试,但失败不会阻止初始化
|
# 尝试进行连接测试,但失败不会阻止初始化
|
||||||
try:
|
try:
|
||||||
info("🔄 正在测试 Mem0 连接...")
|
info("🔄 正在测试 Mem0 连接...")
|
||||||
# 使用短超时的测试
|
|
||||||
await asyncio.wait_for(
|
await asyncio.wait_for(
|
||||||
self.mem0.search("ping", user_id="test", limit=1),
|
self.mem0.search("ping", user_id="test", limit=1),
|
||||||
timeout=10.0
|
timeout=10.0,
|
||||||
)
|
)
|
||||||
info("✅ Mem0 连接测试成功")
|
info("✅ Mem0 连接测试成功")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
warning(f"⚠️ Mem0 连接测试遇到问题(但继续使用): {e}")
|
warning(f"⚠️ Mem0 连接测试遇到问题(但继续使用): {e}")
|
||||||
|
|
||||||
self._initialized = True
|
self._initialized = True
|
||||||
info("🎉 Mem0 初始化完成")
|
info("🎉 Mem0 初始化完成")
|
||||||
|
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
error("❌ Mem0 初始化超时")
|
error("❌ Mem0 初始化超时")
|
||||||
self.mem0 = None
|
self.mem0 = None
|
||||||
@@ -140,11 +147,14 @@ class Mem0Client:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
error(f"❌ Mem0 初始化失败: {e}")
|
error(f"❌ Mem0 初始化失败: {e}")
|
||||||
import traceback
|
import traceback
|
||||||
|
|
||||||
error(f"详细错误信息:\n{traceback.format_exc()}")
|
error(f"详细错误信息:\n{traceback.format_exc()}")
|
||||||
self.mem0 = None
|
self.mem0 = None
|
||||||
self._initialized = False
|
self._initialized = False
|
||||||
|
|
||||||
async def search_memories(self, query: str, user_id: str, limit: int = 5) -> List[str]:
|
async def search_memories(
|
||||||
|
self, query: str, user_id: str, limit: int = 5
|
||||||
|
) -> List[str]:
|
||||||
"""
|
"""
|
||||||
检索相关记忆
|
检索相关记忆
|
||||||
|
|
||||||
@@ -163,7 +173,7 @@ class Mem0Client:
|
|||||||
try:
|
try:
|
||||||
memories = await asyncio.wait_for(
|
memories = await asyncio.wait_for(
|
||||||
self.mem0.search(query, user_id=user_id, limit=limit),
|
self.mem0.search(query, user_id=user_id, limit=limit),
|
||||||
timeout=30.0
|
timeout=30.0,
|
||||||
)
|
)
|
||||||
|
|
||||||
if memories and "results" in memories:
|
if memories and "results" in memories:
|
||||||
@@ -183,17 +193,25 @@ class Mem0Client:
|
|||||||
return []
|
return []
|
||||||
|
|
||||||
async def add_memories(self, messages, user_id):
|
async def add_memories(self, messages, user_id):
|
||||||
if not self.mem0:
|
"""添加记忆"""
|
||||||
return False
|
if not self.mem0:
|
||||||
try:
|
return False
|
||||||
start = time.time()
|
try:
|
||||||
info(f"📝 开始 Mem0 add,消息数: {len(messages)}")
|
start = time.time()
|
||||||
await asyncio.wait_for(
|
info(f"📝 开始 Mem0 add,消息数: {len(messages)}")
|
||||||
self.mem0.add(messages, user_id=user_id, metadata={"type": "conversation"}),
|
await asyncio.wait_for(
|
||||||
timeout=60.0
|
self.mem0.add(
|
||||||
)
|
messages, user_id=user_id, metadata={"type": "conversation"}
|
||||||
info(f"✅ Mem0 add 完成,耗时: {time.time() - start:.2f}s")
|
),
|
||||||
return True
|
timeout=60.0,
|
||||||
except asyncio.TimeoutError:
|
)
|
||||||
error(f"❌ Mem0 记忆添加超时 (60s),已等待 {time.time() - start:.2f}s")
|
info(f"✅ Mem0 add 完成,耗时: {time.time() - start:.2f}s")
|
||||||
return False
|
return True
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
error(
|
||||||
|
f"❌ Mem0 记忆添加超时 (60s),已等待 {time.time() - start:.2f}s"
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
except Exception as e:
|
||||||
|
error(f"❌ Mem0 add 失败: {e}")
|
||||||
|
return False
|
||||||
|
|||||||
@@ -31,25 +31,25 @@ TEST_CASES = [
|
|||||||
"query": "吕布的事迹?",
|
"query": "吕布的事迹?",
|
||||||
"description": "测试快速 RAG 分支"
|
"description": "测试快速 RAG 分支"
|
||||||
},
|
},
|
||||||
# # 测试3: 需要推理的复杂问题 - 应该直接到 React 循环
|
# 测试3: 需要推理的复杂问题 - 应该直接到 React 循环
|
||||||
# {
|
{
|
||||||
# "name": "复杂推理测试",
|
"name": "复杂推理测试",
|
||||||
# "query": "请帮我分析:如果我有10万元,想要在一年内获得15%的收益,有哪些低风险的投资方案?",
|
"query": "请帮我分析:如果我有10万元,想要在一年内获得15%的收益,有哪些低风险的投资方案?",
|
||||||
# "description": "测试 React 循环推理分支"
|
"description": "测试 React 循环推理分支"
|
||||||
# },
|
},
|
||||||
# # 测试4: 需要工具调用的问题
|
# # 测试4: 需要工具调用的问题
|
||||||
# {
|
# {
|
||||||
# "name": "工具调用测试",
|
# "name": "工具调用测试",
|
||||||
# "query": "搜索一下今天的天气怎么样",
|
# "query": "搜索一下今天的天气怎么样",
|
||||||
# "description": "测试工具调用分支"
|
# "description": "测试工具调用分支"
|
||||||
# },
|
# },
|
||||||
# # 测试5: 带记忆的对话
|
# 测试5: 带记忆的对话
|
||||||
# {
|
{
|
||||||
# "name": "记忆测试",
|
"name": "记忆测试",
|
||||||
# "query": "你刚才回答了我什么问题?",
|
"query": "你刚才回答了我什么问题?",
|
||||||
# "description": "测试记忆检索分支",
|
"description": "测试记忆检索分支",
|
||||||
# "thread_id": "test_memory_thread"
|
"thread_id": "test_memory_thread"
|
||||||
# }
|
}
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user