from knowledge_storm import STORMWikiRunner, STORMWikiLMConfigs
from knowledge_storm.lm import LitellmModel
from knowledge_storm.rm import YouRM

# 配置语言模型
lm_configs = STORMWikiLMConfigs()
gpt_4 = LitellmModel(model='gpt-4o', max_tokens=3000, api_key="your_key")
lm_configs.set_article_gen_lm(gpt_4)

# 配置检索模块
rm = YouRM(ydc_api_key="your_key", k=10)

# 创建STORM实例
runner = STORMWikiRunner(engine_args, lm_configs, rm)

# 运行生成
topic = "人工智能的发展历程"
runner.run(topic=topic, do_research=True, do_generate_article=True)

import os
from knowledge_storm import STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs
from knowledge_storm.lm import LitellmModel
from knowledge_storm.rm import YouRM

# 配置语言模型
lm_configs = STORMWikiLMConfigs()
openai_kwargs = {
    'api_key': os.getenv("OPENAI_API_KEY"),
    'temperature': 1.0,
    'top_p': 0.9,
}

# 配置不同组件的模型
gpt_35 = LitellmModel(model='gpt-3.5-turbo', max_tokens=500, **openai_kwargs)
gpt_4 = LitellmModel(model='gpt-4o', max_tokens=3000, **openai_kwargs)

lm_configs.set_conv_simulator_lm(gpt_35)  # 对话模拟器
lm_configs.set_question_asker_lm(gpt_35)  # 问题生成器
lm_configs.set_outline_gen_lm(gpt_4)      # 大纲生成器
lm_configs.set_article_gen_lm(gpt_4)      # 文章生成器

# 配置检索模块
rm = YouRM(ydc_api_key=os.getenv('YDC_API_KEY'), k=10)

# 创建STORM运行器
engine_args = STORMWikiRunnerArguments()
runner = STORMWikiRunner(engine_args, lm_configs, rm)

from knowledge_storm.collaborative_storm.engine import CollaborativeStormLMConfigs, CoStormRunner
from knowledge_storm.lm import LitellmModel
from knowledge_storm.rm import BingSearch

# 配置Co-STORM语言模型
lm_config = CollaborativeStormLMConfigs()
openai_kwargs = {
    "api_key": os.getenv("OPENAI_API_KEY"),
    "temperature": 1.0,
    "top_p": 0.9,
}

# 配置各种AI代理
question_answering_lm = LitellmModel(model='gpt-4o', max_tokens=1000, **openai_kwargs)
discourse_manage_lm = LitellmModel(model='gpt-4o', max_tokens=500, **openai_kwargs)
utterance_polishing_lm = LitellmModel(model='gpt-4o', max_tokens=2000, **openai_kwargs)

lm_config.set_question_answering_lm(question_answering_lm)
lm_config.set_discourse_manage_lm(discourse_manage_lm)
lm_config.set_utterance_polishing_lm(utterance_polishing_lm)

# 配置检索模块
bing_rm = BingSearch(bing_search_api_key=os.environ.get("BING_SEARCH_API_KEY"), k=10)

# 创建Co-STORM运行器
costorm_runner = CoStormRunner(lm_config=lm_config, rm=bing_rm)