import sys import os import uuid import json import zipfile import markdown from dotenv import load_dotenv from openai import OpenAI from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, truncate_filename ) from knowledge_storm.lm import OpenAIModel, DeepSeekModel from knowledge_storm.rm import ( SerperRM ) from dependencies.helper import ( create_manifest, read_txt_file, read_json_file, split_markdown_sections ) load_dotenv() OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") APP_URL = os.getenv("APP_URL") OUTPUT_FOLDER = os.getenv("OUTPUT_FOLDER", "./result") working_dir = "imscc_temp" content_dir = "wiki_content" gpt_35_model_name = "gpt-3.5-turbo" gpt_4_model_name = "gpt-4o" class StormRunner: def __init__(self, topic: str): # self.course_id = course_id self.topic = topic self.article_title = truncate_filename( topic.replace(" ", "_").replace("/", "_") ) self.article_dir = os.path.join(working_dir, self.article_title) self.wiki_content_dir = os.path.join(self.article_dir, content_dir) os.makedirs(working_dir, exist_ok=True) os.makedirs(self.wiki_content_dir, exist_ok=True) file_path = os.path.join(OUTPUT_FOLDER, self.article_title) self.polished_article_path = os.path.join(file_path, "storm_gen_article_polished.txt") self.output_path = os.path.join(OUTPUT_FOLDER, f"{self.article_title}.imscc") self.url_to_info = os.path.join(file_path, "url_to_info.json") if os.getenv('LLM_ENGINE') == "DEEPSEEK": lm_configs = self.deepseek_info() else: lm_configs = self.openai_info() engine_args = STORMWikiRunnerArguments( output_dir=OUTPUT_FOLDER ) rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) self.runner = STORMWikiRunner(engine_args,lm_configs,rm) def openai_info(self): openai_kwargs = { "api_key": os.getenv("OPENAI_API_KEY"), "temperature": 1.0, "top_p": 0.9, } lm_configs = STORMWikiLMConfigs() ModelClass = OpenAIModel conv_simulator_lm = ModelClass( model=gpt_35_model_name, max_tokens=500, **openai_kwargs ) question_asker_lm = ModelClass( model=gpt_35_model_name, max_tokens=500, **openai_kwargs ) outline_gen_lm = ModelClass(model=gpt_4_model_name, max_tokens=400, **openai_kwargs) article_gen_lm = ModelClass(model=gpt_4_model_name, max_tokens=700, **openai_kwargs) article_polish_lm = ModelClass( model=gpt_4_model_name, max_tokens=4000, **openai_kwargs ) lm_configs.set_conv_simulator_lm(conv_simulator_lm) lm_configs.set_question_asker_lm(question_asker_lm) lm_configs.set_outline_gen_lm(outline_gen_lm) lm_configs.set_article_gen_lm(article_gen_lm) lm_configs.set_article_polish_lm(article_polish_lm) model = OpenAIModel(api_key=OPENAI_API_KEY, model=gpt_35_model_name) return lm_configs def deepseek_info(self): lm_configs = STORMWikiLMConfigs() deepseek_kwargs = { "api_key": os.getenv("DEEPSEEK_API_KEY"), "api_base": os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com"), "temperature": 1.0, "top_p": 0.9, } # DeepSeek offers two main models: 'deepseek-chat' for general tasks and 'deepseek-coder' for coding tasks # Users can choose the appropriate model based on their needs conv_simulator_lm = DeepSeekModel( model=os.getenv("DEEPSEEK_MODEL"), max_tokens=500, **deepseek_kwargs ) question_asker_lm = DeepSeekModel( model=os.getenv("DEEPSEEK_MODEL"), max_tokens=500, **deepseek_kwargs ) outline_gen_lm = DeepSeekModel(model=os.getenv("DEEPSEEK_MODEL"), max_tokens=400, **deepseek_kwargs) article_gen_lm = DeepSeekModel(model=os.getenv("DEEPSEEK_MODEL"), max_tokens=700, **deepseek_kwargs) article_polish_lm = DeepSeekModel( model=os.getenv("DEEPSEEK_MODEL"), max_tokens=4000, **deepseek_kwargs ) lm_configs.set_conv_simulator_lm(conv_simulator_lm) lm_configs.set_question_asker_lm(question_asker_lm) lm_configs.set_outline_gen_lm(outline_gen_lm) lm_configs.set_article_gen_lm(article_gen_lm) lm_configs.set_article_polish_lm(article_polish_lm) return lm_configs def process_article(self): # Generate the article result = self.runner.run( topic=self.topic, ) file_path = result imscc_file_name = os.path.join(file_path, f"{self.article_title}.imscc") imscc_file_path = self.generate_imscc() return result def generate_imscc(self): """ Generate an IMSCC package from the article content. """ article_title = self.topic self.save_sections_as_html() # Check if the polished article file exists if not os.path.exists(self.polished_article_path): raise HTTPException(status_code=404, detail="Article not found. Please generate the article first.") self.prepare_article_question() # Create the manifest create_manifest( course_title=article_title, content_dir=self.wiki_content_dir,#os.path.join(self.wiki_content_dir, content_dir), wiki_dir=content_dir, output_file=os.path.join(self.article_dir, "imsmanifest.xml") ) with zipfile.ZipFile(self.output_path, "w", zipfile.ZIP_DEFLATED) as zf: for root, _, files in os.walk(self.wiki_content_dir): for file in files: full_path = os.path.join(root, file) arcname = os.path.relpath(full_path, start=self.wiki_content_dir) zf.write(full_path, arcname) # Clean up (optional) # shutil.rmtree(working_dir) return self.output_path def save_sections_as_html(self): """ Reads a markdown file, splits it into sections, and saves each section as an HTML file. """ # os.makedirs(self.wiki_content_dir, exist_ok=True) # with open(md_path, "r", encoding="utf-8") as f: # content = f.read() content = read_txt_file(self.polished_article_path) sections = split_markdown_sections(content) for section in sections: lines = section.splitlines() if not lines: continue # Skip empty sections title_line = lines[0].strip("# ").strip() filename = f"{title_line.lower().replace(' ', '_')}.html" html = self.markdown_to_html(title_line, section) output_path = os.path.join(self.wiki_content_dir, filename) with open(output_path, "w", encoding="utf-8") as out: out.write(html) print(f"✅ Exported {len(sections)} sections to: {self.wiki_content_dir}") def generate_questions_from_article(self, num_questions: int = 10): article_text = read_txt_file(self.polished_article_path) prompt = f""" Given the following article, generate {num_questions} quiz questions including: - Descriptive - Multiple choice (with 4 options) - Single choice - Fill in the blank - Optional (True/False) Format each question like: Type: Question: Options: (only if applicable) Answer: Article: \"\"\" {article_text} \"\"\" """ openai_client=OpenAI() response = openai_client.chat.completions.create( model="gpt-4", messages=[{"role": "user", "content": prompt}], temperature=0.7 ) return response.choices[0].message.content def prepare_article_question(self): """ Read the polished article content and generate questions from it. """ article_content = read_txt_file(self.polished_article_path) # Generate questions from the article content questions = self.generate_questions_from_article(article_content) question_filename = f"questions_{self.article_title.lower().replace(' ', '_')}.html" html = self.markdown_to_html(self.article_title, questions) output_path = os.path.join(self.wiki_content_dir, question_filename) with open(output_path, "w", encoding="utf-8") as out: out.write(html) def get_article_content(self) -> str: """ Read the polished article content and return it as a string. """ if not os.path.exists(self.polished_article_path): raise HTTPException(status_code=404, detail="Article not found. Please generate the article first.") return read_txt_file(self.polished_article_path) def get_url_to_info(self) -> dict: """ Read the URL to info mapping from the JSON file. """ if not os.path.exists(self.url_to_info): raise HTTPException(status_code=404, detail="URL to info mapping not found. Please generate the article first.") return read_json_file(self.url_to_info) def markdown_to_html(self, title: str, content: str) -> str: body_html = markdown.markdown(content, extensions=['extra', 'codehilite', 'tables', 'toc']) return f""" {title} {body_html} """