panda/bettafish-company
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Different types of base models adapted for each agent.

Browse Source
This commit is contained in:
戒酒的李白
2025-08-23 20:19:57 +08:00
parent f448ddd466
commit 339545f7fb
16 changed files with 1072 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
MediaEngine/agent.py
+6 -1
View File
@@ -9,7 +9,7 @@ import re
from datetime import datetime
from typing import Optional, Dict, Any, List
from .llms import DeepSeekLLM, OpenAILLM, BaseLLM
from .llms import DeepSeekLLM, OpenAILLM, GeminiLLM, BaseLLM
from .nodes import (
ReportStructureNode,
FirstSearchNode,
@@ -67,6 +67,11 @@ class DeepSearchAgent:
api_key=self.config.openai_api_key,
model_name=self.config.openai_model
)
elif self.config.default_llm_provider == "gemini":
return GeminiLLM(
api_key=self.config.gemini_api_key,
model_name=self.config.gemini_model
)
else:
raise ValueError(f"不支持的LLM提供商: {self.config.default_llm_provider}")
MediaEngine/llms/__init__.py
+2 -1
View File
@@ -6,5 +6,6 @@ LLM调用模块
from .base import BaseLLM
from .deepseek import DeepSeekLLM
from .openai_llm import OpenAILLM
from .gemini_llm import GeminiLLM
__all__ = ["BaseLLM", "DeepSeekLLM", "OpenAILLM"]
__all__ = ["BaseLLM", "DeepSeekLLM", "OpenAILLM", "GeminiLLM"]
MediaEngine/llms/gemini_llm.py
+95
View File
@@ -0,0 +1,95 @@
"""
Gemini LLM实现
使用Gemini 2.5-pro中转API进行文本生成
"""
import os
from typing import Optional, Dict, Any
from openai import OpenAI
from .base import BaseLLM
class GeminiLLM(BaseLLM):
"""Gemini LLM实现类"""
def __init__(self, api_key: Optional[str] = None, model_name: Optional[str] = None):
"""
初始化Gemini客户端
Args:
api_key: Gemini API密钥，如果不提供则从环境变量读取
model_name: 模型名称，默认使用gemini-2.5-pro
"""
if api_key is None:
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise ValueError("Gemini API Key未找到！请设置GEMINI_API_KEY环境变量或在初始化时提供")
super().__init__(api_key, model_name)
# 初始化OpenAI客户端，使用Gemini的中转endpoint
self.client = OpenAI(
api_key=self.api_key,
base_url="https://www.chataiapi.com/v1"
)
self.default_model = model_name or self.get_default_model()
def get_default_model(self) -> str:
"""获取默认模型名称"""
return "gemini-2.5-pro"
def invoke(self, system_prompt: str, user_prompt: str, **kwargs) -> str:
"""
调用Gemini API生成回复
Args:
system_prompt: 系统提示词
user_prompt: 用户输入
**kwargs: 其他参数，如temperature、max_tokens等
Returns:
Gemini生成的回复文本
"""
try:
# 构建消息
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
# 设置默认参数
params = {
"model": self.default_model,
"messages": messages,
"temperature": kwargs.get("temperature", 0.7),
"max_tokens": kwargs.get("max_tokens", 4000),
"stream": False
}
# 调用API
response = self.client.chat.completions.create(**params)
# 提取回复内容
if response.choices and response.choices[0].message:
content = response.choices[0].message.content
return self.validate_response(content)
else:
return ""
except Exception as e:
print(f"Gemini API调用错误: {str(e)}")
raise e
def get_model_info(self) -> Dict[str, Any]:
"""
获取当前模型信息
Returns:
模型信息字典
"""
return {
"provider": "Gemini",
"model": self.default_model,
"api_base": "https://www.chataiapi.com/v1"
}
MediaEngine/utils/config.py
+11 -3
View File
@@ -14,12 +14,14 @@ class Config:
# API密钥
deepseek_api_key: Optional[str] = None
openai_api_key: Optional[str] = None
gemini_api_key: Optional[str] = None
bocha_api_key: Optional[str] = None
# 模型配置
default_llm_provider: str = "deepseek" # deepseek 或 openai
default_llm_provider: str = "deepseek" # deepseek、openai 或 gemini
deepseek_model: str = "deepseek-chat"
openai_model: str = "gpt-4o-mini"
gemini_model: str = "gemini-2.5-pro"
# 搜索配置
search_timeout: int = 240
@@ -44,6 +46,10 @@ class Config:
print("错误: OpenAI API Key未设置")
return False
if self.default_llm_provider == "gemini" and not self.gemini_api_key:
print("错误: Gemini API Key未设置")
return False
if not self.bocha_api_key:
print("错误: Bocha API Key未设置")
return False
@@ -65,11 +71,12 @@ class Config:
return cls(
deepseek_api_key=getattr(config_module, "DEEPSEEK_API_KEY", None),
openai_api_key=getattr(config_module, "OPENAI_API_KEY", None),
gemini_api_key=getattr(config_module, "GEMINI_API_KEY", None),
bocha_api_key=getattr(config_module, "BOCHA_API_KEY", None),
default_llm_provider=getattr(config_module, "DEFAULT_LLM_PROVIDER", "deepseek"),
deepseek_model=getattr(config_module, "DEEPSEEK_MODEL", "deepseek-chat"),
openai_model=getattr(config_module, "OPENAI_MODEL", "gpt-4o-mini"),
gemini_model=getattr(config_module, "GEMINI_MODEL", "gemini-2.5-pro"),
search_timeout=getattr(config_module, "SEARCH_TIMEOUT", 240),
max_content_length=getattr(config_module, "SEARCH_CONTENT_MAX_LENGTH", 20000),
max_reflections=getattr(config_module, "MAX_REFLECTIONS", 2),
@@ -92,11 +99,12 @@ class Config:
return cls(
deepseek_api_key=config_dict.get("DEEPSEEK_API_KEY"),
openai_api_key=config_dict.get("OPENAI_API_KEY"),
gemini_api_key=config_dict.get("GEMINI_API_KEY"),
bocha_api_key=config_dict.get("BOCHA_API_KEY"),
default_llm_provider=config_dict.get("DEFAULT_LLM_PROVIDER", "deepseek"),
deepseek_model=config_dict.get("DEEPSEEK_MODEL", "deepseek-chat"),
openai_model=config_dict.get("OPENAI_MODEL", "gpt-4o-mini"),
gemini_model=config_dict.get("GEMINI_MODEL", "gemini-2.5-pro"),
search_timeout=int(config_dict.get("SEARCH_TIMEOUT", "240")),
max_content_length=int(config_dict.get("SEARCH_CONTENT_MAX_LENGTH", "20000")),
max_reflections=int(config_dict.get("MAX_REFLECTIONS", "2")),