"""
NetBrain MCP 提示模板系统
此模块提供了符合MCP协议的提示模板系统，支持：
1. 模板注册和管理
2. 带有占位符的模板渲染
3. 支持字符串和结构化消息格式
4. 与资源系统集成
"""
from typing import Dict, List, Any, Optional, Union, Callable
import logging
import sys
import json
import os
import datetime
import re
from dataclasses import dataclass, asdict, field
# 设置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.StreamHandler(stream=sys.stdout)
]
)
logger = logging.getLogger("mcp_templates")
# 消息对象
@dataclass
class Message:
"""基础消息类"""
role: str
content: str
@dataclass
class UserMessage(Message):
"""用户消息"""
def __init__(self, content: str):
super().__init__(role="user", content=content)
@dataclass
class AssistantMessage(Message):
"""助手消息"""
def __init__(self, content: str):
super().__init__(role="assistant", content=content)
@dataclass
class SystemMessage(Message):
"""系统消息"""
def __init__(self, content: str):
super().__init__(role="system", content=content)
@dataclass
class PromptTemplate:
"""提示模板"""
name: str
description: str
fn: Callable
arguments: List[Dict[str, Any]] = field(default_factory=list)
class TemplateManager:
"""高级提示模板管理器"""
def __init__(self):
self.templates = {}
# 缓存目录
self.templates_dir = os.path.join(os.getcwd(), "templates")
if not os.path.exists(self.templates_dir):
try:
os.makedirs(self.templates_dir)
logger.info(f"已创建模板目录: {self.templates_dir}")
except Exception as e:
logger.warning(f"无法创建模板目录: {str(e)}")
logger.info("模板管理器初始化完成")
def register_template(self, name: str = None, description: str = None):
"""
注册模板装饰器
Args:
name: 模板名称（可选，默认使用函数名）
description: 模板描述（可选）
Returns:
装饰器函数
"""
def decorator(func):
nonlocal name, description
# 如果没有提供名称，使用函数名
template_name = name or func.__name__
# 如果没有提供描述，使用函数文档
template_description = description or (func.__doc__ or "").strip()
if template_name in self.templates:
logger.warning(f"模板已存在: {template_name}，将被覆盖")
# 解析函数参数作为模板参数
import inspect
sig = inspect.signature(func)
args = []
for param_name, param in sig.parameters.items():
# 忽略self和cls参数
if param_name in ('self', 'cls'):
continue
param_desc = ""
if func.__doc__:
# 尝试从文档中提取参数描述
pattern = rf":param {param_name}:\s*(.*?)(?:\n|$)"
match = re.search(pattern, func.__doc__)
if match:
param_desc = match.group(1).strip()
# 确定参数是否必需
required = param.default == inspect.Parameter.empty
args.append({
"name": param_name,
"description": param_desc,
"required": required
})
self.templates[template_name] = PromptTemplate(
name=template_name,
description=template_description,
fn=func,
arguments=args
)
logger.info(f"已注册模板: {template_name}")
return func
return decorator
def get_template(self, name: str) -> Optional[PromptTemplate]:
"""
获取模板
Args:
name: 模板名称
Returns:
模板对象或None
"""
return self.templates.get(name)
def render_template(self, name: str, arguments: Dict[str, Any] = None) -> Optional[Union[str, List[Message]]]:
"""
渲染模板
Args:
name: 模板名称
arguments: 模板参数
Returns:
渲染后的模板内容（字符串或消息列表）
"""
arguments = arguments or {}
template = self.get_template(name)
if not template:
logger.warning(f"模板不存在: {name}")
return None
try:
# 调用模板函数
result = template.fn(**arguments)
logger.info(f"已渲染模板: {name}")
return result
except Exception as e:
logger.error(f"渲染模板失败: {name}, 错误: {str(e)}")
return None
def list_templates(self) -> List[Dict[str, Any]]:
"""
列出所有可用模板
Returns:
模板列表
"""
return [
{
"name": template.name,
"description": template.description,
"arguments": template.arguments
}
for template in self.templates.values()
]
def save_templates(self):
"""保存模板到文件"""
try:
# 将模板序列化并保存
for name, template in self.templates.items():
# 无法直接序列化函数，只保存元数据
template_data = {
"name": template.name,
"description": template.description,
"arguments": template.arguments
}
file_path = os.path.join(self.templates_dir, f"{name}.json")
with open(file_path, 'w', encoding='utf-8') as f:
json.dump(template_data, f, ensure_ascii=False, indent=2)
logger.info(f"已保存所有模板到: {self.templates_dir}")
except Exception as e:
logger.error(f"保存模板失败: {str(e)}")
def load_templates(self):
"""从文件加载模板"""
if not os.path.exists(self.templates_dir):
logger.warning(f"模板目录不存在: {self.templates_dir}")
return
try:
# 加载所有模板文件
for filename in os.listdir(self.templates_dir):
if filename.endswith('.json'):
file_path = os.path.join(self.templates_dir, filename)
with open(file_path, 'r', encoding='utf-8') as f:
template_data = json.load(f)
# 注意：无法加载函数实现，仅加载元数据
logger.info(f"已加载模板元数据: {template_data['name']}")
logger.info(f"已加载所有模板")
except Exception as e:
logger.error(f"加载模板失败: {str(e)}")
# 创建全局模板管理器实例
template_manager = TemplateManager()
# 与资源系统集成的渲染函数
async def render_template_with_resources(
template_name: str,
context: Dict[str, Any],
resource_uris: Dict[str, str],
resource_manager=None
) -> Union[str, List[Message]]:
"""
渲染包含资源引用的模板
Args:
template_name: 模板名称
context: 渲染上下文
resource_uris: 资源URI映射，格式为 {context_key: resource_uri}
resource_manager: 资源管理器实例（可选）
Returns:
渲染后的模板内容
"""
if resource_manager is None:
# 尝试导入资源管理器
try:
from mcp_resources import resource_manager
except ImportError:
logger.error("无法导入资源管理器")
return None
# 扩展上下文，添加资源内容
extended_context = context.copy()
for context_key, uri in resource_uris.items():
try:
# 获取资源内容
resource_content = await resource_manager.get_resource(uri)
extended_context[context_key] = resource_content
except Exception as e:
logger.error(f"获取资源失败: {uri}, 错误: {str(e)}")
extended_context[context_key] = f"错误: 无法加载资源 {uri}"
# 渲染模板
return template_manager.render_template(template_name, extended_context)
# 示例模板
@template_manager.register_template(
name="network_diagnosis",
description="网络诊断提示模板"
)
def network_diagnosis_prompt(device_name: str, issue_description: str, verbose: bool = False) -> str:
"""
创建网络诊断提示
:param device_name: 设备名称
:param issue_description: 问题描述
:param verbose: 是否生成详细输出
"""
base_prompt = f"""
请针对设备 {device_name} 的以下问题进行诊断：
问题描述: {issue_description}
请提供可能的根本原因和解决方案。
"""
if verbose:
base_prompt += """
请详细说明诊断思路和解决步骤，包括：
1. 初步分析和可能原因
2. 建议的排查命令
3. 可能的解决方案，按优先级排序
4. 预防此类问题的最佳实践
"""
return base_prompt
@template_manager.register_template()
def device_configuration(device_type: str, purpose: str) -> List[Message]:
"""
创建设备配置会话
:param device_type: 设备类型（路由器/交换机/防火墙）
:param purpose: 配置目的（新设备/更新/安全强化）
"""
system_message = SystemMessage(
f"你是一名网络工程师，正在帮助配置一台{device_type}，目的是{purpose}。"
)
user_message = UserMessage(
f"我需要为{purpose}配置一台{device_type}，请提供配置步骤和命令。"
)
assistant_message = AssistantMessage(
f"我很乐意帮助您配置{device_type}。让我们开始：\n\n"
f"首先，让我了解一下您的网络环境和具体需求..."
)
return [system_message, user_message, assistant_message]
# 网络设备专用模板
@template_manager.register_template(
name="device_diagnosis",
description="网络设备诊断模板"
)
def device_diagnosis_template(device_info: Dict[str, Any], interfaces_info: str = None, logs: str = None) -> List[Message]:
"""
网络设备诊断提示模板，用于分析设备状态和接口信息
:param device_info: 设备基本信息对象
:param interfaces_info: 接口详细信息
:param logs: 设备日志信息
"""
system_prompt = f"""你是一位专业的网络工程师，负责诊断{device_info.get('vendor', '未知厂商')}的{device_info.get('device_type', '网络设备')}。
请根据提供的设备信息、接口状态和日志，进行专业分析并提供诊断报告。
诊断需要包括：
1. 设备基本状态总结
2. 接口状态分析
3. 潜在问题识别
4. 具体解决方案建议
5. 后续优化建议"""
# 格式化设备信息
device_details = f"""
设备基本信息:
- 名称: {device_info.get('name', 'N/A')}
- IP地址: {device_info.get('ip_address', 'N/A')}
- 型号: {device_info.get('model', 'N/A')}
- 操作系统: {device_info.get('os_version', 'N/A')}
- 状态: {device_info.get('status', 'N/A')}
- 位置: {device_info.get('location', 'N/A')}
- 描述: {device_info.get('description', 'N/A')}
"""
user_prompt = f"""请帮我诊断以下网络设备的状态：
{device_details}
"""
if interfaces_info:
user_prompt += f"""
接口信息:
{interfaces_info}
"""
if logs:
user_prompt += f"""
设备日志:
{logs}
"""
user_prompt += """
请提供详细的诊断报告，包括问题分析和解决方案。"""
return [
SystemMessage(system_prompt),
UserMessage(user_prompt)
]
@template_manager.register_template(
name="config_review",
description="网络设备配置审查模板"
)
def config_review_template(device_info: Dict[str, Any], config_content: str, focus_area: str = None) -> List[Message]:
"""
网络配置审查提示模板，用于审查设备配置并提供改进建议
:param device_info: 设备基本信息对象
:param config_content: 设备配置内容
:param focus_area: 审查重点领域（安全/性能/路由/其他）
"""
# 根据厂商确定专业技能
vendor = device_info.get('vendor', '').lower()
if 'cisco' in vendor:
vendor_expertise = "思科IOS/IOS-XE/NX-OS配置专家"
elif 'huawei' in vendor:
vendor_expertise = "华为VRP配置专家"
elif 'juniper' in vendor:
vendor_expertise = "Juniper JUNOS配置专家"
elif 'h3c' in vendor:
vendor_expertise = "H3C Comware配置专家"
else:
vendor_expertise = "多厂商网络设备配置专家"
system_prompt = f"""你是一名{vendor_expertise}，拥有多年网络配置审查经验。
请对提供的设备配置进行全面审查并提供专业的改进建议。
审查需要包括：
1. 总体配置评估
2. 配置最佳实践符合度
3. 安全风险识别
4. 潜在性能优化点
5. 具体改进建议（包括命令示例）"""
# 格式化设备信息
device_details = f"""
设备信息:
- 名称: {device_info.get('name', 'N/A')}
- 型号: {device_info.get('model', 'N/A')}
- 操作系统: {device_info.get('os_version', 'N/A')}
- 厂商: {device_info.get('vendor', 'N/A')}
"""
user_prompt = f"""请审查以下网络设备的配置：
{device_details}
配置内容:
```
{config_content}
```
"""
if focus_area:
user_prompt += f"""
特别关注以下方面: {focus_area}
"""
user_prompt += """
请提供详细的配置审查报告，包括问题分析和具体改进建议。"""
return [
SystemMessage(system_prompt),
UserMessage(user_prompt)
]
@template_manager.register_template(
name="route_analysis",
description="网络路由分析模板"
)
def route_analysis_template(device_info: Dict[str, Any], routes_output: str, target_network: str = None) -> str:
"""
路由分析提示模板，用于分析路由表并提供优化建议
:param device_info: 设备基本信息对象
:param routes_output: 路由表输出
:param target_network: 目标网络（可选）
"""
prompt = f"""
# 路由分析任务
## 设备信息
- 名称: {device_info.get('name', 'N/A')}
- IP地址: {device_info.get('ip_address', 'N/A')}
- 设备类型: {device_info.get('device_type', 'N/A')}
- 厂商: {device_info.get('vendor', 'N/A')}
- 操作系统: {device_info.get('os_version', 'N/A')}
## 路由表
```
{routes_output}
```
请作为网络工程师分析上述路由表，包括：
1. 默认路由检查
2. 路由类型分布（静态/动态）统计
3. 路由协议使用情况
4. 潜在的路由问题识别
5. 路由优化建议
"""
if target_network:
prompt += f"""
## 特定网络分析
请特别分析到 {target_network} 的路由路径，包括下一跳、出接口、管理距离和度量值。
"""
return prompt
@template_manager.register_template(
name="vlan_config",
description="VLAN配置生成模板"
)
def vlan_config_template(
device_vendor: str,
vlan_id: int,
vlan_name: str,
interfaces: List[str] = None,
ip_address: str = None,
vlan_purpose: str = None
) -> str:
"""
VLAN配置生成提示模板
:param device_vendor: 设备厂商（cisco/huawei/h3c/juniper等）
:param vlan_id: VLAN ID
:param vlan_name: VLAN名称
:param interfaces: 要添加到VLAN的接口列表
:param ip_address: VLAN接口IP地址（可选）
:param vlan_purpose: VLAN用途描述（可选）
"""
# 标准化厂商名称
vendor = device_vendor.lower()
prompt = f"""
# {vlan_name} (VLAN {vlan_id}) 配置生成
请为{device_vendor}设备生成配置VLAN {vlan_id}的完整命令序列，包括：
1. 创建VLAN {vlan_id}
2. 设置VLAN名称为"{vlan_name}"
"""
if vlan_purpose:
prompt += f"3. VLAN用途：{vlan_purpose}\n"
if interfaces:
prompt += f"\n## 需要添加到VLAN {vlan_id}的接口：\n"
for interface in interfaces:
prompt += f"- {interface}\n"
if ip_address:
prompt += f"\n## 配置VLAN接口IP地址：{ip_address}\n"
prompt += f"""
请提供适用于{device_vendor}设备的完整配置命令，包括进入配置模式和退出保存的完整步骤。如果有多种方式，请提供最佳实践方案。
"""
return prompt
@template_manager.register_template(
name="network_topology",
description="网络拓扑分析模板"
)
def network_topology_template(devices: List[Dict[str, Any]], connections: List[Dict[str, Any]]) -> List[Message]:
"""
网络拓扑分析提示模板
:param devices: 设备列表
:param connections: 连接列表
"""
# 格式化设备信息
devices_text = "设备清单:\n"
for i, device in enumerate(devices, 1):
devices_text += f"{i}. {device.get('name', 'Unknown')} ({device.get('ip_address', 'N/A')}) - {device.get('device_type', 'N/A')} - {device.get('vendor', 'N/A')}\n"
# 格式化连接信息
connections_text = "连接清单:\n"
for i, conn in enumerate(connections, 1):
source = conn.get('source_device', 'Unknown')
target = conn.get('target_device', 'Unknown')
src_if = conn.get('source_interface', 'N/A')
tgt_if = conn.get('target_interface', 'N/A')
connections_text += f"{i}. {source} ({src_if}) <--> {target} ({tgt_if})\n"
system_prompt = """你是一名网络拓扑专家，擅长分析网络设计和提供优化建议。
请分析提供的网络拓扑信息，包括设备和连接，然后：
1. 识别网络架构模式（例如核心-汇聚-接入、叶脊等）
2. 评估当前拓扑的优缺点
3. 识别潜在的单点故障
4. 提供拓扑优化建议
5. 讨论可扩展性考虑"""
user_prompt = f"""请分析以下网络拓扑并提供专业评估：
{devices_text}
{connections_text}
请提供详细的拓扑分析报告，包括架构评估、风险点和优化建议。"""
return [
SystemMessage(system_prompt),
UserMessage(user_prompt)
]
@template_manager.register_template(
name="network_security",
description="网络安全评估模板"
)
def network_security_template(
device_info: Dict[str, Any],
config_content: str,
acl_content: str = None,
security_focus: str = None
) -> str:
"""
网络安全评估提示模板
:param device_info: 设备基本信息
:param config_content: 设备配置内容
:param acl_content: 访问控制列表内容（可选）
:param security_focus: 安全评估重点（可选）
"""
prompt = f"""
# 网络安全评估
## 设备信息
- 名称: {device_info.get('name', 'N/A')}
- 类型: {device_info.get('device_type', 'N/A')}
- 厂商: {device_info.get('vendor', 'N/A')}
- 操作系统: {device_info.get('os_version', 'N/A')}
## 设备配置
```
{config_content}
```
"""
if acl_content:
prompt += f"""
## 访问控制列表
```
{acl_content}
```
"""
if security_focus:
prompt += f"""
## 评估重点
{security_focus}
"""
prompt += """
请作为网络安全专家，评估上述配置的安全性，包括：
1. 安全风险识别和严重级别评估
2. 是否符合网络安全最佳实践
3. 身份验证和访问控制评估
4. 加密和安全协议使用情况
5. 详细的安全强化建议（带具体命令）
6. 安全监控和审计建议
请提供全面的安全评估报告和具体改进措施。
"""
return prompt
