139 lines
4.3 KiB
Python
139 lines
4.3 KiB
Python
"""
|
||
A2A Search Agent 配置模块
|
||
|
||
支持用户传入密钥和模型名称，同时支持从环境变量获取
|
||
"""
|
||
import os
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional
|
||
from dotenv import load_dotenv
|
||
|
||
# 加载环境变量
|
||
load_dotenv()
|
||
|
||
|
||
@dataclass
|
||
class LiteLLMConfig:
|
||
"""LiteLLM 配置"""
|
||
# 基础URL - 用户提供的LiteLLM服务地址
|
||
# 优先使用 LLM_BASE_URL（与API格式保持一致），也支持 LITELLM_BASE_URL（向后兼容）
|
||
base_url: str = field(default_factory=lambda: os.getenv(
|
||
"LLM_BASE_URL"
|
||
) or os.getenv(
|
||
"LITELLM_BASE_URL",
|
||
"https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io"
|
||
))
|
||
|
||
# 完整的chat completions端点
|
||
chat_endpoint: str = field(init=False)
|
||
|
||
# API密钥 - 优先使用传入的，否则从环境变量获取
|
||
api_key: Optional[str] = None
|
||
|
||
# 模型名称 - 优先使用传入的，否则从环境变量获取
|
||
model: Optional[str] = None
|
||
|
||
# 请求超时时间(秒)
|
||
timeout: int = 120
|
||
|
||
# 温度参数
|
||
temperature: float = 0.7
|
||
|
||
# 最大token数
|
||
max_tokens: int = 4096
|
||
|
||
def __post_init__(self):
|
||
self.chat_endpoint = f"{self.base_url}/chat/completions"
|
||
|
||
# 从环境变量读取（如果未直接提供）
|
||
# API密钥：优先使用 LITELLM_API_KEY（LiteLLM约定），也支持 LLM_API_KEY（向后兼容）
|
||
if self.api_key is None:
|
||
self.api_key = os.getenv("LITELLM_API_KEY") or os.getenv("LLM_API_KEY")
|
||
if self.model is None:
|
||
# 支持多种环境变量名称：
|
||
# 1. MODEL_NAME - 与API格式保持一致（优先）
|
||
# 2. LLM_MODEL - AKS部署配置使用（必须支持）
|
||
# 3. LITELLM_MODEL - LiteLLM标准约定（向后兼容）
|
||
self.model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL", "gpt-4")
|
||
|
||
def validate(self) -> bool:
|
||
"""验证配置是否完整"""
|
||
if not self.api_key:
|
||
raise ValueError("API密钥未设置! 请设置 LITELLM_API_KEY 或 LLM_API_KEY 环境变量或直接传入 api_key")
|
||
if not self.model:
|
||
raise ValueError("模型名称未设置! 请设置 MODEL_NAME 环境变量或直接传入 model")
|
||
return True
|
||
|
||
|
||
@dataclass
|
||
class AgentConfig:
|
||
"""Agent 配置"""
|
||
# Agent名称
|
||
name: str = "search-agent"
|
||
|
||
# Agent描述
|
||
description: str = "智能AI搜索Agent，基于LiteLLM和A2A协议，能够理解用户查询意图、自动规划搜索策略、从多个来源获取信息，并生成高质量、有来源引用的答案"
|
||
|
||
# Agent版本
|
||
version: str = "1.0.0"
|
||
|
||
# 服务端口
|
||
port: int = 8080
|
||
|
||
# 服务主机
|
||
host: str = "0.0.0.0"
|
||
|
||
# 是否启用流式响应
|
||
enable_streaming: bool = True
|
||
|
||
# 系统提示词
|
||
system_prompt: str = "你是一个智能搜索助手。"
|
||
|
||
|
||
@dataclass
|
||
class A2AConfig:
|
||
"""A2A协议配置"""
|
||
# A2A协议版本
|
||
protocol_version: str = "1.0"
|
||
|
||
# Agent Card配置
|
||
agent_card: dict = field(default_factory=lambda: {
|
||
"name": "search-agent",
|
||
"description": "智能AI搜索Agent，支持A2A协议通信",
|
||
"version": "1.0.0",
|
||
"capabilities": {
|
||
"text": True,
|
||
"streaming": True,
|
||
"push_notifications": False
|
||
},
|
||
"skills": [
|
||
{
|
||
"id": "intelligent-search",
|
||
"name": "智能搜索",
|
||
"description": "理解用户查询意图，自动规划搜索策略，从多个来源获取信息"
|
||
}
|
||
]
|
||
})
|
||
|
||
|
||
def get_config(
|
||
api_key: Optional[str] = None,
|
||
model: Optional[str] = None
|
||
) -> tuple[LiteLLMConfig, AgentConfig, A2AConfig]:
|
||
"""
|
||
获取完整配置
|
||
|
||
Args:
|
||
api_key: LiteLLM API密钥（可选，优先使用，否则从环境变量获取）
|
||
model: 模型名称（可选，优先使用，否则从环境变量获取）
|
||
|
||
Returns:
|
||
(LiteLLMConfig, AgentConfig, A2AConfig) 配置元组
|
||
"""
|
||
litellm_config = LiteLLMConfig(api_key=api_key, model=model)
|
||
agent_config = AgentConfig()
|
||
a2a_config = A2AConfig()
|
||
|
||
return litellm_config, agent_config, a2a_config
|
||
|